Support OAuth M2M for Databricks Apps A2A agents (#29586)

* Add OAuth M2M support for A2A agents targeting Databricks Apps

Databricks App endpoints reject static bearer tokens and require a
short-lived OAuth token minted via the workspace OIDC token endpoint.
A2A agents could previously only authenticate outbound with static_headers
or client header passthrough, so Databricks App agents could not be
registered.

Agents configured with a databricks_oauth block in litellm_params now mint
and cache a client_credentials token and attach it as the outbound
Authorization header on both message/send and message/stream calls,
overriding any statically configured Authorization.

* Add tests covering Databricks App OAuth token error paths

Cover the HTTP status error, transport error, non-object JSON body, and
invalid expires_in fallback branches in the token cache so the failure
handling is locked in by regression tests.

* Harden Databricks App OAuth token cache

Cap the cache TTL at the token's own lifetime so a token whose validity is
shorter than the refresh buffer is never cached and served stale; include a
digest of client_secret in the cache key so a rotated secret mints a fresh
token instead of reusing the old one; and prune the per-key lock when its
cached token is evicted so the lock map stays bounded by the live key set.

* Clear per-key locks on Databricks OAuth cache flush

* fix(a2a/databricks): mint OAuth token via Basic auth header, not unsupported auth= kwarg

litellm's AsyncHTTPHandler.post (what get_async_httpx_client returns) has no
auth parameter, so minting a Databricks App OAuth token raised
"AsyncHTTPHandler.post() got an unexpected keyword argument 'auth'" before any
network call ever left the proxy, breaking the feature end to end. The handler
also calls raise_for_status() internally and re-raises a MaskedHTTPStatusError
(a subclass of httpx.HTTPStatusError), so the explicit raise_for_status() after
post() was dead code.

Build the HTTP Basic Authorization header by hand and pass it via headers, which
is what the Databricks workspace OIDC token endpoint documents for client
authentication. The token-cache tests now model the real handler contract with
create_autospec so the rejected auth= signature is enforced; the previous mocks
accepted any kwargs and silently hid the bug.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

* Prune Databricks OAuth lock on the short-lived-token path

When expires_in is below the refresh buffer the token is intentionally
not cached, so _remove_key never runs for that key and the per-key lock
created by _get_lock leaked permanently. Drop the lock in that branch so
_locks stays bounded by the live key set, and assert the cleanup in the
short-lived-token test

* Gate A2A Databricks OAuth on the databricks_oauth block at the call site

Make the gating explicit where the header is applied so it is clear that only
agents configured with a databricks_oauth block enter the OAuth path; every
other agent is left untouched. Add a regression test asserting a non-Databricks
agent never invokes the token resolver.

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -15,6 +15,10 @@ from fastapi.responses import JSONResponse, StreamingResponse
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.agent_endpoints.databricks_oauth import (
+    DATABRICKS_OAUTH_PARAM,
+    resolve_databricks_app_auth_header,
+)
 from litellm.proxy.agent_endpoints.utils import merge_agent_headers
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.types.utils import all_litellm_params
@@ -677,6 +681,17 @@ async def invoke_agent_a2a(  # noqa: PLR0915
             static_headers=static_headers or None,
         )
 
+        # Databricks App endpoints require a short-lived OAuth M2M token rather
+        # than a static bearer. Only agents explicitly configured with a
+        # ``databricks_oauth`` block get one; every other agent is left untouched.
+        if litellm_params.get(DATABRICKS_OAUTH_PARAM):
+            databricks_auth = await resolve_databricks_app_auth_header(litellm_params)
+            if databricks_auth:
+                agent_extra_headers = {
+                    **(agent_extra_headers or {}),
+                    **databricks_auth,
+                }
+
         # Merge agent-level guardrails into data so post_call_success_hook and
         # _handle_stream_message both pick them up.  A2A agents use model
         # a2a_agent/*, which is not an llm_router deployment, so

litellm/proxy/agent_endpoints/databricks_oauth.py

@@ -0,0 +1,250 @@
+"""
+OAuth M2M (client_credentials) support for A2A agents that target Databricks
+App endpoints.
+
+Databricks Apps reject static bearer tokens; they require a short-lived OAuth
+access token minted from the workspace OIDC token endpoint. When an agent is
+registered with a ``databricks_oauth`` block in its ``litellm_params``, LiteLLM
+fetches that token via the client_credentials grant, caches it until shortly
+before expiry, and attaches it as the outbound ``Authorization`` header on every
+call the proxy makes to the agent.
+
+Config example::
+
+    agents:
+      - agent_name: my-databricks-app
+        agent_card_params:
+          url: https://my-app-1234.aws.databricksapps.com
+        litellm_params:
+          databricks_oauth:
+            client_id: os.environ/DATABRICKS_CLIENT_ID
+            client_secret: os.environ/DATABRICKS_CLIENT_SECRET
+            workspace_url: https://dbc-abc123.cloud.databricks.com
+"""
+
+import asyncio
+import base64
+import hashlib
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import httpx
+
+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.secret_managers.main import get_secret_str
+from litellm.types.llms.custom_http import httpxSpecialProvider
+
+DATABRICKS_OAUTH_PARAM = "databricks_oauth"
+
+_DEFAULT_SCOPE = "all-apis"
+_TOKEN_EXPIRY_BUFFER_SECONDS = 60
+_DEFAULT_TTL_SECONDS = 3600
+
+
+def _resolve_secret(value: Any) -> Optional[str]:
+    """Resolve a config value, expanding ``os.environ/`` references."""
+    if not isinstance(value, str):
+        return None
+    if value.startswith("os.environ/"):
+        return get_secret_str(value)
+    return value
+
+
+def _token_url_from_workspace(workspace_url: str) -> str:
+    """Build the workspace OIDC token endpoint from a workspace URL."""
+    base = workspace_url.strip().rstrip("/")
+    if base.endswith("/serving-endpoints"):
+        base = base[: -len("/serving-endpoints")]
+    return f"{base}/oidc/v1/token"
+
+
+@dataclass(frozen=True)
+class DatabricksAppOAuthConfig:
+    client_id: str
+    client_secret: str
+    token_url: str
+    scope: str
+
+    @property
+    def cache_key(self) -> str:
+        # Include a digest of the secret so a rotated client_secret yields a new
+        # key and forces a fresh token instead of serving the stale one.
+        secret_digest = hashlib.sha256(self.client_secret.encode()).hexdigest()[:16]
+        return f"{self.token_url}|{self.client_id}|{self.scope}|{secret_digest}"
+
+
+def parse_databricks_oauth_config(
+    litellm_params: Optional[Dict[str, Any]],
+) -> Optional[DatabricksAppOAuthConfig]:
+    """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.
+    """
+    if not litellm_params:
+        return None
+
+    raw = litellm_params.get(DATABRICKS_OAUTH_PARAM)
+    if raw is None:
+        return None
+    if not isinstance(raw, dict):
+        raise ValueError(
+            f"'{DATABRICKS_OAUTH_PARAM}' must be a mapping of OAuth settings, "
+            f"got {type(raw).__name__}"
+        )
+
+    client_id = _resolve_secret(raw.get("client_id"))
+    client_secret = _resolve_secret(raw.get("client_secret"))
+    workspace_url = _resolve_secret(raw.get("workspace_url"))
+
+    missing = [
+        name
+        for name, value in (
+            ("client_id", client_id),
+            ("client_secret", client_secret),
+            ("workspace_url", workspace_url),
+        )
+        if not value
+    ]
+    if missing:
+        raise ValueError(
+            f"Databricks App OAuth config is missing required field(s): "
+            f"{', '.join(missing)}"
+        )
+
+    scope = _resolve_secret(raw.get("scope")) or _DEFAULT_SCOPE
+
+    return DatabricksAppOAuthConfig(
+        client_id=client_id,  # type: ignore[arg-type]
+        client_secret=client_secret,  # type: ignore[arg-type]
+        token_url=_token_url_from_workspace(workspace_url),  # type: ignore[arg-type]
+        scope=scope,
+    )
+
+
+class DatabricksAppOAuthTokenCache(InMemoryCache):
+    """In-memory cache for Databricks App OAuth client_credentials tokens.
+
+    Keyed by token endpoint + client_id + scope so distinct agents and service
+    principals never share a token. A per-key ``asyncio.Lock`` collapses
+    concurrent fetches into a single token request.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(default_ttl=_DEFAULT_TTL_SECONDS)
+        self._locks: Dict[str, asyncio.Lock] = {}
+
+    def _get_lock(self, cache_key: str) -> asyncio.Lock:
+        return self._locks.setdefault(cache_key, asyncio.Lock())
+
+    def _remove_key(self, key: str) -> None:
+        # Drop the per-key lock alongside the cached token so ``_locks`` stays
+        # bounded by the live key set rather than growing for every key ever seen.
+        super()._remove_key(key)
+        self._locks.pop(key, None)
+
+    def flush_cache(self) -> None:
+        super().flush_cache()
+        self._locks.clear()
+
+    async def async_get_token(self, config: DatabricksAppOAuthConfig) -> str:
+        cache_key = config.cache_key
+
+        cached = self.get_cache(cache_key)
+        if cached is not None:
+            return cached
+
+        async with self._get_lock(cache_key):
+            cached = self.get_cache(cache_key)
+            if cached is not None:
+                return cached
+
+            token, ttl = await self._fetch_token(config)
+            # ttl == 0 means the token's own lifetime is shorter than the
+            # refresh buffer; skip caching so we never hand out a stale token,
+            # and drop the lock we just created since no cached entry will ever
+            # trigger _remove_key to clean it up.
+            if ttl > 0:
+                self.set_cache(cache_key, token, ttl=ttl)
+            else:
+                self._locks.pop(cache_key, None)
+            return token
+
+    async def _fetch_token(self, config: DatabricksAppOAuthConfig) -> Tuple[str, int]:
+        client = get_async_httpx_client(llm_provider=httpxSpecialProvider.A2A)
+
+        verbose_logger.debug(
+            "Fetching Databricks App OAuth token from %s", config.token_url
+        )
+
+        basic_auth = base64.b64encode(
+            f"{config.client_id}:{config.client_secret}".encode()
+        ).decode()
+        try:
+            response = await client.post(
+                config.token_url,
+                data={
+                    "grant_type": "client_credentials",
+                    "scope": config.scope,
+                },
+                headers={
+                    "Authorization": f"Basic {basic_auth}",
+                    "Content-Type": "application/x-www-form-urlencoded",
+                },
+            )
+        except httpx.HTTPStatusError as exc:
+            raise ValueError(
+                "Databricks App OAuth token request failed with status "
+                f"{exc.response.status_code}"
+            ) from exc
+        except httpx.HTTPError as exc:
+            raise ValueError(
+                f"Databricks App OAuth token request failed: {exc}"
+            ) from exc
+
+        body = response.json()
+        if not isinstance(body, dict):
+            raise ValueError(
+                "Databricks App OAuth token response returned non-object JSON "
+                f"(got {type(body).__name__})"
+            )
+
+        access_token = body.get("access_token")
+        if not access_token:
+            raise ValueError(
+                "Databricks App OAuth token response missing 'access_token'"
+            )
+
+        raw_expires_in = body.get("expires_in")
+        try:
+            expires_in = (
+                int(raw_expires_in)
+                if raw_expires_in is not None
+                else _DEFAULT_TTL_SECONDS
+            )
+        except (TypeError, ValueError):
+            expires_in = _DEFAULT_TTL_SECONDS
+
+        ttl = max(expires_in - _TOKEN_EXPIRY_BUFFER_SECONDS, 0)
+        return access_token, ttl
+
+
+databricks_app_oauth_token_cache = DatabricksAppOAuthTokenCache()
+
+
+async def resolve_databricks_app_auth_header(
+    litellm_params: Optional[Dict[str, Any]],
+) -> Optional[Dict[str, str]]:
+    """Return ``{"Authorization": "Bearer <token>"}`` for a Databricks App agent.
+
+    Returns ``None`` when the agent is not configured for Databricks App OAuth.
+    """
+    config = parse_databricks_oauth_config(litellm_params)
+    if config is None:
+        return None
+
+    token = await databricks_app_oauth_token_cache.async_get_token(config)
+    return {"Authorization": f"Bearer {token}"}

tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py

@@ -14,7 +14,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
-
 # ---------------------------------------------------------------------------
 # Helper: build a minimal mock agent
 # ---------------------------------------------------------------------------
@@ -307,6 +306,97 @@ async def test_convention_unrelated_prefix_not_forwarded():
     assert headers is None
 
 
+# ---------------------------------------------------------------------------
+# Databricks App OAuth M2M injection
+# ---------------------------------------------------------------------------
+
+
+def _mock_databricks_token_client(access_token="dbx-oauth-token"):
+    response = MagicMock()
+    response.raise_for_status = MagicMock()
+    response.json = MagicMock(
+        return_value={"access_token": access_token, "expires_in": 3600}
+    )
+    client = MagicMock()
+    client.post = AsyncMock(return_value=response)
+    return client
+
+
+@pytest.mark.asyncio
+async def test_databricks_oauth_header_injected():
+    """A databricks_oauth block mints an outbound Bearer Authorization header."""
+    from litellm.proxy.agent_endpoints import databricks_oauth
+
+    databricks_oauth.databricks_app_oauth_token_cache.flush_cache()
+
+    mock_agent = _make_mock_agent()
+    mock_agent.litellm_params = {
+        "databricks_oauth": {
+            "client_id": "cid",
+            "client_secret": "secret",
+            "workspace_url": "https://dbc.cloud.databricks.com",
+        }
+    }
+    mock_request = _make_mock_request()
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=_mock_databricks_token_client("minted-token"),
+    ):
+        mock_asend = await _invoke(mock_agent, mock_request, None)
+
+    headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
+    assert headers is not None
+    assert headers.get("Authorization") == "Bearer minted-token"
+
+
+@pytest.mark.asyncio
+async def test_databricks_oauth_overrides_static_authorization():
+    """The minted OAuth token wins over a statically configured Authorization."""
+    from litellm.proxy.agent_endpoints import databricks_oauth
+
+    databricks_oauth.databricks_app_oauth_token_cache.flush_cache()
+
+    mock_agent = _make_mock_agent(static_headers={"Authorization": "Bearer static-pat"})
+    mock_agent.litellm_params = {
+        "databricks_oauth": {
+            "client_id": "cid",
+            "client_secret": "secret",
+            "workspace_url": "https://dbc.cloud.databricks.com",
+        }
+    }
+    mock_request = _make_mock_request()
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=_mock_databricks_token_client("oauth-wins"),
+    ):
+        mock_asend = await _invoke(mock_agent, mock_request, None)
+
+    headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
+    assert headers is not None
+    assert headers.get("Authorization") == "Bearer oauth-wins"
+
+
+@pytest.mark.asyncio
+async def test_non_databricks_agent_skips_oauth_resolution():
+    """Agents without a databricks_oauth block never enter the OAuth path."""
+    mock_agent = _make_mock_agent(static_headers={"x-custom": "v"})
+    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:
+        mock_asend = await _invoke(mock_agent, mock_request, None)
+
+    mock_resolve.assert_not_called()
+    headers = mock_asend.call_args.kwargs.get("agent_extra_headers")
+    assert headers == {"x-custom": "v"}
+    assert "Authorization" not in headers
+
+
 # ---------------------------------------------------------------------------
 # Direct unit test for the merge utility
 # ---------------------------------------------------------------------------

tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py

@@ -0,0 +1,496 @@
+"""
+Unit tests for Databricks App OAuth M2M support for A2A agents.
+
+Covers config parsing (including os.environ/ resolution and validation),
+workspace token-URL construction, client_credentials token fetching, caching
+with expiry buffering, and the public ``resolve_databricks_app_auth_header``
+helper.
+"""
+
+import base64
+from unittest.mock import MagicMock, create_autospec, patch
+
+import httpx
+import pytest
+
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.proxy.agent_endpoints.databricks_oauth import (
+    DatabricksAppOAuthConfig,
+    DatabricksAppOAuthTokenCache,
+    parse_databricks_oauth_config,
+    resolve_databricks_app_auth_header,
+)
+
+
+def _expected_basic_auth(client_id: str, client_secret: str) -> str:
+    token = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
+    return f"Basic {token}"
+
+
+def _mock_http_handler(access_token="tok-abc", expires_in=3600, post_error=None):
+    """Return a mock that mirrors litellm's ``AsyncHTTPHandler`` contract.
+
+    Two properties of the real handler matter for these tests and were the
+    source of a runtime bug the original suite missed:
+
+    1. ``post`` does not accept an ``auth`` kwarg. ``create_autospec`` enforces
+       the real signature, so reintroducing HTTP Basic via ``auth=`` fails with
+       ``TypeError`` instead of silently passing.
+    2. ``post`` calls ``raise_for_status`` internally and raises
+       ``httpx.HTTPStatusError`` itself on non-2xx; callers never inspect the
+       returned response's status. Error-path tests therefore raise from
+       ``post`` rather than from ``response.raise_for_status``.
+    """
+    handler = create_autospec(AsyncHTTPHandler, instance=True)
+    if post_error is not None:
+        handler.post.side_effect = post_error
+    else:
+        response = MagicMock()
+        response.json = MagicMock(
+            return_value={"access_token": access_token, "expires_in": expires_in}
+        )
+        handler.post.return_value = response
+    return handler
+
+
+# ---------------------------------------------------------------------------
+# Config parsing
+# ---------------------------------------------------------------------------
+
+
+def test_parse_returns_none_without_block():
+    assert parse_databricks_oauth_config(None) is None
+    assert parse_databricks_oauth_config({}) is None
+    assert parse_databricks_oauth_config({"other": "value"}) is None
+
+
+def test_parse_builds_config_and_token_url():
+    config = parse_databricks_oauth_config(
+        {
+            "databricks_oauth": {
+                "client_id": "cid",
+                "client_secret": "secret",
+                "workspace_url": "https://dbc-abc.cloud.databricks.com",
+            }
+        }
+    )
+    assert config == DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="secret",
+        token_url="https://dbc-abc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+
+
+def test_parse_strips_serving_endpoints_and_trailing_slash():
+    config = parse_databricks_oauth_config(
+        {
+            "databricks_oauth": {
+                "client_id": "cid",
+                "client_secret": "secret",
+                "workspace_url": "https://dbc-abc.cloud.databricks.com/serving-endpoints/",
+            }
+        }
+    )
+    assert config is not None
+    assert config.token_url == "https://dbc-abc.cloud.databricks.com/oidc/v1/token"
+
+
+def test_parse_custom_scope():
+    config = parse_databricks_oauth_config(
+        {
+            "databricks_oauth": {
+                "client_id": "cid",
+                "client_secret": "secret",
+                "workspace_url": "https://dbc-abc.cloud.databricks.com",
+                "scope": "custom-scope",
+            }
+        }
+    )
+    assert config is not None
+    assert config.scope == "custom-scope"
+
+
+@pytest.mark.parametrize(
+    "missing_field", ["client_id", "client_secret", "workspace_url"]
+)
+def test_parse_raises_on_missing_field(missing_field):
+    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})
+
+
+def test_parse_raises_on_non_mapping_block():
+    with pytest.raises(ValueError, match="mapping"):
+        parse_databricks_oauth_config({"databricks_oauth": "not-a-dict"})
+
+
+def test_parse_resolves_os_environ_references(monkeypatch):
+    monkeypatch.setenv("MY_DBX_CLIENT_ID", "env-cid")
+    monkeypatch.setenv("MY_DBX_SECRET", "env-secret")
+    config = parse_databricks_oauth_config(
+        {
+            "databricks_oauth": {
+                "client_id": "os.environ/MY_DBX_CLIENT_ID",
+                "client_secret": "os.environ/MY_DBX_SECRET",
+                "workspace_url": "https://dbc-abc.cloud.databricks.com",
+            }
+        }
+    )
+    assert config is not None
+    assert config.client_id == "env-cid"
+    assert config.client_secret == "env-secret"
+
+
+# ---------------------------------------------------------------------------
+# Token fetching + caching
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_fetch_token_posts_client_credentials_with_basic_auth():
+    cache = DatabricksAppOAuthTokenCache()
+    config = DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    client = _mock_http_handler(access_token="tok-1")
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        token = await cache.async_get_token(config)
+
+    assert token == "tok-1"
+    client.post.assert_awaited_once()
+    call = client.post.call_args
+    assert call.args[0] == config.token_url
+    assert call.kwargs["data"] == {
+        "grant_type": "client_credentials",
+        "scope": "all-apis",
+    }
+    # 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 "auth" not in call.kwargs
+
+
+@pytest.mark.asyncio
+async def test_token_is_cached_across_calls():
+    cache = DatabricksAppOAuthTokenCache()
+    config = DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    client = _mock_http_handler(access_token="tok-cached")
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        first = await cache.async_get_token(config)
+        second = await cache.async_get_token(config)
+
+    assert first == second == "tok-cached"
+    client.post.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_distinct_clients_do_not_share_token():
+    cache = DatabricksAppOAuthTokenCache()
+    config_a = DatabricksAppOAuthConfig(
+        client_id="cid-a",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    config_b = DatabricksAppOAuthConfig(
+        client_id="cid-b",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+
+    clients = [_mock_http_handler("tok-a"), _mock_http_handler("tok-b")]
+
+    def _next_client(*args, **kwargs):
+        return clients.pop(0)
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        side_effect=_next_client,
+    ):
+        token_a = await cache.async_get_token(config_a)
+        token_b = await cache.async_get_token(config_b)
+
+    assert token_a == "tok-a"
+    assert token_b == "tok-b"
+
+
+@pytest.mark.asyncio
+async def test_ttl_applies_expiry_buffer():
+    cache = DatabricksAppOAuthTokenCache()
+    config = DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    client = _mock_http_handler(access_token="tok", expires_in=600)
+
+    captured = {}
+    real_set = cache.set_cache
+
+    def _spy_set(key, value, **kwargs):
+        captured["ttl"] = kwargs.get("ttl")
+        return real_set(key, value, **kwargs)
+
+    with (
+        patch(
+            "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+            return_value=client,
+        ),
+        patch.object(cache, "set_cache", side_effect=_spy_set),
+    ):
+        await cache.async_get_token(config)
+
+    assert captured["ttl"] == 600 - 60
+
+
+@pytest.mark.asyncio
+async def test_missing_access_token_raises():
+    cache = DatabricksAppOAuthTokenCache()
+    config = DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    client = _mock_http_handler()
+    client.post.return_value.json.return_value = {"not_a_token": "x"}
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        with pytest.raises(ValueError, match="access_token"):
+            await cache.async_get_token(config)
+
+
+def _config():
+    return DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+
+
+@pytest.mark.asyncio
+async def test_http_status_error_raises_value_error():
+    cache = DatabricksAppOAuthTokenCache()
+    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
+        )
+    )
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        with pytest.raises(ValueError, match="status 401"):
+            await cache.async_get_token(_config())
+
+
+@pytest.mark.asyncio
+async def test_transport_error_raises_value_error():
+    cache = DatabricksAppOAuthTokenCache()
+    client = _mock_http_handler(post_error=httpx.ConnectError("boom"))
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        with pytest.raises(ValueError, match="token request failed"):
+            await cache.async_get_token(_config())
+
+
+@pytest.mark.asyncio
+async def test_non_object_json_body_raises():
+    cache = DatabricksAppOAuthTokenCache()
+    client = _mock_http_handler()
+    client.post.return_value.json.return_value = ["not", "an", "object"]
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        with pytest.raises(ValueError, match="non-object JSON"):
+            await cache.async_get_token(_config())
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("expires_in", [None, "not-a-number"])
+async def test_invalid_expires_in_falls_back_to_default_ttl(expires_in):
+    cache = DatabricksAppOAuthTokenCache()
+    client = _mock_http_handler(expires_in=expires_in)
+
+    captured = {}
+    real_set = cache.set_cache
+
+    def _spy_set(key, value, **kwargs):
+        captured["ttl"] = kwargs.get("ttl")
+        return real_set(key, value, **kwargs)
+
+    with (
+        patch(
+            "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+            return_value=client,
+        ),
+        patch.object(cache, "set_cache", side_effect=_spy_set),
+    ):
+        await cache.async_get_token(_config())
+
+    # default TTL (3600) minus the 60s expiry buffer
+    assert captured["ttl"] == 3600 - 60
+
+
+@pytest.mark.asyncio
+async def test_short_lived_token_not_cached():
+    """A token whose lifetime is below the refresh buffer is never cached and
+    leaves no per-key lock behind."""
+    cache = DatabricksAppOAuthTokenCache()
+    config = _config()
+    client = _mock_http_handler(access_token="short", expires_in=30)
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        await cache.async_get_token(config)
+        await cache.async_get_token(config)
+
+    assert cache.get_cache(config.cache_key) is None
+    assert config.cache_key not in cache._locks
+    assert client.post.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_rotated_secret_forces_new_token():
+    """Rotating client_secret changes the cache key so a fresh token is minted."""
+    cache = DatabricksAppOAuthTokenCache()
+    old = DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="old-secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    rotated = DatabricksAppOAuthConfig(
+        client_id="cid",
+        client_secret="new-secret",
+        token_url="https://dbc.cloud.databricks.com/oidc/v1/token",
+        scope="all-apis",
+    )
+    assert old.cache_key != rotated.cache_key
+
+    clients = [_mock_http_handler("old-token"), _mock_http_handler("new-token")]
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        side_effect=lambda *a, **k: clients.pop(0),
+    ):
+        assert await cache.async_get_token(old) == "old-token"
+        assert await cache.async_get_token(rotated) == "new-token"
+
+
+@pytest.mark.asyncio
+async def test_lock_pruned_when_token_evicted():
+    """The per-key lock is removed when its cached token is deleted/evicted."""
+    cache = DatabricksAppOAuthTokenCache()
+    config = _config()
+    client = _mock_http_handler("tok")
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        await cache.async_get_token(config)
+
+    assert config.cache_key in cache._locks
+
+    cache.delete_cache(config.cache_key)
+
+    assert config.cache_key not in cache._locks
+
+
+@pytest.mark.asyncio
+async def test_flush_cache_clears_locks():
+    """flush_cache drops the per-key locks alongside the cached tokens."""
+    cache = DatabricksAppOAuthTokenCache()
+    config = _config()
+    client = _mock_http_handler("tok")
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        await cache.async_get_token(config)
+
+    assert config.cache_key in cache._locks
+
+    cache.flush_cache()
+
+    assert cache._locks == {}
+    assert cache.get_cache(config.cache_key) is None
+
+
+# ---------------------------------------------------------------------------
+# Public helper
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_resolve_returns_none_when_not_configured():
+    assert await resolve_databricks_app_auth_header(None) is None
+    assert await resolve_databricks_app_auth_header({"foo": "bar"}) is None
+
+
+@pytest.mark.asyncio
+async def test_resolve_returns_bearer_header():
+    from litellm.proxy.agent_endpoints.databricks_oauth import (
+        databricks_app_oauth_token_cache,
+    )
+
+    databricks_app_oauth_token_cache.flush_cache()
+
+    litellm_params = {
+        "databricks_oauth": {
+            "client_id": "resolve-cid",
+            "client_secret": "secret",
+            "workspace_url": "https://resolve.cloud.databricks.com",
+        }
+    }
+    client = _mock_http_handler(access_token="resolved-token")
+
+    with patch(
+        "litellm.proxy.agent_endpoints.databricks_oauth.get_async_httpx_client",
+        return_value=client,
+    ):
+        header = await resolve_databricks_app_auth_header(litellm_params)
+
+    assert header == {"Authorization": "Bearer resolved-token"}