From f3d24ff9f24a018a98577f3777ce11ec457ff521 Mon Sep 17 00:00:00 2001
From: PRABHU KIRAN VANDRANKI <72809214+VANDRANKI@users.noreply.github.com>
Date: Tue, 19 May 2026 20:34:31 -0400
Subject: [PATCH] fix(a2a): forward agent_extra_headers through completion
 bridge for custom_llm_provider agents
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When an A2A agent uses `custom_llm_provider` (e.g. LangGraph), `asend_message` and
`asend_message_streaming` route through `_send_message_via_completion_bridge` instead
of the native A2A client path.  The native path correctly merged `agent_extra_headers`
into the HTTP client before making the call, but the completion-bridge path silently
dropped them — so convention-based headers (`x-a2a-{alias}-{target}`) and
admin-allowlisted `extra_headers` never reached the downstream backend.

Fix:
- Add `agent_extra_headers` parameter to `_send_message_via_completion_bridge`.
- Before calling `A2ACompletionBridgeHandler.handle_non_streaming / handle_streaming`,
  shallow-copy `litellm_params` and merge `agent_extra_headers` into its `extra_headers`
  key so they flow through to `litellm.acompletion`.
- Pass `agent_extra_headers` from `asend_message` into the bridge call.
- Apply the same merge in the streaming branch of `asend_message_streaming`.

Fixes #28267
---
 litellm/a2a_protocol/main.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py
index 3ad5485dea1..269ce2d027b 100644
--- a/litellm/a2a_protocol/main.py
+++ b/litellm/a2a_protocol/main.py
@@ -132,6 +132,7 @@ async def _send_message_via_completion_bridge(
     custom_llm_provider: str,
     api_base: Optional[str],
     litellm_params: Dict[str, Any],
+    agent_extra_headers: Optional[Dict[str, str]] = None,
 ) -> LiteLLMSendMessageResponse:
     """
     Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
@@ -152,6 +153,13 @@ async def _send_message_via_completion_bridge(
         else dict(request.params)
     )
 
+    if agent_extra_headers:
+        merged = dict(litellm_params)
+        existing = dict(merged.get("extra_headers") or {})
+        existing.update(agent_extra_headers)
+        merged["extra_headers"] = existing
+        litellm_params = merged
+
     response_dict = await A2ACompletionBridgeHandler.handle_non_streaming(
         request_id=str(request.id),
         params=params,
@@ -281,6 +289,7 @@ async def asend_message(
             custom_llm_provider=custom_llm_provider,
             api_base=api_base,
             litellm_params=litellm_params,
+            agent_extra_headers=agent_extra_headers,
         )
 
     # Standard A2A client flow
@@ -509,6 +518,13 @@ async def asend_message_streaming(  # noqa: PLR0915
             else dict(request.params)
         )
 
+        if agent_extra_headers:
+            merged = dict(litellm_params)
+            existing = dict(merged.get("extra_headers") or {})
+            existing.update(agent_extra_headers)
+            merged["extra_headers"] = existing
+            litellm_params = merged
+
         async for chunk in A2ACompletionBridgeHandler.handle_streaming(
             request_id=str(request.id),
             params=params,