mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 00:48:01 +00:00
Mateo Wang
778a7f752d
* 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>