Support OAuth M2M for Databricks Apps A2A agents

[mateo-berri]

Relevant issues

Linear ticket

Resolves LIT-3549

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Screenshots / Proof of Fix

Register a Databricks App agent in the proxy config, then call it through the gateway. The proxy mints an OAuth token from the workspace and forwards it to the App.

agents:
  - agent_name: my-databricks-app
    agent_card_params:
      name: my-databricks-app
      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

curl -s http://localhost:4000/a2a/my-databricks-app \
  -H "Authorization: Bearer $LITELLM_KEY" \
  -H "Content-Type: application/json" \
  -d '{
        "jsonrpc": "2.0",
        "id": "1",
        "method": "message/send",
        "params": {
          "message": {
            "role": "user",
            "parts": [{"kind": "text", "text": "ping"}],
            "messageId": "m1"
          }
        }
      }'

The outbound request to the Databricks App carries a freshly minted Authorization: Bearer <oauth-token> rather than the inbound LiteLLM key, and the token is cached until shortly before it expires so repeat calls reuse it.

Type

🆕 New Feature

Changes

Databricks App endpoints require a short-lived OAuth token minted from the workspace OIDC token endpoint instead of a static bearer; the existing A2A outbound auth only supported static_headers and client header passthrough, so Databricks App agents could not be registered.

This wires OAuth M2M (client_credentials) into the A2A agent gateway. When an agent is configured with a databricks_oauth block in its litellm_params (client_id, client_secret, workspace_url, optional scope), LiteLLM mints a token via the workspace oidc/v1/token endpoint, caches it keyed by token endpoint plus client_id plus scope until shortly before expiry, and attaches it as the outbound Authorization header on both message/send and message/stream. The minted token overrides any statically configured Authorization, which is what Databricks Apps require. Credentials support os.environ/ references so secrets stay out of config. The token acquisition mirrors the MCP client_credentials flow and reuses the same workspace token endpoint that the Databricks model provider already uses.

No DB schema change is needed since the config lives inside the existing litellm_params JSON blob.

---

Generated by Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 99.11504% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/agent_endpoints/databricks_oauth.py	99.07%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[mateo-berri]

@greptileai

[greptile-apps]

Greptile Summary

This PR wires OAuth M2M (client_credentials) into the A2A agent gateway so that agents configured with a databricks_oauth block in litellm_params automatically receive a short-lived Bearer token on every outbound call instead of a static credential.

• databricks_oauth.py: new module that parses the OAuth config (including os.environ/ resolution), fetches a token from the workspace OIDC endpoint with HTTP Basic auth, and caches it with TTL minus a 60 s expiry buffer. A SHA-256 digest of client_secret is included in the cache key, per-key asyncio.Lock objects are pruned via an overridden _remove_key, and tokens whose lifetime is shorter than the buffer are never cached (lock popped immediately). All concerns raised in the previous review round have been addressed.
• a2a_endpoints.py: one-line injection after header construction, applied before both message/send and message/stream dispatch so both code paths receive the minted token.
• Tests: test_databricks_oauth.py (unit, comprehensive) and two new cases in test_agent_headers.py (integration). All use mocks; no real network calls.

Confidence Score: 5/5

Safe to merge. The change is well-scoped, strictly additive, and touches only the new Databricks OAuth path — agents without a databricks_oauth block are unaffected.

All issues raised in the previous review round have been correctly addressed: the lock dictionary is pruned via an overridden _remove_key, the client_secret digest is included in the cache key so rotated credentials get a fresh token, and short-lived tokens bypass the cache without leaving an orphaned lock entry. The implementation is thoroughly tested with mocks, covering config parsing, TTL buffering, lock eviction, and all error paths.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/agent_endpoints/databricks_oauth.py	New module implementing OAuth M2M (client_credentials) for Databricks App A2A agents. Correctly handles token caching with per-key asyncio locks, expiry buffering, secret digest in cache key, lock cleanup via _remove_key override, and short-lived token edge case. Previous review concerns have all been addressed.
litellm/proxy/agent_endpoints/a2a_endpoints.py	Small addition to invoke_agent_a2a: calls resolve_databricks_app_auth_header and merges the resulting Bearer token into agent_extra_headers before both message/send and message/stream dispatch.
tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py	Comprehensive unit tests covering config parsing, env-var resolution, token fetching with Basic auth, caching/TTL, short-lived token bypass, lock eviction, and error paths.
tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py	Two new integration-style tests verifying OAuth header injection and that the minted token overrides a statically configured Authorization header.

