Support OAuth M2M for Databricks Apps A2A agents #29586
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
|
|
||
|
Comment on lines
+1
to
+38
Contributor
There was a problem hiding this comment.
The team rule flags Databricks-specific logic placed outside Rule Used: What: Avoid writing provider-specific code outside... (source) Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time! |
||
| 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()) | ||
|
greptile-apps[bot] marked this conversation as resolved.
greptile-apps[bot] marked this conversation as resolved.
|
||
|
|
||
| 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) | ||
|
greptile-apps[bot] marked this conversation as resolved.
|
||
| 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}"} | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
llms/The team rule asks that provider-specific logic live inside the
litellm/llms/directory. All Databricks-specific OAuth logic here (_token_url_from_workspace,DatabricksAppOAuthConfig, token fetching) is tied to one provider and would be better housed underlitellm/llms/databricks/(or a shared auth utility beneath it) so it remains discoverable alongside the existing Databricks LLM provider and stays consistent with the codebase convention.Rule Used: What: Avoid writing provider-specific code outside... (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!