diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306000000_add_byok_mcp_jwt_key_mapping/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306000000_add_byok_mcp_jwt_key_mapping/migration.sql new file mode 100644 index 00000000000..33c081b6a1f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260306000000_add_byok_mcp_jwt_key_mapping/migration.sql @@ -0,0 +1,49 @@ +-- Re-add spec_path (was added in 20260220, dropped in 20260224, re-added in schema) +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "spec_path" TEXT; + +-- Add BYOK MCP fields +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "tool_name_to_display_name" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "tool_name_to_description" JSONB DEFAULT '{}'; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "is_byok" BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "byok_description" TEXT[] DEFAULT ARRAY[]::TEXT[]; +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "byok_api_key_help_url" TEXT; + +-- CreateTable +CREATE TABLE "LiteLLM_MCPUserCredentials" ( + "id" TEXT NOT NULL, + "user_id" TEXT NOT NULL, + "server_id" TEXT NOT NULL, + "credential_b64" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_MCPUserCredentials_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_MCPUserCredentials_user_id_server_id_key" ON "LiteLLM_MCPUserCredentials"("user_id", "server_id"); + +-- CreateTable +CREATE TABLE "LiteLLM_JWTKeyMapping" ( + "id" TEXT NOT NULL, + "jwt_claim_name" TEXT NOT NULL, + "jwt_claim_value" TEXT NOT NULL, + "token" TEXT NOT NULL, + "description" TEXT, + "is_active" BOOLEAN NOT NULL DEFAULT true, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_JWTKeyMapping_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value"); + +-- CreateIndex +CREATE INDEX "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value", "is_active"); + +-- AddForeignKey +ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE RESTRICT ON UPDATE CASCADE; diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 5c063839304..cde932d6aac 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -1644,6 +1644,37 @@ if MCP_AVAILABLE: }, ) + async def _resolve_byok_auth( + mcp_server: MCPServer, + user_api_key_auth: Optional[UserAPIKeyAuth], + mcp_auth_header: Optional[str], + ) -> Optional[str]: + """Resolve BYOK credential for a server, returning the auth header to use.""" + if not mcp_server.is_byok: + return mcp_auth_header + if not mcp_auth_header: + byok_cred = await _get_byok_credential(mcp_server, user_api_key_auth) + if byok_cred is None: + raise HTTPException( + status_code=401, + detail={ + "error": "byok_auth_required", + "server_id": mcp_server.server_id, + "server_name": mcp_server.server_name or mcp_server.name, + "message": ( + "No stored credential found for this BYOK server. " + "Complete the OAuth authorization flow to provide your API key." + ), + }, + headers={ + "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' + }, + ) + return byok_cred + # External auth header supplied; still enforce user-identity check. + await _check_byok_credential(mcp_server, user_api_key_auth) + return mcp_auth_header + async def execute_mcp_tool( name: str, arguments: Dict[str, Any], @@ -1733,30 +1764,10 @@ if MCP_AVAILABLE: "mcp_tool_call_metadata" ] = standard_logging_mcp_tool_call - # BYOK: retrieve the stored per-user credential. A single DB call - # both checks existence and fetches the value, avoiding a double query. - if mcp_server.is_byok and not mcp_auth_header: - byok_cred = await _get_byok_credential(mcp_server, user_api_key_auth) - if byok_cred is None: - raise HTTPException( - status_code=401, - detail={ - "error": "byok_auth_required", - "server_id": mcp_server.server_id, - "server_name": mcp_server.server_name or mcp_server.name, - "message": ( - "No stored credential found for this BYOK server. " - "Complete the OAuth authorization flow to provide your API key." - ), - }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, - ) - mcp_auth_header = byok_cred - elif mcp_server.is_byok: - # External auth header supplied; still enforce user-identity check. - await _check_byok_credential(mcp_server, user_api_key_auth) + # BYOK: retrieve or validate the per-user credential + mcp_auth_header = await _resolve_byok_auth( + mcp_server, user_api_key_auth, mcp_auth_header + ) # Check if tool exists in local registry first (for OpenAPI-based tools) # These tools are registered with their prefixed names diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index b48db72a536..f7a4cec301b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -81,7 +81,6 @@ if MCP_AVAILABLE: delete_user_credential, get_all_mcp_servers_for_user, get_mcp_server, - get_user_credential, store_user_credential, update_mcp_server, ) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 55ca2452e7b..4854be02a17 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1458,6 +1458,12 @@ async def test_add_update_server_with_alias(): mock_mcp_server.available_on_public_internet = True mock_mcp_server.created_at = None mock_mcp_server.updated_at = None + # BYOK fields + mock_mcp_server.tool_name_to_display_name = None + mock_mcp_server.tool_name_to_description = None + mock_mcp_server.is_byok = False + mock_mcp_server.byok_description = [] + mock_mcp_server.byok_api_key_help_url = None # Add server to manager await test_manager.add_server(mock_mcp_server) @@ -1504,6 +1510,12 @@ async def test_add_update_server_without_alias(): mock_mcp_server.available_on_public_internet = True mock_mcp_server.created_at = None mock_mcp_server.updated_at = None + # BYOK fields + mock_mcp_server.tool_name_to_display_name = None + mock_mcp_server.tool_name_to_description = None + mock_mcp_server.is_byok = False + mock_mcp_server.byok_description = [] + mock_mcp_server.byok_api_key_help_url = None # Add server to manager await test_manager.add_server(mock_mcp_server) @@ -1550,6 +1562,12 @@ async def test_add_update_server_fallback_to_server_id(): mock_mcp_server.available_on_public_internet = True mock_mcp_server.created_at = None mock_mcp_server.updated_at = None + # BYOK fields + mock_mcp_server.tool_name_to_display_name = None + mock_mcp_server.tool_name_to_description = None + mock_mcp_server.is_byok = False + mock_mcp_server.byok_description = [] + mock_mcp_server.byok_api_key_help_url = None # Add server to manager await test_manager.add_server(mock_mcp_server) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 89543e42956..c309ea1bf68 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -1044,6 +1044,19 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True} ) + # Initialize jwt_handler with litellm_jwtauth so user_api_key_auth can + # access jwt_handler.litellm_jwtauth before auth_builder is called + from litellm.proxy.auth.handle_jwt import JWTHandler + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.caching import DualCache + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth(), + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler) + # Mock enterprise license check and JWTAuthManager.auth_builder # License check must be mocked to avoid environment variable pollution # in parallel test execution diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index a7391666cda..abda3b1b250 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -477,7 +477,6 @@ async def test_check_byok_credential_missing_credential(): with patch( "litellm.proxy._experimental.mcp_server.db.get_user_credential", new=AsyncMock(return_value=None), - ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): with pytest.raises(HTTPException) as exc_info: await _check_byok_credential(server, user_auth) @@ -511,7 +510,6 @@ async def test_check_byok_credential_has_credential(): with patch( "litellm.proxy._experimental.mcp_server.db.get_user_credential", new=AsyncMock(return_value="some-credential-value"), - ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): # Should not raise await _check_byok_credential(server, user_auth)