scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

feat(mcp): startup backfill stamping oauth2_flow on legacy null rows (#32290)

Browse source 
* feat(mcp): startup backfill stamping oauth2_flow on legacy null rows

Rows created before the write-side stamps carry a null oauth2_flow and rely on
read-time field-shape inference, which cannot tell a DCR-registered interactive
server (client creds + token_url, no persisted authorization_url) from an M2M
server unless endpoint discovery succeeds first; on a transient discovery failure
those servers flip to client_credentials for that registry load

The backfill classifies each null oauth2 row once, at rest, ordered by signal
strength: per-user token rows (only the interactive flow mints them, so this is
definitive and catches the DCR-trap cohort), then a persisted authorization_url,
then a persisted registration_url (DCR implies interactive; this covers
registered-but-never-signed-in rows), then the M2M credential shape mirroring the
legacy inference, else the interactive default that matches how
needs_user_oauth_token treats a null flow. Every stamp is logged with the rule
that fired and written with updated_by=oauth2_flow_backfill for auditability

Runs in _init_mcp_servers_in_db before the registry load so the first build of
the boot classifies from the column, is isolated so a failure cannot block server
loading, and is idempotent: a healed fleet exits after one indexed query. This
unblocks deleting the read-time inference for DB rows in the follow-up

Third step of the oauth2_flow persistence sequence, after #32283 and #32288

* fix(mcp): backfill leaves the ambiguous M2M shape unstamped instead of guessing client_credentials

The credential shape (client_id + client_secret + token_url, no interactive signal) is
shared by real M2M servers and DCR-registered interactive servers nobody has signed
into: the DCR persist writes creds and token_url but not authorization_url or
registration_url. Stamping client_credentials from that shape permanently mislabeled
the interactive cohort, and once explicit the value is authoritative, so per-user
traffic would run on the proxy's stored client credential with no discovery rescue
and no backstop (it only guards null rows)

The backfill now stamps only what it can prove. Interactive signals keep stamping
authorization_code; the ambiguous shape is left null with an actionable warning naming
the server and the fix (set oauth2_flow via the dashboard or PUT /v1/mcp/server). A
true M2M row keeps working per-request through the security backstop while the warning
nags; an interactive row keeps its Authorize button (null renders interactive), and one
completed sign-in creates the per-user token that stamps it authorization_code at the
next boot. Mirrors the config-level rule: M2M is asserted by a human, never guessed

Raised by review on the PR

* perf(mcp): batch the backfill stamps into one update_many per flow value

The per-row update loop issued one DB round-trip per legacy row at startup; rows
sharing a stamped value now go out as a single update_many, so the DB cost is
constant in fleet size. Per-row logging keeps the rule that fired for each server

Raised by review on the PR

* fix(mcp): backfill stamps only rows still null at write time and counts only real OAuth token rows as sign-in proof

Two review findings. The batched update_many matched on server_id alone, so an
explicit oauth2_flow set between the backfill's read and its write (an admin PUT or
a sign-in's DCR stamp landing in the boot window) would be overwritten with the
inferred value; the where clause now also requires oauth2_flow to still be null, so
an explicit value can never be clobbered under any interleaving

And the per_user_tokens rule counted any LiteLLM_MCPUserCredentials row as proof of
an interactive sign-in, but that table doubles as BYOK storage for user-supplied API
keys; a BYOK-flavored row would have stamped an M2M-shaped server authorization_code.
The rule now counts only rows whose payload decodes as a type oauth2 token via the
existing _decode_oauth_payload discriminator, so bare keys, undecodable rows, and
stale leftovers from a BYOK-to-oauth2 auth switch prove nothing

Raised by review on the PR
This commit is contained in:
tin-berri 2026-07-06 18:42:08 -07:00 • committed by GitHub
parent 43b0a25f07
commit 4e3ebbb164
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 400 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

155
litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py Normal file
View file

@ -0,0 +1,155 @@
"""Startup backfill for oauth2 MCP server rows persisted before oauth2_flow was written.
Rows created before the write-side stamps (DCR persist, UI create, REST create) carry a
null ``oauth2_flow`` and rely on read-time field-shape inference, which cannot tell a
DCR-registered interactive server from an M2M server unless endpoint discovery succeeds
first. This backfill classifies each null row once, at rest, using signals inference
never had, and persists the result so the read path never has to infer again.
Signal order, strongest first:
1. Per-user OAuth token rows exist for the server: only the interactive flow mints
per-user tokens, so this is definitive and immune to the discovery trap. BYOK API
keys share the same table (``LiteLLM_MCPUserCredentials``), so only rows whose
payload decodes as a ``type: oauth2`` token count as proof; bare keys and
undecodable rows prove nothing about the flow.
2. ``authorization_url`` persisted: interactive needs a user-facing authorization
endpoint; M2M (RFC 6749 section 4.4) never has one.
3. ``registration_url`` persisted: dynamic client registration (RFC 7591) exists to mint
clients for the interactive flow; M2M servers are configured with static credentials.
4. ``token_url`` plus decryptable ``client_id`` and ``client_secret``: ambiguous, left
unstamped. The shape is shared by M2M servers and DCR-registered interactive servers
whose authorization endpoint lives only in discovery (registered but never signed
in), so stamping client_credentials here could permanently route per-user traffic
through the proxy's stored client credential. The row keeps working through the
request-time backstop and a warning names it with the one-line fix (set oauth2_flow
via the dashboard or ``PUT /v1/mcp/server``); a completed interactive sign-in also
heals it via rule 1 at the next boot.
5. Anything else is interactive: matching how ``needs_user_oauth_token`` treats a null
flow, so the stamp never changes runtime routing for rows no rule recognizes.
The backfill never stamps client_credentials: M2M is asserted by a human (config
requires it, the API accepts it, the dashboard sets it), mirroring the config-level
validation error. Runs before the first registry load on every boot and is idempotent:
a healed fleet has no null rows and the backfill exits after one query.
"""
import json
from collections import Counter
from typing import Any, Literal, Optional
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials
from litellm.proxy.utils import PrismaClient
from litellm.types.mcp import MCPCredentials
OAuth2Flow = Literal["client_credentials", "authorization_code"]
BackfillRule = Literal[
"per_user_tokens",
"authorization_url",
"registration_url",
"ambiguous_m2m_shape",
"interactive_default",
]
_BACKFILL_AUDIT_ACTOR = "oauth2_flow_backfill"
def _decrypted_credentials(raw_credentials: Any) -> Optional[MCPCredentials]:
if raw_credentials is None:
return None
if isinstance(raw_credentials, str):
try:
parsed = json.loads(raw_credentials)
except (ValueError, TypeError):
return None
else:
parsed = raw_credentials
if not isinstance(parsed, dict):
return None
return decrypt_credentials(credentials=dict(parsed))
def classify_null_flow_row(
*,
has_per_user_tokens: bool,
authorization_url: Optional[str],
registration_url: Optional[str],
token_url: Optional[str],
credentials: Optional[MCPCredentials],
) -> tuple[Optional[OAuth2Flow], BackfillRule]:
if has_per_user_tokens:
return "authorization_code", "per_user_tokens"
if authorization_url:
return "authorization_code", "authorization_url"
if registration_url:
return "authorization_code", "registration_url"
if token_url and credentials and credentials.get("client_id") and credentials.get("client_secret"):
return None, "ambiguous_m2m_shape"
return "authorization_code", "interactive_default"
async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]:
"""Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable
ones, warn on the ambiguous ones, and return counts per rule."""
null_rows: list[Any] = await prisma_client.db.litellm_mcpservertable.find_many(
where={"auth_type": "oauth2", "oauth2_flow": None},
)
if not null_rows:
return {}
server_ids = [row.server_id for row in null_rows]
token_rows: list[Any] = await prisma_client.db.litellm_mcpusercredentials.find_many(
where={"server_id": {"in": server_ids}},
)
server_ids_with_oauth_tokens: set[str] = {
token_row.server_id for token_row in token_rows if _decode_oauth_payload(token_row.credential_b64) is not None
}
classified = tuple(
(
row,
classify_null_flow_row(
has_per_user_tokens=row.server_id in server_ids_with_oauth_tokens,
authorization_url=row.authorization_url,
registration_url=row.registration_url,
token_url=row.token_url,
credentials=_decrypted_credentials(row.credentials),
),
)
for row in null_rows
)
for row, (flow, rule) in classified:
if flow is None:
verbose_proxy_logger.warning(
"oauth2_flow backfill: server_id=%s is ambiguous (client credentials + token_url, "
"no interactive signal); left unstamped. Set oauth2_flow explicitly via the "
"dashboard or PUT /v1/mcp/server: client_credentials if this server is M2M, or "
"complete an interactive sign-in and it will be stamped authorization_code at the "
"next boot.",
row.server_id,
)
else:
verbose_proxy_logger.info(
"oauth2_flow backfill: server_id=%s stamped %s (rule=%s)",
row.server_id,
flow,
rule,
)
stamped_flows = {flow for _, (flow, _) in classified if flow is not None}
for stamped_flow in stamped_flows:
server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow]
await prisma_client.db.litellm_mcpservertable.update_many(
where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None},
data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR},
)
counts: dict[BackfillRule, int] = dict(Counter(rule for _, (_, rule) in classified))
verbose_proxy_logger.info(
"oauth2_flow backfill: processed %d oauth2 server row(s): %s",
len(null_rows),
counts,
)
return counts

