Merge b9c53038fc into eb0e3f8c18

litellm/llms/databricks/common_utils.py

@@ -13,6 +13,7 @@ Authentication priority:
 import os
 import re
 from typing import Any, Final, Literal
+from urllib.parse import urlsplit, urlunsplit
 
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 
@@ -224,11 +225,8 @@ class DatabricksBase:
         """
         import requests
 
-        # Extract workspace URL from api_base
-        workspace_url = api_base.rstrip("/")
-        if "/serving-endpoints" in workspace_url:
-            workspace_url = workspace_url.replace("/serving-endpoints", "")
-
+        api_base_parts: Final = urlsplit(api_base)
+        workspace_url: Final = urlunsplit((api_base_parts.scheme, api_base_parts.netloc, "", "", ""))
         token_url: Final = f"{workspace_url}/oidc/v1/token"
 
         try:

tests/test_litellm/llms/databricks/test_databricks_partner_integration.py

@@ -245,6 +245,24 @@ class TestOAuthM2M:
             assert "/serving-endpoints" not in call_url
             assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
 
+    def test_oauth_m2m_strips_ai_gateway_path(self):
+        """OAuth M2M derives the token URL from the workspace origin."""
+        databricks_base = DatabricksBase()
+
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"access_token": "token"}
+
+        with patch("requests.post", return_value=mock_response) as mock_post:
+            databricks_base._get_oauth_m2m_token(
+                api_base="https://adb-123.azuredatabricks.net/ai-gateway/mlflow/v1",
+                client_id="id",
+                client_secret="secret",
+            )
+
+            call_url = mock_post.call_args[0][0]
+            assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
+
 
 class TestValidateEnvironmentWithOAuth:
     """Test OAuth M2M is used when credentials are available."""