_{Reviews (7): Last reviewed commit: "Prune Databricks OAuth lock on the short..." | Re-trigger Greptile}

[greptile-apps]

Provider-specific code placed outside 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 under litellm/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!

[greptile-apps]

Greptile Summary

This PR wires OAuth M2M (client_credentials) support into the A2A agent gateway so agents that target Databricks App endpoints can mint short-lived workspace tokens instead of forwarding the incoming LiteLLM bearer token.

• Adds databricks_oauth.py: parses a databricks_oauth block from litellm_params, resolves os.environ/ secrets, fetches a token from the workspace OIDC endpoint, and caches it with an expiry buffer. Double-checked locking prevents thundering-herd on cache miss.
• Integrates the new helper into invoke_agent_a2a so the minted token wins over any statically configured Authorization header, matching Databricks Apps' auth requirements.
• Comprehensive mock-only unit tests cover config validation, TTL arithmetic, error paths, and the override behavior.

Confidence Score: 4/5

Safe to merge with minor follow-up items; the core token-minting and caching logic is correct and the tests are thorough mock-only tests.

The double-checked locking, expiry buffering, and error propagation are all correctly implemented. Two non-blocking concerns are worth a follow-up: the _locks dictionary accumulates entries forever because eviction only clears cache_dict/ttl_dict, and client_secret is absent from the cache key, meaning a rotated credential won't take effect until the cached token expires naturally.

litellm/proxy/agent_endpoints/databricks_oauth.py — the token cache subclass needs attention for the lock-eviction gap and the cache key composition.

Important Files Changed

Filename	Overview
litellm/proxy/agent_endpoints/databricks_oauth.py	New file: implements OAuth M2M (client_credentials) token acquisition and caching for Databricks App A2A agents. The double-checked locking pattern and expiry-buffered TTL are sound, but the _locks dict is never pruned and client_secret is absent from the cache key.
litellm/proxy/agent_endpoints/a2a_endpoints.py	Minimal change: adds a call to resolve_databricks_app_auth_header after existing header assembly and merges the result into agent_extra_headers. Errors propagate to the existing except Exception handler, producing a JSON-RPC 500 response. No structural changes.
tests/test_litellm/proxy/agent_endpoints/test_databricks_oauth.py	New test file: comprehensive mock-only unit tests covering config parsing, secret resolution, TTL/buffer arithmetic, error paths (HTTP 4xx, transport error, malformed body), and the public header helper. No real network calls.
tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py	Adds two integration-style tests verifying OAuth header injection and that the minted token overrides a statically configured Authorization header. Existing tests are unchanged; new tests use proper mocking.

_{Reviews (2): Last reviewed commit: "Add tests covering Databricks App OAuth ..." | Re-trigger Greptile}

[greptile-apps]

Provider-specific code outside llms/

The team rule flags Databricks-specific logic placed outside litellm/llms/. The OAuth flow here is specific to the Databricks workspace OIDC endpoint. If future providers need similar M2M flows for A2A agents, each will add a new top-level file in agent_endpoints/ instead of being co-located with its provider. Consider whether this belongs under litellm/llms/databricks/ or a shared litellm/proxy/agent_endpoints/oauth/ abstraction.

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!

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[mateo-berri]

@greptileai

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ mateo-berri
❌ cursoragent
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[mateo-berri]

@greptileai

[mateo-berri]

QA Proof

[Sameerlite]

There is no gating condition to this. So if i am trying to use any other agent. Will it also get added with this auth?

[mateo-berri]

Good catch on the readability. Other agents are not affected: the gating is on the presence of a databricks_oauth block in the agent's litellm_params. resolve_databricks_app_auth_header returns None for any agent without that block, so the header is never added

That gating was hidden inside the resolver though, so I made it explicit at the call site in 89c559b:

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}

Now an agent only enters the OAuth path when it is explicitly configured for Databricks; every other agent skips it entirely. I also added a regression test (test_non_databricks_agent_skips_oauth_resolution) asserting the resolver is never invoked for a non-Databricks agent

[mateo-berri]

@Sameerlite re-requesting review. Your gating concern is addressed: the OAuth path is now explicitly guarded at the call site in a2a_endpoints.py, so only agents that carry a databricks_oauth block in their litellm_params ever mint a token; every other agent skips it entirely and its headers are 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}

This block runs before both the message/send and message/stream dispatch, so both paths are covered. Added a regression test test_non_databricks_agent_skips_oauth_resolution that asserts the resolver is never invoked for a non-Databricks agent and no Authorization header is added; the full test_agent_headers.py suite (16 tests) passes locally

[Sameerlite]

LGTM