scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(a2a): preserve agent url, fix streaming chunk envelope, and protect forwarded metadata

Browse source 
- Streaming chunk: move final out of the message object into the
  result envelope per the A2A spec.
- Agent card merge: keep upstream url on the stored card so the
  runtime invocation path can locate the upstream backend; the public
  well-known endpoint already rewrites this field to the proxy URL
  before exposing it to clients.
- Completion bridge: apply A2A forward metadata after merging
  litellm_params so an agent-configured extra_body cannot
  overwrite the forwarded metadata.

Co-authored-by: Yassin Kortam <yassin@berri.ai>
This commit is contained in:
Cursor Agent 2026-05-26 09:32:50 +00:00
parent 79ffe24c6f
commit 20129d4405
No known key found for this signature in database
5 changed files with 57 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -100,11 +100,6 @@ class A2ACompletionBridgeHandler:
"api_base": api_base,
"stream": False,
}
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
@ -112,6 +107,14 @@ class A2ACompletionBridgeHandler:
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so an
# agent-configured ``extra_body`` does not overwrite the forwarded
# A2A metadata; the helper merges into any existing ``extra_body``.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# Call litellm.acompletion
response = await litellm.acompletion(**completion_params)
@ -212,11 +215,6 @@ class A2ACompletionBridgeHandler:
"api_base": api_base,
"stream": True,
}
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
@ -224,6 +222,14 @@ class A2ACompletionBridgeHandler:
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so an
# agent-configured ``extra_body`` does not overwrite the forwarded
# A2A metadata; the helper merges into any existing ``extra_body``.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)

14
litellm/a2a_protocol/litellm_completion_bridge/transformation.py
View file

@ -327,15 +327,19 @@ class A2ACompletionBridgeTransformation:
if not content and not is_final:
return None
# Build A2A streaming chunk (legacy format)
# Build A2A streaming chunk (legacy format). ``final`` is an
# envelope-level streaming property per the A2A spec and must live
# alongside ``message`` in ``result``, not inside the message object.
a2a_chunk = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
"message": {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
},
"final": is_final,
},
}

26
litellm/a2a_protocol/providers/litellm_completion/handler.py
View file

@ -96,11 +96,6 @@ class A2ACompletionBridgeHandler:
"api_base": api_base,
"stream": False,
}
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
@ -108,6 +103,14 @@ class A2ACompletionBridgeHandler:
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so an
# agent-configured ``extra_body`` does not overwrite the forwarded
# A2A metadata; the helper merges into any existing ``extra_body``.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# Call litellm.acompletion
response = await litellm.acompletion(**completion_params)
@ -211,11 +214,6 @@ class A2ACompletionBridgeHandler:
"api_base": api_base,
"stream": True,
}
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
@ -223,6 +221,14 @@ class A2ACompletionBridgeHandler:
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so an
# agent-configured ``extra_body`` does not overwrite the forwarded
# A2A metadata; the helper merges into any existing ``extra_body``.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)

12
litellm/proxy/a2a/agent_card.py
View file

@ -63,6 +63,12 @@ _ALLOWED_TOP_LEVEL_KEYS = {
"security",
"supportsAuthenticatedExtendedCard",
"signatures",
# ``url`` is retained on the stored card because the runtime A2A invocation
# path (``a2a_endpoints.py``) reads ``agent.agent_card_params['url']`` to
# locate the upstream backend. The public ``/.well-known/agent-card.json``
# endpoint rewrites this field to the proxy URL before serving it to
# clients, so retaining it here does not leak the upstream to A2A callers.
"url",
}
_DEFAULT_SKILLS: List[Dict[str, Any]] = [
@ -121,8 +127,10 @@ def merge_agent_card(
"""
base: Dict[str, Any] = deepcopy(dict(upstream_card)) if upstream_card else {}
# Strip the upstream URL so clients don't accidentally bypass the proxy.
base.pop("url", None)
# Keep the upstream ``url`` on the stored card: the runtime A2A
# invocation path reads it from ``agent_card_params`` to know where to
# proxy requests. The public well-known endpoint rewrites this field
# to the proxy URL before exposing the card to clients.
base["protocolVersion"] = LITELLM_A2A_PROTOCOL_VERSION

8
tests/test_litellm/proxy/a2a/test_agent_card.py
View file

@ -39,11 +39,15 @@ def _full_upstream_card() -> dict:
}
def test_strips_top_level_url():
def test_preserves_top_level_url_for_runtime_invocation():
# The runtime A2A invocation path reads ``agent_card_params['url']`` to
# know where to proxy requests, so the merge must keep the upstream URL
# on the stored card. The public well-known endpoint rewrites this field
# to the proxy URL before exposing it to clients.
merged = merge_agent_card(
_full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE
)
assert "url" not in merged
assert merged["url"] == "http://internal:9999/"
def test_overrides_protocol_version():