feat(agents): support Databricks Agent Platform agents in the agent gateway (#45681)

* feat(agents): support Databricks Agent Platform agents in the agent gateway

Add a databricks_agent provider that speaks the MLflow ResponsesAgent wire
(Responses-shaped input, output_text items, output_text.delta SSE) to Model
Serving endpoints (/invocations or the unified /serving-endpoints/responses)
and Databricks Apps (/responses), with a PAT or OAuth M2M client credentials

Register it through the agents config block or the Admin UI agent catalog and
call it through /a2a/<agent> (message/send, message/stream) or model:
a2a/<agent> chat completions. The chat route now runs every completion-bridge
agent through its own provider, and attaches the minted Databricks OAuth token
for both bridge and URL agents. Consumed OAuth settings never reach the
provider call

* fix(agents): harden databricks_agent input, output, and OAuth handling

Non-text content parts answer 400 unless drop_params is set, a reply without
an output list answers 502, flat OAuth fields resolve an os.environ api_base,
the bridge debug log no longer echoes the request model, routing tests stub
only the HTTP boundary, and the catalog description fits the dropdown

* fix(agents): merge agent and client extra_headers on the a2a chat route, and reject an empty databricks_oauth block

* fix(agents): bridge only databricks_agent agents on the a2a chat route, keep the card url over a client api_base, and require workspace_url for a Databricks App

* refactor(agents): read the a2a chat bridge providers from a constants set instead of branching on a provider name

* fix(databricks_agent): let the configured credential win, reject malformed client extra_headers, and keep a missing endpoint name empty

* test(agents): drive the flat OAuth a2a send through the real bridge

* fix(agents): keep an empty databricks_oauth block ignored on url agents, as before this change

* fix(agents): accept bare App hosts, env service principals, and os.environ refs on Databricks agents

* fix(databricks_agent): address review caveats on chat bridging, params, and auth

Map max_tokens and max_completion_tokens to max_output_tokens and pass
temperature, top_p, and user, so /v1/messages works without drop_params.
Send A2A message metadata only with string values, answering 400 or
dropping the rest under drop_params. Join message items with a blank line
in both stream and non-stream replies. Cache the env service principal
token until shortly before it expires. Keep databricks_agent replies out
of the response cache.

On the chat bridge, let request fields win over the agent's stored params
while pinning the agent's api_base and api_key, and apply the agent's
static_headers between client and agent extra_headers, as /a2a does.
Resolve os.environ/ only for databricks_agent api_base and api_key instead
of every bridge provider. Ignore an unusable databricks_oauth block on URL
agents on chat, as at the merge base, instead of answering 400.

* test(internal_context): list the databricks agent token cache as process-local

* fix(databricks_agent): keep the agent's endpoint, input, and target out of client control

extra_body model and input answer 400, or are dropped under drop_params.
A client api_base or base_url never replaces the agent's card URL on the
bridge. Chat to a URL agent keeps its merge-base call, and an incomplete
databricks_oauth block on a URL agent keeps its merge-base ValueError.

* test(databricks): expect the merge-base 400 for a client api_base on agent chat

* fix(databricks): keep agent-owned inputs and identity headers out of a chat caller's reach

* fix(databricks): bound the agent-over-caller merge to two levels so it needs no recursion

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>

litellm/a2a_protocol/litellm_completion_bridge/handler.py

@@ -11,6 +11,7 @@ A2A Streaming Events (in order):
 """
 
 from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
+from types import MappingProxyType
 from typing import Any, Final
 
 import litellm
@@ -21,6 +22,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
     A2AStreamingContext,
 )
 from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
+from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS
 from litellm.interactions.agents.utils import merge_agent_headers
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.utils import ModelResponse
@@ -44,6 +46,32 @@ _AGENT_ONLY_PARAMS: Final = frozenset(
 )
 
 
+def bridge_model_name(litellm_params: Mapping[str, object]) -> str:
+    """The litellm model string a bridge agent completes with: ``<custom_llm_provider>/<model>``.
+
+    A provider bridged on chat names a real upstream endpoint with ``model``, so a missing one stays empty for that
+    provider to reject or ignore instead of becoming a made-up ``agent`` endpoint.
+    """
+    custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
+    default_model: Final = "" if custom_llm_provider in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS else "agent"
+    configured_model: Final = litellm_params.get("model")
+    model: Final = configured_model if isinstance(configured_model, str) else default_model
+    if not custom_llm_provider or model.startswith(f"{custom_llm_provider}/"):
+        return model
+    return f"{custom_llm_provider}/{model}"
+
+
+def agent_completion_kwargs(litellm_params: Mapping[str, object]) -> Mapping[str, object]:
+    """The agent's litellm_params that are valid litellm.acompletion kwargs (api_key, api_base, ...)."""
+    return MappingProxyType(
+        {
+            key: value
+            for key, value in litellm_params.items()
+            if key not in ("model", "custom_llm_provider") and key not in _AGENT_ONLY_PARAMS
+        }
+    )
+
+
 class A2ACompletionBridgeHandler:
     """
     Static methods for handling A2A requests via LiteLLM completion.
@@ -64,16 +92,7 @@ class A2ACompletionBridgeHandler:
         # Transform A2A message to OpenAI format
         openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
 
-        # Get completion params
-        custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
-        model: Final[str] = litellm_params.get("model", "agent")
-
-        # Build full model string if provider specified
-        # Skip prepending if model already starts with the provider prefix
-        if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
-            full_model = f"{custom_llm_provider}/{model}"
-        else:
-            full_model = model
+        full_model: Final = bridge_model_name(litellm_params)
 
         if stream:
             verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
@@ -87,13 +106,7 @@ class A2ACompletionBridgeHandler:
             "api_base": api_base,
             "stream": stream,
         }
-        # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
-        litellm_params_to_add: Final = {
-            k: v
-            for k, v in litellm_params.items()
-            if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
-        }
-        completion_params.update(litellm_params_to_add)
+        completion_params.update(agent_completion_kwargs(litellm_params))
         # Apply forward metadata AFTER the litellm_params merge so the helper
         # sees any agent-owner-configured ``extra_body.metadata`` and can keep
         # those keys authoritative over the client-supplied A2A metadata.

litellm/constants.py

@@ -17,7 +17,7 @@ OAUTH_TOKEN_EXCHANGE_CACHE_SAFETY_MARGIN_SECONDS: Final = 60
 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_PROFILE: Final = "jwt_bearer_obo"
 MICROSOFT_365_COPILOT_DEFAULT_TOKEN_EXCHANGE_SCOPE: Final = "https://graph.microsoft.com/.default"
 # Replies depend on the caller's delegated identity, which response-cache keys do not include.
-RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot"})
+RESPONSE_CACHE_EXCLUDED_PROVIDERS: Final = frozenset({"microsoft_365_copilot", "databricks_agent"})
 MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
 GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
 GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
@@ -845,6 +845,8 @@ PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
     }
 )
 
+A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS: Final = frozenset({"databricks_agent"})
+
 LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [
     "openai",
     "azure",

litellm/llms/databricks/agent/responses_output.py

@@ -0,0 +1,102 @@
+from collections.abc import Mapping, Sequence
+from typing import Final
+
+from pydantic import BaseModel, ConfigDict
+
+from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.llms.databricks.common_utils import DatabricksException
+from litellm.types.utils import GenericStreamingChunk, Usage
+
+
+class _ContentPart(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    type: str = ""
+    text: str | None = None
+
+
+class OutputItem(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    type: str = ""
+    id: str | None = None
+    content: tuple[_ContentPart, ...] | None = None
+
+
+class AgentUsage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int | None = None
+
+    def to_usage(self) -> Usage:
+        return Usage(
+            prompt_tokens=self.input_tokens,
+            completion_tokens=self.output_tokens,
+            total_tokens=self.total_tokens if self.total_tokens is not None else self.input_tokens + self.output_tokens,
+        )
+
+
+class AgentResponse(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    output: tuple[OutputItem, ...]
+    custom_outputs: object = None
+    usage: AgentUsage | None = None
+
+
+class _StreamEvent(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    type: str = ""
+    item_id: str | None = None
+    delta: str | None = None
+    item: OutputItem | None = None
+    message: object = None
+
+
+_MESSAGE_SEPARATOR: Final = "\n\n"
+
+
+def output_text(items: Sequence[OutputItem] | None) -> str:
+    texts: Final = (_message_text(item) for item in items or ())
+    return _MESSAGE_SEPARATOR.join(text for text in texts if text)
+
+
+def _message_text(item: OutputItem) -> str:
+    if item.type != "message":
+        return ""
+    return "".join(part.text or "" for part in item.content or () if part.type == "output_text")
+
+
+def _text_chunk(text: str) -> GenericStreamingChunk:
+    return GenericStreamingChunk(text=text, is_finished=False, finish_reason="", usage=None, index=0, tool_use=None)
+
+
+class DatabricksAgentResponsesIterator(BaseModelResponseIterator):
+    _streamed_item_ids: frozenset[str] = frozenset()
+    _emitted_text: bool = False
+    _last_text_item_id: str | None = None
+
+    def chunk_parser(self, chunk: Mapping[str, object]) -> GenericStreamingChunk:
+        event: Final = _StreamEvent.model_validate(chunk)
+        if event.type == "response.output_text.delta":
+            if event.item_id is not None:
+                self._streamed_item_ids = self._streamed_item_ids | {event.item_id}
+            return self._item_text_chunk(event.delta or "", event.item_id)
+        if event.type == "response.output_item.done":
+            if event.item is None or event.item.id in self._streamed_item_ids:
+                return _text_chunk("")
+            return self._item_text_chunk(output_text((event.item,)), event.item.id)
+        if event.type == "error":
+            raise DatabricksException(status_code=500, message=f"Databricks agent stream error: {event.message}")
+        return _text_chunk("")
+
+    def _item_text_chunk(self, text: str, item_id: str | None) -> GenericStreamingChunk:
+        if not text:
+            return _text_chunk("")
+        separator: Final = _MESSAGE_SEPARATOR if self._emitted_text and item_id != self._last_text_item_id else ""
+        self._emitted_text = True
+        self._last_text_item_id = item_id
+        return _text_chunk(f"{separator}{text}")

litellm/llms/databricks/agent/transformation.py

@@ -0,0 +1,467 @@
+"""Databricks ResponsesAgent endpoints.
+
+Model Serving: POST {workspace}/serving-endpoints/{endpoint}/invocations
+Model Serving, unified: POST {workspace}/serving-endpoints/responses with the endpoint as the body "model"
+Databricks Apps: POST {app}/responses
+"""
+
+import hashlib
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
+from typing import TYPE_CHECKING, Final, Protocol
+from urllib.parse import quote, urlsplit
+
+import httpx
+from pydantic import BaseModel, ConfigDict, StrictStr, TypeAdapter, ValidationError
+
+import litellm
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.litellm_core_utils.prompt_templates.common_utils import (
+    convert_content_list_to_str,
+)
+from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
+from litellm.llms.databricks.agent.responses_output import (
+    AgentResponse,
+    DatabricksAgentResponsesIterator,
+    output_text,
+)
+from litellm.llms.databricks.common_utils import DatabricksBase, DatabricksException
+from litellm.secret_managers.main import get_secret_str
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import Choices, Message, ModelResponse, Usage
+
+if TYPE_CHECKING:
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
+
+PROVIDER_NAME: Final = "databricks_agent"
+_MODEL_PREFIX: Final = f"{PROVIDER_NAME}/"
+_INVOCATIONS_SUFFIX: Final = "/invocations"
+_APP_RESPONSES_SUFFIX: Final = "/responses"
+_ENDPOINT_URL_SUFFIXES: Final = (_INVOCATIONS_SUFFIX, _APP_RESPONSES_SUFFIX)
+_SERVING_ENDPOINTS_SEGMENT: Final = "/serving-endpoints/"
+_UNIFIED_RESPONSES_SUFFIX: Final = "/serving-endpoints/responses"
+_APP_HOST_SUFFIX: Final = ".databricksapps.com"
+_TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60
+
+
+class _TokenCache(Protocol):
+    def get_cache(self, key: str) -> object: ...
+
+    def set_cache(self, key: str, value: str, *, ttl: int) -> None: ...
+
+
+def _new_token_cache() -> _TokenCache:
+    return InMemoryCache(max_size_in_memory=64)
+
+
+_ENV_M2M_TOKENS: Final = _new_token_cache()
+_INPUT_ROLES: Final = frozenset({"system", "developer", "user", "assistant"})
+_BODY_PARAMS: Final = (
+    "max_output_tokens",
+    "temperature",
+    "top_p",
+    "user",
+    "custom_inputs",
+    "context",
+    "databricks_options",
+)
+_SAMPLING_PARAMS: Final = ("temperature", "top_p", "user")
+_OUTPUT_TOKEN_PARAMS: Final = ("max_completion_tokens", "max_tokens")
+_STRING_METADATA: Final = TypeAdapter(Mapping[str, StrictStr])
+_METADATA: Final = TypeAdapter(Mapping[str, object])
+_PINNED_BODY_FIELDS: Final = ("model", "input")
+
+
+def endpoint_name_from_model(model: str) -> str:
+    return model.removeprefix(_MODEL_PREFIX).strip()
+
+
+def _normalized_base(api_base: str) -> str:
+    return api_base.strip().rstrip("/")
+
+
+def _is_unified_responses_url(api_base: str) -> bool:
+    return _normalized_base(api_base).endswith(_UNIFIED_RESPONSES_SUFFIX)
+
+
+def _is_databricks_app_host(api_base: str) -> bool:
+    return (urlsplit(_normalized_base(api_base)).hostname or "").endswith(_APP_HOST_SUFFIX)
+
+
+def is_databricks_app_url(api_base: str) -> bool:
+    if _is_unified_responses_url(api_base):
+        return False
+    return _is_databricks_app_host(api_base) or _normalized_base(api_base).endswith(_APP_RESPONSES_SUFFIX)
+
+
+def workspace_url_from_api_base(api_base: str | None) -> str | None:
+    if not api_base or not api_base.strip() or is_databricks_app_url(api_base):
+        return None
+    parts: Final = urlsplit(_normalized_base(api_base))
+    if not parts.scheme or not parts.netloc:
+        return None
+    return f"{parts.scheme}://{parts.netloc}"
+
+
+def _required_endpoint_name(model: str) -> str:
+    endpoint: Final = endpoint_name_from_model(model)
+    if not endpoint:
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                "model must name the Databricks serving endpoint (model: databricks_agent/<endpoint>) "
+                "unless api_base is the endpoint's /invocations URL or a Databricks App /responses URL"
+            ),
+        )
+    return endpoint
+
+
+def _app_endpoint_url(app_base: str) -> str:
+    return f"{app_base}{_APP_RESPONSES_SUFFIX}" if urlsplit(app_base).path in ("", "/") else app_base
+
+
+def resolve_endpoint_url(api_base: str, model: str) -> str:
+    base: Final = _normalized_base(api_base)
+    if _is_unified_responses_url(base):
+        _required_endpoint_name(model)
+        return base
+    if base.endswith(_ENDPOINT_URL_SUFFIXES):
+        return base
+    if _is_databricks_app_host(base):
+        return _app_endpoint_url(base)
+    if _SERVING_ENDPOINTS_SEGMENT in base:
+        return f"{base}{_INVOCATIONS_SUFFIX}"
+    workspace: Final = base.removesuffix("/serving-endpoints")
+    return f"{workspace}/serving-endpoints/{quote(_required_endpoint_name(model), safe='')}/invocations"
+
+
+def _body_model_field(api_base: object, model: str) -> Mapping[str, str]:
+    if not isinstance(api_base, str) or not _is_unified_responses_url(api_base):
+        return {}
+    return {"model": _required_endpoint_name(model)}
+
+
+class _InputContentPart(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    type: str = "unknown"
+
+
+class _InputMessage(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    content: str | tuple[_InputContentPart, ...] | None = None
+    tool_calls: tuple[object, ...] | None = None
+
+
+def _drop_params_enabled(litellm_params: Mapping[str, object]) -> bool:
+    return litellm.drop_params is True or litellm_params.get("drop_params") is True
+
+
+def _input_item(message: AllMessageValues, drop_unsupported: bool) -> Mapping[str, str] | None:
+    role: Final = message["role"]
+    if role not in _INPUT_ROLES:
+        if drop_unsupported:
+            return None
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                f"message role '{role}' is not supported by a Databricks ResponsesAgent; "
+                "send system, developer, user, or assistant messages, or set drop_params: true to drop it"
+            ),
+        )
+    unsupported: Final = _unsupported_message_parts(message)
+    if unsupported and not drop_unsupported:
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                f"databricks_agent sends text content only; remove {', '.join(unsupported)} from the "
+                "message, or set drop_params: true to drop them"
+            ),
+        )
+    content: Final = convert_content_list_to_str(message)
+    if unsupported and not content:
+        return None
+    return {"role": role, "content": content}
+
+
+def _input_items(messages: Sequence[AllMessageValues], drop_unsupported: bool) -> Sequence[Mapping[str, str]]:
+    items: Final = (_input_item(message, drop_unsupported) for message in messages)
+    return [item for item in items if item is not None]
+
+
+def _unsupported_message_parts(message: AllMessageValues) -> tuple[str, ...]:
+    parsed: Final = _InputMessage.model_validate(message)
+    content_parts: Final = parsed.content if isinstance(parsed.content, tuple) else ()
+    content_types: Final[frozenset[str]] = frozenset(part.type for part in content_parts) - {"text"}
+    tool_calls: Final[frozenset[str]] = frozenset({"tool_calls"}) if parsed.tool_calls else frozenset()
+    return tuple(sorted(content_types | tool_calls))
+
+
+def _has_authorization(headers: Mapping[str, object]) -> bool:
+    return any(name.lower() == "authorization" for name in headers)
+
+
+def _without_authorization(headers: Mapping[str, object]) -> Mapping[str, object]:
+    return {name: value for name, value in headers.items() if name.lower() != "authorization"}
+
+
+def _env_oauth_m2m_token(api_base: str | None) -> str | None:
+    client_id: Final = get_secret_str("DATABRICKS_CLIENT_ID")
+    client_secret: Final = get_secret_str("DATABRICKS_CLIENT_SECRET")
+    if not (client_id and client_secret):
+        return None
+    workspace: Final = get_secret_str("DATABRICKS_HOST") or workspace_url_from_api_base(api_base)
+    if workspace is None:
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                "DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET are set but the OAuth token endpoint is unknown: "
+                "set DATABRICKS_HOST to the workspace URL (a Databricks App URL is not a workspace)"
+            ),
+        )
+    return _cached_m2m_token(workspace, client_id, client_secret)
+
+
+def _cached_m2m_token(workspace: str, client_id: str, client_secret: str) -> str:
+    secret_digest: Final = hashlib.sha256(client_secret.encode()).hexdigest()[:16]
+    cache_key: Final = f"{workspace}|{client_id}|{secret_digest}"
+    cached: Final[object] = _ENV_M2M_TOKENS.get_cache(cache_key)
+    if isinstance(cached, str):
+        return cached
+    token: Final = DatabricksBase().get_oauth_m2m_token_and_lifetime(workspace, client_id, client_secret)
+    ttl: Final = token.expires_in - _TOKEN_EXPIRY_BUFFER_SECONDS
+    if ttl > 0:
+        _ENV_M2M_TOKENS.set_cache(cache_key, token.access_token, ttl=ttl)
+    return token.access_token
+
+
+def _resolved_env_reference(value: str | None) -> str | None:
+    if value is not None and value.startswith("os.environ/"):
+        return get_secret_str(value)
+    return value
+
+
+def _max_output_tokens(non_default_params: Mapping[str, object]) -> Mapping[str, object]:
+    limits: Final = (non_default_params[key] for key in _OUTPUT_TOKEN_PARAMS if non_default_params.get(key) is not None)
+    limit: Final = next(limits, None)
+    return {} if limit is None else {"max_output_tokens": limit}
+
+
+def _string_metadata_entries(metadata: object) -> Mapping[str, str]:
+    try:
+        entries: Final = _METADATA.validate_python(metadata)
+    except ValidationError:
+        return {}
+    return {key: value for key, value in entries.items() if isinstance(value, str)}
+
+
+def _has_string_values_only(metadata: object) -> bool:
+    try:
+        _STRING_METADATA.validate_python(metadata)
+    except ValidationError:
+        return False
+    return True
+
+
+def _without_pinned_fields(extra_body: Mapping[str, object], drop_params: bool) -> Mapping[str, object]:
+    pinned: Final = tuple(key for key in _PINNED_BODY_FIELDS if key in extra_body)
+    if not pinned:
+        return extra_body
+    if not drop_params:
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                f"extra_body cannot set {', '.join(pinned)} on a Databricks ResponsesAgent request, "
+                "since the agent's endpoint and the messages set them; remove them, "
+                "or set drop_params: true to drop them"
+            ),
+        )
+    return {key: value for key, value in extra_body.items() if key not in _PINNED_BODY_FIELDS}
+
+
+def _with_string_metadata(extra_body: Mapping[str, object], drop_params: bool) -> Mapping[str, object]:
+    if "metadata" not in extra_body or _has_string_values_only(extra_body["metadata"]):
+        return extra_body
+    if not drop_params:
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                "metadata sent to a Databricks ResponsesAgent must map names to string values; "
+                "send string values only, or set drop_params: true to drop the others"
+            ),
+        )
+    kept: Final = _string_metadata_entries(extra_body["metadata"])
+    rest: Final = {key: value for key, value in extra_body.items() if key != "metadata"}
+    return {**rest, **({"metadata": kept} if kept else {})}
+
+
+class DatabricksAgentConfig(BaseConfig):
+    @staticmethod
+    def resolve_api_base_and_key(api_base: str | None, api_key: str | None) -> tuple[str | None, str | None]:
+        return (
+            _resolved_env_reference(api_base) or get_secret_str("DATABRICKS_API_BASE"),
+            _resolved_env_reference(api_key),
+        )
+
+    def get_supported_openai_params(self, model: str) -> list[str]:  # mutable-ok: mirrors override contract
+        return ["stream", *_OUTPUT_TOKEN_PARAMS, *_SAMPLING_PARAMS]
+
+    def map_openai_params(
+        self,
+        non_default_params: Mapping[str, object],
+        optional_params: Mapping[str, object],
+        model: str,
+        drop_params: bool,
+    ) -> dict[str, object]:  # mutable-ok: mirrors override contract
+        sampling: Final = {key: non_default_params[key] for key in _SAMPLING_PARAMS if key in non_default_params}
+        return {**optional_params, **sampling, **_max_output_tokens(non_default_params)}
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, object],
+        model: str,
+        messages: Sequence[AllMessageValues],
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> dict[str, object]:  # mutable-ok: mirrors override contract
+        bearer: Final = api_key or _env_oauth_m2m_token(api_base) or get_secret_str("DATABRICKS_API_KEY")
+        if bearer:
+            return {
+                **_without_authorization(headers),
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {bearer}",
+            }
+        if _has_authorization(headers):
+            return {**headers, "Content-Type": "application/json"}
+        raise DatabricksException(
+            status_code=400,
+            message=(
+                "Missing Databricks credentials: set api_key to a personal access token, set the "
+                "DATABRICKS_API_KEY environment variable, set DATABRICKS_CLIENT_ID and DATABRICKS_CLIENT_SECRET "
+                "for OAuth M2M, or register the agent with OAuth client credentials"
+            ),
+        )
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        api_key: str | None,
+        model: str,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        stream: bool | None = None,
+    ) -> str:
+        if not api_base:
+            raise DatabricksException(
+                status_code=400,
+                message=(
+                    "api_base is required for databricks_agent: the workspace URL, the Databricks App URL, "
+                    "or the full endpoint URL (set it or the DATABRICKS_API_BASE environment variable)"
+                ),
+            )
+        return resolve_endpoint_url(api_base, model)
+
+    def transform_request(
+        self,
+        model: str,
+        messages: Sequence[AllMessageValues],
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        headers: Mapping[str, object],
+    ) -> dict[str, object]:  # mutable-ok: mirrors override contract
+        body_params: Final = {key: optional_params[key] for key in _BODY_PARAMS if key in optional_params}
+        stream: Final = {"stream": True} if optional_params.get("stream") is True else {}
+        return {
+            **_body_model_field(litellm_params.get("api_base"), model),
+            "input": _input_items(messages, _drop_params_enabled(litellm_params)),
+            **stream,
+            **body_params,
+        }
+
+    def transform_extra_body(
+        self,
+        extra_body: Mapping[str, object],
+        request: Mapping[str, object],
+        model: str,
+        litellm_params: Mapping[str, object],
+    ) -> Mapping[str, object]:
+        drop_params: Final = _drop_params_enabled(litellm_params)
+        return _with_string_metadata(_without_pinned_fields(extra_body, drop_params), drop_params)
+
+    def transform_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: ModelResponse,
+        logging_obj: "LiteLLMLoggingObj",
+        request_data: Mapping[str, object],
+        messages: Sequence[AllMessageValues],
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        encoding: "Tokenizer | None",
+        api_key: str | None = None,
+        json_mode: bool | None = None,
+    ) -> ModelResponse:
+        logging_obj.post_call(
+            input=messages,
+            api_key="",
+            original_response=raw_response.text,
+            additional_args={"complete_input_dict": request_data},
+        )
+        try:
+            agent_response: Final = AgentResponse.model_validate_json(raw_response.content)
+        except ValidationError as e:
+            raise DatabricksException(
+                status_code=502,
+                message=f"Databricks agent returned a response that is not a ResponsesAgent response: {e}",
+            )
+
+        content: Final = output_text(agent_response.output)
+        message: Final = Message(
+            content=content,
+            role="assistant",
+            provider_specific_fields=(
+                {"custom_outputs": agent_response.custom_outputs} if agent_response.custom_outputs is not None else None
+            ),
+        )
+        return ModelResponse(
+            id=model_response.id,
+            created=model_response.created,
+            model=model,
+            choices=[Choices(finish_reason="stop", index=0, message=message)],
+            usage=agent_response.usage.to_usage()
+            if agent_response.usage is not None
+            else self._estimated_usage(model, messages, content),
+        )
+
+    @staticmethod
+    def _estimated_usage(model: str, messages: Sequence[AllMessageValues], content: str) -> Usage:
+        from litellm.utils import token_counter
+
+        prompt_tokens: Final = token_counter(model=model, messages=list(messages))
+        completion_tokens: Final = token_counter(model=model, text=content, count_response_tokens=True)
+        return Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+        )
+
+    def get_model_response_iterator(
+        self,
+        streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse,
+        sync_stream: bool,
+        json_mode: bool | None = False,
+    ) -> DatabricksAgentResponsesIterator:
+        return DatabricksAgentResponsesIterator(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
+
+    def get_error_class(
+        self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers
+    ) -> BaseLLMException:
+        return DatabricksException(status_code=status_code, message=error_message)

litellm/llms/databricks/common_utils.py

@@ -15,13 +15,24 @@ import re
 from typing import Final, Literal
 from urllib.parse import urlsplit, urlunsplit
 
+from pydantic import BaseModel, ConfigDict
+
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 
+_DEFAULT_M2M_TOKEN_LIFETIME_SECONDS: Final = 3600
+
 
 class DatabricksException(BaseLLMException):
     pass
 
 
+class OAuthM2MToken(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    access_token: str
+    expires_in: int = _DEFAULT_M2M_TOKEN_LIFETIME_SECONDS
+
+
 class DatabricksBase:
     """
     Base class for Databricks integration with authentication,
@@ -206,12 +217,20 @@ class DatabricksBase:
             return f"{normalized_api_base}/ai-gateway/mlflow/v1"
         return api_base
 
-    def _get_oauth_m2m_token(
+    def get_oauth_m2m_token(
         self,
         api_base: str,
         client_id: str,
         client_secret: str,
     ) -> str:
+        return self.get_oauth_m2m_token_and_lifetime(api_base, client_id, client_secret).access_token
+
+    def get_oauth_m2m_token_and_lifetime(
+        self,
+        api_base: str,
+        client_id: str,
+        client_secret: str,
+    ) -> OAuthM2MToken:
         """
         Obtain an OAuth M2M access token using client credentials flow.
 
@@ -224,7 +243,7 @@ class DatabricksBase:
             client_secret: OAuth client secret
 
         Returns:
-            Access token string
+            The access token and its lifetime in seconds
 
         Raises:
             DatabricksException: If token request fails
@@ -258,8 +277,7 @@ class DatabricksBase:
                 message=f"OAuth M2M token request failed: {response.text}",
             )
 
-        token_data: Final = response.json()
-        return token_data["access_token"]
+        return OAuthM2MToken.model_validate(response.json())
 
     def _get_databricks_credentials(
         self, api_key: str | None, api_base: str | None, headers: dict | None
@@ -346,7 +364,7 @@ class DatabricksBase:
         if client_id and client_secret and api_base:
             # Use OAuth M2M flow (preferred for production)
             verbose_logger.debug("Using OAuth M2M authentication for Databricks")
-            access_token: Final = self._get_oauth_m2m_token(api_base, client_id, client_secret)
+            access_token: Final = self.get_oauth_m2m_token(api_base, client_id, client_secret)
             headers = headers or {}
             headers["Authorization"] = f"Bearer {access_token}"
             headers["Content-Type"] = "application/json"

litellm/main.py

@@ -5143,6 +5143,33 @@ def _complete_langflow(ctx: CompletionDispatchContext) -> _CompletionDispatchRes
     )
 
 
+def _complete_databricks_agent(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
+    from litellm.llms.databricks.agent.transformation import DatabricksAgentConfig
+
+    api_base, api_key = DatabricksAgentConfig.resolve_api_base_and_key(
+        api_base=ctx.api_base or litellm.api_base,
+        api_key=ctx.api_key or litellm.api_key,
+    )
+    return base_llm_http_handler.completion(
+        model=ctx.model,
+        stream=ctx.stream,
+        messages=ctx.messages,
+        acompletion=ctx.acompletion,
+        api_base=api_base,
+        model_response=ctx.model_response,
+        optional_params=ctx.optional_params,
+        litellm_params={**ctx.litellm_params, "api_base": api_base},
+        shared_session=ctx.shared_session,
+        custom_llm_provider=ctx.custom_llm_provider,
+        timeout=ctx.timeout,
+        headers=ctx.headers or litellm.headers,
+        encoding=_get_encoding(),
+        api_key=api_key,
+        logging_obj=ctx.logging,
+        client=_dispatch_client_http(ctx),
+    )
+
+
 @tracer.wrap()
 @client
 def completion(
@@ -6101,6 +6128,9 @@ def completion(
             # LangFlow - Visual AI Agent Platform
             response = _complete_langflow(_dispatch_ctx)
 
+        elif custom_llm_provider == "databricks_agent":
+            response = _complete_databricks_agent(_dispatch_ctx)  # rebind-ok: dispatch chain binds response per branch
+
         else:
             raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)
         return response

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -34,8 +34,9 @@ from litellm.proxy.a2a.version_convert import (
     normalize_stream_event,
 )
 from litellm.proxy.agent_endpoints.databricks_oauth import (
-    DATABRICKS_OAUTH_PARAM,
-    resolve_databricks_app_auth_header,
+    AgentBackendAuth,
+    has_databricks_oauth,
+    resolve_databricks_backend_auth,
 )
 from litellm.proxy.agent_endpoints.utils import merge_agent_headers
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@@ -166,13 +167,17 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str,
     )
 
 
-async def _resolve_backend_auth_header(
+async def _resolve_backend_auth(
     litellm_params: dict[str, object],
     custom_llm_provider: object,
-) -> Mapping[str, str] | None:
-    if litellm_params.get(DATABRICKS_OAUTH_PARAM):
-        return await resolve_databricks_app_auth_header(litellm_params)
-    return await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider)
+) -> AgentBackendAuth:
+    databricks_auth: Final = await resolve_databricks_backend_auth(litellm_params)
+    if has_databricks_oauth(litellm_params):
+        return databricks_auth
+    return AgentBackendAuth(
+        litellm_params=databricks_auth.litellm_params,
+        header=await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider),
+    )
 
 
 def _forwarding_headers(
@@ -819,6 +824,7 @@ async def invoke_agent_a2a(
                     if header_name:
                         dynamic_headers[header_name] = val
 
+        backend_auth: Final = await _resolve_backend_auth(litellm_params, custom_llm_provider)
         agent_extra_headers: Final = _forwarding_headers(
             caller_identity=caller_identity,
             request_data=data,
@@ -826,8 +832,9 @@ async def invoke_agent_a2a(
                 dynamic_headers=dynamic_headers or None,
                 static_headers=static_headers or None,
             ),
-            backend_auth_header=await _resolve_backend_auth_header(litellm_params, custom_llm_provider),
+            backend_auth_header=backend_auth.header,
         )
+        forwarded_litellm_params: Final = backend_auth.litellm_params
 
         # Merge agent-level guardrails into data so post_call_success_hook and
         # _handle_stream_message both pick them up.  A2A agents use model
@@ -872,7 +879,7 @@ async def invoke_agent_a2a(
                 model=f"a2a_agent/{agent_name}",
                 request=a2a_request,
                 api_base=agent_url,
-                litellm_params=litellm_params,
+                litellm_params=forwarded_litellm_params,
                 agent_id=agent.agent_id,
                 metadata=data.get("metadata", {}),
                 proxy_server_request=data.get("proxy_server_request"),
@@ -912,7 +919,7 @@ async def invoke_agent_a2a(
                 api_base=agent_url,
                 request_id=request_id if request_id is not None else "",
                 params=params,
-                litellm_params=litellm_params,
+                litellm_params=forwarded_litellm_params,
                 agent_id=agent.agent_id,
                 metadata=data.get("metadata", {}),
                 proxy_server_request=data.get("proxy_server_request"),

litellm/proxy/agent_endpoints/a2a_routing.py

@@ -5,17 +5,115 @@ Handles routing for A2A agents (models with "a2a/<agent-name>" prefix).
 Looks up agents in the registry and injects their API base URL.
 """
 
+from collections.abc import Mapping
 from typing import Any, Final
 
 from fastapi import HTTPException
+from pydantic import TypeAdapter, ValidationError
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm.a2a_protocol.litellm_completion_bridge.handler import (
+    agent_completion_kwargs,
+    bridge_model_name,
+)
+from litellm.constants import A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.agent_endpoints.databricks_oauth import (
+    resolve_databricks_backend_auth,
+)
+from litellm.proxy.agent_endpoints.utils import merge_agent_headers
+from litellm.types.agents import AgentResponse
+
+_HEADERS: Final = TypeAdapter(Mapping[str, str])
+_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object])
+_AGENT_PINNED_KWARGS: Final = ("api_base", "api_key")
+_AGENT_LAYERED_KWARGS: Final = ("custom_inputs", "extra_body")
+_CLIENT_TARGET_KWARGS: Final = ("api_base", "base_url")
+_PROXY_HEADER_PREFIX: Final = "x-litellm-"
+
+
+def _agent_extra_headers(agent_kwargs: Mapping[str, object]) -> Mapping[str, str] | None:
+    extra_headers: Final = agent_kwargs.get("extra_headers")
+    return None if extra_headers is None else _HEADERS.validate_python(extra_headers)
+
+
+def _drop_params_enabled(data: Mapping[str, object], litellm_params: Mapping[str, object]) -> bool:
+    return litellm.drop_params is True or data.get("drop_params") is True or litellm_params.get("drop_params") is True
+
+
+def _client_extra_headers(data: Mapping[str, object], litellm_params: Mapping[str, object]) -> Mapping[str, str] | None:
+    extra_headers: Final = data.get("extra_headers")
+    if not extra_headers:
+        return None
+    try:
+        headers: Final = _HEADERS.validate_python(extra_headers)
+    except ValidationError as e:
+        if _drop_params_enabled(data, litellm_params):
+            return None
+        raise HTTPException(
+            status_code=400,
+            detail=(
+                "extra_headers must be an object mapping header names to string values, "
+                "or set drop_params: true to drop it"
+            ),
+        ) from e
+    return {name: value for name, value in headers.items() if not name.lower().startswith(_PROXY_HEADER_PREFIX)}
+
+
+def _mappings(client_value: object, agent_value: object) -> tuple[Mapping[str, object], Mapping[str, object]] | None:
+    try:
+        return _OBJECT_MAPPING.validate_python(client_value), _OBJECT_MAPPING.validate_python(agent_value)
+    except ValidationError:
+        return None
+
+
+def _agent_entries_over_client(client_value: object, agent_value: object) -> object:
+    mappings: Final = _mappings(client_value, agent_value)
+    return agent_value if mappings is None else {**mappings[0], **mappings[1]}
+
+
+def _agent_keys_over_client(client_value: object, agent_value: object) -> object:
+    mappings: Final = _mappings(client_value, agent_value)
+    if mappings is None:
+        return agent_value
+    client, agent = mappings
+    return {**client, **{key: _agent_entries_over_client(client.get(key), value) for key, value in agent.items()}}
+
+
+def _with_headers(data: Mapping[str, object], extra_headers: Mapping[str, str] | None) -> Mapping[str, object]:
+    request: Final = {key: value for key, value in data.items() if key != "extra_headers"}
+    return {**request, **({"extra_headers": extra_headers} if extra_headers else {})}
+
+
+def _bridge_request_data(
+    data: Mapping[str, object],
+    litellm_params: Mapping[str, object],
+    static_headers: Mapping[str, str] | None,
+) -> Mapping[str, object]:
+    agent_kwargs: Final = agent_completion_kwargs(litellm_params)
+    pinned: Final = {key: agent_kwargs[key] for key in _AGENT_PINNED_KWARGS if key in agent_kwargs}
+    agent_headers: Final = merge_agent_headers(
+        dynamic_headers=static_headers,
+        static_headers=_agent_extra_headers(agent_kwargs),
+    )
+    extra_headers: Final = merge_agent_headers(
+        dynamic_headers=_client_extra_headers(data, litellm_params),
+        static_headers=agent_headers,
+    )
+    layered: Final = {
+        key: _agent_keys_over_client(data[key], agent_kwargs[key])
+        for key in _AGENT_LAYERED_KWARGS
+        if key in data and key in agent_kwargs
+    }
+    return _with_headers(
+        {**agent_kwargs, **data, **layered, **pinned, "model": bridge_model_name(litellm_params)},
+        extra_headers,
+    )
 
 
 async def route_a2a_agent_request(
-    data: dict,
+    data: dict[str, object],  # mutable-ok: the URL agent path writes api_base into the caller's request
     route_type: str,
     user_api_key_dict: UserAPIKeyAuth | None = None,
 ) -> Any | None:
