From 20129d44055ae27f166ce8fa7710787314da25d2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 26 May 2026 09:32:50 +0000 Subject: [PATCH] fix(a2a): preserve agent url, fix streaming chunk envelope, and protect forwarded metadata - 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 --- .../litellm_completion_bridge/handler.py | 26 ++++++++++++------- .../transformation.py | 14 ++++++---- .../providers/litellm_completion/handler.py | 26 ++++++++++++------- litellm/proxy/a2a/agent_card.py | 12 +++++++-- .../test_litellm/proxy/a2a/test_agent_card.py | 8 ++++-- 5 files changed, 57 insertions(+), 29 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 38ab94fe56c..36ad4f2444b 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -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) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index f5b8b40f95a..3cab71558b6 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -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, }, } diff --git a/litellm/a2a_protocol/providers/litellm_completion/handler.py b/litellm/a2a_protocol/providers/litellm_completion/handler.py index 08d24014fdc..f7e806f86e5 100644 --- a/litellm/a2a_protocol/providers/litellm_completion/handler.py +++ b/litellm/a2a_protocol/providers/litellm_completion/handler.py @@ -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) diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index f4a6862997c..6cc9670bc09 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -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 diff --git a/tests/test_litellm/proxy/a2a/test_agent_card.py b/tests/test_litellm/proxy/a2a/test_agent_card.py index 4c9901ccb55..46961cf21f9 100644 --- a/tests/test_litellm/proxy/a2a/test_agent_card.py +++ b/tests/test_litellm/proxy/a2a/test_agent_card.py @@ -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():