From 3f79222350854f422ca6a52bf55e14faf0a94409 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Fri, 5 Jun 2026 15:22:52 +0200 Subject: [PATCH] fix(proxy): persist oauth2_flow on MCP server registration (#29690) --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/_types.py | 2 + litellm/proxy/schema.prisma | 1 + schema.prisma | 1 + .../test_mcp_management_endpoints.py | 52 +++++++++++++++++++ 6 files changed, 59 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260604120000_add_oauth2_flow_to_mcp_servers/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260604120000_add_oauth2_flow_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260604120000_add_oauth2_flow_to_mcp_servers/migration.sql new file mode 100644 index 0000000000..fee6926d96 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260604120000_add_oauth2_flow_to_mcp_servers/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth2_flow" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index c4754ef611..6715f80464 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -322,6 +322,7 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + oauth2_flow String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) delegate_auth_to_upstream Boolean @default(false) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6c36c813e7..6de8f0a3cd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1376,6 +1376,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None allow_all_keys: bool = False available_on_public_internet: bool = True delegate_auth_to_upstream: bool = False @@ -1449,6 +1450,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): authorization_url: Optional[str] = None token_url: Optional[str] = None registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None allow_all_keys: bool = False available_on_public_internet: bool = True delegate_auth_to_upstream: bool = False diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c4754ef611..6715f80464 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -322,6 +322,7 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + oauth2_flow String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) delegate_auth_to_upstream Boolean @default(false) diff --git a/schema.prisma b/schema.prisma index c4754ef611..6715f80464 100644 --- a/schema.prisma +++ b/schema.prisma @@ -322,6 +322,7 @@ model LiteLLM_MCPServerTable { authorization_url String? token_url String? registration_url String? + oauth2_flow String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) delegate_auth_to_upstream Boolean @default(false) diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index a5c8320a5c..7e044442e2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -3317,3 +3317,55 @@ def test_sanitize_mcp_server_for_non_admin_clears_credential_fields(): # server without exposing secrets. assert sanitized.server_id == server.server_id assert sanitized.alias == server.alias + + +def test_oauth2_flow_accepted_on_create_request(): + """NewMCPServerRequest carries oauth2_flow through to the persisted dict.""" + from litellm.proxy._experimental.mcp_server.db import _prepare_mcp_server_data + + payload = NewMCPServerRequest( + server_name="m2m-server", + url="https://example.com/mcp", + transport="http", + auth_type="oauth2", + token_url="https://idp.example.com/oauth/token", + oauth2_flow="client_credentials", + ) + data_dict = _prepare_mcp_server_data(payload) + assert data_dict["oauth2_flow"] == "client_credentials" + + +def test_oauth2_flow_round_trips_on_update_and_response_models(): + """oauth2_flow survives UpdateMCPServerRequest and the LiteLLM_MCPServerTable + response model. Before the fix these models dropped the field (no attribute), + which is why a persisted value never round-tripped.""" + from litellm.proxy._types import ( + LiteLLM_MCPServerTable, + UpdateMCPServerRequest, + ) + + update = UpdateMCPServerRequest( + server_id="srv-1", oauth2_flow="client_credentials" + ) + assert update.oauth2_flow == "client_credentials" + + row = LiteLLM_MCPServerTable( + server_id="srv-1", + transport="http", + oauth2_flow="client_credentials", + ) + assert row.oauth2_flow == "client_credentials" + + +def test_oauth2_flow_defaults_to_none_when_omitted(): + """Omitting oauth2_flow is valid and resolves to None (runtime infers it).""" + from litellm.proxy._types import ( + LiteLLM_MCPServerTable, + UpdateMCPServerRequest, + ) + + assert UpdateMCPServerRequest(server_id="srv-1").oauth2_flow is None + assert ( + LiteLLM_MCPServerTable(server_id="srv-1", transport="http").oauth2_flow + is None + )