11
litellm/proxy/proxy_server.py
View file

@ -6366,6 +6366,17 @@ class ProxyConfig:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.oauth2_flow_backfill import (
backfill_null_oauth2_flows,
)
try:
if prisma_client is not None:
await backfill_null_oauth2_flows(prisma_client)
except Exception as e: # noqa: BLE001
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_mcp_servers_in_db backfill - {}".format(str(e))
)
try:
await global_mcp_server_manager.reload_servers_from_database()

234
tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_flow_backfill.py Normal file
View file

@ -0,0 +1,234 @@
"""
Tests for the startup oauth2_flow backfill.
Legacy oauth2 rows with a null oauth2_flow are classified once, at rest, using
signals read-time inference never had (per-user token rows first), and the
result is persisted so the read path never infers again. The signal order is
the spec, and so is the refusal to stamp client_credentials: the M2M credential
shape is shared by DCR-registered interactive servers whose authorization
endpoint lives only in discovery, so ambiguous rows are left unstamped for a
human to assert rather than being permanently mislabeled M2M.
"""
import base64
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy._experimental.mcp_server.oauth2_flow_backfill import (
backfill_null_oauth2_flows,
classify_null_flow_row,
)
def test_classify_per_user_tokens_beat_m2m_shape():
"""The DCR trap row: creds + token_url, no authorization_url, but a user has
signed in. Tokens are definitive; the M2M shape must not win."""
flow, rule = classify_null_flow_row(
has_per_user_tokens=True,
authorization_url=None,
registration_url=None,
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
assert flow == "authorization_code"
assert rule == "per_user_tokens"
def test_classify_authorization_url_beats_m2m_shape():
flow, rule = classify_null_flow_row(
has_per_user_tokens=False,
authorization_url="https://idp.example.com/authorize",
registration_url=None,
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
assert flow == "authorization_code"
assert rule == "authorization_url"
def test_classify_registration_url_beats_m2m_shape():
"""A registration endpoint means DCR, and DCR exists to mint interactive
clients; an abandoned-DCR row (no sign-in yet) must not be stamped M2M."""
flow, rule = classify_null_flow_row(
has_per_user_tokens=False,
authorization_url=None,
registration_url="https://idp.example.com/register",
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
assert flow == "authorization_code"
assert rule == "registration_url"
def test_classify_m2m_shape_is_ambiguous_and_unstamped():
"""The M2M shape alone must never stamp client_credentials: a DCR-registered
interactive server that nobody signed into yet has the identical shape, and a
wrong M2M stamp would permanently route its per-user traffic through the
proxy's stored client credential."""
flow, rule = classify_null_flow_row(
has_per_user_tokens=False,
authorization_url=None,
registration_url=None,
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
assert flow is None
assert rule == "ambiguous_m2m_shape"
def test_classify_partial_credentials_default_interactive():
"""token_url without a full credential pair is not the M2M shape."""
flow, rule = classify_null_flow_row(
has_per_user_tokens=False,
authorization_url=None,
registration_url=None,
token_url="https://idp.example.com/token",
credentials={"client_id": "cid"},
)
assert flow == "authorization_code"
assert rule == "interactive_default"
def test_classify_bare_row_default_interactive():
flow, rule = classify_null_flow_row(
has_per_user_tokens=False,
authorization_url=None,
registration_url=None,
token_url=None,
credentials=None,
)
assert flow == "authorization_code"
assert rule == "interactive_default"
def _row(server_id, *, authorization_url=None, registration_url=None, token_url=None, credentials=None):
return SimpleNamespace(
server_id=server_id,
authorization_url=authorization_url,
registration_url=registration_url,
token_url=token_url,
credentials=credentials,
)
def _oauth_token_row(server_id):
payload = json.dumps({"type": "oauth2", "access_token": "tok", "connected_at": "2026-07-01T00:00:00Z"})
return SimpleNamespace(
server_id=server_id,
user_id="u1",
credential_b64=base64.urlsafe_b64encode(payload.encode()).decode(),
)
def _byok_key_row(server_id):
return SimpleNamespace(
server_id=server_id,
user_id="u1",
credential_b64=base64.urlsafe_b64encode(b"sk-user-supplied-upstream-key").decode(),
)
def _mock_prisma(null_rows, token_rows):
mock_prisma = MagicMock()
mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=null_rows)
mock_prisma.db.litellm_mcpservertable.update_many = AsyncMock(return_value=MagicMock())
mock_prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=token_rows)
return mock_prisma
@pytest.mark.asyncio
async def test_backfill_only_targets_null_flow_oauth2_rows():
"""The where clause is the guard that explicit and non-oauth2 rows are never touched."""
mock_prisma = _mock_prisma([], [])
counts = await backfill_null_oauth2_flows(mock_prisma)
assert counts == {}
mock_prisma.db.litellm_mcpservertable.find_many.assert_awaited_once_with(
where={"auth_type": "oauth2", "oauth2_flow": None},
)
mock_prisma.db.litellm_mcpusercredentials.find_many.assert_not_awaited()
mock_prisma.db.litellm_mcpservertable.update_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_backfill_stamps_rows_and_reports_rule_counts():
dcr_trap_row = _row(
"signed_in_dcr",
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
m2m_row = _row(
"legacy_m2m",
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
interactive_row = _row("legacy_interactive", authorization_url="https://idp.example.com/authorize")
mock_prisma = _mock_prisma(
[dcr_trap_row, m2m_row, interactive_row],
[_oauth_token_row("signed_in_dcr")],
)
counts = await backfill_null_oauth2_flows(mock_prisma)
assert counts == {"per_user_tokens": 1, "ambiguous_m2m_shape": 1, "authorization_url": 1}
mock_prisma.db.litellm_mcpservertable.update_many.assert_awaited_once()
call = mock_prisma.db.litellm_mcpservertable.update_many.await_args
assert sorted(call.kwargs["where"]["server_id"]["in"]) == ["legacy_interactive", "signed_in_dcr"]
assert "oauth2_flow" in call.kwargs["where"] and call.kwargs["where"]["oauth2_flow"] is None
assert call.kwargs["data"] == {"oauth2_flow": "authorization_code", "updated_by": "oauth2_flow_backfill"}
@pytest.mark.asyncio
async def test_backfill_handles_json_string_credentials():
"""JSON-string credential blobs must decode: the M2M shape is recognized (and
therefore deliberately left unstamped) rather than misread as credential-less."""
m2m_row = _row(
"json_creds_m2m",
token_url="https://idp.example.com/token",
credentials='{"client_id": "cid", "client_secret": "csecret"}',
)
mock_prisma = _mock_prisma([m2m_row], [])
counts = await backfill_null_oauth2_flows(mock_prisma)
assert counts == {"ambiguous_m2m_shape": 1}
mock_prisma.db.litellm_mcpservertable.update_many.assert_not_awaited()
@pytest.mark.asyncio
async def test_backfill_treats_undecodable_credentials_as_absent():
row = _row(
"corrupt_creds",
token_url="https://idp.example.com/token",
credentials="not-json",
)
mock_prisma = _mock_prisma([row], [])
counts = await backfill_null_oauth2_flows(mock_prisma)
assert counts == {"interactive_default": 1}
@pytest.mark.asyncio
async def test_backfill_byok_key_rows_are_not_sign_in_proof():
"""BYOK API keys live in the same table as per-user OAuth tokens; a bare key row
must not satisfy the per_user_tokens rule, or a BYOK-flavored M2M-shaped server
would be permanently stamped authorization_code. Only rows whose payload decodes
as a type oauth2 token count."""
byok_shaped_row = _row(
"byok_m2m_shape",
token_url="https://idp.example.com/token",
credentials={"client_id": "cid", "client_secret": "csecret"},
)
mock_prisma = _mock_prisma([byok_shaped_row], [_byok_key_row("byok_m2m_shape")])
counts = await backfill_null_oauth2_flows(mock_prisma)
assert counts == {"ambiguous_m2m_shape": 1}
mock_prisma.db.litellm_mcpservertable.update_many.assert_not_awaited()