fix(a2a): forward agent_extra_headers through completion bridge for custom_llm_provider agents

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
This commit is contained in:
PRABHU KIRAN VANDRANKI 2026-05-19 20:34:31 -04:00
parent a72414a061
commit f3d24ff9f2

View file

@ -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,