diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 9f1403d432..7b2f75e1cf 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/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 diff --git a/litellm/proxy/agent_endpoints/databricks_oauth.py b/litellm/proxy/agent_endpoints/databricks_oauth.py new file mode 100644 index 0000000000..1c1f5a2b4c --- /dev/null +++ b/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 "}`` 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}"} diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py index 93ba9dc922..fa530e0975 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py b/tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py new file mode 100644 index 0000000000..39f51ee040 --- /dev/null +++ b/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"}