@@ -68,6 +166,12 @@ async def route_a2a_agent_request(
                 detail=f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.",
             )
 
+    litellm_params: Final = agent.litellm_params or {}
+
+    if litellm_params.get("custom_llm_provider") in A2A_CHAT_COMPLETION_BRIDGE_PROVIDERS:
+        bridge_auth: Final = await resolve_databricks_backend_auth(litellm_params)
+        return _call_route(route_type, _bridge_agent_request(data, agent, bridge_auth.litellm_params))
+
     # Get API base URL from agent config
     if not agent.agent_card_params or "url" not in agent.agent_card_params:
         verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name)
@@ -78,4 +182,19 @@ async def route_a2a_agent_request(
     data["api_base"] = agent.agent_card_params["url"]
     verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"])
 
-    return getattr(litellm, f"{route_type}")(**data)
+    return _call_route(route_type, data)
+
+
+def _bridge_agent_request(
+    data: Mapping[str, object],
+    agent: AgentResponse,
+    litellm_params: Mapping[str, object],
+) -> Mapping[str, object]:
+    card_url: Final[object] = (agent.agent_card_params or {}).get("url")
+    api_base: Final = {"api_base": card_url} if card_url else {}
+    request: Final = {key: value for key, value in data.items() if key not in _CLIENT_TARGET_KWARGS}
+    return _bridge_request_data({**request, **api_base}, litellm_params, agent.static_headers)
+
+
+def _call_route(route_type: str, request_data: Mapping[str, object]) -> object:
+    return getattr(litellm, route_type)(**request_data)

litellm/proxy/agent_endpoints/databricks_oauth.py

@@ -30,18 +30,24 @@ from dataclasses import dataclass
 from typing import Final
 
 import httpx
+from pydantic import TypeAdapter, ValidationError
 
+import litellm
 from litellm._logging import verbose_logger
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.llms.databricks.agent.transformation import is_databricks_app_url, workspace_url_from_api_base
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.custom_http import httpxSpecialProvider
 
 DATABRICKS_OAUTH_PARAM: Final = "databricks_oauth"
+DATABRICKS_AGENT_PROVIDER: Final = "databricks_agent"
+_FLAT_OAUTH_FIELDS: Final = ("client_id", "client_secret", "workspace_url", "scope")
 
 _DEFAULT_SCOPE: Final = "all-apis"
 _TOKEN_EXPIRY_BUFFER_SECONDS: Final = 60
 _DEFAULT_TTL_SECONDS: Final = 3600
+_OAUTH_BLOCK: Final = TypeAdapter(Mapping[str, object])
 
 
 def _resolve_secret(value: object) -> str | None:
@@ -75,41 +81,118 @@ class DatabricksAppOAuthConfig:
         return f"{self.token_url}|{self.client_id}|{self.scope}|{secret_digest}"
 
 
+def _workspace_origin(url: object) -> str | None:
+    return workspace_url_from_api_base(url) if isinstance(url, str) else None
+
+
+@dataclass(frozen=True, slots=True)
+class DatabricksOAuthConfigError:
+    message: str
+
+
+def _config_error(message: str) -> litellm.BadRequestError:
+    return litellm.BadRequestError(message=message, model="a2a", llm_provider=DATABRICKS_AGENT_PROVIDER)
+
+
+def _is_databricks_agent(litellm_params: Mapping[str, object]) -> bool:
+    return litellm_params.get("custom_llm_provider") == DATABRICKS_AGENT_PROVIDER
+
+
+def _flat_databricks_agent_oauth(litellm_params: Mapping[str, object]) -> Mapping[str, object] | None:
+    """The Admin UI catalog cannot nest fields, so a ``databricks_agent`` spells its OAuth block as flat
+    ``client_id`` / ``client_secret`` / ``workspace_url`` / ``scope``. ``workspace_url`` falls back to the
+    ``api_base`` origin for Model Serving; a Databricks App host is never a token endpoint, so an App must set it."""
+    if not _is_databricks_agent(litellm_params):
+        return None
+    if not litellm_params.get("client_id") and not litellm_params.get("client_secret"):
+        return None
+    return {
+        "client_id": litellm_params.get("client_id"),
+        "client_secret": litellm_params.get("client_secret"),
+        "workspace_url": litellm_params.get("workspace_url")
+        or _workspace_origin(_resolve_secret(litellm_params.get("api_base"))),
+        "scope": litellm_params.get("scope"),
+    }
+
+
+def _raw_oauth_block(litellm_params: Mapping[str, object]) -> object | None:
+    nested: Final = litellm_params.get(DATABRICKS_OAUTH_PARAM)
+    if not _is_databricks_agent(litellm_params):
+        return nested or None
+    return nested if nested is not None else _flat_databricks_agent_oauth(litellm_params)
+
+
+def has_databricks_oauth(litellm_params: Mapping[str, object]) -> bool:
+    return _raw_oauth_block(litellm_params) is not None
+
+
+def without_databricks_oauth_params(
+    litellm_params: Mapping[str, object],
+) -> dict[str, object]:  # mutable-ok: asend_message takes litellm_params as a dict
+    consumed: Final = (
+        frozenset({DATABRICKS_OAUTH_PARAM, *_FLAT_OAUTH_FIELDS})
+        if _is_databricks_agent(litellm_params)
+        else frozenset({DATABRICKS_OAUTH_PARAM})
+    )
+    return {key: value for key, value in litellm_params.items() if key not in consumed}
+
+
 def parse_databricks_oauth_config(
     litellm_params: Mapping[str, object] | None,
 ) -> DatabricksAppOAuthConfig | None:
     """Build a Databricks App OAuth config from an agent's ``litellm_params``.
 
-    Returns ``None`` when the agent has no ``databricks_oauth`` block. Raises
-    ``ValueError`` when the block is present but incomplete, so misconfiguration
-    surfaces loudly instead of silently sending an unauthenticated request.
+    Returns ``None`` when the agent has no ``databricks_oauth`` block. A block that is present but incomplete
+    raises instead of sending an unauthenticated request: a 400 ``litellm.BadRequestError`` for a
+    ``databricks_agent``, and the ``ValueError`` URL agents raised before ``databricks_agent`` existed.
     """
+    config: Final = databricks_oauth_config(litellm_params)
+    if not isinstance(config, DatabricksOAuthConfigError):
+        return config
+    if litellm_params and _is_databricks_agent(litellm_params):
+        raise _config_error(config.message)
+    raise ValueError(config.message)
+
+
+def databricks_oauth_config(
+    litellm_params: Mapping[str, object] | None,
+) -> DatabricksAppOAuthConfig | DatabricksOAuthConfigError | None:
     if not litellm_params:
         return None
 
-    raw: Final = litellm_params.get(DATABRICKS_OAUTH_PARAM)
+    raw: Final = _raw_oauth_block(litellm_params)
     if raw is None:
         return None
-    if not isinstance(raw, dict):
-        raise ValueError(f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, got {type(raw).__name__}")
-
-    client_id: Final = _resolve_secret(raw.get("client_id"))
-    client_secret: Final = _resolve_secret(raw.get("client_secret"))
-    workspace_url: Final = _resolve_secret(raw.get("workspace_url"))
-
-    missing: Final = [
-        name
-        for name, value in (
-            ("client_id", client_id),
-            ("client_secret", client_secret),
-            ("workspace_url", workspace_url),
+    try:
+        block: Final = _OAUTH_BLOCK.validate_python(raw)
+    except ValidationError:
+        return DatabricksOAuthConfigError(
+            f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, got {type(raw).__name__}"
         )
-        if not value
-    ]
-    if missing:
-        raise ValueError(f"Databricks App OAuth config is missing required field(s): {', '.join(missing)}")
 
-    scope: Final = _resolve_secret(raw.get("scope")) or _DEFAULT_SCOPE
+    client_id: Final = _resolve_secret(block.get("client_id"))
+    client_secret: Final = _resolve_secret(block.get("client_secret"))
+    workspace_url: Final = _resolve_secret(block.get("workspace_url"))
+
+    if not (client_id and client_secret and workspace_url):
+        missing: Final = tuple(
+            name
+            for name, value in (
+                ("client_id", client_id),
+                ("client_secret", client_secret),
+                ("workspace_url", workspace_url),
+            )
+            if not value
+        )
+        return DatabricksOAuthConfigError(
+            f"Databricks App OAuth config is missing required field(s): {', '.join(missing)}"
+        )
+    if is_databricks_app_url(workspace_url):
+        return DatabricksOAuthConfigError(
+            "workspace_url must be the Databricks workspace URL, not a Databricks App URL"
+        )
+
+    scope: Final = _resolve_secret(block.get("scope")) or _DEFAULT_SCOPE
 
     return DatabricksAppOAuthConfig(
         client_id=client_id,
@@ -215,6 +298,13 @@ class DatabricksAppOAuthTokenCache(InMemoryCache):
 databricks_app_oauth_token_cache: Final = DatabricksAppOAuthTokenCache()
 
 
+async def _resolve_token(litellm_params: Mapping[str, object] | None) -> str | None:
+    config: Final = parse_databricks_oauth_config(litellm_params)
+    if config is None:
+        return None
+    return await databricks_app_oauth_token_cache.async_get_token(config)
+
+
 async def resolve_databricks_app_auth_header(
     litellm_params: Mapping[str, object] | None,
 ) -> dict[str, str] | None:
@@ -222,9 +312,29 @@ async def resolve_databricks_app_auth_header(
 
     Returns ``None`` when the agent is not configured for Databricks App OAuth.
     """
-    config: Final = parse_databricks_oauth_config(litellm_params)
-    if config is None:
-        return None
+    token: Final = await _resolve_token(litellm_params)
+    return None if token is None else {"Authorization": f"Bearer {token}"}
 
-    token: Final = await databricks_app_oauth_token_cache.async_get_token(config)
-    return {"Authorization": f"Bearer {token}"}
+
+@dataclass(frozen=True, slots=True)
+class AgentBackendAuth:
+    litellm_params: dict[str, object]  # mutable-ok: asend_message takes litellm_params as a dict
+    header: Mapping[str, str] | None
+
+
+async def resolve_databricks_backend_auth(litellm_params: Mapping[str, object]) -> AgentBackendAuth:
+    """A ``databricks_agent`` completes with the minted token as its ``api_key``, so the token outranks a personal
+    access token and a client ``Authorization`` header; any other agent sends it as the outbound header."""
+    return await _backend_auth(litellm_params, parse_databricks_oauth_config(litellm_params))
+
+
+async def _backend_auth(
+    litellm_params: Mapping[str, object], config: DatabricksAppOAuthConfig | None
+) -> AgentBackendAuth:
+    forwarded: Final = without_databricks_oauth_params(litellm_params)
+    token: Final = None if config is None else await databricks_app_oauth_token_cache.async_get_token(config)
+    if token is None:
+        return AgentBackendAuth(litellm_params=forwarded, header=None)
+    if _is_databricks_agent(litellm_params):
+        return AgentBackendAuth(litellm_params={**forwarded, "api_key": token}, header=None)
+    return AgentBackendAuth(litellm_params=forwarded, header={"Authorization": f"Bearer {token}"})

litellm/proxy/public_endpoints/agent_create_fields.json

@@ -305,6 +305,77 @@
     "litellm_params_template": {
       "custom_llm_provider": "watsonx_orchestrate"
     }
+  },
+  {
+    "agent_type": "databricks_agent",
+    "agent_type_display_name": "Databricks Agent",
+    "description": "Connect to Databricks Agent Platform ResponsesAgents on Model Serving or Databricks Apps",
+    "logo_url": "/ui/assets/logos/databricks.svg",
+    "credential_fields": [
+      {
+        "key": "api_base",
+        "label": "Databricks URL",
+        "placeholder": "https://<workspace>.cloud.databricks.com",
+        "tooltip": "The workspace URL for a Model Serving endpoint, or a full URL: https://<workspace>/serving-endpoints/<endpoint>/invocations, https://<workspace>/serving-endpoints/responses, or https://<app>.databricksapps.com/responses",
+        "required": true,
+        "field_type": "text",
+        "default_value": null,
+        "include_in_litellm_params": true
+      },
+      {
+        "key": "model",
+        "label": "Serving endpoint name",
+        "placeholder": "my-agent-endpoint",
+        "tooltip": "The Model Serving endpoint that hosts the agent. Leave empty when the Databricks URL is an endpoint /invocations URL or a Databricks App /responses URL",
+        "required": false,
+        "field_type": "text",
+        "default_value": null,
+        "include_in_litellm_params": true
+      },
+      {
+        "key": "api_key",
+        "label": "Personal access token",
+        "placeholder": null,
+        "tooltip": "A Databricks PAT. Works for Model Serving only; Databricks Apps accept OAuth client credentials only",
+        "required": false,
+        "field_type": "password",
+        "default_value": null,
+        "include_in_litellm_params": true
+      },
+      {
+        "key": "client_id",
+        "label": "OAuth client ID",
+        "placeholder": null,
+        "tooltip": "Service principal application ID for OAuth M2M (client credentials). Required for Databricks Apps",
+        "required": false,
+        "field_type": "text",
+        "default_value": null,
+        "include_in_litellm_params": true
+      },
+      {
+        "key": "client_secret",
+        "label": "OAuth client secret",
+        "placeholder": null,
+        "tooltip": "Service principal OAuth secret",
+        "required": false,
+        "field_type": "password",
+        "default_value": null,
+        "include_in_litellm_params": true
+      },
+      {
+        "key": "workspace_url",
+        "label": "Workspace URL (OAuth)",
+        "placeholder": "https://<workspace>.cloud.databricks.com",
+        "tooltip": "The workspace that mints the OAuth token. Defaults to the Databricks URL origin; set it for a Databricks App, whose host is not the workspace",
+        "required": false,
+        "field_type": "text",
+        "default_value": null,
+        "include_in_litellm_params": true
+      }
+    ],
+    "litellm_params_template": {
+      "custom_llm_provider": "databricks_agent"
+    }
   }
 ]
 

litellm/types/utils.py

@@ -4305,6 +4305,7 @@ class LlmProviders(str, Enum):
     TRITON = "triton"
     PREDIBASE = "predibase"
     DATABRICKS = "databricks"
+    DATABRICKS_AGENT = "databricks_agent"
     EMPOWER = "empower"
     GITHUB = "github"
     RAGFLOW = "ragflow"

litellm/utils.py

@@ -8777,6 +8777,10 @@ class ProviderConfigManager:
                 lambda: ProviderConfigManager._get_langflow_config(),
                 False,
             ),
+            LlmProviders.DATABRICKS_AGENT: (
+                lambda: ProviderConfigManager._get_databricks_agent_config(),
+                False,
+            ),
             LlmProviders.GDC: (
                 lambda: litellm.GDCGeminiConfig(),
                 False,
@@ -8879,6 +8883,12 @@ class ProviderConfigManager:
 
         return LangFlowConfig()
 
+    @staticmethod
+    def _get_databricks_agent_config() -> BaseConfig:
+        from litellm.llms.databricks.agent.transformation import DatabricksAgentConfig
+
+        return DatabricksAgentConfig()
+
     @staticmethod
     def get_provider_chat_config(
         model: str,

tests/integration/providers/test_databricks_agent_wire.py

@@ -0,0 +1,984 @@
+import asyncio
+import base64
+import json
+import re
+import threading
+import uuid
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+from urllib.parse import parse_qs
+
+import httpx
+import openai
+import pytest
+import yaml
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.process import graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_ENDPOINT: Final = "support-agent"
+_INVOCATIONS: Final = f"/serving-endpoints/{_ENDPOINT}/invocations"
+_UNIFIED: Final = "/serving-endpoints/responses"
+_APP: Final = "/responses"
+_TOKEN: Final = "/oidc/v1/token"
+_PAT: Final = "synthetic-databricks-pat"
+_CLIENT_ID: Final = "synthetic-databricks-sp"
+_CLIENT_SECRET: Final = "synthetic-databricks-sp-secret"
+_MINTED: Final = "synthetic-databricks-minted-token"
+_REPLY: Final = "the export job is waiting on a lock"
+_PROMPT: Final = "Why does the export job hang?"
+_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 11, "output_tokens": 3, "total_tokens": 14}
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
+_OWNED_PROXY_BUDGET: Final = int(2 * graceful_stop_seconds() + 120)
+
+
+def _basic(client_id: str, client_secret: str) -> str:
+    return "Basic " + base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
+
+
+def _agent_reply(text: str = _REPLY) -> bytes:
+    return json.dumps(
+        {
+            "output": [
+                {
+                    "type": "message",
+                    "id": "msg-1",
+                    "role": "assistant",
+                    "content": [{"type": "output_text", "text": text}],
+                }
+            ],
+            "usage": _USAGE,
+        }
+    ).encode()
+
+
+def _agent_stream(text: str = _REPLY) -> bytes:
+    head, tail = text[: len(text) // 2], text[len(text) // 2 :]
+    events: Final = (
+        {"type": "response.output_text.delta", "item_id": "msg-1", "delta": head},
+        {"type": "response.output_text.delta", "item_id": "msg-1", "delta": tail},
+        {
+            "type": "response.output_item.done",
+            "item": {"type": "message", "id": "msg-1", "content": [{"type": "output_text", "text": text}]},
+        },
+    )
+    return "".join(f"data: {json.dumps(event)}\n\n" for event in events).encode()
+
+
+def _token_reply() -> Reply:
+    return Reply(body=json.dumps({"access_token": _MINTED, "token_type": "Bearer", "expires_in": 3600}).encode())
+
+
+@dataclass(frozen=True, slots=True)
+class _Databricks:
+    token_status: Callable[[], int] = lambda: 200
+    endpoint_reply: Callable[[Request], Reply] | None = None
+
+    def __call__(self, request: Request) -> Reply:
+        if request.target == _TOKEN:
+            status: Final = self.token_status()
+            if status != 200:
+                return Reply(status=status, body=b'{"error":"invalid_client"}')
+            return _token_reply()
+        if self.endpoint_reply is not None:
+            return self.endpoint_reply(request)
+        if _JSON_OBJECT.validate_json(request.body).get("stream") is True:
+            return Reply(body=_agent_stream(), content_type="text/event-stream")
+        return Reply(body=_agent_reply())
+
+
+def _calls_to(requests: Sequence[Request], target: str) -> tuple[Request, ...]:
+    return tuple(request for request in requests if request.target == target)
+
+
+def _body(request: Request) -> dict[str, JsonValue]:
+    return _JSON_OBJECT.validate_json(request.body)
+
+
+def _sdk_base_url(gateway: Gateway) -> str:
+    return f"{str(gateway.client.base_url).rstrip('/')}/v1"
+
+
+def _register(
+    gateway: Gateway,
+    cleanups: Callable[[Callable[[], None]], object],
+    card_url: str,
+    static_headers: dict[str, JsonValue] | None = None,
+    **params: JsonValue,
+) -> str:
+    name: Final = f"dbx-agent-{uuid.uuid4().hex[:12]}"
+    created: Final = gateway.request(
+        "POST",
+        "/v1/agents",
+        {
+            "agent_name": name,
+            "agent_card_params": {
+                "protocolVersion": "0.3",
+                "name": name,
+                "description": "Databricks Agent agent",
+                "url": card_url,
+                "version": "1.0.0",
+                "defaultInputModes": ["text"],
+                "defaultOutputModes": ["text"],
+                "capabilities": {"streaming": True},
+                "skills": [],
+            },
+            "litellm_params": {"custom_llm_provider": "databricks_agent", **params},
+            **({"static_headers": static_headers} if static_headers else {}),
+        },
+    )
+    assert created.status_code == 200, created.text
+    agent_id: Final = created.json()["agent_id"]
+
+    def delete() -> None:
+        deleted: Final = gateway.request("DELETE", f"/v1/agents/{agent_id}")
+        assert deleted.status_code == 200, deleted.text
+
+    cleanups(delete)
+    return name
+
+
+@contextmanager
+def _agent(
+    gateway: Gateway,
+    databricks: _Databricks,
+    *,
+    url_path: str = "",
+    static_headers: dict[str, JsonValue] | None = None,
+    **params: JsonValue,
+) -> Iterator[tuple[Wire, str]]:
+    with wire_server(databricks) as wire, gateway.scenario() as scenario:
+        params_with_base: Final = {"api_base": wire.url + url_path, **params}
+        yield (
+            wire,
+            _register(gateway, scenario.cleanups.callback, wire.url + url_path, static_headers, **params_with_base),
+        )
+
+
+def _chat(gateway: Gateway, name: str, **extra: JsonValue) -> httpx.Response:
+    return gateway.request(
+        "POST",
+        "/v1/chat/completions",
+        {"model": f"a2a/{name}", "messages": [{"role": "user", "content": _PROMPT}], **extra},
+    )
+
+
+def _a2a(
+    gateway: Gateway, name: str, method: str, marker: str, metadata: dict[str, JsonValue] | None = None
+) -> httpx.Response:
+    message: Final[dict[str, JsonValue]] = {
+        "role": "user",
+        "messageId": marker,
+        "parts": [{"kind": "text", "text": _PROMPT}],
+        **({"metadata": metadata} if metadata is not None else {}),
+    }
+    return gateway.request(
+        "POST",
+        f"/a2a/{name}",
+        {"jsonrpc": "2.0", "id": marker, "method": method, "params": {"message": message}},
+    )
+
+
+def _sse(text: str) -> tuple[dict[str, JsonValue], ...]:
+    return tuple(
+        _JSON_OBJECT.validate_json(line.removeprefix("data: "))
+        for line in text.splitlines()
+        if line.startswith("data: ") and line != "data: [DONE]"
+    )
+
+
+def _spend_status(request_id: str) -> str:
+    rows: Final = eventually(
+        lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)),
+        lambda found: len(found) == 1,
+        seconds=70,
+    )
+    return str(rows[0]["status"])
+
+
+def _assert_invocation(request: Request, bearer: str, *, stream: bool = False) -> None:
+    assert request.method == "POST" and request.target == _INVOCATIONS
+    assert request.headers["authorization"] == f"Bearer {bearer}"
+    assert _body(request) == {
+        "input": [{"role": "user", "content": _PROMPT}],
+        **({"stream": True} if stream else {}),
+    }
+
+
+def _assert_token_request(request: Request) -> None:
+    assert request.method == "POST"
+    assert request.headers["authorization"] == _basic(_CLIENT_ID, _CLIENT_SECRET)
+    assert parse_qs(request.body.decode()) == {"grant_type": ["client_credentials"], "scope": ["all-apis"]}
+
+
+def test_chat_on_a_pat_agent_reaches_the_serving_endpoint_and_bills_the_agent_usage(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == 200, response.text
+        payload: Final = response.json()
+        assert payload["choices"][0]["message"] == {"content": _REPLY, "role": "assistant"}
+        assert payload["choices"][0]["finish_reason"] == "stop"
+        assert payload["usage"]["prompt_tokens"] == 11 and payload["usage"]["completion_tokens"] == 3
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        _assert_invocation(received[0], _PAT)
+        assert _spend_status(payload["id"]) == "success"
+
+
+def test_async_chat_stream_on_a_pat_agent_sends_the_reply_once(gateway: Gateway) -> None:
+    async def stream(name: str) -> str:
+        client: Final = openai.AsyncOpenAI(base_url=_sdk_base_url(gateway), api_key=gateway.key, max_retries=0)
+        chunks: Final = await client.chat.completions.create(
+            model=f"a2a/{name}", messages=[{"role": "user", "content": _PROMPT}], stream=True
+        )
+        return "".join([chunk.choices[0].delta.content or "" async for chunk in chunks if chunk.choices])
+
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        assert asyncio.run(stream(name)) == _REPLY
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        _assert_invocation(received[0], _PAT, stream=True)
+
+
+def test_sync_sdk_chat_on_a_pat_agent_returns_the_agent_reply(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        client: Final = openai.OpenAI(base_url=_sdk_base_url(gateway), api_key=gateway.key, max_retries=0)
+        completion: Final = client.chat.completions.create(
+            model=f"a2a/{name}", messages=[{"role": "user", "content": _PROMPT}]
+        )
+        assert completion.choices[0].message.content == _REPLY
+        assert len(_calls_to(wire.drain(), _INVOCATIONS)) == 1
+
+
+def test_responses_on_a_pat_agent_returns_the_reply_in_output(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = gateway.request("POST", "/v1/responses", {"model": f"a2a/{name}", "input": _PROMPT})
+        assert response.status_code == 200, response.text
+        texts: Final = tuple(
+            part["text"]
+            for item in response.json()["output"]
+            if item["type"] == "message"
+            for part in item["content"]
+            if part["type"] == "output_text"
+        )
+        assert texts == (_REPLY,)
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        _assert_invocation(received[0], _PAT)
+
+
+@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop-params"))
+def test_messages_on_a_pat_agent_sends_max_tokens_as_max_output_tokens(gateway: Gateway, drop_params: bool) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = gateway.request(
+            "POST",
+            "/v1/messages",
+            {
+                "model": f"a2a/{name}",
+                "max_tokens": 64,
+                "messages": [{"role": "user", "content": _PROMPT}],
+                **({"drop_params": True} if drop_params else {}),
+            },
+        )
+        assert response.status_code == 200, response.text
+        assert response.json()["content"] == [{"type": "text", "text": _REPLY}]
+        invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+        assert len(invocations) == 1, invocations
+        assert _body(invocations[0]) == {"input": [{"role": "user", "content": _PROMPT}], "max_output_tokens": 64}
+
+
+def test_chat_sampling_params_reach_the_agent_in_its_own_shape(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(gateway, name, max_completion_tokens=32, temperature=0.2, top_p=0.9, user="u-1")
+        assert response.status_code == 200, response.text
+        invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+        assert len(invocations) == 1, invocations
+        assert _body(invocations[0]) == {
+            "input": [{"role": "user", "content": _PROMPT}],
+            "max_output_tokens": 32,
+            "temperature": 0.2,
+            "top_p": 0.9,
+            "user": "u-1",
+        }
+
+
+def test_a_reply_with_several_message_items_reads_the_same_streamed_or_not(gateway: Gateway) -> None:
+    first, second = "Let me look at the export queue.", "The export job is waiting on a lock."
+
+    def message(item_id: str, text: str) -> dict[str, JsonValue]:
+        return {"type": "message", "id": item_id, "content": [{"type": "output_text", "text": text}]}
+
+    def reply(request: Request) -> Reply:
+        if _body(request).get("stream") is not True:
+            output: Final = (message("msg-1", first), {"type": "function_call", "id": "fc-1"}, message("msg-2", second))
+            return Reply(body=json.dumps({"output": output, "usage": _USAGE}).encode())
+        events: Final = (
+            {"type": "response.output_text.delta", "item_id": "msg-1", "delta": first},
+            {"type": "response.output_item.done", "item": {"type": "function_call", "id": "fc-1"}},
+            {"type": "response.output_text.delta", "item_id": "msg-2", "delta": second},
+            {"type": "response.output_item.done", "item": message("msg-2", second)},
+        )
+        return Reply(
+            body="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(),
+            content_type="text/event-stream",
+        )
+
+    with _agent(gateway, _Databricks(endpoint_reply=reply), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        client: Final = openai.OpenAI(base_url=_sdk_base_url(gateway), api_key=gateway.key, max_retries=0)
+        messages: Final = [{"role": "user", "content": _PROMPT}]
+        whole: Final = client.chat.completions.create(model=f"a2a/{name}", messages=messages)
+        streamed: Final = client.chat.completions.create(model=f"a2a/{name}", messages=messages, stream=True)
+        expected: Final = f"{first}\n\n{second}"
+        assert whole.choices[0].message.content == expected
+        assert "".join(chunk.choices[0].delta.content or "" for chunk in streamed if chunk.choices) == expected
+        assert len(_calls_to(wire.drain(), _INVOCATIONS)) == 2
+
+
+def test_identical_chat_requests_on_an_agent_each_reach_the_agent(gateway: Gateway) -> None:
+    replies: Final = iter(("first answer", "second answer"))
+    databricks: Final = _Databricks(endpoint_reply=lambda _: Reply(body=_agent_reply(next(replies))))
+    with _agent(gateway, databricks, model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        answers: Final = tuple(_chat(gateway, name) for _ in range(2))
+        assert [answer.json()["choices"][0]["message"]["content"] for answer in answers] == [
+            "first answer",
+            "second answer",
+        ], [answer.text for answer in answers]
+        assert len(_calls_to(wire.drain(), _INVOCATIONS)) == 2
+
+
+def test_static_headers_sit_between_client_and_agent_extra_headers_on_chat(gateway: Gateway) -> None:
+    with _agent(
+        gateway,
+        _Databricks(),
+        static_headers={"X-Static": "static", "X-Shared": "from-static", "X-Both": "from-static"},
+        model=_ENDPOINT,
+        api_key=_PAT,
+        extra_headers={"x-both": "from-agent"},
+    ) as (wire, name):
+        response: Final = _chat(gateway, name, extra_headers={"x-shared": "from-client", "X-Trace": "trace-1"})
+        assert response.status_code == 200, response.text
+        invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+        assert len(invocations) == 1, invocations
+        headers: Final = invocations[0].headers
+        assert (headers["x-static"], headers["x-shared"], headers["x-both"]) == ("static", "from-static", "from-agent")
+        assert headers["x-trace"] == "trace-1"
+        _assert_invocation(invocations[0], _PAT)
+
+
+def test_chat_never_forwards_client_litellm_identity_headers_to_the_agent(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(
+            gateway,
+            name,
+            extra_headers={"X-LiteLLM-User-Id": "spoofed-user", "x-litellm-team-id": "spoofed-team", "X-Trace": "t-1"},
+        )
+        assert response.status_code == 200, response.text
+        invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+        assert len(invocations) == 1, invocations
+        headers: Final = invocations[0].headers
+        assert "x-litellm-user-id" not in headers and "x-litellm-team-id" not in headers, headers
+        assert headers["x-trace"] == "t-1"
+        _assert_invocation(invocations[0], _PAT)
+
+
+def test_the_agent_custom_inputs_and_extra_body_keys_win_over_the_callers_on_chat(gateway: Gateway) -> None:
+    with _agent(
+        gateway,
+        _Databricks(),
+        model=_ENDPOINT,
+        api_key=_PAT,
+        custom_inputs={"tenant": "owner"},
+        extra_body={"metadata": {"tenant": "owner"}},
+    ) as (wire, name):
+        response: Final = _chat(
+            gateway,
+            name,
+            custom_inputs={"tenant": "caller", "session": "s-1"},
+            extra_body={"metadata": {"tenant": "caller", "turn": "2"}},
+        )
+        assert response.status_code == 200, response.text
+        invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+        assert len(invocations) == 1, invocations
+        assert _body(invocations[0]) == {
+            "input": [{"role": "user", "content": _PROMPT}],
+            "custom_inputs": {"tenant": "owner", "session": "s-1"},
+            "metadata": {"tenant": "owner", "turn": "2"},
+        }
+
+
+@pytest.mark.parametrize("drop_params", (False, True), ids=("strict", "drop-params"))
+def test_a2a_message_metadata_reaches_the_agent_only_as_string_values(gateway: Gateway, drop_params: bool) -> None:
+    params: Final[dict[str, JsonValue]] = {"drop_params": True} if drop_params else {}
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT, **params) as (wire, name):
+        strings: Final = _a2a(gateway, name, "message/send", f"meta-{uuid.uuid4().hex}", {"session": "s-1"})
+        assert strings.status_code == 200 and "error" not in strings.json(), strings.text
+        assert _body(_calls_to(wire.drain(), _INVOCATIONS)[0])["metadata"] == {"session": "s-1"}
+        mixed: Final = _a2a(gateway, name, "message/send", f"meta-{uuid.uuid4().hex}", {"session": "s-1", "turn": 1})
+        sent: Final = _calls_to(wire.drain(), _INVOCATIONS)
+        if drop_params:
+            assert mixed.status_code == 200 and "error" not in mixed.json(), mixed.text
+            assert [_body(request)["metadata"] for request in sent] == [{"session": "s-1"}]
+        else:
+            assert "string values" in mixed.text, mixed.text
+            assert sent == ()
+
+
+def test_a2a_message_send_on_a_pat_agent_answers_with_the_agent_text(gateway: Gateway) -> None:
+    marker: Final = f"send-{uuid.uuid4().hex}"
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _a2a(gateway, name, "message/send", marker)
+        assert response.status_code == 200, response.text
+        body: Final = response.json()
+        assert body["id"] == marker and "error" not in body, body
+        assert body["result"]["parts"] == [{"kind": "text", "text": _REPLY}]
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        _assert_invocation(received[0], _PAT)
+
+
+def test_a2a_message_stream_on_a_pat_agent_streams_the_text_and_completes(gateway: Gateway) -> None:
+    marker: Final = f"stream-{uuid.uuid4().hex}"
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _a2a(gateway, name, "message/stream", marker)
+        assert response.status_code == 200, response.text
+        events: Final = tuple(event["result"] for event in _sse(response.text))
+        artifacts: Final = tuple(
+            part["text"]
+            for event in events
+            if event["kind"] == "artifact-update"
+            for part in event["artifact"]["parts"]
+        )
+        assert artifacts == (_REPLY,)
+        assert events[-1]["kind"] == "status-update" and events[-1]["final"] is True
+        assert events[-1]["status"]["state"] == "completed"
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        _assert_invocation(received[0], _PAT, stream=True)
+
+
+def test_the_unified_responses_url_names_the_endpoint_in_the_body(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), url_path=_UNIFIED, model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == 200, response.text
+        assert response.json()["choices"][0]["message"]["content"] == _REPLY
+        received: Final = wire.drain()
+        assert [(request.method, request.target) for request in received] == [("POST", _UNIFIED)]
+        assert received[0].headers["authorization"] == f"Bearer {_PAT}"
+        assert _body(received[0]) == {"model": _ENDPOINT, "input": [{"role": "user", "content": _PROMPT}]}
+
+
+def test_an_app_agent_mints_its_token_at_the_workspace_and_calls_the_app(gateway: Gateway) -> None:
+    with (
+        wire_server(_Databricks()) as workspace,
+        _agent(
+            gateway,
+            _Databricks(),
+            url_path=_APP,
+            client_id=_CLIENT_ID,
+            client_secret=_CLIENT_SECRET,
+            workspace_url=workspace.url,
+        ) as (app, name),
+    ):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == 200, response.text
+        assert response.json()["choices"][0]["message"]["content"] == _REPLY
+        minted: Final = workspace.drain()
+        assert [request.target for request in minted] == [_TOKEN]
+        _assert_token_request(minted[0])
+        called: Final = app.drain()
+        assert [(request.method, request.target) for request in called] == [("POST", _APP)]
+        assert called[0].headers["authorization"] == f"Bearer {_MINTED}"
+        assert _body(called[0]) == {"input": [{"role": "user", "content": _PROMPT}]}
+
+
+def test_a_nested_oauth_block_mints_once_for_chat_and_a2a_calls_on_one_agent(gateway: Gateway) -> None:
+    def oauth(url: str) -> dict[str, JsonValue]:
+        return {"client_id": _CLIENT_ID, "client_secret": _CLIENT_SECRET, "workspace_url": url}
+
+    with wire_server(_Databricks()) as wire, gateway.scenario() as scenario:
+        name: Final = _register(
+            gateway,
+            scenario.cleanups.callback,
+            wire.url,
+            model=_ENDPOINT,
+            api_base=wire.url,
+            databricks_oauth=oauth(wire.url),
+        )
+        for _ in range(3):
+            response = _chat(gateway, name)
+            assert response.status_code == 200, response.text
+        sent: Final = _a2a(gateway, name, "message/send", f"send-{uuid.uuid4().hex}")
+        assert sent.status_code == 200 and "error" not in sent.json(), sent.text
+        received: Final = wire.drain()
+        tokens: Final = _calls_to(received, _TOKEN)
+        assert len(tokens) <= 2, [request.target for request in received]
+        _assert_token_request(tokens[0])
+        invocations: Final = _calls_to(received, _INVOCATIONS)
+        assert len(invocations) == 4
+        for invocation in invocations:
+            _assert_invocation(invocation, _MINTED)
+
+
+def test_the_minted_token_outranks_the_agent_pat_and_a_client_authorization(gateway: Gateway) -> None:
+    with _agent(
+        gateway,
+        _Databricks(),
+        model=_ENDPOINT,
+        api_key=_PAT,
+        client_id=_CLIENT_ID,
+        client_secret=_CLIENT_SECRET,
+    ) as (wire, name):
+        response: Final = _chat(gateway, name, extra_headers={"Authorization": "Bearer client-supplied"})
+        assert response.status_code == 200, response.text
+        received: Final = wire.drain()
+        assert [request.target for request in received] == [_TOKEN, _INVOCATIONS]
+        _assert_invocation(received[1], _MINTED)
+
+
+def test_the_agent_pat_outranks_a_client_authorization(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(gateway, name, extra_headers={"Authorization": "Bearer client-supplied"})
+        assert response.status_code == 200, response.text
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        _assert_invocation(received[0], _PAT)
+
+
+def test_empty_oauth_fields_beside_a_pat_keep_the_pat_path(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT, client_id="", client_secret="") as (wire, name):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == 200, response.text
+        received: Final = wire.drain()
+        assert [request.target for request in received] == [_INVOCATIONS]
+        _assert_invocation(received[0], _PAT)
+
+
+def test_custom_inputs_reach_the_agent_beside_the_input(gateway: Gateway) -> None:
+    custom_inputs: Final[dict[str, JsonValue]] = {"session_id": "s-1", "tier": "gold"}
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(gateway, name, custom_inputs=custom_inputs)
+        assert response.status_code == 200, response.text
+        received: Final = wire.drain()
+        assert len(received) == 1, received
+        assert _body(received[0]) == {"input": [{"role": "user", "content": _PROMPT}], "custom_inputs": custom_inputs}
+
+
+def test_two_agents_on_one_endpoint_name_never_answer_from_each_other(gateway: Gateway) -> None:
+    def answering(text: str) -> _Databricks:
+        return _Databricks(endpoint_reply=lambda _: Reply(body=_agent_reply(text)))
+
+    with (
+        _agent(gateway, answering("from workspace one"), model=_ENDPOINT, api_key=_PAT) as (first_wire, first),
+        _agent(gateway, answering("from workspace two"), model=_ENDPOINT, api_key=_PAT) as (second_wire, second),
+    ):
+        responses: Final = tuple(_chat(gateway, name) for name in (first, second, first))
+        assert all(response.status_code == 200 for response in responses), [r.text for r in responses]
+        answers: Final = tuple(response.json()["choices"][0]["message"]["content"] for response in responses)
+        assert answers == ("from workspace one", "from workspace two", "from workspace one")
+        assert len(_calls_to(first_wire.drain(), _INVOCATIONS)) == 2
+        assert len(_calls_to(second_wire.drain(), _INVOCATIONS)) == 1
+
+
+def test_a_workspace_agent_with_no_endpoint_is_refused_before_any_call(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), api_key=_PAT) as (wire, name):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == 400, response.text
+        assert "model must name the Databricks serving endpoint" in response.text
+        assert wire.drain() == ()
+
+
+def test_an_app_agent_without_a_workspace_url_is_refused_before_its_secret_leaves(gateway: Gateway) -> None:
+    with _agent(gateway, _Databricks(), url_path=_APP, client_id=_CLIENT_ID, client_secret=_CLIENT_SECRET) as (
+        wire,
+        name,
+    ):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == 400, response.text
+        assert "workspace_url" in response.text
+        assert wire.drain() == ()
+
+
+@pytest.mark.parametrize(("drop_params", "status"), ((False, 400), (True, 200)), ids=("strict", "drop-params"))
+def test_malformed_client_extra_headers_follow_drop_params(gateway: Gateway, drop_params: bool, status: int) -> None:
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(
+            gateway, name, extra_headers=["Authorization", "x"], **({"drop_params": True} if drop_params else {})
+        )
+        assert response.status_code == status, response.text
+        if drop_params:
+            received: Final = wire.drain()
+            assert len(received) == 1, received
+            _assert_invocation(received[0], _PAT)
+        else:
+            assert "extra_headers must be an object" in response.text
+            assert wire.drain() == ()
+
+
+@pytest.mark.parametrize(("drop_params", "status"), ((False, 400), (True, 200)), ids=("strict", "drop-params"))
+def test_an_image_part_follows_drop_params(gateway: Gateway, drop_params: bool, status: int) -> None:
+    content: Final[list[JsonValue]] = [
+        {"type": "text", "text": _PROMPT},
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
+    ]
+    with _agent(gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {
+                "model": f"a2a/{name}",
+                "messages": [{"role": "user", "content": content}],
+                **({"drop_params": True} if drop_params else {}),
+            },
+        )
+        assert response.status_code == status, response.text
+        if drop_params:
+            received: Final = wire.drain()
+            assert len(received) == 1, received
+            _assert_invocation(received[0], _PAT)
+        else:
+            assert "image_url" in response.text
+            assert wire.drain() == ()
+
+
+@pytest.mark.parametrize(
+    ("reply", "status"),
+    (
+        (Reply(status=401, body=b'{"error_code":"UNAUTHENTICATED","message":"Invalid token"}'), 401),
+        (Reply(status=500, body=b'{"error_code":"INTERNAL_ERROR","message":"agent crashed"}'), 500),
+        (Reply(body=b"{}"), 502),
+    ),
+    ids=("upstream-401", "upstream-500", "not-a-responses-agent"),
+)
+def test_a_failing_agent_keeps_its_status_and_unrelated_models_keep_serving(
+    gateway: Gateway, reply: Reply, status: int
+) -> None:
+    with (
+        _agent(gateway, _Databricks(endpoint_reply=lambda _: reply), model=_ENDPOINT, api_key=_PAT) as (wire, name),
+        gateway.scenario() as scenario,
+    ):
+        response: Final = _chat(gateway, name)
+        assert response.status_code == status, response.text
+        assert len(_calls_to(wire.drain(), _INVOCATIONS)) >= 1
+        unrelated: Final = gateway.chat(scenario.model())
+        assert unrelated["choices"], unrelated
+
+
+def test_a_refused_token_request_answers_an_error_without_calling_the_agent(gateway: Gateway) -> None:
+    with _agent(
+        gateway,
+        _Databricks(token_status=lambda: 400),
+        model=_ENDPOINT,
+        client_id=_CLIENT_ID,
+        client_secret=_CLIENT_SECRET,
+    ) as (wire, name):
+        response: Final = _chat(gateway, name)
+        assert response.status_code >= 400, response.text
+        received: Final = wire.drain()
+        assert [request.target for request in received] == [_TOKEN]
+
+
+def test_reading_an_agent_back_redacts_its_pat_and_client_secret(gateway: Gateway) -> None:
+    with _agent(
+        gateway, _Databricks(), model=_ENDPOINT, api_key=_PAT, client_id=_CLIENT_ID, client_secret=_CLIENT_SECRET
+    ) as (_, name):
+        listed: Final = gateway.request("GET", "/v1/agents")
+        assert listed.status_code == 200, listed.text
+        text: Final = listed.text
+        assert name in text
+        assert _PAT not in text and _CLIENT_SECRET not in text
+
+
+@pytest.mark.parametrize(
+    "oauth_fields",
+    (None, (), ("client_id",), ("client_id", "client_secret", "workspace_url")),
+    ids=(
+        "plain",
+        "with-empty-databricks-oauth",
+        "with-partial-databricks-oauth",
+        "with-complete-databricks-oauth",
+    ),
+)
+def test_a_url_agent_keeps_its_a2a_call_to_the_card_url(gateway: Gateway, oauth_fields: tuple[str, ...] | None) -> None:
+    marker: Final = f"url-agent-{uuid.uuid4().hex}"
+
+    def peer(request: Request) -> Reply:
+        if request.method == "GET":
+            return Reply(status=404)
+        outbound: Final = _body(request)
+        return Reply(
+            body=json.dumps(
+                {
+                    "jsonrpc": "2.0",
+                    "id": outbound["id"],
+                    "result": {
+                        "kind": "message",
+                        "role": "agent",
+                        "messageId": marker,
+                        "parts": [{"kind": "text", "text": _REPLY}],
+                    },
+                }
+            ).encode()
+        )
+
+    with wire_server(peer) as wire, gateway.scenario() as scenario:
+        complete_block: Final = {"client_id": _CLIENT_ID, "client_secret": _CLIENT_SECRET, "workspace_url": wire.url}
+        oauth_block: Final = None if oauth_fields is None else {field: complete_block[field] for field in oauth_fields}
+        created: Final = gateway.request(
+            "POST",
+            "/v1/agents",
+            {
+                "agent_name": marker,
+                "agent_card_params": {
+                    "protocolVersion": "0.3",
+                    "name": marker,
+                    "description": "URL agent",
+                    "url": wire.url + "/",
+                    "version": "1.0.0",
+                    "defaultInputModes": ["text"],
+                    "defaultOutputModes": ["text"],
+                    "capabilities": {"streaming": False},
+                    "skills": [],
+                },
+                **({"litellm_params": {"databricks_oauth": oauth_block}} if oauth_block is not None else {}),
+            },
+        )
+        assert created.status_code == 200, created.text
+        agent_id: Final = created.json()["agent_id"]
+        scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}"))
+        response: Final = _chat(gateway, marker, extra_headers={"Authorization": "Bearer client-token"})
+        assert response.status_code == 200, response.text
+        assert response.json()["choices"][0]["message"]["content"] == _REPLY
+        posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+        assert [request.target for request in posts] == ["/"]
+        assert _body(posts[0])["method"] == "message/send"
+        assert posts[0].headers["authorization"] == "Bearer client-token"
+
+
+def test_a_url_agent_with_an_incomplete_databricks_oauth_answers_a2a_with_an_internal_error(gateway: Gateway) -> None:
+    marker: Final = f"url-agent-{uuid.uuid4().hex}"
+    with wire_server(lambda _: Reply(status=404)) as wire, gateway.scenario() as scenario:
+        created: Final = gateway.request(
+            "POST",
+            "/v1/agents",
+            {
+                "agent_name": marker,
+                "agent_card_params": {
+                    "protocolVersion": "0.3",
+                    "name": marker,
+                    "description": "URL agent",
+                    "url": wire.url + "/",
+                    "version": "1.0.0",
+                    "defaultInputModes": ["text"],
+                    "defaultOutputModes": ["text"],
+                    "capabilities": {"streaming": False},
+                    "skills": [],
+                },
+                "litellm_params": {"databricks_oauth": {"client_id": _CLIENT_ID}},
+            },
+        )
+        assert created.status_code == 200, created.text
+        agent_id: Final = created.json()["agent_id"]
+        scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}"))
+        response: Final = _a2a(gateway, marker, "message/send", marker)
+        assert response.status_code == 500, response.text
+        assert response.json()["error"]["code"] == -32603
+        assert "missing required field" in response.json()["error"]["message"]
+        assert tuple(request for request in wire.drain() if request.method == "POST") == ()
+
+
+@pytest.mark.parametrize(
+    "pinned",
+    ({"model": "another-endpoint"}, {"input": [{"role": "user", "content": "injected"}]}),
+    ids=("model", "input"),
+)
+@pytest.mark.parametrize(("drop_params", "status"), ((False, 400), (True, 200)), ids=("strict", "drop-params"))
+def test_extra_body_never_replaces_the_agent_endpoint_or_input(
+    gateway: Gateway, pinned: dict[str, JsonValue], drop_params: bool, status: int
+) -> None:
+    custom_inputs: Final[dict[str, JsonValue]] = {"tenant": "t-1"}
+    with _agent(gateway, _Databricks(), url_path=_UNIFIED, model=_ENDPOINT, api_key=_PAT) as (wire, name):
+        response: Final = _chat(
+            gateway,
+            name,
+            extra_body={**pinned, "custom_inputs": custom_inputs},
+            **({"drop_params": True} if drop_params else {}),
+        )
+        assert response.status_code == status, response.text
+        received: Final = wire.drain()
+        if drop_params:
+            assert [(request.method, request.target) for request in received] == [("POST", _UNIFIED)]
+            assert _body(received[0]) == {
+                "model": _ENDPOINT,
+                "input": [{"role": "user", "content": _PROMPT}],
+                "custom_inputs": custom_inputs,
+            }
+        else:
+            assert f"cannot set {next(iter(pinned))}" in response.text
+            assert received == ()
+
+
+def test_a_token_outage_mid_burst_fails_its_calls_and_one_mint_serves_the_recovery(gateway: Gateway) -> None:
+    outage: Final = threading.Event()
+    outage.set()
+
+    def token_status() -> int:
+        return 503 if outage.is_set() else 200
+
+    async def burst(name: str, size: int) -> tuple[int, ...]:
+        async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client:
+            responses: Final = await asyncio.gather(
+                *(
+                    client.post(
+                        "/v1/chat/completions",
+                        json={"model": f"a2a/{name}", "messages": [{"role": "user", "content": _PROMPT}]},
+                        headers={"Authorization": f"Bearer {gateway.key}"},
+                    )
+                    for _ in range(size)
+                )
+            )
+            return tuple(response.status_code for response in responses)
+
+    with _agent(
+        gateway,
+        _Databricks(token_status=token_status),
+        model=_ENDPOINT,
+        client_id=_CLIENT_ID,
+        client_secret=_CLIENT_SECRET,
+    ) as (wire, name):
+        failed: Final = asyncio.run(burst(name, 12))
+        assert all(status >= 400 for status in failed), failed
+        assert _calls_to(wire.drain(), _INVOCATIONS) == ()
+        outage.clear()
+        recovered: Final = asyncio.run(burst(name, 12))
+        assert recovered == (200,) * 12, recovered
+        received: Final = wire.drain()
+        assert len(_calls_to(received, _TOKEN)) <= 2
+        for invocation in _calls_to(received, _INVOCATIONS):
+            _assert_invocation(invocation, _MINTED)
+        assert len(_calls_to(received, _INVOCATIONS)) == 12
+
+
+def _owned_config(
+    directory: Path, wire: Wire, api_key: str | None = _PAT, *, allow_client_side_credentials: bool = False
+) -> Path:
+    base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    general_settings: Final = _JSON_OBJECT.validate_python(base.get("general_settings") or {})
+    name: Final = "dbx-owned-agent"
+    config: Final[dict[str, JsonValue]] = {
+        **base,
+        "general_settings": {
+            **general_settings,
+            **({"allow_client_side_credentials": True} if allow_client_side_credentials else {}),
+        },
+        "agents": [
+            {
+                "agent_name": name,
+                "agent_card_params": {
+                    "protocolVersion": "0.3",
+                    "name": name,
+                    "description": "Databricks Agent agent",
+                    "url": wire.url,
+                    "version": "1.0.0",
+                    "defaultInputModes": ["text"],
+                    "defaultOutputModes": ["text"],
+                    "capabilities": {"streaming": True},
+                    "skills": [],
+                },
+                "litellm_params": {
+                    "custom_llm_provider": "databricks_agent",
+                    "model": _ENDPOINT,
+                    "api_base": wire.url,
+                    **({"api_key": api_key} if api_key is not None else {}),
+                },
+            }
+        ],
+    }
+    path: Final = directory / "databricks-agent-owned.yaml"
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+@pytest.mark.timeout(_OWNED_PROXY_BUDGET)
+def test_a_client_api_base_or_base_url_never_takes_the_agent_pat_off_its_card_url(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with wire_server(_Databricks()) as wire, wire_server(lambda _: Reply(body=_agent_reply())) as collector:
+        config: Final = _owned_config(tmp_path, wire, allow_client_side_credentials=True)
+        with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
+            rejected: Final = _chat(owned.gateway, "dbx-owned-agent", api_base=collector.url)
+            served: Final = _chat(owned.gateway, "dbx-owned-agent", base_url=collector.url)
+            assert rejected.status_code == 400, rejected.text
+            assert served.status_code == 200, served.text
+            assert collector.drain() == ()
+            invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+            assert len(invocations) == 1
+            _assert_invocation(invocations[0], _PAT)
+
+
+@pytest.mark.timeout(_OWNED_PROXY_BUDGET)
+def test_a_killed_worker_leaves_the_sibling_serving_the_databricks_agent(gateway: Gateway, tmp_path: Path) -> None:
+    with wire_server(_Databricks()) as wire:
+        with owned_proxy_process(gateway, tmp_path, {}, config=_owned_config(tmp_path, wire), workers=2) as owned:
+            candidate: Final = owned.gateway
+            workers: Final = eventually(
+                lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
+                lambda pids: len(pids) == 2,
+                seconds=graceful_stop_seconds() + 60,
+            )
+            first: Final = _chat(candidate, "dbx-owned-agent")
+            assert first.status_code == 200, first.text
+            replacement: Final = owned.kill_worker(workers[0])
+            assert replacement not in workers
+            served: Final = tuple(_chat(candidate, "dbx-owned-agent") for _ in range(6))
+            assert all(response.status_code == 200 for response in served), [r.text for r in served]
+            assert {response.json()["choices"][0]["message"]["content"] for response in served} == {_REPLY}
+            invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+            assert len(invocations) == 7
+            for invocation in invocations:
+                _assert_invocation(invocation, _PAT)
+
+
+@pytest.mark.timeout(_OWNED_PROXY_BUDGET)
+def test_the_env_service_principal_mints_once_over_an_env_pat_and_env_references_resolve_on_an_api_created_agent(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with wire_server(_Databricks()) as wire:
+        environment: Final = {
+            "DATABRICKS_CLIENT_ID": _CLIENT_ID,
+            "DATABRICKS_CLIENT_SECRET": _CLIENT_SECRET,
+            "DATABRICKS_HOST": wire.url,
+            "DATABRICKS_API_KEY": "env-pat",
+            "DBX_AGENT_HOST": wire.url,
+            "DBX_AGENT_PAT": _PAT,
+        }
+        with owned_proxy_process(gateway, tmp_path, environment, config=_owned_config(tmp_path, wire, None)) as owned:
+            candidate: Final = owned.gateway
+            served: Final = tuple(_chat(candidate, "dbx-owned-agent") for _ in range(3))
+            assert all(response.status_code == 200 for response in served), [r.text for r in served]
+            received: Final = wire.drain()
+            assert len(_calls_to(received, _TOKEN)) == 1, received
+            _assert_token_request(_calls_to(received, _TOKEN)[0])
+            for invocation in _calls_to(received, _INVOCATIONS):
+                _assert_invocation(invocation, _MINTED)
+            assert len(_calls_to(received, _INVOCATIONS)) == 3
+
+            with candidate.scenario() as scenario:
+                name: Final = _register(
+                    candidate,
+                    scenario.cleanups.callback,
+                    wire.url,
+                    model=_ENDPOINT,
+                    api_base="os.environ/DBX_AGENT_HOST",
+                    api_key="os.environ/DBX_AGENT_PAT",
+                )
+                response: Final = _chat(candidate, name)
+                assert response.status_code == 200, response.text
+                invocations: Final = _calls_to(wire.drain(), _INVOCATIONS)
+                assert len(invocations) == 1, invocations
+                _assert_invocation(invocations[0], _PAT)

tests/unit/a2a_protocol/litellm_completion_bridge/test_handler.py

@@ -0,0 +1,151 @@
+import json
+from typing import Final
+
+import httpx
+import pytest
+import respx
+
+import litellm
+from litellm.a2a_protocol.litellm_completion_bridge.handler import (
+    A2A_USER_API_KEY_HASH_PARAM,
+    A2ACompletionBridgeHandler,
+    agent_completion_kwargs,
+    bridge_model_name,
+)
+
+WORKSPACE: Final = "https://adb-1.azuredatabricks.net"
+INVOCATIONS_URL: Final = f"{WORKSPACE}/serving-endpoints/my-agent/invocations"
+APP_URL: Final = "https://my-app-1.azure.databricksapps.com/responses"
+MESSAGE_ITEM: Final = {
+    "type": "message",
+    "id": "msg_1",
+    "role": "assistant",
+    "content": [{"type": "output_text", "text": "pong"}],
+}
+
+
+@pytest.fixture
+def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+
+
+def _a2a_params() -> dict[str, object]:
+    return {"message": {"role": "user", "parts": [{"kind": "text", "text": "ping"}], "messageId": "m-1"}}
+
+
+@pytest.mark.parametrize(
+    ("litellm_params", "expected"),
+    [
+        ({"custom_llm_provider": "databricks_agent", "model": "my-agent"}, "databricks_agent/my-agent"),
+        (
+            {"custom_llm_provider": "databricks_agent", "model": "databricks_agent/my-agent"},
+            "databricks_agent/my-agent",
+        ),
+        ({"custom_llm_provider": "databricks_agent"}, "databricks_agent/"),
+        ({"custom_llm_provider": "databricks_agent", "model": None}, "databricks_agent/"),
+        ({"custom_llm_provider": "databricks_agent", "model": ""}, "databricks_agent/"),
+        ({"custom_llm_provider": "langgraph"}, "langgraph/agent"),
+        ({"custom_llm_provider": "langgraph", "model": None}, "langgraph/agent"),
+        ({"model": "gpt-5"}, "gpt-5"),
+    ],
+)
+def test_bridge_model_name(litellm_params: dict[str, object], expected: str) -> None:
+    assert bridge_model_name(litellm_params) == expected
+
+
+def test_agent_completion_kwargs_drops_model_provider_and_agent_only_keys() -> None:
+    kwargs = agent_completion_kwargs(
+        {
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": WORKSPACE,
+            "api_key": "pat-1",
+            "is_public": True,
+            "agent_name": "dbx",
+            "agent_id": "id-1",
+            "agent_card_params": {"url": WORKSPACE},
+            A2A_USER_API_KEY_HASH_PARAM: "hash",
+        }
+    )
+    assert dict(kwargs) == {"api_base": WORKSPACE, "api_key": "pat-1"}
+
+
+def test_agent_completion_kwargs_never_expands_proxy_env_secrets(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("PROXY_ONLY_SECRET", "proxy-secret-value")
+    kwargs = agent_completion_kwargs(
+        {
+            "custom_llm_provider": "langgraph",
+            "api_base": "https://agent.example.com",
+            "api_key": "os.environ/PROXY_ONLY_SECRET",
+        }
+    )
+    assert dict(kwargs) == {"api_base": "https://agent.example.com", "api_key": "os.environ/PROXY_ONLY_SECRET"}
+
+
+@respx.mock
+async def test_databricks_agent_message_send_through_the_bridge(httpx_transport: None) -> None:
+    route = respx.post(INVOCATIONS_URL).mock(
+        return_value=httpx.Response(200, json={"object": "response", "output": [MESSAGE_ITEM]})
+    )
+    response = await A2ACompletionBridgeHandler.handle_non_streaming(
+        request_id="req-1",
+        params=_a2a_params(),
+        litellm_params={
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": WORKSPACE,
+            "api_key": "pat-1",
+        },
+        api_base=None,
+    )
+    sent = route.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer pat-1"
+    assert json.loads(sent.content) == {"input": [{"role": "user", "content": "ping"}]}
+    result = response["result"]
+    assert result["kind"] == "message"
+    assert result["parts"] == [{"kind": "text", "text": "pong"}]
+
+
+@respx.mock
+async def test_databricks_app_credential_wins_over_a_forwarded_client_authorization(httpx_transport: None) -> None:
+    route = respx.post(APP_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    await A2ACompletionBridgeHandler.handle_non_streaming(
+        request_id="req-1",
+        params=_a2a_params(),
+        litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "minted-oauth"},
+        api_base=APP_URL,
+        agent_extra_headers={"authorization": "Bearer client-token", "X-LiteLLM-User-Id": "u-1"},
+    )
+    sent = route.calls.last.request
+    assert sent.headers.get_list("Authorization") == ["Bearer minted-oauth"]
+    assert sent.headers["X-LiteLLM-User-Id"] == "u-1"
+
+
+@respx.mock
+async def test_databricks_agent_message_stream_through_the_bridge(httpx_transport: None) -> None:
+    events = [
+        {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "po"},
+        {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "ng"},
+        {"type": "response.output_item.done", "item": MESSAGE_ITEM},
+    ]
+    body = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n"
+    respx.post(INVOCATIONS_URL).mock(
+        return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body)
+    )
+    chunks = [
+        chunk
+        async for chunk in A2ACompletionBridgeHandler.handle_streaming(
+            request_id="req-1",
+            params=_a2a_params(),
+            litellm_params={
+                "custom_llm_provider": "databricks_agent",
+                "model": "my-agent",
+                "api_base": WORKSPACE,
+                "api_key": "pat-1",
+            },
+        )
+    ]
+    kinds = [chunk["result"]["kind"] for chunk in chunks]
+    assert kinds == ["task", "status-update", "artifact-update", "status-update"]
+    assert chunks[2]["result"]["artifact"]["parts"] == [{"kind": "text", "text": "pong"}]
+    assert chunks[3]["result"]["status"]["state"] == "completed"

tests/unit/caching/test_caching_handler.py

@@ -172,6 +172,50 @@ async def test_async_copilot_requests_bypass_response_cache(copilot_response_cac
     _assert_copilot_cache_empty(copilot_response_cache)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stream", [False, True])
+async def test_databricks_agent_replies_are_never_served_from_the_response_cache(
+    copilot_response_cache: Cache, stream: bool
+) -> None:
+    workspace: Final = "https://adb-cache.azuredatabricks.net"
+    replies: Final = iter(("first agent answer", "second agent answer"))
+    requests: Final[list[httpx.Request]] = []
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        text: Final = next(replies)
+        if json.loads(request.content).get("stream"):
+            frame: Final = {"type": "response.output_text.delta", "item_id": "m", "delta": text}
+            return httpx.Response(200, content=f"data: {json.dumps(frame)}\n\n".encode(), request=request)
+        message: Final = {"type": "message", "id": "m", "content": [{"type": "output_text", "text": text}]}
+        return httpx.Response(200, json={"output": [message]}, request=request)
+
+    async def answer() -> str:
+        response: Final = await litellm.acompletion(
+            model="databricks_agent/support-agent",
+            messages=[{"role": "user", "content": "same question"}],
+            api_base=workspace,
+            api_key="dbx-pat",
+            client=client,
+            stream=stream,
+            caching=True,
+        )
+        if isinstance(response, CustomStreamWrapper):
+            return "".join([chunk.choices[0].delta.content or "" async for chunk in response])
+        return cast(str, response.choices[0].message.content)  # cast-ok: a ModelResponse's message content
+
+    client: Final = AsyncHTTPHandler(transport=_CopilotCacheAsyncTransport(respond))
+    async with client.client:
+        first: Final = await answer()
+        await asyncio.gather(*_PENDING_CACHE_WRITES)
+        second: Final = await answer()
+        await asyncio.gather(*_PENDING_CACHE_WRITES)
+
+    assert (first, second) == ("first agent answer", "second agent answer")
+    assert len(requests) == 2
+    _assert_copilot_cache_empty(copilot_response_cache)
+
+
 @pytest.mark.asyncio
 async def test_existing_provider_response_cache_still_hits(copilot_response_cache: Cache) -> None:
     messages: Final = [{"role": "user", "content": "cached provider control"}]

tests/unit/llms/databricks/agent/test_databricks_agent_transformation.py

@@ -0,0 +1,725 @@
+import json
+from collections.abc import Mapping
+from typing import Final
+from unittest.mock import MagicMock
+from uuid import uuid4
+
+import httpx
+import pytest
+import responses
+import respx
+
+import litellm
+from litellm.llms.databricks.agent.responses_output import (
+    AgentResponse,
+    DatabricksAgentResponsesIterator,
+    output_text,
+)
+from litellm.llms.databricks.agent.transformation import (
+    DatabricksAgentConfig,
+    is_databricks_app_url,
+    resolve_endpoint_url,
+    workspace_url_from_api_base,
+)
+from litellm.llms.databricks.common_utils import DatabricksException
+from litellm.types.utils import ModelResponse
+
+WORKSPACE: Final = "https://adb-1.azuredatabricks.net"
+TOKEN_URL: Final = f"{WORKSPACE}/oidc/v1/token"
+INVOCATIONS_URL: Final = f"{WORKSPACE}/serving-endpoints/my-agent/invocations"
+UNIFIED_URL: Final = f"{WORKSPACE}/serving-endpoints/responses"
+APP_HOST: Final = "https://my-app-1.azure.databricksapps.com"
+APP_URL: Final = f"{APP_HOST}/responses"
+MESSAGE_ITEM: Final = {
+    "type": "message",
+    "id": "msg_1",
+    "role": "assistant",
+    "content": [{"type": "output_text", "text": "pong", "annotations": []}],
+}
+_TOOL_CALL: Final = {
+    "id": "call_1",
+    "type": "function",
+    "function": {"name": "get_weather", "arguments": "{}"},
+}
+
+
+@pytest.fixture
+def env_oauth_m2m(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("DATABRICKS_HOST", raising=False)
+    monkeypatch.setenv("DATABRICKS_CLIENT_ID", "sp-id")
+    monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", f"sp-secret-{uuid4().hex}")
+
+
+@pytest.fixture
+def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+
+
+def _sse(events: list[dict[str, object]]) -> str:
+    return "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n"
+
+
+@pytest.mark.parametrize(
+    ("api_base", "model", "expected"),
+    [
+        (WORKSPACE, "my-agent", INVOCATIONS_URL),
+        (f"{WORKSPACE}/", "databricks_agent/my-agent", INVOCATIONS_URL),
+        (f"{WORKSPACE}/serving-endpoints", "my-agent", INVOCATIONS_URL),
+        (INVOCATIONS_URL, "anything", INVOCATIONS_URL),
+        (f"{WORKSPACE}/serving-endpoints/my-agent", "", INVOCATIONS_URL),
+        (f"{APP_URL}/", "agent", APP_URL),
+        (APP_HOST, "", APP_URL),
+        (f"{APP_HOST}/", "agent", APP_URL),
+        (f"{APP_HOST}/api/agent", "", f"{APP_HOST}/api/agent"),
+        (f"{UNIFIED_URL}/", "databricks_agent/my-agent", UNIFIED_URL),
+        (WORKSPACE, "team/agent v2", f"{WORKSPACE}/serving-endpoints/team%2Fagent%20v2/invocations"),
+    ],
+)
+def test_resolve_endpoint_url(api_base: str, model: str, expected: str) -> None:
+    assert resolve_endpoint_url(api_base, model) == expected
+
+
+@pytest.mark.parametrize("app_base", [APP_HOST, f"{APP_HOST}/", APP_URL, "https://other.example/responses"])
+def test_an_app_url_is_never_a_workspace(app_base: str) -> None:
+    assert is_databricks_app_url(app_base)
+    assert workspace_url_from_api_base(app_base) is None
+
+
+@pytest.mark.parametrize("api_base", [WORKSPACE, f"{WORKSPACE}/serving-endpoints", INVOCATIONS_URL, UNIFIED_URL])
+def test_model_serving_urls_resolve_to_the_workspace(api_base: str) -> None:
+    assert not is_databricks_app_url(api_base)
+    assert workspace_url_from_api_base(api_base) == WORKSPACE
+
+
+def test_resolve_endpoint_url_needs_an_endpoint_name_for_a_workspace_url() -> None:
+    with pytest.raises(DatabricksException, match="model must name the Databricks serving endpoint") as info:
+        resolve_endpoint_url(WORKSPACE, "databricks_agent/")
+    assert info.value.status_code == 400
+
+
+def test_unified_responses_url_needs_an_endpoint_name_for_the_body_model() -> None:
+    with pytest.raises(DatabricksException, match="model must name the Databricks serving endpoint") as info:
+        resolve_endpoint_url(UNIFIED_URL, "databricks_agent/")
+    assert info.value.status_code == 400
+
+
+def test_get_complete_url_requires_api_base() -> None:
+    with pytest.raises(DatabricksException, match="api_base is required for databricks_agent"):
+        DatabricksAgentConfig().get_complete_url(
+            api_base=None, api_key=None, model="my-agent", optional_params={}, litellm_params={}
+        )
+
+
+def test_transform_request_maps_messages_to_responses_input_with_passthrough_fields() -> None:
+    body = DatabricksAgentConfig().transform_request(
+        model="my-agent",
+        messages=[
+            {"role": "system", "content": "be terse"},
+            {"role": "user", "content": [{"type": "text", "text": "ping"}, {"type": "text", "text": "pong?"}]},
+            {"role": "assistant", "content": "pong"},
+        ],
+        optional_params={"custom_inputs": {"tenant": "t-1"}, "context": {"conversation_id": "c-1"}, "ignored": 1},
+        litellm_params={},
+        headers={},
+    )
+    assert body == {
+        "input": [
+            {"role": "system", "content": "be terse"},
+            {"role": "user", "content": "pingpong?"},
+            {"role": "assistant", "content": "pong"},
+        ],
+        "custom_inputs": {"tenant": "t-1"},
+        "context": {"conversation_id": "c-1"},
+    }
+
+
+def test_transform_request_sets_stream_only_when_streaming() -> None:
+    config = DatabricksAgentConfig()
+    messages = [{"role": "user", "content": "ping"}]
+    assert "stream" not in config.transform_request("m", messages, {"stream": False}, {}, {})
+    assert config.transform_request("m", messages, {"stream": True}, {}, {})["stream"] is True
+
+
+def test_transform_request_rejects_roles_a_responses_agent_cannot_take(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "drop_params", False)
+    with pytest.raises(DatabricksException, match="role 'tool' is not supported") as info:
+        DatabricksAgentConfig().transform_request(
+            model="my-agent",
+            messages=[{"role": "tool", "content": "42", "tool_call_id": "call_1"}],
+            optional_params={},
+            litellm_params={},
+            headers={},
+        )
+    assert info.value.status_code == 400
+
+
+def test_transform_request_rejects_assistant_tool_calls(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "drop_params", False)
+    with pytest.raises(DatabricksException, match="remove tool_calls from the message") as info:
+        DatabricksAgentConfig().transform_request(
+            model="my-agent",
+            messages=[
+                {"role": "user", "content": "weather?"},
+                {"role": "assistant", "content": None, "tool_calls": [_TOOL_CALL]},
+            ],
+            optional_params={},
+            litellm_params={},
+            headers={},
+        )
+    assert info.value.status_code == 400
+
+
+def test_transform_request_drops_tool_turns_and_empty_messages_under_drop_params(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "drop_params", False)
+    body = DatabricksAgentConfig().transform_request(
+        model="my-agent",
+        messages=[
+            {"role": "user", "content": "weather?"},
+            {"role": "assistant", "content": None, "tool_calls": [_TOOL_CALL]},
+            {"role": "tool", "content": "sunny", "tool_call_id": "call_1"},
+            {"role": "assistant", "content": "Checking the weather", "tool_calls": [_TOOL_CALL]},
+            {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}]},
+            {"role": "user", "content": "and tomorrow?"},
+        ],
+        optional_params={},
+        litellm_params={"drop_params": True},
+        headers={},
+    )
+    assert body == {
+        "input": [
+            {"role": "user", "content": "weather?"},
+            {"role": "assistant", "content": "Checking the weather"},
+            {"role": "user", "content": "and tomorrow?"},
+        ]
+    }
+
+
+@pytest.mark.parametrize(
+    "part",
+    [
+        {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
+        {"type": "file", "file": {"file_id": "file-1"}},
+    ],
+)
+def test_transform_request_rejects_content_a_responses_agent_cannot_take(
+    part: dict[str, object], monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "drop_params", False)
+    with pytest.raises(DatabricksException, match=str(part["type"])) as info:
+        DatabricksAgentConfig().transform_request(
+            model="my-agent",
+            messages=[{"role": "user", "content": [{"type": "text", "text": "describe this"}, part]}],
+            optional_params={},
+            litellm_params={},
+            headers={},
+        )
+    assert info.value.status_code == 400
+
+
+@pytest.mark.parametrize(
+    ("global_drop_params", "litellm_params"),
+    [(True, {}), (False, {"drop_params": True})],
+)
+def test_transform_request_drops_non_text_content_under_drop_params(
+    global_drop_params: bool, litellm_params: dict[str, object], monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "drop_params", global_drop_params)
+    body = DatabricksAgentConfig().transform_request(
+        model="my-agent",
+        messages=[
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "describe this"},
+                    {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
+                ],
+            }
+        ],
+        optional_params={},
+        litellm_params=litellm_params,
+        headers={},
+    )
+    assert body == {"input": [{"role": "user", "content": "describe this"}]}
+
+
+def test_transform_request_joins_text_content_parts() -> None:
+    body = DatabricksAgentConfig().transform_request(
+        model="my-agent",
+        messages=[{"role": "user", "content": [{"type": "text", "text": "pi"}, {"type": "text", "text": "ng"}]}],
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )
+    assert body == {"input": [{"role": "user", "content": "ping"}]}
+
+
+def test_validate_environment_lets_the_api_key_replace_a_header_authorization() -> None:
+    headers = DatabricksAgentConfig().validate_environment(
+        headers={"authorization": "Bearer client-token", "X-LiteLLM-User-Id": "u1"},
+        model="my-agent",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key="pat-1",
+    )
+    assert headers == {"X-LiteLLM-User-Id": "u1", "Content-Type": "application/json", "Authorization": "Bearer pat-1"}
+
+
+def test_validate_environment_uses_a_header_authorization_when_no_api_key_is_set() -> None:
+    headers = DatabricksAgentConfig().validate_environment(
+        headers={"authorization": "Bearer client-token"},
+        model="my-agent",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key=None,
+    )
+    assert headers == {"authorization": "Bearer client-token", "Content-Type": "application/json"}
+
+
+def test_validate_environment_uses_the_pat_as_bearer() -> None:
+    headers = DatabricksAgentConfig().validate_environment(
+        headers={}, model="my-agent", messages=[], optional_params={}, litellm_params={}, api_key="pat-1"
+    )
+    assert headers == {"Content-Type": "application/json", "Authorization": "Bearer pat-1"}
+
+
+def test_validate_environment_without_any_credential_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("DATABRICKS_CLIENT_ID", raising=False)
+    monkeypatch.delenv("DATABRICKS_CLIENT_SECRET", raising=False)
+    with pytest.raises(DatabricksException, match="Missing Databricks credentials") as info:
+        DatabricksAgentConfig().validate_environment(
+            headers={}, model="my-agent", messages=[], optional_params={}, litellm_params={}, api_key=None
+        )
+    assert info.value.status_code == 400
+
+
+def _validate_environment(api_key: str | None, api_base: str | None) -> dict[str, object]:
+    return DatabricksAgentConfig().validate_environment(
+        headers={"authorization": "Bearer client-token"},
+        model="my-agent",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key=api_key,
+        api_base=api_base,
+    )
+
+
+@responses.activate
+@pytest.mark.parametrize("api_base", [WORKSPACE, f"{WORKSPACE}/serving-endpoints", INVOCATIONS_URL])
+def test_validate_environment_mints_an_oauth_m2m_token_from_the_env_service_principal(
+    env_oauth_m2m: None, api_base: str
+) -> None:
+    token = responses.post(TOKEN_URL, json={"access_token": "m2m-token", "expires_in": 3600})
+    headers = _validate_environment(api_key=None, api_base=api_base)
+    assert headers == {"Content-Type": "application/json", "Authorization": "Bearer m2m-token"}
+    assert token.calls[0].request.headers["Authorization"].startswith("Basic ")
+    assert "grant_type=client_credentials" in str(token.calls[0].request.body)
+
+
+@responses.activate
+def test_validate_environment_mints_the_env_token_from_databricks_host_for_an_app(
+    env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("DATABRICKS_HOST", f"{WORKSPACE}/")
+    responses.post(TOKEN_URL, json={"access_token": "m2m-token"})
+    assert _validate_environment(api_key=None, api_base=APP_HOST)["Authorization"] == "Bearer m2m-token"
+
+
+@responses.activate
+def test_validate_environment_never_posts_the_env_secret_to_an_app_host(env_oauth_m2m: None) -> None:
+    with pytest.raises(DatabricksException, match="set DATABRICKS_HOST") as info:
+        _validate_environment(api_key=None, api_base=APP_URL)
+    assert info.value.status_code == 400
+    assert not responses.calls
+
+
+@responses.activate
+def test_validate_environment_prefers_the_api_key_over_the_env_service_principal(env_oauth_m2m: None) -> None:
+    assert _validate_environment(api_key="pat-1", api_base=WORKSPACE)["Authorization"] == "Bearer pat-1"
+    assert not responses.calls
+
+
+def test_output_text_joins_message_items_and_skips_tool_trace_items() -> None:
+    output = [
+        {"type": "function_call", "name": "lookup", "arguments": "{}", "call_id": "c1"},
+        {"type": "function_call_output", "call_id": "c1", "output": "secret trace"},
+        {
+            "type": "message",
+            "content": [{"type": "output_text", "text": "Let me check."}, {"type": "refusal", "refusal": "no"}],
+        },
+        {"type": "message", "content": [{"type": "output_text", "text": "It is sunny."}]},
+    ]
+    assert output_text(AgentResponse.model_validate({"output": output}).output) == "Let me check.\n\nIt is sunny."
+    assert output_text(AgentResponse.model_validate({"output": output[:2]}).output) == ""
+    assert output_text(None) == ""
+
+
+def test_transform_response_reads_usage_and_custom_outputs() -> None:
+    raw = httpx.Response(
+        200,
+        json={
+            "object": "response",
+            "output": [MESSAGE_ITEM],
+            "usage": {"input_tokens": 7, "output_tokens": 3},
+            "custom_outputs": {"trace_id": "tr-1"},
+        },
+        request=httpx.Request("POST", INVOCATIONS_URL),
+    )
+    logging_obj = MagicMock()
+    response = DatabricksAgentConfig().transform_response(
+        model="my-agent",
+        raw_response=raw,
+        model_response=ModelResponse(),
+        logging_obj=logging_obj,
+        request_data={},
+        messages=[{"role": "user", "content": "ping"}],
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+    assert logging_obj.post_call.call_args.kwargs["original_response"] == raw.text
+    assert response.choices[0].message.content == "pong"
+    assert response.choices[0].message.provider_specific_fields == {"custom_outputs": {"trace_id": "tr-1"}}
+    assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (7, 3, 10)
+
+
+def test_transform_response_estimates_usage_when_the_agent_reports_none() -> None:
+    raw = httpx.Response(
+        200, json={"object": "response", "output": [MESSAGE_ITEM]}, request=httpx.Request("POST", INVOCATIONS_URL)
+    )
+    response = DatabricksAgentConfig().transform_response(
+        model="my-agent",
+        raw_response=raw,
+        model_response=ModelResponse(),
+        logging_obj=MagicMock(),
+        request_data={},
+        messages=[{"role": "user", "content": "ping"}],
+        optional_params={},
+        litellm_params={},
+        encoding=None,
+    )
+    assert response.choices[0].message.provider_specific_fields is None
+    assert response.usage.prompt_tokens > 0
+    assert response.usage.completion_tokens > 0
+    assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens
+
+
+@pytest.mark.parametrize("body", ["<html>login</html>", "[]", "{}", '{"output": "pong"}'])
+def test_transform_response_rejects_a_body_that_is_not_a_responses_agent_response(body: str) -> None:
+    raw = httpx.Response(200, text=body, request=httpx.Request("POST", APP_URL))
+    with pytest.raises(DatabricksException, match="not a ResponsesAgent response") as info:
+        DatabricksAgentConfig().transform_response(
+            model="my-agent",
+            raw_response=raw,
+            model_response=ModelResponse(),
+            logging_obj=MagicMock(),
+            request_data={},
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            encoding=None,
+        )
+    assert info.value.status_code == 502
+
+
+def test_stream_iterator_emits_deltas_once_and_whole_items_only_when_nothing_streamed() -> None:
+    iterator = DatabricksAgentResponsesIterator(streaming_response=iter(()), sync_stream=True)
+    chunks = [
+        iterator.chunk_parser({"type": "response.output_text.delta", "item_id": "msg_1", "delta": "po"}),
+        iterator.chunk_parser({"type": "response.output_text.delta", "item_id": "msg_1", "delta": "ng"}),
+        iterator.chunk_parser({"type": "response.output_item.done", "item": MESSAGE_ITEM}),
+        iterator.chunk_parser({"type": "response.output_item.done", "item": {**MESSAGE_ITEM, "id": "msg_2"}}),
+        iterator.chunk_parser({"type": "response.output_item.done", "item": {"type": "function_call", "id": "fc_1"}}),
+    ]
+    assert [chunk["text"] for chunk in chunks] == ["po", "ng", "", "\n\npong", ""]
+    assert all(chunk["is_finished"] is False for chunk in chunks)
+
+
+def test_stream_iterator_separates_consecutive_message_items_like_the_non_streaming_reply() -> None:
+    iterator = DatabricksAgentResponsesIterator(streaming_response=iter(()), sync_stream=True)
+    events: list[dict[str, object]] = [
+        {"type": "response.output_text.delta", "item_id": "msg_1", "delta": ""},
+        {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Let me check."},
+        {"type": "response.output_item.done", "item": {"type": "function_call", "id": "fc_1"}},
+        {"type": "response.output_text.delta", "item_id": "msg_2", "delta": "It is"},
+        {"type": "response.output_text.delta", "item_id": "msg_2", "delta": " sunny."},
+    ]
+    streamed = "".join(iterator.chunk_parser(event)["text"] for event in events)
+    whole = output_text(
+        AgentResponse.model_validate(
+            {
+                "output": [
+                    {"type": "message", "content": [{"type": "output_text", "text": "Let me check."}]},
+                    {"type": "function_call", "id": "fc_1"},
+                    {"type": "message", "content": [{"type": "output_text", "text": "It is sunny."}]},
+                ]
+            }
+        ).output
+    )
+    assert streamed == whole == "Let me check.\n\nIt is sunny."
+
+
+def test_stream_iterator_surfaces_agent_errors() -> None:
+    iterator = DatabricksAgentResponsesIterator(streaming_response=iter(()), sync_stream=True)
+    with pytest.raises(DatabricksException, match="tool failed"):
+        iterator.chunk_parser({"type": "error", "message": "tool failed"})
+
+
+@respx.mock
+async def test_acompletion_posts_responses_input_to_model_serving(httpx_transport: None) -> None:
+    route = respx.post(INVOCATIONS_URL).mock(
+        return_value=httpx.Response(200, json={"object": "response", "output": [MESSAGE_ITEM]})
+    )
+    response = await litellm.acompletion(
+        model="databricks_agent/my-agent",
+        messages=[{"role": "user", "content": "ping"}],
+        api_base=WORKSPACE,
+        api_key="pat-1",
+        custom_inputs={"tenant": "t-1"},
+    )
+    sent = route.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer pat-1"
+    assert json.loads(sent.content) == {
+        "input": [{"role": "user", "content": "ping"}],
+        "custom_inputs": {"tenant": "t-1"},
+    }
+    assert response.choices[0].message.content == "pong"
+    assert response.model == "my-agent"
+
+
+@respx.mock
+async def test_acompletion_names_the_endpoint_in_the_body_for_the_unified_responses_url(
+    httpx_transport: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("DATABRICKS_API_BASE", UNIFIED_URL)
+    monkeypatch.setenv("DATABRICKS_API_KEY", "pat-env")
+    route = respx.post(UNIFIED_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    response = await litellm.acompletion(
+        model="databricks_agent/my-agent",
+        messages=[{"role": "user", "content": "ping"}],
+    )
+    sent = route.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer pat-env"
+    assert json.loads(sent.content) == {"model": "my-agent", "input": [{"role": "user", "content": "ping"}]}
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+async def test_acompletion_forwards_extra_headers_and_keeps_the_minted_token(httpx_transport: None) -> None:
+    route = respx.post(APP_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    await litellm.acompletion(
+        model="databricks_agent/agent",
+        messages=[{"role": "user", "content": "ping"}],
+        api_base=APP_URL,
+        extra_headers={"Authorization": "Bearer minted-oauth", "X-LiteLLM-Trace-Id": "tr-1"},
+    )
+    sent = route.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer minted-oauth"
+    assert sent.headers["X-LiteLLM-Trace-Id"] == "tr-1"
+
+
+@respx.mock
+async def test_acompletion_streams_output_text_deltas(httpx_transport: None) -> None:
+    events: list[dict[str, object]] = [
+        {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "po"},
+        {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "ng"},
+        {"type": "response.output_item.done", "item": MESSAGE_ITEM},
+    ]
+    route = respx.post(INVOCATIONS_URL).mock(
+        return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_sse(events))
+    )
+    stream = await litellm.acompletion(
+        model="databricks_agent/my-agent",
+        messages=[{"role": "user", "content": "ping"}],
+        api_base=WORKSPACE,
+        api_key="pat-1",
+        stream=True,
+    )
+    deltas = [chunk.choices[0].delta.content async for chunk in stream if chunk.choices[0].delta.content]
+    assert json.loads(route.calls.last.request.content)["stream"] is True
+    assert "".join(deltas) == "pong"
+
+
+@pytest.mark.parametrize("env_pat", (None, "pat-env"), ids=("alone", "beside-an-env-pat"))
+@responses.activate
+@respx.mock
+async def test_acompletion_authenticates_with_the_env_service_principal(
+    httpx_transport: None, env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch, env_pat: str | None
+) -> None:
+    if env_pat is None:
+        monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
+    else:
+        monkeypatch.setenv("DATABRICKS_API_KEY", env_pat)
+    responses.post(TOKEN_URL, json={"access_token": "m2m-token"})
+    route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    response = await litellm.acompletion(
+        model="databricks_agent/my-agent",
+        messages=[{"role": "user", "content": "ping"}],
+        api_base=WORKSPACE,
+    )
+    assert route.calls.last.request.headers["Authorization"] == "Bearer m2m-token"
+    assert response.choices[0].message.content == "pong"
+
+
+async def test_unsupported_openai_params_follow_drop_params(httpx_transport: None) -> None:
+    messages = [{"role": "user", "content": "ping"}]
+    with pytest.raises(litellm.UnsupportedParamsError):
+        await litellm.acompletion(
+            model="databricks_agent/my-agent",
+            messages=messages,
+            api_base=WORKSPACE,
+            api_key="pat-1",
+            frequency_penalty=0.2,
+        )
+    with respx.mock:
+        route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+        await litellm.acompletion(
+            model="databricks_agent/my-agent",
+            messages=messages,
+            api_base=WORKSPACE,
+            api_key="pat-1",
+            frequency_penalty=0.2,
+            drop_params=True,
+        )
+    assert "frequency_penalty" not in json.loads(route.calls.last.request.content)
+
+
+@respx.mock
+@pytest.mark.parametrize("token_limit", [{"max_tokens": 64}, {"max_completion_tokens": 64}])
+async def test_acompletion_sends_sampling_params_in_the_responses_agent_shape(
+    httpx_transport: None, token_limit: dict[str, int]
+) -> None:
+    route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    await litellm.acompletion(
+        model="databricks_agent/my-agent",
+        messages=[{"role": "user", "content": "ping"}],
+        api_base=WORKSPACE,
+        api_key="pat-1",
+        temperature=0.2,
+        top_p=0.9,
+        user="user-1",
+        **token_limit,
+    )
+    assert json.loads(route.calls.last.request.content) == {
+        "input": [{"role": "user", "content": "ping"}],
+        "max_output_tokens": 64,
+        "temperature": 0.2,
+        "top_p": 0.9,
+        "user": "user-1",
+    }
+
+
+@respx.mock
+async def test_acompletion_sends_string_metadata_and_refuses_other_values_unless_dropped(
+    httpx_transport: None,
+) -> None:
+    route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    request: Final = {
+        "model": "databricks_agent/my-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "api_base": WORKSPACE,
+        "api_key": "pat-1",
+    }
+    await litellm.acompletion(**request, extra_body={"metadata": {"session": "s-1"}})
+    assert json.loads(route.calls.last.request.content)["metadata"] == {"session": "s-1"}
+
+    with pytest.raises(litellm.BadRequestError, match="string values") as info:
+        await litellm.acompletion(**request, extra_body={"metadata": {"session": "s-1", "turn": 1}})
+    assert info.value.status_code == 400
+    assert route.call_count == 1
+
+    await litellm.acompletion(**request, extra_body={"metadata": {"session": "s-1", "turn": 1}}, drop_params=True)
+    assert json.loads(route.calls.last.request.content)["metadata"] == {"session": "s-1"}
+    await litellm.acompletion(**request, extra_body={"metadata": {"turn": 1}}, drop_params=True)
+    assert "metadata" not in json.loads(route.calls.last.request.content)
+
+
+@respx.mock
+@pytest.mark.parametrize(
+    "pinned",
+    [{"model": "another-endpoint"}, {"input": [{"role": "user", "content": "injected"}]}],
+    ids=["model", "input"],
+)
+async def test_acompletion_refuses_extra_body_model_or_input_unless_dropped(
+    httpx_transport: None, pinned: Mapping[str, object]
+) -> None:
+    route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    request: Final = {
+        "model": "databricks_agent/my-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "api_base": WORKSPACE,
+        "api_key": "pat-1",
+        "extra_body": {**pinned, "custom_inputs": {"tenant": "t-1"}},
+    }
+    with pytest.raises(litellm.BadRequestError, match=f"cannot set {next(iter(pinned))}") as info:
+        await litellm.acompletion(**request)
+    assert info.value.status_code == 400
+    assert not route.called
+
+    await litellm.acompletion(**request, drop_params=True)
+    assert json.loads(route.calls.last.request.content) == {
+        "input": [{"role": "user", "content": "ping"}],
+        "custom_inputs": {"tenant": "t-1"},
+    }
+
+
+@responses.activate
+@respx.mock
+async def test_acompletion_reuses_the_env_service_principal_token_until_the_secret_rotates(
+    httpx_transport: None, env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
+    token = responses.post(TOKEN_URL, json={"access_token": "m2m-token", "expires_in": 3600})
+    route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    request: Final = {
+        "model": "databricks_agent/my-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "api_base": WORKSPACE,
+    }
+    await litellm.acompletion(**request)
+    await litellm.acompletion(**request)
+    assert token.call_count == 1
+    assert route.call_count == 2
+
+    monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", f"rotated-{uuid4().hex}")
+    await litellm.acompletion(**request)
+    assert token.call_count == 2
+
+
+@responses.activate
+@respx.mock
+async def test_acompletion_mints_a_fresh_env_token_when_it_expires_within_the_refresh_buffer(
+    httpx_transport: None, env_oauth_m2m: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
+    token = responses.post(TOKEN_URL, json={"access_token": "m2m-token", "expires_in": 30})
+    respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    request: Final = {
+        "model": "databricks_agent/my-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "api_base": WORKSPACE,
+    }
+    await litellm.acompletion(**request)
+    await litellm.acompletion(**request)
+    assert token.call_count == 2
+
+
+@respx.mock
+async def test_acompletion_resolves_os_environ_api_base_and_api_key(
+    httpx_transport: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("MY_DBX_HOST", WORKSPACE)
+    monkeypatch.setenv("MY_DBX_PAT", "pat-from-env")
+    route = respx.post(INVOCATIONS_URL).mock(return_value=httpx.Response(200, json={"output": [MESSAGE_ITEM]}))
+    response = await litellm.acompletion(
+        model="databricks_agent/my-agent",
+        messages=[{"role": "user", "content": "ping"}],
+        api_base="os.environ/MY_DBX_HOST",
+        api_key="os.environ/MY_DBX_PAT",
+    )
+    assert route.calls.last.request.headers["Authorization"] == "Bearer pat-from-env"
+    assert response.choices[0].message.content == "pong"

tests/unit/llms/databricks/test_databricks_partner_integration.py

@@ -197,7 +197,7 @@ class TestOAuthM2M:
         mock_response.json.return_value = {"access_token": "test-access-token"}
 
         with patch("requests.post", return_value=mock_response) as mock_post:
-            token = databricks_base._get_oauth_m2m_token(
+            token = databricks_base.get_oauth_m2m_token(
                 api_base="https://adb-123.azuredatabricks.net/serving-endpoints",
                 client_id="test-client-id",
                 client_secret="test-client-secret",
@@ -219,7 +219,7 @@ class TestOAuthM2M:
 
         with patch("requests.post", return_value=mock_response):
             with pytest.raises(DatabricksException) as exc_info:
-                databricks_base._get_oauth_m2m_token(
+                databricks_base.get_oauth_m2m_token(
                     api_base="https://adb-123.azuredatabricks.net",
                     client_id="bad-client-id",
                     client_secret="bad-secret",
@@ -235,7 +235,7 @@ class TestOAuthM2M:
         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(
+            databricks_base.get_oauth_m2m_token(
                 api_base="https://adb-123.azuredatabricks.net/serving-endpoints",
                 client_id="id",
                 client_secret="secret",
@@ -254,7 +254,7 @@ class TestOAuthM2M:
         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(
+            databricks_base.get_oauth_m2m_token(
                 api_base="https://adb-123.azuredatabricks.net/ai-gateway/mlflow/v1",
                 client_id="id",
                 client_secret="secret",
@@ -278,7 +278,7 @@ class TestValidateEnvironmentWithOAuth:
         databricks_base = DatabricksBase()
 
         with patch.object(
-            databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
+            databricks_base, "get_oauth_m2m_token", return_value="oauth-token"
         ) as mock_oauth:
             api_base, headers = databricks_base.databricks_validate_environment(
                 api_key=None,
@@ -546,7 +546,7 @@ class TestAuthenticationPriority:
         databricks_base = DatabricksBase()
 
         with patch.object(
-            databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
+            databricks_base, "get_oauth_m2m_token", return_value="oauth-token"
         ) as mock_oauth:
             api_base, headers = databricks_base.databricks_validate_environment(
                 api_key=None,  # No PAT provided - OAuth should be used
@@ -570,7 +570,7 @@ class TestAuthenticationPriority:
 
         # Mock the OAuth call - it will be attempted but PAT should override
         with patch.object(
-            databricks_base, "_get_oauth_m2m_token", return_value="oauth-token"
+            databricks_base, "get_oauth_m2m_token", return_value="oauth-token"
         ):
             api_base, headers = databricks_base.databricks_validate_environment(
                 api_key="dapi-explicit-pat",

tests/unit/proxy/agent_endpoints/test_agent_headers.py

@@ -9,10 +9,17 @@ Tests cover:
 - merge_agent_headers utility
 """
 
+import base64
+import json
 import sys
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+import respx
+
+import litellm
 
 # ---------------------------------------------------------------------------
 # Helper: build a minimal mock agent
@@ -378,6 +385,72 @@ async def test_databricks_oauth_overrides_static_authorization():
     assert headers.get("Authorization") == "Bearer oauth-wins"
 
 
+@respx.mock
+@pytest.mark.asyncio
+async def test_databricks_agent_flat_oauth_send_reaches_the_endpoint_with_the_minted_token(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
+    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
+    from litellm.proxy.agent_endpoints.databricks_oauth import databricks_app_oauth_token_cache
+    from litellm.types.agents import AgentResponse
+
+    workspace: Final = "https://adb-1.azuredatabricks.net"
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    databricks_app_oauth_token_cache.flush_cache()
+    registry: Final = AgentRegistry()
+    registry.register_agent(
+        AgentResponse(
+            agent_id="dbx-agent-id",
+            agent_name="dbx-agent",
+            agent_card_params={"url": workspace},
+            litellm_params={
+                "custom_llm_provider": "databricks_agent",
+                "model": "my-agent",
+                "api_base": workspace,
+                "client_id": "cid",
+                "client_secret": "secret",
+            },
+        )
+    )
+    token: Final = respx.post(f"{workspace}/oidc/v1/token").mock(
+        return_value=httpx.Response(200, json={"access_token": "flat-minted", "expires_in": 3600})
+    )
+    endpoint: Final = respx.post(f"{workspace}/serving-endpoints/my-agent/invocations").mock(
+        return_value=httpx.Response(
+            200,
+            json={
+                "output": [
+                    {
+                        "type": "message",
+                        "id": "msg_1",
+                        "role": "assistant",
+                        "content": [{"type": "output_text", "text": "pong"}],
+                    }
+                ]
+            },
+        )
+    )
+
+    with (
+        patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry),
+    ):
+        response = await invoke_agent_a2a(
+            agent_id="dbx-agent",
+            request=_make_mock_request(),
+            fastapi_response=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1"),
+        )
+
+    basic: Final = base64.b64encode(b"cid:secret").decode()
+    assert token.calls.last.request.headers["Authorization"] == f"Basic {basic}"
+    sent: Final = endpoint.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer flat-minted"
+    assert json.loads(sent.content) == {"input": [{"role": "user", "content": "Hello"}]}
+    assert json.loads(response.body)["result"]["parts"][0]["text"] == "pong"
+
+
 @pytest.mark.asyncio
 async def test_non_databricks_agent_skips_oauth_resolution():
     """Agents without a databricks_oauth block never enter the OAuth path."""
@@ -385,13 +458,10 @@ async def test_non_databricks_agent_skips_oauth_resolution():
     mock_agent.litellm_params = {"require_trace_id_on_calls_to_agent": False}
     mock_request = _make_mock_request()
 
-    with patch(
-        "litellm.proxy.agent_endpoints.a2a_endpoints.resolve_databricks_app_auth_header",
-        new_callable=AsyncMock,
-    ) as mock_resolve:
+    with patch("litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client") as mock_token_client:
         mock_asend = await _invoke(mock_agent, mock_request, None)
 
-    mock_resolve.assert_not_called()
+    mock_token_client.assert_not_called()
     headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
     assert headers == {"x-custom": "v", "X-LiteLLM-User-Id": "u1"}
     assert "Authorization" not in headers

tests/unit/proxy/agent_endpoints/test_databricks_oauth.py

@@ -8,17 +8,22 @@ helper.
 """
 
 import base64
+from typing import Final
 from unittest.mock import MagicMock, create_autospec, patch
 
 import httpx
 import pytest
 
+import litellm
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.proxy.agent_endpoints.databricks_oauth import (
     DatabricksAppOAuthConfig,
     DatabricksAppOAuthTokenCache,
+    has_databricks_oauth,
     parse_databricks_oauth_config,
     resolve_databricks_app_auth_header,
+    resolve_databricks_backend_auth,
+    without_databricks_oauth_params,
 )
 
 
@@ -46,9 +51,7 @@ def _mock_http_handler(access_token="tok-abc", expires_in=3600, post_error=None)
         handler.post.side_effect = post_error
     else:
         response = MagicMock()
-        response.json = MagicMock(
-            return_value={"access_token": access_token, "expires_in": expires_in}
-        )
+        response.json = MagicMock(return_value={"access_token": access_token, "expires_in": expires_in})
         handler.post.return_value = response
     return handler
 
@@ -111,23 +114,49 @@ def test_parse_custom_scope():
     assert config.scope == "custom-scope"
 
 
-@pytest.mark.parametrize(
-    "missing_field", ["client_id", "client_secret", "workspace_url"]
+AGENT_KINDS: Final = pytest.mark.parametrize(
+    ("agent_params", "error"),
+    [({}, ValueError), ({"custom_llm_provider": "databricks_agent"}, litellm.BadRequestError)],
+    ids=["url-agent", "databricks-agent"],
 )
-def test_parse_raises_on_missing_field(missing_field):
+
+
+@AGENT_KINDS
+@pytest.mark.parametrize("missing_field", ["client_id", "client_secret", "workspace_url"])
+def test_parse_raises_on_missing_field(missing_field, agent_params, error):
     block = {
         "client_id": "cid",
         "client_secret": "secret",
         "workspace_url": "https://dbc-abc.cloud.databricks.com",
     }
     block.pop(missing_field)
-    with pytest.raises(ValueError, match=missing_field):
-        parse_databricks_oauth_config({"databricks_oauth": block})
+    with pytest.raises(error, match=missing_field):
+        parse_databricks_oauth_config({**agent_params, "databricks_oauth": block})
 
 
-def test_parse_raises_on_non_mapping_block():
-    with pytest.raises(ValueError, match="mapping"):
-        parse_databricks_oauth_config({"databricks_oauth": "not-a-dict"})
+@AGENT_KINDS
+def test_parse_raises_on_non_mapping_block(agent_params, error):
+    with pytest.raises(error, match="mapping"):
+        parse_databricks_oauth_config({**agent_params, "databricks_oauth": "not-a-dict"})
+
+
+@pytest.mark.parametrize(
+    ("block", "error"),
+    [({}, "missing required field"), ("", "mapping")],
+)
+def test_parse_raises_on_an_empty_nested_block(block, error):
+    params = {"custom_llm_provider": "databricks_agent", "databricks_oauth": block}
+    assert has_databricks_oauth(params)
+    with pytest.raises(litellm.BadRequestError, match=error):
+        parse_databricks_oauth_config(params)
+
+
+@pytest.mark.parametrize("block", [{}, ""])
+@pytest.mark.asyncio
+async def test_an_empty_nested_block_on_a_url_agent_is_ignored_as_before(block):
+    auth: Final = await resolve_databricks_backend_auth({"databricks_oauth": block, "api_key": "hop-key"})
+    assert auth.header is None
+    assert auth.litellm_params == {"api_key": "hop-key"}
 
 
 def test_parse_resolves_os_environ_references(monkeypatch):
@@ -179,9 +208,7 @@ async def test_fetch_token_posts_client_credentials_with_basic_auth():
     }
     # Databricks authenticates the client with HTTP Basic; it must be sent as a
     # header because litellm's AsyncHTTPHandler.post has no ``auth`` parameter.
-    assert call.kwargs["headers"]["Authorization"] == _expected_basic_auth(
-        "cid", "secret"
-    )
+    assert call.kwargs["headers"]["Authorization"] == _expected_basic_auth("cid", "secret")
     assert "auth" not in call.kwargs
 
 
@@ -304,9 +331,7 @@ async def test_http_status_error_raises_value_error():
     request = httpx.Request("POST", _config().token_url)
     error_response = httpx.Response(status_code=401, request=request)
     client = _mock_http_handler(
-        post_error=httpx.HTTPStatusError(
-            "unauthorized", request=request, response=error_response
-        )
+        post_error=httpx.HTTPStatusError("unauthorized", request=request, response=error_response)
     )
 
     with patch(
@@ -494,3 +519,169 @@ async def test_resolve_returns_bearer_header():
         header = await resolve_databricks_app_auth_header(litellm_params)
 
     assert header == {"Authorization": "Bearer resolved-token"}
+
+
+def test_parse_builds_config_from_flat_databricks_agent_fields():
+    config = parse_databricks_oauth_config(
+        {
+            "custom_llm_provider": "databricks_agent",
+            "api_base": "https://adb-1.azuredatabricks.net/serving-endpoints",
+            "model": "my-agent",
+            "client_id": "sp-id",
+            "client_secret": "sp-secret",
+        }
+    )
+    assert config == DatabricksAppOAuthConfig(
+        client_id="sp-id",
+        client_secret="sp-secret",
+        token_url="https://adb-1.azuredatabricks.net/oidc/v1/token",
+        scope="all-apis",
+    )
+
+
+def test_parse_flat_fields_take_the_workspace_from_an_os_environ_api_base(monkeypatch):
+    monkeypatch.setenv("MY_DBX_HOST", "https://adb-3.azuredatabricks.net/serving-endpoints")
+    config = parse_databricks_oauth_config(
+        {
+            "custom_llm_provider": "databricks_agent",
+            "api_base": "os.environ/MY_DBX_HOST",
+            "client_id": "sp-id",
+            "client_secret": "sp-secret",
+        }
+    )
+    assert config is not None
+    assert config.token_url == "https://adb-3.azuredatabricks.net/oidc/v1/token"
+
+
+def test_parse_flat_fields_prefer_an_explicit_workspace_url_for_an_app():
+    config = parse_databricks_oauth_config(
+        {
+            "custom_llm_provider": "databricks_agent",
+            "api_base": "https://my-app-1.azure.databricksapps.com/responses",
+            "client_id": "sp-id",
+            "client_secret": "sp-secret",
+            "workspace_url": "https://adb-1.azuredatabricks.net",
+            "scope": "custom-scope",
+        }
+    )
+    assert config is not None
+    assert config.token_url == "https://adb-1.azuredatabricks.net/oidc/v1/token"
+    assert config.scope == "custom-scope"
+
+
+@pytest.mark.parametrize(
+    "app_base",
+    ["https://my-app-1.azure.databricksapps.com/responses", "https://my-app-1.azure.databricksapps.com"],
+)
+def test_parse_flat_fields_of_an_app_without_a_workspace_url_raise_instead_of_using_the_app_host(app_base):
+    with pytest.raises(litellm.BadRequestError, match="workspace_url") as info:
+        parse_databricks_oauth_config(
+            {
+                "custom_llm_provider": "databricks_agent",
+                "api_base": app_base,
+                "client_id": "sp-id",
+                "client_secret": "sp-secret",
+            }
+        )
+    assert info.value.status_code == 400
+
+
+@pytest.mark.parametrize(
+    "workspace_url",
+    ["https://my-app-1.azure.databricksapps.com", "https://my-app-1.azure.databricksapps.com/responses"],
+)
+def test_parse_rejects_an_app_url_typed_as_the_workspace_url(workspace_url):
+    with pytest.raises(litellm.BadRequestError, match="not a Databricks App URL") as info:
+        parse_databricks_oauth_config(
+            {
+                "custom_llm_provider": "databricks_agent",
+                "api_base": "https://my-app-1.azure.databricksapps.com/responses",
+                "client_id": "sp-id",
+                "client_secret": "sp-secret",
+                "workspace_url": workspace_url,
+            }
+        )
+    assert info.value.status_code == 400
+
+
+def test_parse_flat_fields_with_only_a_client_id_raise_on_the_missing_secret():
+    with pytest.raises(litellm.BadRequestError, match="client_secret"):
+        parse_databricks_oauth_config(
+            {
+                "custom_llm_provider": "databricks_agent",
+                "api_base": "https://adb-1.azuredatabricks.net",
+                "client_id": "sp",
+            }
+        )
+
+
+def test_parse_ignores_flat_client_fields_of_other_providers():
+    assert (
+        parse_databricks_oauth_config(
+            {
+                "custom_llm_provider": "azure_ai",
+                "client_id": "entra-id",
+                "client_secret": "entra-secret",
+                "tenant_id": "t",
+            }
+        )
+        is None
+    )
+    assert parse_databricks_oauth_config({"custom_llm_provider": "databricks_agent", "api_key": "pat"}) is None
+
+
+def test_parse_nested_block_wins_over_flat_fields():
+    config = parse_databricks_oauth_config(
+        {
+            "custom_llm_provider": "databricks_agent",
+            "client_id": "flat-id",
+            "client_secret": "flat-secret",
+            "databricks_oauth": {
+                "client_id": "nested-id",
+                "client_secret": "nested-secret",
+                "workspace_url": "https://adb-2.azuredatabricks.net",
+            },
+        }
+    )
+    assert config is not None
+    assert (config.client_id, config.token_url) == ("nested-id", "https://adb-2.azuredatabricks.net/oidc/v1/token")
+
+
+def test_without_oauth_params_drops_every_consumed_key_of_a_databricks_agent():
+    forwarded = without_databricks_oauth_params(
+        {
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": "https://adb-1.azuredatabricks.net",
+            "client_id": "sp-id",
+            "client_secret": "sp-secret",
+            "workspace_url": "https://adb-1.azuredatabricks.net",
+            "scope": "all-apis",
+            "databricks_oauth": {"client_id": "nested-id", "client_secret": "nested-secret"},
+            "custom_inputs": {"tenant": "t-1"},
+        }
+    )
+    assert forwarded == {
+        "custom_llm_provider": "databricks_agent",
+        "model": "my-agent",
+        "api_base": "https://adb-1.azuredatabricks.net",
+        "custom_inputs": {"tenant": "t-1"},
+    }
+
+
+def test_without_oauth_params_keeps_other_providers_client_credentials():
+    forwarded = without_databricks_oauth_params(
+        {
+            "custom_llm_provider": "azure_ai",
+            "client_id": "entra-id",
+            "client_secret": "entra-secret",
+            "tenant_id": "t",
+            "databricks_oauth": {"client_id": "sp"},
+        }
+    )
+    assert forwarded == {
+        "custom_llm_provider": "azure_ai",
+        "client_id": "entra-id",
+        "client_secret": "entra-secret",
+        "tenant_id": "t",
+    }

tests/unit/proxy/public_endpoints/test_public_endpoints.py

@@ -507,6 +507,7 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset(
         "charity_engine",
         "chutes",
         "darkbloom",
+        "databricks_agent",
         "gdc",
         "helicone",
         "inception",

tests/unit/proxy/test_route_a2a_models.py

@@ -4,14 +4,22 @@ Test A2A model routing in proxy.
 Maps to: litellm/proxy/agent_endpoints/a2a_routing.py
 """
 
-
-
+import base64
+import json
+from collections.abc import Iterator
 from unittest.mock import AsyncMock, Mock, patch
 
+import httpx
 import pytest
+import respx
+from fastapi import HTTPException
 
+import litellm
 from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request
+from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
+from litellm.proxy.agent_endpoints.databricks_oauth import databricks_app_oauth_token_cache
 from litellm.proxy.route_llm_request import route_request
+from litellm.types.agents import AgentResponse
 
 
 @pytest.mark.asyncio
@@ -183,3 +191,477 @@ async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_re
     assert call_kwargs["model"] == f"a2a/{agent_name}"
     assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com"
     prisma_client.db.litellm_agentstable.find_unique.assert_awaited()
+
+
+WORKSPACE = "https://adb-1.azuredatabricks.net"
+APP_URL = "https://my-app-1.azure.databricksapps.com/responses"
+AGENT_REPLY = {
+    "output": [
+        {"type": "message", "id": "msg_1", "role": "assistant", "content": [{"type": "output_text", "text": "pong"}]}
+    ]
+}
+A2A_REPLY = {
+    "jsonrpc": "2.0",
+    "id": "1",
+    "result": {"kind": "message", "role": "agent", "messageId": "m-1", "parts": [{"kind": "text", "text": "pong"}]},
+}
+
+
+@pytest.fixture
+def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    databricks_app_oauth_token_cache.flush_cache()
+    yield
+    databricks_app_oauth_token_cache.flush_cache()
+
+
+def _registry_with(agent: AgentResponse) -> AgentRegistry:
+    registry = AgentRegistry()
+    registry.register_agent(agent)
+    return registry
+
+
+def _token_route(workspace: str) -> respx.Route:
+    return respx.post(f"{workspace}/oidc/v1/token").mock(
+        return_value=httpx.Response(200, json={"access_token": "minted-oauth", "expires_in": 3600})
+    )
+
+
+@respx.mock
+async def test_route_a2a_bridge_agent_calls_the_endpoint_with_the_agent_params(httpx_transport: None) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-agent-id",
+        agent_name="dbx-agent",
+        agent_card_params={"url": WORKSPACE},
+        litellm_params={
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": WORKSPACE,
+            "api_key": "pat-1",
+            "custom_inputs": {"tenant": "t-1"},
+            "is_public": True,
+        },
+    )
+    endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock(
+        return_value=httpx.Response(200, json=AGENT_REPLY)
+    )
+    data = {"model": "a2a/dbx-agent", "messages": [{"role": "user", "content": "ping"}]}
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    sent = endpoint.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer pat-1"
+    assert json.loads(sent.content) == {
+        "input": [{"role": "user", "content": "ping"}],
+        "custom_inputs": {"tenant": "t-1"},
+    }
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+async def test_route_a2a_bridge_agent_merges_client_static_and_agent_extra_headers(httpx_transport: None) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-agent-id",
+        agent_name="dbx-agent",
+        agent_card_params={"url": WORKSPACE},
+        static_headers={"X-Static": "static", "x-shared": "from-static", "X-Both": "from-static"},
+        litellm_params={
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": WORKSPACE,
+            "api_key": "pat-1",
+            "extra_headers": {"X-Agent": "agent", "x-both": "from-agent"},
+        },
+    )
+    endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock(
+        return_value=httpx.Response(200, json=AGENT_REPLY)
+    )
+    data = {
+        "model": "a2a/dbx-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "extra_headers": {"X-Trace": "trace-1", "X-Shared": "from-client", "X-LiteLLM-User-Id": "spoofed-user"},
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    sent = endpoint.calls.last.request
+    assert "X-LiteLLM-User-Id" not in sent.headers
+    assert sent.headers["X-Trace"] == "trace-1"
+    assert sent.headers["X-Static"] == "static"
+    assert sent.headers["X-Agent"] == "agent"
+    assert sent.headers.get_list("X-Shared") == ["from-static"]
+    assert sent.headers.get_list("X-Both") == ["from-agent"]
+    assert sent.headers["Authorization"] == "Bearer pat-1"
+
+
+async def test_route_a2a_bridge_agent_lets_the_request_win_except_for_the_agent_credentials() -> None:
+    agent = AgentResponse(
+        agent_id="dbx-agent-id",
+        agent_name="dbx-agent",
+        agent_card_params={"url": WORKSPACE},
+        litellm_params={
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": WORKSPACE,
+            "api_key": "pat-1",
+            "metadata": {"tags": ["agent-tag"]},
+            "custom_inputs": {"tenant": "t-1"},
+            "timeout": 30,
+        },
+    )
+    proxy_metadata = {"user_api_key_hash": "key-hash", "user_api_key_team_id": "team-1"}
+    data = {
+        "model": "a2a/dbx-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "metadata": proxy_metadata,
+        "api_base": "https://elsewhere.example.com",
+        "api_key": "client-key",
+    }
+    completion = Mock(return_value="reply")
+
+    with (
+        patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)),
+        patch.object(litellm, "acompletion", completion),
+    ):
+        await route_a2a_agent_request(data=data, route_type="acompletion")
+
+    sent = completion.call_args.kwargs
+    assert sent["metadata"] == proxy_metadata
+    assert (sent["api_base"], sent["api_key"]) == (WORKSPACE, "pat-1")
+    assert (sent["custom_inputs"], sent["timeout"]) == ({"tenant": "t-1"}, 30)
+    assert sent["model"] == "databricks_agent/my-agent"
+
+
+@respx.mock
+async def test_route_a2a_bridge_agent_keeps_its_custom_inputs_and_extra_body_keys_over_the_callers(
+    httpx_transport: None,
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-agent-id",
+        agent_name="dbx-agent",
+        agent_card_params={"url": WORKSPACE},
+        litellm_params={
+            "custom_llm_provider": "databricks_agent",
+            "model": "my-agent",
+            "api_base": WORKSPACE,
+            "api_key": "pat-1",
+            "custom_inputs": {"tenant": "owner"},
+            "extra_body": {"metadata": {"tenant": "owner"}},
+        },
+    )
+    endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock(
+        return_value=httpx.Response(200, json=AGENT_REPLY)
+    )
+    data = {
+        "model": "a2a/dbx-agent",
+        "messages": [{"role": "user", "content": "ping"}],
+        "custom_inputs": {"tenant": "caller", "session": "s-1"},
+        "extra_body": {"metadata": {"tenant": "caller", "turn": "2"}},
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert json.loads(endpoint.calls.last.request.content) == {
+        "input": [{"role": "user", "content": "ping"}],
+        "custom_inputs": {"tenant": "owner", "session": "s-1"},
+        "metadata": {"tenant": "owner", "turn": "2"},
+    }
+
+
+@respx.mock
+async def test_route_a2a_bridge_agent_with_an_empty_oauth_block_fails_before_calling_the_app(
+    httpx_transport: None,
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-app-id",
+        agent_name="dbx-app",
+        agent_card_params={"url": APP_URL},
+        litellm_params={"custom_llm_provider": "databricks_agent", "databricks_oauth": {}},
+    )
+    app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {"model": "a2a/dbx-app", "messages": [{"role": "user", "content": "ping"}]}
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        with pytest.raises(litellm.BadRequestError, match="missing required field") as info:
+            await route_a2a_agent_request(data=data, route_type="acompletion")
+
+    assert info.value.status_code == 400
+    assert not app.called
+
+
+@respx.mock
+@pytest.mark.parametrize(
+    "oauth_params",
+    [
+        {"databricks_oauth": {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE}},
+        {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE, "scope": "all-apis"},
+    ],
+)
+async def test_route_a2a_bridge_agent_with_databricks_oauth_sends_the_minted_token(
+    httpx_transport: None, oauth_params: dict[str, object]
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-app-id",
+        agent_name="dbx-app",
+        agent_card_params={"url": APP_URL},
+        litellm_params={"custom_llm_provider": "databricks_agent", **oauth_params},
+    )
+    token = _token_route(WORKSPACE)
+    app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {
+        "model": "a2a/dbx-app",
+        "messages": [{"role": "user", "content": "ping"}],
+        "extra_headers": {"X-Tenant": "t"},
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert token.calls.last.request.headers["Authorization"] == f"Basic {base64.b64encode(b'sp:s').decode()}"
+    sent = app.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer minted-oauth"
+    assert sent.headers["X-Tenant"] == "t"
+    assert json.loads(sent.content) == {"input": [{"role": "user", "content": "ping"}]}
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+@pytest.mark.parametrize(
+    "oauth_block",
+    [
+        {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE},
+        {"client_id": "sp"},
+        "not-a-mapping",
+        {"client_id": "sp", "client_secret": "s", "workspace_url": "https://my-app-1.azure.databricksapps.com"},
+    ],
+    ids=["complete", "partial", "not-a-mapping", "app-host-as-workspace"],
+)
+async def test_route_a2a_url_agent_chat_keeps_the_client_authorization_over_its_databricks_oauth(
+    httpx_transport: None, oauth_block: object
+) -> None:
+    agent_url = "https://my-a2a-app.azure.databricksapps.com"
+    agent = AgentResponse(
+        agent_id="a2a-app-id",
+        agent_name="a2a-app",
+        agent_card_params={"url": agent_url},
+        litellm_params={"databricks_oauth": oauth_block},
+    )
+    token = _token_route(WORKSPACE)
+    app = respx.post(agent_url).mock(return_value=httpx.Response(200, json=A2A_REPLY))
+    data = {
+        "model": "a2a/a2a-app",
+        "messages": [{"role": "user", "content": "ping"}],
+        "extra_headers": {"Authorization": "Bearer client-token"},
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    sent = app.calls.last.request
+    assert sent.headers["Authorization"] == "Bearer client-token"
+    assert json.loads(sent.content)["method"] == "message/send"
+    assert not token.called
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+async def test_route_a2a_agent_with_another_provider_keeps_calling_its_card_url_over_a2a(
+    httpx_transport: None,
+) -> None:
+    agent_url = "https://pydantic-agent.example.com"
+    agent = AgentResponse(
+        agent_id="pydantic-agent-id",
+        agent_name="pydantic-agent",
+        agent_card_params={"url": agent_url},
+        litellm_params={"custom_llm_provider": "pydantic_ai_agents", "model": "pydantic-agent"},
+    )
+    agent_endpoint = respx.post(agent_url).mock(return_value=httpx.Response(200, json=A2A_REPLY))
+    data = {"model": "a2a/pydantic-agent", "messages": [{"role": "user", "content": "ping"}]}
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert json.loads(agent_endpoint.calls.last.request.content)["method"] == "message/send"
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+@pytest.mark.parametrize("client_target", ["api_base", "base_url"])
+async def test_route_a2a_bridge_agent_ignores_a_client_target_over_its_card_url(
+    httpx_transport: None, client_target: str
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-agent-id",
+        agent_name="dbx-agent",
+        agent_card_params={"url": WORKSPACE},
+        litellm_params={"custom_llm_provider": "databricks_agent", "model": "my-agent", "api_key": "pat-1"},
+    )
+    endpoint = respx.post(f"{WORKSPACE}/serving-endpoints/my-agent/invocations").mock(
+        return_value=httpx.Response(200, json=AGENT_REPLY)
+    )
+    data = {
+        "model": "a2a/dbx-agent",
+        client_target: "https://collector.example.com",
+        "messages": [{"role": "user", "content": "ping"}],
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert endpoint.call_count == 1
+    assert endpoint.calls.last.request.headers["Authorization"] == "Bearer pat-1"
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+@pytest.mark.parametrize("endpoint_name", [{"model": ""}, {"model": None}, {}])
+async def test_route_a2a_bridge_agent_without_an_endpoint_name_answers_400_before_calling_databricks(
+    httpx_transport: None, endpoint_name: dict[str, object]
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-agent-id",
+        agent_name="dbx-agent",
+        agent_card_params={"url": WORKSPACE},
+        litellm_params={
+            "custom_llm_provider": "databricks_agent",
+            "api_base": WORKSPACE,
+            "api_key": "p",
+            **endpoint_name,
+        },
+    )
+    workspace = respx.route(host="adb-1.azuredatabricks.net").mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {"model": "a2a/dbx-agent", "messages": [{"role": "user", "content": "ping"}]}
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        with pytest.raises(litellm.BadRequestError, match="must name the Databricks serving endpoint"):
+            await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert not workspace.called
+
+
+@respx.mock
+async def test_route_a2a_bridge_app_agent_without_an_endpoint_name_calls_the_app(httpx_transport: None) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-app-id",
+        agent_name="dbx-app",
+        agent_card_params={"url": APP_URL},
+        litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "p"},
+    )
+    app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {"model": "a2a/dbx-app", "messages": [{"role": "user", "content": "ping"}]}
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert app.called
+    assert response.choices[0].message.content == "pong"
+
+
+@respx.mock
+@pytest.mark.parametrize(
+    ("credential", "env_api_key", "expected_authorization"),
+    [
+        ({"api_key": "pat-1"}, None, "Bearer pat-1"),
+        ({}, "env-pat", "Bearer env-pat"),
+        ({"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE}, "env-pat", "Bearer minted-oauth"),
+        (
+            {"client_id": "sp", "client_secret": "s", "workspace_url": WORKSPACE, "api_key": "pat-1"},
+            None,
+            "Bearer minted-oauth",
+        ),
+        ({}, None, "Bearer client-token"),
+    ],
+)
+async def test_route_a2a_bridge_agent_credential_wins_over_a_client_authorization(
+    httpx_transport: None,
+    monkeypatch: pytest.MonkeyPatch,
+    credential: dict[str, object],
+    env_api_key: str | None,
+    expected_authorization: str,
+) -> None:
+    if env_api_key is None:
+        monkeypatch.delenv("DATABRICKS_API_KEY", raising=False)
+    else:
+        monkeypatch.setenv("DATABRICKS_API_KEY", env_api_key)
+    agent = AgentResponse(
+        agent_id="dbx-app-id",
+        agent_name="dbx-app",
+        agent_card_params={"url": APP_URL},
+        litellm_params={"custom_llm_provider": "databricks_agent", **credential},
+    )
+    _token_route(WORKSPACE)
+    app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {
+        "model": "a2a/dbx-app",
+        "messages": [{"role": "user", "content": "ping"}],
+        "extra_headers": {"authorization": "Bearer client-token", "X-Tenant": "t"},
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    sent = app.calls.last.request
+    assert sent.headers.get_list("Authorization") == [expected_authorization]
+    assert sent.headers["X-Tenant"] == "t"
+
+
+@respx.mock
+@pytest.mark.parametrize("extra_headers", [5, ["X-A"], {"X-N": 5}, "x"])
+async def test_route_a2a_bridge_agent_answers_400_for_malformed_client_extra_headers(
+    httpx_transport: None, extra_headers: object
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-app-id",
+        agent_name="dbx-app",
+        agent_card_params={"url": APP_URL},
+        litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "p"},
+    )
+    app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {"model": "a2a/dbx-app", "messages": [{"role": "user", "content": "ping"}], "extra_headers": extra_headers}
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        with pytest.raises(HTTPException) as info:
+            await route_a2a_agent_request(data=data, route_type="acompletion")
+
+    assert info.value.status_code == 400
+    assert "extra_headers" in str(info.value.detail)
+    assert not app.called
+
+
+@respx.mock
+@pytest.mark.parametrize(
+    ("extra_headers", "drop_params"),
+    [
+        (5, {"drop_params": True}),
+        ({"X-N": 5}, {"drop_params": True}),
+        ("", {}),
+        ({}, {}),
+        (None, {}),
+    ],
+)
+async def test_route_a2a_bridge_agent_calls_without_client_extra_headers_it_drops_or_finds_empty(
+    httpx_transport: None, extra_headers: object, drop_params: dict[str, object]
+) -> None:
+    agent = AgentResponse(
+        agent_id="dbx-app-id",
+        agent_name="dbx-app",
+        agent_card_params={"url": APP_URL},
+        litellm_params={"custom_llm_provider": "databricks_agent", "api_key": "p"},
+    )
+    app = respx.post(APP_URL).mock(return_value=httpx.Response(200, json=AGENT_REPLY))
+    data = {
+        "model": "a2a/dbx-app",
+        "messages": [{"role": "user", "content": "ping"}],
+        "extra_headers": extra_headers,
+        **drop_params,
+    }
+
+    with patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", _registry_with(agent)):
+        response = await (await route_a2a_agent_request(data=data, route_type="acompletion"))
+
+    assert app.calls.last.request.headers["Authorization"] == "Bearer p"
+    assert "X-N" not in app.calls.last.request.headers
+    assert response.choices[0].message.content == "pong"

tests/unit/test_internal_context.py

@@ -54,6 +54,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset(
         "litellm/llms/azure/common_utils.py",
         "litellm/llms/bedrock/base_aws_llm.py",
         "litellm/llms/custom_httpx/http_handler.py",
+        "litellm/llms/databricks/agent/transformation.py",
         "litellm/llms/gigachat/authenticator.py",
         "litellm/llms/github_copilot/per_user_auth.py",
         "litellm/llms/litellm_proxy/skills/handler.py",

ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.test.ts

@@ -96,6 +96,21 @@ describe("detectAgentType", () => {
 
     expect(detectAgentType(agent)).toBe("bedrock_agentcore");
   });
+
+  it("detects databricks_agent agents from custom_llm_provider so the edit form shows the Databricks fields", () => {
+    const agent = {
+      agent_id: "agent-2",
+      agent_name: "dbx-agent",
+      agent_card_params: { url: "https://adb-1.azuredatabricks.net" },
+      litellm_params: {
+        custom_llm_provider: "databricks_agent",
+        model: "my-agent-endpoint",
+        api_base: "https://adb-1.azuredatabricks.net",
+      },
+    } as unknown as Agent;
+
+    expect(detectAgentType(agent)).toBe("databricks_agent");
+  });
 });
 
 describe("API agent metadata validation", () => {

ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_type_utils.ts

@@ -12,6 +12,7 @@ export const detectAgentType = (agent: Agent): string => {
 
   // Check by custom_llm_provider first
   if (customProvider === "langflow") return "langflow";
+  if (customProvider === "databricks_agent") return "databricks_agent";
   if (customProvider === "langgraph") return "langgraph";
   if (customProvider === "azure_ai") return "azure_ai_foundry";
   if (customProvider === "bedrock") return "bedrock_agentcore";