mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
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:
parent
a72414a061
commit
f3d24ff9f2
1 changed files with 16 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue