From b9c53038fc0a150116052cfc742a222a1678ffe0 Mon Sep 17 00:00:00 2001 From: ump45nose <52391318+ump45nose@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:25:02 +0800 Subject: [PATCH] fix(databricks): derive OAuth URL from workspace origin --- litellm/llms/databricks/common_utils.py | 8 +++----- .../test_databricks_partner_integration.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b8dc98f2582..7695b1cb35e 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/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: diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index 139990021b4..b72a5b8a02e 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -249,6 +249,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."""