From 29874191177d15086b866165b193f83bddc281d4 Mon Sep 17 00:00:00 2001 From: Jose Maria Vilaplana Date: Fri, 17 Apr 2026 17:16:29 +0200 Subject: [PATCH] fix(databricks): respect custom User-Agent set via extra_headers Currently, `databricks_validate_environment()` unconditionally overwrites the User-Agent header, discarding any value the caller explicitly set via `extra_headers`. This makes it impossible to pass a custom User-Agent string to Databricks endpoints. This makes the assignment conditional: if a User-Agent is already present in the headers dict (e.g. set via extra_headers), it is preserved. Otherwise, the default user agent is built as before. Databricks is the only provider whose `validate_environment` actively sets User-Agent. Other providers only add Authorization and Content-Type, leaving User-Agent untouched when set by the caller. This fix aligns the Databricks behavior with that pattern. --- litellm/llms/databricks/common_utils.py | 6 ++-- .../test_databricks_partner_integration.py | 34 +++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 608f29a03a7..afbf7c5df81 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -387,8 +387,10 @@ class DatabricksBase: if api_key is not None: headers["Authorization"] = f"Bearer {api_key}" - # Set User-Agent with optional custom prefix - headers["User-Agent"] = self._build_user_agent(custom_user_agent) + # Set User-Agent with optional custom prefix, but don't overwrite + # if the caller explicitly set one (e.g. via extra_headers) + if "User-Agent" not in headers: + headers["User-Agent"] = self._build_user_agent(custom_user_agent) # Debug logging with redaction (never log actual tokens) verbose_logger.debug( 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 800066ac5bf..aa7901598dd 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -334,6 +334,40 @@ class TestValidateEnvironmentUserAgent: assert headers["User-Agent"].startswith("mycompany_litellm/") + def test_extra_headers_user_agent_preserved(self, monkeypatch): + """User-Agent set via headers dict (extra_headers) is not overwritten.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers={"User-Agent": "My-Custom-Agent/1.0"}, + ) + + assert headers["User-Agent"] == "My-Custom-Agent/1.0" + + def test_default_user_agent_when_not_in_headers(self, monkeypatch): + """Default User-Agent is set when headers dict has no User-Agent.""" + monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False) + monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False) + + databricks_base = DatabricksBase() + + api_base, headers = databricks_base.databricks_validate_environment( + api_key="test-key", + api_base="https://adb-123.net/serving-endpoints", + endpoint_type="chat_completions", + custom_endpoint=False, + headers={}, + ) + + assert headers["User-Agent"].startswith("litellm/") + class TestSDKPartnerTelemetry: """Test that SDK partner telemetry is registered."""