mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #34990 from BerriAI/litellm_fix_mcp_oauth_discovered_issuer_anchoring
fix(mcp): never write discovery results to the row, heal already-stamped rows, and retry failed discovery with backoff
This commit is contained in:
tin-berri
GitHub
commit
732364e260
8 changed files with 768 additions and 536 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -190,7 +190,7 @@ const OAuthFormFields: React.FC<OAuthFormFieldsProps> = ({
|
|||
label={
|
||||
<FieldLabel
|
||||
label="Issuer (optional)"
|
||||
tooltip="OAuth 2.0 authorization server issuer (RFC 8414). Auto-discovered from the upstream on first connect; set it explicitly to pin the trust anchor so token and scope discovery is fetched from and validated against this issuer (RFC 8414 §3.3) instead of anything the resource advertises."
|
||||
tooltip="OAuth 2.0 authorization server issuer (RFC 8414). Leave empty to discover endpoints from the upstream resource; set it to pin the trust anchor, which makes this issuer's document the only endpoint source (RFC 8414 §3.3), overriding the Authorization/Token/Registration URLs above and failing closed if its metadata cannot be fetched."
|
||||
/>
|
||||
}
|
||||
name="issuer"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue