From c3d599dd2705c4544c5fbb8a30fac924aa0af248 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 00:23:31 +0900 Subject: [PATCH 01/20] fix: honor registered A2A providers --- litellm/proxy/agent_endpoints/a2a_routing.py | 151 +++++++++++++++++- .../agent_endpoints/model_list_helpers.py | 11 +- .../test_model_list_helpers.py | 31 +++- .../proxy/test_route_a2a_models.py | 114 ++++++++++++- 4 files changed, 294 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 8a795214750..0373033ee10 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -5,20 +5,134 @@ Handles routing for A2A agents (models with "a2a/" prefix). Looks up agents in the registry and injects their API base URL. """ -from typing import Any, Final +from __future__ import annotations + +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final +from uuid import uuid4 from fastapi import HTTPException +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm.llms.a2a.common_utils import A2AError, convert_messages_to_prompt, extract_text_from_a2a_response from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Choices, Message, ModelResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str]) +_MESSAGES_ADAPTER: Final = TypeAdapter(list[AllMessageValues]) + + +class _A2ATextPart(TypedDict): + kind: ReadOnly[str] + text: ReadOnly[str] + + +class _A2AMessage(TypedDict): + role: ReadOnly[str] + parts: ReadOnly[tuple[_A2ATextPart, ...]] + messageId: ReadOnly[str] + + +class _A2AParams(TypedDict): + message: ReadOnly[_A2AMessage] + + +async def _route_registered_provider( + data: Mapping[str, object], + model_name: str, + api_base: str, + litellm_params: Mapping[str, object], +) -> ModelResponse | CustomStreamWrapper: + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator + + raw_messages: Final = data.get("messages") + messages: Final = _MESSAGES_ADAPTER.validate_python(raw_messages) + stream: Final = data.get("stream") is True + request_id: Final = str(uuid4()) + params: Final[_A2AParams] = { + "message": { + "role": "user", + "parts": ({"kind": "text", "text": convert_messages_to_prompt(messages)},), + "messageId": str(uuid4()), + } + } + provider_params: Final = _OBJECT_DICT_ADAPTER.validate_python(litellm_params) + bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python(params) + configured_headers: Final = litellm_params.get("extra_headers") or litellm_params.get("headers") + agent_extra_headers: Final = ( + _HEADERS_ADAPTER.validate_python(configured_headers) if isinstance(configured_headers, dict) else None + ) + + if stream: + streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming( + request_id=request_id, + params=bridge_params, + litellm_params=provider_params, + api_base=api_base, + agent_extra_headers=agent_extra_headers, + ) + completion_stream: Final = A2AModelResponseIterator( + streaming_response=streaming_response, + sync_stream=False, + model=model_name, + ) + logging_obj: Final = data.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + raise TypeError("litellm_logging_obj is required for streaming A2A requests") + return CustomStreamWrapper( + completion_stream=completion_stream, + model=model_name, + custom_llm_provider="a2a", + logging_obj=logging_obj, + stream_options=data.get("stream_options"), + ) + + response: Final = await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=request_id, + params=bridge_params, + litellm_params=provider_params, + api_base=api_base, + agent_extra_headers=agent_extra_headers, + ) + error_value: Final = response.get("error") + if isinstance(error_value, dict): + error: Final = _OBJECT_DICT_ADAPTER.validate_python(error_value) + error_message: Final = error.get("message") + raise A2AError( + status_code=500, + message=f"A2A error: {error_message if isinstance(error_message, str) else 'Unknown error'}", + ) + + text: Final = extract_text_from_a2a_response(response) + model_response: Final = ModelResponse( + id=str(response.get("id") or request_id), + model=model_name, + choices=[ # mutable-ok: ModelResponse requires a choices list + Choices(finish_reason="stop", index=0, message=Message(content=text, role="assistant")) + ], + ) + return model_response async def route_a2a_agent_request( - data: dict, + data: Mapping[str, object], route_type: str, user_api_key_dict: UserAPIKeyAuth | None = None, -) -> Any | None: +) -> Awaitable[object] | None: """ Route A2A agent requests directly to litellm with injected API base. @@ -69,13 +183,34 @@ async def route_a2a_agent_request( ) # Get API base URL from agent config - if not agent.agent_card_params or "url" not in agent.agent_card_params: + agent_card_params: Final = agent.agent_card_params + agent_url: Final = agent_card_params.get("url") if agent_card_params else None + if not isinstance(agent_url, str) or not agent_url: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) - # Inject API base and route to litellm - data["api_base"] = agent.agent_card_params["url"] - verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, data["api_base"]) + registered_params_value: Final = agent.litellm_params + registered_provider_value: Final = ( + registered_params_value.get("custom_llm_provider") if registered_params_value else None + ) + registered_provider: Final = registered_provider_value if isinstance(registered_provider_value, str) else None + configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None + api_base: Final = configured_api_base if isinstance(configured_api_base, str) else agent_url + if ( + registered_provider + and registered_provider != "a2a" + and route_type == "acompletion" + and registered_params_value is not None + ): + verbose_proxy_logger.debug("[A2A] Routing %s through %s", model_name, registered_provider) + return _route_registered_provider( + data=data, + model_name=model_name, + api_base=api_base, + litellm_params=registered_params_value, + ) - return getattr(litellm, f"{route_type}")(**data) + completion_data: Final = MappingProxyType({**data, "api_base": api_base}) + verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, api_base) + return getattr(litellm, f"{route_type}")(**completion_data) # pyright: ignore[reportAny] # dynamic SDK route diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 4a88644bae5..624136c020b 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -4,12 +4,18 @@ Helper functions for appending A2A agents to model lists. Used by proxy model endpoints to make agents appear in UI alongside models. """ +from typing import Final + +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) +_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) + async def append_agents_to_model_group( model_groups: list[ModelGroupInfoProxy], @@ -70,12 +76,15 @@ async def append_agents_to_model_info( for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: + agent_params = agent.litellm_params + provider_value = agent_params.get("custom_llm_provider") if agent_params else None + custom_llm_provider = provider_value if isinstance(provider_value, str) else "a2a" models.append( { "model_name": f"a2a/{agent.agent_name}", "litellm_params": { "model": f"a2a/{agent.agent_name}", - "custom_llm_provider": "a2a", + "custom_llm_provider": custom_llm_provider, }, "model_info": { "id": agent.agent_id, diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py index 939ab1cab40..b82cd1dd231 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -4,8 +4,6 @@ Test appending A2A agents to model lists. Maps to: litellm/proxy/agent_endpoints/model_list_helpers.py """ - - from unittest.mock import AsyncMock, Mock, patch import pytest @@ -109,3 +107,32 @@ async def test_append_agents_to_model_info(): assert result[0]["litellm_params"]["custom_llm_provider"] == "a2a" assert result[0]["model_info"]["id"] == "agent-123" assert result[0]["model_info"]["mode"] == "chat" + + +@pytest.mark.asyncio +async def test_append_agents_to_model_info_preserves_registered_provider(): + agent = AgentResponse( + agent_id="agent-123", + agent_name="test-agent", + agent_card_params={"url": "http://example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + registry = Mock() + registry.get_agent_by_id = Mock(return_value=agent) + + with ( + patch( # test-quality-ok: access resolution is outside model-list assembly + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))), + ), + patch( # test-quality-ok: registry output drives model-list assembly + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ), + ): + result = await append_agents_to_model_info( + models=[], + user_api_key_dict=Mock(spec=UserAPIKeyAuth), + ) + + assert result[0]["litellm_params"]["custom_llm_provider"] == "pydantic_ai_agents" diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 35308474949..6273fdc0df0 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -4,8 +4,6 @@ Test A2A model routing in proxy. Maps to: litellm/proxy/agent_endpoints/a2a_routing.py """ - - from unittest.mock import AsyncMock, Mock, patch import pytest @@ -72,6 +70,118 @@ async def test_route_a2a_model_bypasses_router(): assert call_kwargs["api_base"] == "http://agent.example.com" +@pytest.mark.asyncio +async def test_route_a2a_model_uses_registered_provider(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + data = { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + } + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": "Hello back"}], + "messageId": "message-id", + }, + } + + with ( + patch( # test-quality-ok: registry lookup is the routing seam + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( # test-quality-ok: access control is outside this routing test + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( # test-quality-ok: provider dispatch is the tested seam + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + patch( # test-quality-ok: generic dispatch must stay unused + "litellm.acompletion", AsyncMock() + ) as generic_completion, + ): + call = await route_a2a_agent_request(data, "acompletion") + response = await call + + bridge.assert_awaited_once() + generic_completion.assert_not_called() + assert response.choices[0].message.content == "Hello back" + + +@pytest.mark.asyncio +async def test_route_a2a_stream_uses_registered_provider(): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + logging_obj = Mock(spec=Logging) + data = { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + "stream": True, + "litellm_logging_obj": logging_obj, + } + provider_stream = object() + completion_stream = object() + wrapper = object() + + with ( + patch( # test-quality-ok: registry lookup is the routing seam + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( # test-quality-ok: access control is outside this routing test + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( # test-quality-ok: provider dispatch is the tested seam + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_streaming", + Mock(return_value=provider_stream), + ) as bridge, + patch( # test-quality-ok: iterator wiring is the tested seam + "litellm.llms.a2a.chat.streaming_iterator.A2AModelResponseIterator", + Mock(return_value=completion_stream), + ), + patch( # test-quality-ok: wrapper wiring is the tested seam + "litellm.litellm_core_utils.streaming_handler.CustomStreamWrapper", + Mock(return_value=wrapper), + ) as stream_wrapper, + patch( # test-quality-ok: generic dispatch must stay unused + "litellm.acompletion", AsyncMock() + ) as generic_completion, + ): + call = await route_a2a_agent_request(data, "acompletion") + response = await call + + bridge.assert_called_once() + stream_wrapper.assert_called_once_with( + completion_stream=completion_stream, + model="a2a/test-agent", + custom_llm_provider="a2a", + logging_obj=logging_obj, + stream_options=None, + ) + generic_completion.assert_not_called() + assert response is wrapper + + @pytest.mark.asyncio async def test_route_non_a2a_model_raises_error_if_not_in_router(): """Test that non-a2a models that aren't in router raise an error""" From 02c5c74b37bd912b5d924078b1b8356aa4504afa Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 14:18:12 +0900 Subject: [PATCH 02/20] fix: preserve registered A2A request context --- litellm/llms/a2a/chat/streaming_iterator.py | 7 + litellm/proxy/agent_endpoints/a2a_routing.py | 134 ++++++++++++++++-- .../agent_endpoints/model_list_helpers.py | 19 +-- .../a2a/chat/test_a2a_streaming_iterator.py | 26 ++++ .../test_model_list_helpers.py | 26 ++++ .../proxy/test_route_a2a_models.py | 60 +++++++- 6 files changed, 253 insertions(+), 19 deletions(-) create mode 100644 tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 1983c18a6b3..0cbff95b9d3 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -83,6 +83,13 @@ class A2AModelResponseIterator(BaseModelResponseIterator): tool_use=None, ) + def _handle_string_chunk( + self, str_line: str | dict + ) -> GenericStreamingChunk | ModelResponseStream: + if isinstance(str_line, dict): + return self.chunk_parser(chunk=str_line) + return super()._handle_string_chunk(str_line=str_line) + def _get_finish_reason(self, chunk: dict) -> str | None: """Extract finish reason from A2A chunk""" result: Final = chunk.get("result", {}) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 0373033ee10..2963240f41a 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -7,6 +7,7 @@ Looks up agents in the registry and injects their API base URL. from __future__ import annotations +import asyncio from collections.abc import Awaitable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -18,6 +19,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm.interactions.agents.utils import merge_agent_headers from litellm.llms.a2a.common_utils import A2AError, convert_messages_to_prompt, extract_text_from_a2a_response from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.types.llms.openai import AllMessageValues @@ -29,6 +31,41 @@ if TYPE_CHECKING: _OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str]) _MESSAGES_ADAPTER: Final = TypeAdapter(list[AllMessageValues]) +_FORWARDED_REQUEST_PARAMS: Final = frozenset( + { + "audio", + "frequency_penalty", + "functions", + "function_call", + "include_server_side_tool_invocations", + "logit_bias", + "logprobs", + "guardrails", + "max_completion_tokens", + "max_tokens", + "modalities", + "n", + "parallel_tool_calls", + "prediction", + "presence_penalty", + "reasoning_effort", + "response_format", + "seed", + "service_tier", + "stop", + "store", + "temperature", + "thinking", + "timeout", + "tool_choice", + "tools", + "top_logprobs", + "top_p", + "user", + "verbosity", + "web_search_options", + } +) class _A2ATextPart(TypedDict): @@ -49,8 +86,9 @@ class _A2AParams(TypedDict): async def _route_registered_provider( data: Mapping[str, object], model_name: str, - api_base: str, + api_base: str | None, litellm_params: Mapping[str, object], + static_headers: Mapping[str, str] | None, ) -> ModelResponse | CustomStreamWrapper: from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, @@ -70,12 +108,25 @@ async def _route_registered_provider( "messageId": str(uuid4()), } } - provider_params: Final = _OBJECT_DICT_ADAPTER.validate_python(litellm_params) + provider_params: Final = { + **_OBJECT_DICT_ADAPTER.validate_python(litellm_params), + **{ + key: data[key] + for key in _FORWARDED_REQUEST_PARAMS + if key in data and data[key] is not None + }, + } bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python(params) configured_headers: Final = litellm_params.get("extra_headers") or litellm_params.get("headers") - agent_extra_headers: Final = ( + configured_headers_dict: Final = ( _HEADERS_ADAPTER.validate_python(configured_headers) if isinstance(configured_headers, dict) else None ) + agent_extra_headers: Final = merge_agent_headers( + dynamic_headers=configured_headers_dict, + static_headers=static_headers, + ) + if agent_extra_headers: + provider_params["extra_headers"] = agent_extra_headers if stream: streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming( @@ -125,9 +176,63 @@ async def _route_registered_provider( Choices(finish_reason="stop", index=0, message=Message(content=text, role="assistant")) ], ) + usage: Final = response.get("usage") + if usage is not None: + setattr(model_response, "usage", usage) + + logging_obj: Final = data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + + def _enqueue_logging() -> None: + asyncio.create_task( + logging_obj.dispatch_success_handlers( + model_response, + cache_hit=False, + prefer_async_handlers=True, + ) + ) + + logging_obj._enqueue_deferred_logging = _enqueue_logging + return model_response +def _merge_agent_guardrails( + data: Mapping[str, object], + agent_guardrails: object, +) -> Mapping[str, object]: + if not agent_guardrails: + return data + + configured_guardrails: list[object] = ( + agent_guardrails if isinstance(agent_guardrails, list) else [agent_guardrails] + ) + metadata = data.get("metadata") + metadata_guardrails = metadata.get("guardrails") if isinstance(metadata, dict) else None + root_guardrails = data.get("guardrails") + existing_guardrails: list[object] = [] + for value in (metadata_guardrails, root_guardrails): + if isinstance(value, list): + existing_guardrails.extend(value) + elif value: + existing_guardrails.append(value) + + merged_guardrails = existing_guardrails + [ + guardrail for guardrail in configured_guardrails if guardrail not in existing_guardrails + ] + if isinstance(data, dict): + data["guardrails"] = merged_guardrails + if isinstance(metadata, dict): + metadata["guardrails"] = merged_guardrails + return data + + merged_data = dict(data) + merged_data["guardrails"] = merged_guardrails + if isinstance(metadata, dict): + merged_data["metadata"] = {**metadata, "guardrails": merged_guardrails} + return merged_data + + async def route_a2a_agent_request( data: Mapping[str, object], route_type: str, @@ -185,11 +290,6 @@ async def route_a2a_agent_request( # Get API base URL from agent config agent_card_params: Final = agent.agent_card_params agent_url: Final = agent_card_params.get("url") if agent_card_params else None - if not isinstance(agent_url, str) or not agent_url: - verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) - route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) - raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) - registered_params_value: Final = agent.litellm_params registered_provider_value: Final = ( registered_params_value.get("custom_llm_provider") if registered_params_value else None @@ -197,6 +297,19 @@ async def route_a2a_agent_request( registered_provider: Final = registered_provider_value if isinstance(registered_provider_value, str) else None configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None api_base: Final = configured_api_base if isinstance(configured_api_base, str) else agent_url + registered_model: Final = registered_params_value.get("model") if registered_params_value else None + cardless_provider: Final = ( + registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model + ) + if (not isinstance(agent_url, str) or not agent_url) and not cardless_provider: + verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) + route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) + raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) + + routed_data: Final = _merge_agent_guardrails( + data=data, + agent_guardrails=registered_params_value.get("guardrails") if registered_params_value else None, + ) if ( registered_provider and registered_provider != "a2a" @@ -205,12 +318,13 @@ async def route_a2a_agent_request( ): verbose_proxy_logger.debug("[A2A] Routing %s through %s", model_name, registered_provider) return _route_registered_provider( - data=data, + data=routed_data, model_name=model_name, api_base=api_base, litellm_params=registered_params_value, + static_headers=agent.static_headers, ) - completion_data: Final = MappingProxyType({**data, "api_base": api_base}) + completion_data: Final = MappingProxyType({**routed_data, "api_base": api_base}) verbose_proxy_logger.debug("[A2A] Routing %s to %s", model_name, api_base) return getattr(litellm, f"{route_type}")(**completion_data) # pyright: ignore[reportAny] # dynamic SDK route diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 624136c020b..67d721f9faa 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -6,16 +6,12 @@ Used by proxy model endpoints to make agents appear in UI alongside models. from typing import Final -from pydantic import TypeAdapter - from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) -_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) - async def append_agents_to_model_group( model_groups: list[ModelGroupInfoProxy], @@ -39,11 +35,16 @@ async def append_agents_to_model_group( for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: + agent_params: Final = agent.litellm_params + provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None + custom_llm_provider: Final = ( + provider_value if isinstance(provider_value, str) else "a2a" + ) model_groups.append( ModelGroupInfoProxy( model_group=f"a2a/{agent.agent_name}", mode="chat", - providers=["a2a"], + providers=[custom_llm_provider], ) ) case _: @@ -76,9 +77,11 @@ async def append_agents_to_model_info( for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) if agent is not None: - agent_params = agent.litellm_params - provider_value = agent_params.get("custom_llm_provider") if agent_params else None - custom_llm_provider = provider_value if isinstance(provider_value, str) else "a2a" + agent_params: Final = agent.litellm_params + provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None + custom_llm_provider: Final = ( + provider_value if isinstance(provider_value, str) else "a2a" + ) models.append( { "model_name": f"a2a/{agent.agent_name}", diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py new file mode 100644 index 00000000000..579d6f8efef --- /dev/null +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -0,0 +1,26 @@ +"""Tests for the A2A chat streaming iterator.""" + +import pytest + +from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator + + +@pytest.mark.asyncio +async def test_async_iterator_accepts_decoded_a2a_events(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "artifact-update", + "artifact": {"parts": [{"kind": "text", "text": "Hello"}]}, + }, + } + + iterator = A2AModelResponseIterator( + streaming_response=_events(), + sync_stream=False, + ) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["text"] == "Hello" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py index b82cd1dd231..6ce22ed1a3f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -64,6 +64,32 @@ async def test_append_agents_to_model_group(): assert result[0].providers == ["a2a"] +@pytest.mark.asyncio +async def test_append_agents_to_model_group_preserves_registered_provider(): + agent = AgentResponse( + agent_id="agent-123", + agent_name="test-agent", + agent_card_params={"url": "http://example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + registry = Mock() + registry.get_agent_by_id = Mock(return_value=agent) + + with ( + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))), + ), + patch("litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry), + ): + result = await append_agents_to_model_group( + model_groups=[], + user_api_key_dict=Mock(spec=UserAPIKeyAuth), + ) + + assert result[0].providers == ["pydantic_ai_agents"] + + @pytest.mark.asyncio async def test_append_agents_to_model_info(): """Test agents are converted to model info format with a2a/ prefix""" diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 6273fdc0df0..1dd2edc8bf7 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -78,11 +78,20 @@ async def test_route_a2a_model_uses_registered_provider(): agent_id="test-agent-id", agent_name="test-agent", agent_card_params={"url": "http://agent.example.com"}, - litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + litellm_params={ + "custom_llm_provider": "pydantic_ai_agents", + "guardrails": ["agent-guardrail"], + }, + static_headers={"Authorization": "Bearer static"}, ) data = { "model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}], + "guardrails": ["request-guardrail"], + "max_tokens": 32, + "temperature": 0.2, + "timeout": 12.0, + "tools": [{"type": "function", "function": {"name": "lookup"}}], } bridge_response = { "jsonrpc": "2.0", @@ -118,6 +127,55 @@ async def test_route_a2a_model_uses_registered_provider(): bridge.assert_awaited_once() generic_completion.assert_not_called() assert response.choices[0].message.content == "Hello back" + bridge_kwargs = bridge.await_args.kwargs + assert bridge_kwargs["litellm_params"]["max_tokens"] == 32 + assert bridge_kwargs["litellm_params"]["temperature"] == 0.2 + assert bridge_kwargs["litellm_params"]["timeout"] == 12.0 + assert bridge_kwargs["litellm_params"]["tools"] == data["tools"] + assert bridge_kwargs["litellm_params"]["guardrails"] == ["request-guardrail", "agent-guardrail"] + assert bridge_kwargs["litellm_params"]["extra_headers"] == {"Authorization": "Bearer static"} + + +@pytest.mark.asyncio +async def test_route_a2a_cardless_bedrock_agentcore_uses_registered_model(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={}, + litellm_params={ + "custom_llm_provider": "bedrock", + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123:runtime/test", + }, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["api_base"] is None @pytest.mark.asyncio From 41bec1c2d0a752551d83fdfee86bb1bece8c17e6 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:13:57 +0900 Subject: [PATCH 03/20] fix: forward A2A headers and errors --- litellm/llms/a2a/chat/streaming_iterator.py | 11 ++++- litellm/proxy/agent_endpoints/a2a_routing.py | 47 ++++++++++++++++++- .../a2a/chat/test_a2a_streaming_iterator.py | 12 +++++ .../proxy/test_route_a2a_models.py | 13 ++++- 4 files changed, 80 insertions(+), 3 deletions(-) diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 0cbff95b9d3..1c3fa0c6c92 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -7,7 +7,7 @@ from typing import Final from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.utils import GenericStreamingChunk, ModelResponseStream -from ..common_utils import extract_text_from_a2a_response +from ..common_utils import A2AError, extract_text_from_a2a_response class A2AModelResponseIterator(BaseModelResponseIterator): @@ -56,6 +56,15 @@ class A2AModelResponseIterator(BaseModelResponseIterator): } } """ + if "error" in chunk: + error_value: Final = chunk["error"] + error_message: Final = ( + error_value.get("message") + if isinstance(error_value, dict) and isinstance(error_value.get("message"), str) + else str(error_value) + ) + raise A2AError(status_code=500, message=f"A2A error: {error_message}") + try: # Extract text from A2A response text: Final = extract_text_from_a2a_response(chunk) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 2963240f41a..bd8b222ce4f 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -89,6 +89,7 @@ async def _route_registered_provider( api_base: str | None, litellm_params: Mapping[str, object], static_headers: Mapping[str, str] | None, + dynamic_headers: Mapping[str, str] | None = None, ) -> ModelResponse | CustomStreamWrapper: from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, @@ -122,7 +123,10 @@ async def _route_registered_provider( _HEADERS_ADAPTER.validate_python(configured_headers) if isinstance(configured_headers, dict) else None ) agent_extra_headers: Final = merge_agent_headers( - dynamic_headers=configured_headers_dict, + dynamic_headers=merge_agent_headers( + dynamic_headers=dynamic_headers, + static_headers=configured_headers_dict, + ), static_headers=static_headers, ) if agent_extra_headers: @@ -233,6 +237,40 @@ def _merge_agent_guardrails( return merged_data +def _get_agent_dynamic_headers( + data: Mapping[str, object], + agent_id: str, + agent_name: str, + extra_headers: list[str] | None, +) -> dict[str, str]: + proxy_request: Final = data.get("proxy_server_request") + raw_headers: object = proxy_request.get("headers") if isinstance(proxy_request, Mapping) else None + if not isinstance(raw_headers, Mapping): + metadata: Final = data.get("metadata") + raw_headers = metadata.get("headers") if isinstance(metadata, Mapping) else None + normalized_headers: Final = ( + {str(key).lower(): str(value) for key, value in raw_headers.items()} + if isinstance(raw_headers, Mapping) + else {} + ) + + dynamic_headers: dict[str, str] = {} + for header_name in extra_headers or []: + header_name_str: Final = str(header_name) + value: Final = normalized_headers.get(header_name_str.lower()) + if value is not None: + dynamic_headers[header_name_str] = value + + for alias in (agent_id.lower(), agent_name.lower()): + prefix: Final = f"x-a2a-{alias}-" + for key, value in normalized_headers.items(): + if key.startswith(prefix): + header_name: Final = key[len(prefix) :] + if header_name: + dynamic_headers[header_name] = value + return dynamic_headers + + async def route_a2a_agent_request( data: Mapping[str, object], route_type: str, @@ -310,6 +348,12 @@ async def route_a2a_agent_request( data=data, agent_guardrails=registered_params_value.get("guardrails") if registered_params_value else None, ) + registered_dynamic_headers: Final = _get_agent_dynamic_headers( + data=routed_data, + agent_id=agent.agent_id, + agent_name=agent.agent_name, + extra_headers=agent.extra_headers, + ) if ( registered_provider and registered_provider != "a2a" @@ -323,6 +367,7 @@ async def route_a2a_agent_request( api_base=api_base, litellm_params=registered_params_value, static_headers=agent.static_headers, + dynamic_headers=registered_dynamic_headers, ) completion_data: Final = MappingProxyType({**routed_data, "api_base": api_base}) diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index 579d6f8efef..b301b2f3c6e 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -3,6 +3,7 @@ import pytest from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator +from litellm.llms.a2a.common_utils import A2AError @pytest.mark.asyncio @@ -24,3 +25,14 @@ async def test_async_iterator_accepts_decoded_a2a_events(): chunk = await iterator.__aiter__().__anext__() assert chunk["text"] == "Hello" + + +@pytest.mark.asyncio +async def test_async_iterator_propagates_jsonrpc_errors(): + async def _events(): + yield {"jsonrpc": "2.0", "error": {"code": -32000, "message": "agent failed"}} + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + with pytest.raises(A2AError, match="agent failed"): + await iterator.__aiter__().__anext__() diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 1dd2edc8bf7..d96ee37dbdf 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -83,6 +83,7 @@ async def test_route_a2a_model_uses_registered_provider(): "guardrails": ["agent-guardrail"], }, static_headers={"Authorization": "Bearer static"}, + extra_headers=["X-Tenant"], ) data = { "model": "a2a/test-agent", @@ -92,6 +93,12 @@ async def test_route_a2a_model_uses_registered_provider(): "temperature": 0.2, "timeout": 12.0, "tools": [{"type": "function", "function": {"name": "lookup"}}], + "proxy_server_request": { + "headers": { + "x-tenant": "tenant-1", + "x-a2a-test-agent-x-run": "run-1", + } + }, } bridge_response = { "jsonrpc": "2.0", @@ -133,7 +140,11 @@ async def test_route_a2a_model_uses_registered_provider(): assert bridge_kwargs["litellm_params"]["timeout"] == 12.0 assert bridge_kwargs["litellm_params"]["tools"] == data["tools"] assert bridge_kwargs["litellm_params"]["guardrails"] == ["request-guardrail", "agent-guardrail"] - assert bridge_kwargs["litellm_params"]["extra_headers"] == {"Authorization": "Bearer static"} + assert bridge_kwargs["litellm_params"]["extra_headers"] == { + "X-Tenant": "tenant-1", + "x-run": "run-1", + "Authorization": "Bearer static", + } @pytest.mark.asyncio From 07cea813c6a59bfe04f533e0c38679ae5b3a4490 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 17:26:08 +0900 Subject: [PATCH 04/20] fix: harden registered A2A routing --- .../transformation.py | 21 ++ litellm/proxy/agent_endpoints/a2a_routing.py | 129 +++++++-- litellm/proxy/common_request_processing.py | 6 + .../proxy/test_route_a2a_models.py | 249 +++++++++++++++++- 4 files changed, 384 insertions(+), 21 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 15cf77708f9..1ac90b3d294 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -173,6 +173,23 @@ class A2ACompletionBridgeTransformation: if hasattr(choice, "message") and choice.message: content = choice.message.content or "" + tool_calls: list[Any] | None = None + finish_reason: str | None = None + if hasattr(response, "choices") and response.choices: + choice = response.choices[0] + finish_reason = getattr(choice, "finish_reason", None) + message = getattr(choice, "message", None) + raw_tool_calls = getattr(message, "tool_calls", None) + if raw_tool_calls: + tool_calls = [ + call.model_dump(exclude_none=True) + if hasattr(call, "model_dump") + else call.dict(exclude_none=True) + if hasattr(call, "dict") + else call + for call in raw_tool_calls + ] + # Build A2A message a2a_message: Final = { "kind": "message", @@ -180,6 +197,10 @@ class A2ACompletionBridgeTransformation: "parts": [{"kind": "text", "text": content}], "messageId": uuid4().hex, } + if tool_calls: + a2a_message["tool_calls"] = tool_calls + if finish_reason: + a2a_message["finish_reason"] = finish_reason # Build A2A response a2a_response: Final = { diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index bd8b222ce4f..6256d9acccc 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -23,7 +23,7 @@ from litellm.interactions.agents.utils import merge_agent_headers from litellm.llms.a2a.common_utils import A2AError, convert_messages_to_prompt, extract_text_from_a2a_response from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.types.utils import Choices, CustomPricingLiteLLMParams, Message, ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper @@ -66,6 +66,24 @@ _FORWARDED_REQUEST_PARAMS: Final = frozenset( "web_search_options", } ) +_A2A_PRICING_PARAMS: Final = frozenset({"cost_per_query", "response_cost"}) | frozenset( + CustomPricingLiteLLMParams.model_fields +) + + +def _get_agent_request_headers(data: Mapping[str, object]) -> dict[str, str]: + proxy_request: Final = data.get("proxy_server_request") + raw_headers: object = proxy_request.get("headers") if isinstance(proxy_request, Mapping) else None + if not isinstance(raw_headers, Mapping): + metadata: Final = data.get("metadata") + if not isinstance(metadata, Mapping): + metadata = data.get("litellm_metadata") + raw_headers = metadata.get("headers") if isinstance(metadata, Mapping) else None + return ( + {str(key).lower(): str(value) for key, value in raw_headers.items()} + if isinstance(raw_headers, Mapping) + else {} + ) class _A2ATextPart(TypedDict): @@ -132,6 +150,18 @@ async def _route_registered_provider( if agent_extra_headers: provider_params["extra_headers"] = agent_extra_headers + logging_obj: Final = data.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + pricing_params = { + key: litellm_params[key] + for key in _A2A_PRICING_PARAMS + if key in litellm_params and litellm_params[key] is not None + } + if pricing_params: + logging_obj.litellm_params.update(pricing_params) + logging_obj.model_call_details["litellm_params"].update(pricing_params) + logging_obj.custom_pricing = True + if stream: streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming( request_id=request_id, @@ -145,7 +175,6 @@ async def _route_registered_provider( sync_stream=False, model=model_name, ) - logging_obj: Final = data.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): raise TypeError("litellm_logging_obj is required for streaming A2A requests") return CustomStreamWrapper( @@ -172,19 +201,35 @@ async def _route_registered_provider( message=f"A2A error: {error_message if isinstance(error_message, str) else 'Unknown error'}", ) + result: Final = response.get("result") + result_dict: Final = result if isinstance(result, Mapping) else {} + nested_message: Final = result_dict.get("message") + response_message: Final = nested_message if isinstance(nested_message, Mapping) else result_dict + tool_calls: Final = response_message.get("tool_calls") + normalized_tool_calls: Final = tool_calls if isinstance(tool_calls, list) else None + finish_reason: Final = response_message.get("finish_reason") text: Final = extract_text_from_a2a_response(response) model_response: Final = ModelResponse( id=str(response.get("id") or request_id), model=model_name, choices=[ # mutable-ok: ModelResponse requires a choices list - Choices(finish_reason="stop", index=0, message=Message(content=text, role="assistant")) + Choices( + finish_reason=( + finish_reason + if isinstance(finish_reason, str) + else "tool_calls" + if normalized_tool_calls + else "stop" + ), + index=0, + message=Message(content=text, role="assistant", tool_calls=normalized_tool_calls), + ) ], ) usage: Final = response.get("usage") if usage is not None: setattr(model_response, "usage", usage) - logging_obj: Final = data.get("litellm_logging_obj") if isinstance(logging_obj, Logging): def _enqueue_logging() -> None: @@ -211,7 +256,8 @@ def _merge_agent_guardrails( configured_guardrails: list[object] = ( agent_guardrails if isinstance(agent_guardrails, list) else [agent_guardrails] ) - metadata = data.get("metadata") + metadata_key: Final = "litellm_metadata" if "litellm_metadata" in data else "metadata" + metadata = data.get(metadata_key) metadata_guardrails = metadata.get("guardrails") if isinstance(metadata, dict) else None root_guardrails = data.get("guardrails") existing_guardrails: list[object] = [] @@ -227,36 +273,42 @@ def _merge_agent_guardrails( if isinstance(data, dict): data["guardrails"] = merged_guardrails if isinstance(metadata, dict): - metadata["guardrails"] = merged_guardrails + data[metadata_key] = {**metadata, "guardrails": merged_guardrails} return data merged_data = dict(data) merged_data["guardrails"] = merged_guardrails if isinstance(metadata, dict): - merged_data["metadata"] = {**metadata, "guardrails": merged_guardrails} + merged_data[metadata_key] = {**metadata, "guardrails": merged_guardrails} return merged_data +async def merge_a2a_agent_guardrails_before_hooks(data: Mapping[str, object]) -> Mapping[str, object]: + model_name: Final = data.get("model") + if not isinstance(model_name, str) or not model_name.startswith("a2a/"): + return data + + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent = await get_agent_with_read_through(model_name[4:]) + if agent is None or not agent.litellm_params: + return data + return _merge_agent_guardrails(data, agent.litellm_params.get("guardrails")) + + def _get_agent_dynamic_headers( data: Mapping[str, object], agent_id: str, agent_name: str, extra_headers: list[str] | None, ) -> dict[str, str]: - proxy_request: Final = data.get("proxy_server_request") - raw_headers: object = proxy_request.get("headers") if isinstance(proxy_request, Mapping) else None - if not isinstance(raw_headers, Mapping): - metadata: Final = data.get("metadata") - raw_headers = metadata.get("headers") if isinstance(metadata, Mapping) else None - normalized_headers: Final = ( - {str(key).lower(): str(value) for key, value in raw_headers.items()} - if isinstance(raw_headers, Mapping) - else {} - ) + normalized_headers: Final = _get_agent_request_headers(data) dynamic_headers: dict[str, str] = {} for header_name in extra_headers or []: header_name_str: Final = str(header_name) + if header_name_str.lower().startswith("x-litellm-"): + continue value: Final = normalized_headers.get(header_name_str.lower()) if value is not None: dynamic_headers[header_name_str] = value @@ -266,11 +318,32 @@ def _get_agent_dynamic_headers( for key, value in normalized_headers.items(): if key.startswith(prefix): header_name: Final = key[len(prefix) :] - if header_name: + if header_name and not header_name.lower().startswith("x-litellm-"): dynamic_headers[header_name] = value return dynamic_headers +def _get_agent_identity_headers(user_api_key_dict: UserAPIKeyAuth | None) -> dict[str, str]: + if user_api_key_dict is None: + return {} + headers: dict[str, str] = {} + if user_api_key_dict.user_id: + headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id + if user_api_key_dict.team_id: + headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id + return headers + + +def _enforce_inbound_trace_id(data: Mapping[str, object], agent_id: str) -> None: + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + if not get_chain_id_from_headers(_get_agent_request_headers(data)): + raise HTTPException( + status_code=400, + detail=f"Agent '{agent_id}' requires x-litellm-trace-id header on all inbound requests.", + ) + + async def route_a2a_agent_request( data: Mapping[str, object], route_type: str, @@ -325,6 +398,9 @@ async def route_a2a_agent_request( detail=f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.", ) + if (agent.litellm_params or {}).get("require_trace_id_on_calls_to_agent"): + _enforce_inbound_trace_id(data, agent.agent_id) + # Get API base URL from agent config agent_card_params: Final = agent.agent_card_params agent_url: Final = agent_card_params.get("url") if agent_card_params else None @@ -336,7 +412,7 @@ async def route_a2a_agent_request( configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None api_base: Final = configured_api_base if isinstance(configured_api_base, str) else agent_url registered_model: Final = registered_params_value.get("model") if registered_params_value else None - cardless_provider: Final = ( + cardless_provider: Final = registered_provider == "watsonx_orchestrate" or ( registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model ) if (not isinstance(agent_url, str) or not agent_url) and not cardless_provider: @@ -354,6 +430,19 @@ async def route_a2a_agent_request( agent_name=agent.agent_name, extra_headers=agent.extra_headers, ) + registered_static_headers: Mapping[str, str] | None = agent.static_headers + if registered_params_value and registered_params_value.get("databricks_oauth"): + from litellm.proxy.agent_endpoints.databricks_oauth import resolve_databricks_app_auth_header + + databricks_headers = await resolve_databricks_app_auth_header(dict(registered_params_value)) + registered_static_headers = merge_agent_headers( + dynamic_headers=registered_static_headers, + static_headers=databricks_headers, + ) + registered_static_headers = merge_agent_headers( + dynamic_headers=registered_static_headers, + static_headers=_get_agent_identity_headers(user_api_key_dict), + ) if ( registered_provider and registered_provider != "a2a" @@ -366,7 +455,7 @@ async def route_a2a_agent_request( model_name=model_name, api_base=api_base, litellm_params=registered_params_value, - static_headers=agent.static_headers, + static_headers=registered_static_headers, dynamic_headers=registered_dynamic_headers, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index dbbf9cb673e..f7b12357af4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1845,6 +1845,12 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_logging_obj"] = logging_obj + from litellm.proxy.agent_endpoints.a2a_routing import ( + merge_a2a_agent_guardrails_before_hooks, + ) + + self.data = await merge_a2a_agent_guardrails_before_hooks(self.data) + # Merge model-level guardrails before pre_call_hook so DB/UI-configured # guardrails actually execute on pre_call. Without this, guardrails set # via litellm_params.guardrails are only honored on post_call paths diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index d96ee37dbdf..a30b7d422cf 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -7,8 +7,13 @@ Maps to: litellm/proxy/agent_endpoints/a2a_routing.py from unittest.mock import AsyncMock, Mock, patch import pytest +from fastapi import HTTPException -from litellm.proxy.agent_endpoints.a2a_routing import route_a2a_agent_request +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.a2a_routing import ( + merge_a2a_agent_guardrails_before_hooks, + route_a2a_agent_request, +) from litellm.proxy.route_llm_request import route_request @@ -189,6 +194,248 @@ async def test_route_a2a_cardless_bedrock_agentcore_uses_registered_model(): assert bridge.await_args.kwargs["api_base"] is None +@pytest.mark.asyncio +async def test_route_a2a_cardless_watsonx_orchestrate_uses_registered_model(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={}, + litellm_params={ + "custom_llm_provider": "watsonx_orchestrate", + "model": "agent", + "cp4d_host": "https://wxo.example.com", + "instance_id": "instance", + }, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["api_base"] is None + + +@pytest.mark.asyncio +async def test_route_a2a_registered_provider_preserves_identity_headers(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + "proxy_server_request": { + "headers": { + "x-a2a-test-agent-x-litellm-user-id": "attacker", + "x-a2a-test-agent-x-litellm-team-id": "attacker-team", + } + }, + }, + "acompletion", + user_api_key_dict=UserAPIKeyAuth(user_id="trusted-user", team_id="trusted-team"), + ) + await call + + headers = bridge.await_args.kwargs["agent_extra_headers"] + assert headers["X-LiteLLM-User-Id"] == "trusted-user" + assert headers["X-LiteLLM-Team-Id"] == "trusted-team" + assert "x-litellm-user-id" not in {key.lower() for key in headers if key != "X-LiteLLM-User-Id"} + + +@pytest.mark.asyncio +async def test_route_a2a_requires_inbound_trace_id(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={ + "custom_llm_provider": "pydantic_ai_agents", + "require_trace_id_on_calls_to_agent": True, + }, + ) + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + ): + with pytest.raises(HTTPException, match="requires x-litellm-trace-id"): + await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + + +@pytest.mark.asyncio +async def test_route_a2a_resolves_databricks_oauth_headers(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "databricks", "databricks_oauth": {"client_id": "id"}}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.proxy.agent_endpoints.databricks_oauth.resolve_databricks_app_auth_header", + AsyncMock(return_value={"Authorization": "Bearer minted"}), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["agent_extra_headers"]["Authorization"] == "Bearer minted" + + +@pytest.mark.asyncio +async def test_registered_provider_response_preserves_tool_calls(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": { + "kind": "message", + "parts": [], + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + "finish_reason": "tool_calls", + }, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ), + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + response = await call + + assert response.choices[0].finish_reason == "tool_calls" + assert response.choices[0].message.tool_calls[0].id == "call-1" + + +@pytest.mark.asyncio +async def test_a2a_agent_guardrails_merge_before_hooks(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"guardrails": ["agent-guardrail"]}, + ) + with patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ): + merged = await merge_a2a_agent_guardrails_before_hooks( + {"model": "a2a/test-agent", "guardrails": ["request-guardrail"]} + ) + + assert merged["guardrails"] == ["request-guardrail", "agent-guardrail"] + + @pytest.mark.asyncio async def test_route_a2a_stream_uses_registered_provider(): from litellm.litellm_core_utils.litellm_logging import Logging From 32cb4c076b740d97c89e944e1b4b652573b1fdc4 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 18:17:58 +0900 Subject: [PATCH 05/20] fix: close registered A2A review gaps --- .../litellm_completion_bridge/handler.py | 21 ++++- .../transformation.py | 4 + litellm/llms/a2a/chat/streaming_iterator.py | 23 +++-- litellm/proxy/agent_endpoints/a2a_routing.py | 90 ++++++++++++++++--- .../agent_endpoints/model_list_helpers.py | 8 +- litellm/proxy/common_request_processing.py | 14 ++- .../a2a/chat/test_a2a_streaming_iterator.py | 21 +++++ .../proxy/test_route_a2a_models.py | 52 +++++++++++ 8 files changed, 200 insertions(+), 33 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index a62a2b0c724..4a3f7d608ce 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -59,7 +59,12 @@ class A2ACompletionBridgeHandler: message: Final = params.get("message", {}) # Transform A2A message to OpenAI format - openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) + supplied_messages: Final = params.get("messages") + openai_messages: Final = ( + supplied_messages + if isinstance(supplied_messages, list) + else A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) + ) # Get completion params custom_llm_provider: Final = litellm_params.get("custom_llm_provider") @@ -149,10 +154,12 @@ class A2ACompletionBridgeHandler: if a2a_provider_config is not None: verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider) + provider_params: Final = {key: value for key, value in params.items() if key != "messages"} return await a2a_provider_config.handle_non_streaming( request_id=request_id, - params=params, + params=provider_params, api_base=api_base, + timeout=litellm_params.get("timeout") or 60.0, litellm_params=litellm_params, agent_extra_headers=agent_extra_headers, ) @@ -218,10 +225,12 @@ class A2ACompletionBridgeHandler: if a2a_provider_config is not None: verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider) + provider_params: Final = {key: value for key, value in params.items() if key != "messages"} async for chunk in a2a_provider_config.handle_streaming( request_id=request_id, - params=params, + params=provider_params, api_base=api_base, + timeout=litellm_params.get("timeout") or 60.0, litellm_params=litellm_params, agent_extra_headers=agent_extra_headers, ): @@ -261,6 +270,7 @@ class A2ACompletionBridgeHandler: # 3. Accumulate content and emit artifact update accumulated_text = "" + accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas chunk_count = 0 async for chunk in response: chunk_count += 1 @@ -271,6 +281,9 @@ class A2ACompletionBridgeHandler: choice = chunk.choices[0] if hasattr(choice, "delta") and choice.delta: content = choice.delta.content or "" + tool_calls = getattr(choice.delta, "tool_calls", None) + if isinstance(tool_calls, (list, tuple)): + accumulated_tool_calls.extend(tool_calls) if content: accumulated_text += content @@ -289,6 +302,8 @@ class A2ACompletionBridgeHandler: state="completed", final=True, ) + if accumulated_tool_calls: + completed_event["result"]["tool_calls"] = accumulated_tool_calls yield completed_event verbose_logger.info( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 1ac90b3d294..c87b1367377 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -202,12 +202,16 @@ class A2ACompletionBridgeTransformation: if finish_reason: a2a_message["finish_reason"] = finish_reason + usage: Final = getattr(response, "usage", None) + # Build A2A response a2a_response: Final = { "jsonrpc": "2.0", "id": request_id, "result": a2a_message, } + if usage is not None: + a2a_response["usage"] = usage.model_dump(exclude_none=True) if hasattr(usage, "model_dump") else usage verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(content)) diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 1c3fa0c6c92..9da95139c4e 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -71,15 +71,16 @@ class A2AModelResponseIterator(BaseModelResponseIterator): # Determine finish reason finish_reason: Final = self._get_finish_reason(chunk) + tool_calls: Final = self._get_tool_calls(chunk) # Return generic streaming chunk return GenericStreamingChunk( text=text, - is_finished=bool(finish_reason), - finish_reason=finish_reason or "", + is_finished=bool(finish_reason or tool_calls), + finish_reason=finish_reason or ("tool_calls" if tool_calls else ""), usage=None, index=0, - tool_use=None, + tool_use=tool_calls, ) except Exception: # Return empty chunk on parse error @@ -92,9 +93,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): tool_use=None, ) - def _handle_string_chunk( - self, str_line: str | dict - ) -> GenericStreamingChunk | ModelResponseStream: + def _handle_string_chunk(self, str_line: str | dict) -> GenericStreamingChunk | ModelResponseStream: if isinstance(str_line, dict): return self.chunk_parser(chunk=str_line) return super()._handle_string_chunk(str_line=str_line) @@ -118,3 +117,15 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return "stop" return None + + def _get_tool_calls(self, chunk: dict) -> list[dict] | None: + result: Final = chunk.get("result", {}) + if not isinstance(result, dict): + return None + tool_calls = result.get("tool_calls") + if isinstance(tool_calls, list): + return tool_calls + message = result.get("message") + if isinstance(message, dict) and isinstance(message.get("tool_calls"), list): + return message["tool_calls"] + return None diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 6256d9acccc..69c3aebf4f6 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -80,9 +80,7 @@ def _get_agent_request_headers(data: Mapping[str, object]) -> dict[str, str]: metadata = data.get("litellm_metadata") raw_headers = metadata.get("headers") if isinstance(metadata, Mapping) else None return ( - {str(key).lower(): str(value) for key, value in raw_headers.items()} - if isinstance(raw_headers, Mapping) - else {} + {str(key).lower(): str(value) for key, value in raw_headers.items()} if isinstance(raw_headers, Mapping) else {} ) @@ -95,10 +93,12 @@ class _A2AMessage(TypedDict): role: ReadOnly[str] parts: ReadOnly[tuple[_A2ATextPart, ...]] messageId: ReadOnly[str] + contextId: ReadOnly[str | None] class _A2AParams(TypedDict): message: ReadOnly[_A2AMessage] + messages: ReadOnly[list[AllMessageValues]] async def _route_registered_provider( @@ -120,20 +120,27 @@ async def _route_registered_provider( messages: Final = _MESSAGES_ADAPTER.validate_python(raw_messages) stream: Final = data.get("stream") is True request_id: Final = str(uuid4()) + raw_session_id: Final = data.get("litellm_session_id") + metadata: Final = data.get("metadata") + session_id: Final = ( + raw_session_id + if isinstance(raw_session_id, str) + else metadata.get("session_id") + if isinstance(metadata, Mapping) and isinstance(metadata.get("session_id"), str) + else None + ) params: Final[_A2AParams] = { "message": { "role": "user", "parts": ({"kind": "text", "text": convert_messages_to_prompt(messages)},), "messageId": str(uuid4()), - } + "contextId": session_id, + }, + "messages": messages, } provider_params: Final = { **_OBJECT_DICT_ADAPTER.validate_python(litellm_params), - **{ - key: data[key] - for key in _FORWARDED_REQUEST_PARAMS - if key in data and data[key] is not None - }, + **{key: data[key] for key in _FORWARDED_REQUEST_PARAMS if key in data and data[key] is not None}, } bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python(params) configured_headers: Final = litellm_params.get("extra_headers") or litellm_params.get("headers") @@ -161,6 +168,7 @@ async def _route_registered_provider( logging_obj.litellm_params.update(pricing_params) logging_obj.model_call_details["litellm_params"].update(pricing_params) logging_obj.custom_pricing = True + provider_params["no-log"] = True if stream: streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming( @@ -226,9 +234,12 @@ async def _route_registered_provider( ) ], ) - usage: Final = response.get("usage") + raw_usage: Final = response.get("usage") + usage: Final = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage if usage is not None: setattr(model_response, "usage", usage) + if isinstance(logging_obj, Logging): + logging_obj.model_call_details["usage"] = usage if isinstance(logging_obj, Logging): @@ -253,9 +264,7 @@ def _merge_agent_guardrails( if not agent_guardrails: return data - configured_guardrails: list[object] = ( - agent_guardrails if isinstance(agent_guardrails, list) else [agent_guardrails] - ) + configured_guardrails: list[object] = agent_guardrails if isinstance(agent_guardrails, list) else [agent_guardrails] metadata_key: Final = "litellm_metadata" if "litellm_metadata" in data else "metadata" metadata = data.get(metadata_key) metadata_guardrails = metadata.get("guardrails") if isinstance(metadata, dict) else None @@ -296,6 +305,49 @@ async def merge_a2a_agent_guardrails_before_hooks(data: Mapping[str, object]) -> return _merge_agent_guardrails(data, agent.litellm_params.get("guardrails")) +async def authorize_a2a_agent_before_hooks( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth | None, +) -> Mapping[str, object]: + model_name: Final = data.get("model") + if not isinstance(model_name, str) or not model_name.startswith("a2a/"): + return data + + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentRequestHandler + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent = await get_agent_with_read_through(model_name[4:]) + if agent is None: + return data + + is_admin: Final = user_api_key_dict is not None and ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + if not is_admin: + is_allowed: Final = await AgentRequestHandler.is_agent_allowed( + agent_id=agent.agent_id, + user_api_key_auth=user_api_key_dict, + ) + if not is_allowed: + raise HTTPException( + status_code=403, + detail=f"Agent '{agent.agent_name}' is not allowed for your key/team. Contact proxy admin for access.", + ) + + if (agent.litellm_params or {}).get("require_trace_id_on_calls_to_agent"): + _enforce_inbound_trace_id(data, agent.agent_id) + + if isinstance(data, dict): + data["agent_id"] = agent.agent_id + metadata = data.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + data["metadata"] = metadata + metadata["agent_id"] = agent.agent_id + return data + + def _get_agent_dynamic_headers( data: Mapping[str, object], agent_id: str, @@ -409,6 +461,16 @@ async def route_a2a_agent_request( registered_params_value.get("custom_llm_provider") if registered_params_value else None ) registered_provider: Final = registered_provider_value if isinstance(registered_provider_value, str) else None + from litellm.a2a_protocol.litellm_completion_bridge.handler import A2A_USER_API_KEY_HASH_PARAM + + registered_params_for_route: Final[Mapping[str, object]] = ( + { + **registered_params_value, + A2A_USER_API_KEY_HASH_PARAM: user_api_key_dict.api_key, + } + if registered_params_value and user_api_key_dict is not None and user_api_key_dict.api_key + else registered_params_value or {} + ) configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None api_base: Final = configured_api_base if isinstance(configured_api_base, str) else agent_url registered_model: Final = registered_params_value.get("model") if registered_params_value else None @@ -454,7 +516,7 @@ async def route_a2a_agent_request( data=routed_data, model_name=model_name, api_base=api_base, - litellm_params=registered_params_value, + litellm_params=registered_params_for_route, static_headers=registered_static_headers, dynamic_headers=registered_dynamic_headers, ) diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 67d721f9faa..3fa62d0275c 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -37,9 +37,7 @@ async def append_agents_to_model_group( if agent is not None: agent_params: Final = agent.litellm_params provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None - custom_llm_provider: Final = ( - provider_value if isinstance(provider_value, str) else "a2a" - ) + custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else "a2a" model_groups.append( ModelGroupInfoProxy( model_group=f"a2a/{agent.agent_name}", @@ -79,9 +77,7 @@ async def append_agents_to_model_info( if agent is not None: agent_params: Final = agent.litellm_params provider_value: Final = agent_params.get("custom_llm_provider") if agent_params else None - custom_llm_provider: Final = ( - provider_value if isinstance(provider_value, str) else "a2a" - ) + custom_llm_provider: Final = provider_value if isinstance(provider_value, str) else "a2a" models.append( { "model_name": f"a2a/{agent.agent_name}", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f7b12357af4..16a10a16e70 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1836,6 +1836,16 @@ class ProxyBaseLLMRequestProcessing: ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call ## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse. + from litellm.proxy.agent_endpoints.a2a_routing import ( + authorize_a2a_agent_before_hooks, + merge_a2a_agent_guardrails_before_hooks, + ) + + self.data = await authorize_a2a_agent_before_hooks( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) + logging_obj, self.data = litellm.utils.function_setup( original_function=route_type, rules_obj=litellm.utils.Rules(), @@ -1845,10 +1855,6 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_logging_obj"] = logging_obj - from litellm.proxy.agent_endpoints.a2a_routing import ( - merge_a2a_agent_guardrails_before_hooks, - ) - self.data = await merge_a2a_agent_guardrails_before_hooks(self.data) # Merge model-level guardrails before pre_call_hook so DB/UI-configured diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index b301b2f3c6e..f436fe27a57 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -27,6 +27,27 @@ async def test_async_iterator_accepts_decoded_a2a_events(): assert chunk["text"] == "Hello" +@pytest.mark.asyncio +async def test_async_iterator_preserves_tool_calls(): + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + + async def _events(): + yield {"jsonrpc": "2.0", "result": {"tool_calls": tool_calls}} + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["tool_use"] == tool_calls + assert chunk["finish_reason"] == "tool_calls" + + @pytest.mark.asyncio async def test_async_iterator_propagates_jsonrpc_errors(): async def _events(): diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index a30b7d422cf..8cbeaf6f6d1 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -290,6 +290,58 @@ async def test_route_a2a_registered_provider_preserves_identity_headers(): assert "x-litellm-user-id" not in {key.lower() for key in headers if key != "X-LiteLLM-User-Id"} +@pytest.mark.asyncio +async def test_route_a2a_registered_provider_preserves_messages_and_session(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import A2A_USER_API_KEY_HASH_PARAM + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "langflow", "model": "flow"}, + ) + data = { + "model": "a2a/test-agent", + "messages": [ + {"role": "system", "content": "Be concise"}, + {"role": "user", "content": "Hello"}, + ], + "litellm_session_id": "session-1", + } + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + data, + "acompletion", + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + ) + await call + + bridge_kwargs = bridge.await_args.kwargs + assert bridge_kwargs["params"]["messages"] == data["messages"] + assert bridge_kwargs["params"]["message"]["contextId"] == "session-1" + assert bridge_kwargs["litellm_params"][A2A_USER_API_KEY_HASH_PARAM] == "hashed-key" + + @pytest.mark.asyncio async def test_route_a2a_requires_inbound_trace_id(): from litellm.types.agents import AgentResponse From d2f05d1cf1388f2eefcebd3f936309f9a1268970 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 22:47:40 +0900 Subject: [PATCH 06/20] fix: address A2A review feedback --- .../litellm_completion_bridge/handler.py | 53 ++++++------ .../transformation.py | 75 ++++++++-------- litellm/llms/a2a/chat/streaming_iterator.py | 13 +-- litellm/proxy/agent_endpoints/a2a_routing.py | 66 +++++++++++--- .../test_completion_bridge_streaming.py | 20 +++-- .../a2a/chat/test_a2a_streaming_iterator.py | 2 +- .../proxy/test_route_a2a_models.py | 86 +++++++++++++++++++ 7 files changed, 228 insertions(+), 87 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 4a3f7d608ce..a5c4463da8d 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -155,14 +155,16 @@ class A2ACompletionBridgeHandler: verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider) provider_params: Final = {key: value for key, value in params.items() if key != "messages"} - return await a2a_provider_config.handle_non_streaming( - request_id=request_id, - params=provider_params, - api_base=api_base, - timeout=litellm_params.get("timeout") or 60.0, - litellm_params=litellm_params, - agent_extra_headers=agent_extra_headers, - ) + provider_kwargs: Final[dict[str, Any]] = { + "request_id": request_id, + "params": provider_params, + "api_base": api_base, + "litellm_params": litellm_params, + "agent_extra_headers": agent_extra_headers, + } + if litellm_params.get("timeout") is not None: + provider_kwargs["timeout"] = litellm_params["timeout"] + return await a2a_provider_config.handle_non_streaming(**provider_kwargs) completion_params: Final = A2ACompletionBridgeHandler._build_completion_params( params=params, @@ -226,14 +228,16 @@ class A2ACompletionBridgeHandler: verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider) provider_params: Final = {key: value for key, value in params.items() if key != "messages"} - async for chunk in a2a_provider_config.handle_streaming( - request_id=request_id, - params=provider_params, - api_base=api_base, - timeout=litellm_params.get("timeout") or 60.0, - litellm_params=litellm_params, - agent_extra_headers=agent_extra_headers, - ): + provider_kwargs: Final[dict[str, Any]] = { + "request_id": request_id, + "params": provider_params, + "api_base": api_base, + "litellm_params": litellm_params, + "agent_extra_headers": agent_extra_headers, + } + if litellm_params.get("timeout") is not None: + provider_kwargs["timeout"] = litellm_params["timeout"] + async for chunk in a2a_provider_config.handle_streaming(**provider_kwargs): yield chunk return @@ -268,8 +272,7 @@ class A2ACompletionBridgeHandler: # Call litellm.acompletion with streaming response: Final = await A2ACompletionBridgeHandler._acompletion(completion_params) - # 3. Accumulate content and emit artifact update - accumulated_text = "" + # 3. Forward content as artifact updates accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas chunk_count = 0 async for chunk in response: @@ -286,15 +289,11 @@ class A2ACompletionBridgeHandler: accumulated_tool_calls.extend(tool_calls) if content: - accumulated_text += content - - # Emit artifact update with accumulated content - if accumulated_text: - artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=accumulated_text, - ) - yield artifact_event + artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=content, + ) + yield artifact_event # 4. Emit final status update (kind: "status-update", status: "completed", final: true) completed_event: Final = A2ACompletionBridgeTransformation.create_status_update_event( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index c87b1367377..c24af243c83 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -166,41 +166,44 @@ class A2ACompletionBridgeTransformation: Returns: A2A SendMessageResponse dict """ - # Extract content from response - content = "" - if hasattr(response, "choices") and response.choices: - choice: Final = response.choices[0] - if hasattr(choice, "message") and choice.message: - content = choice.message.content or "" + serialized_choices: list[dict[str, Any]] = [] + raw_choices: Final = getattr(response, "choices", None) + if raw_choices: + for choice in raw_choices: + content: Final = ( + getattr(getattr(choice, "message", None), "content", None) or "" + ) + message: Final = { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": content}], + "messageId": uuid4().hex, + } + raw_tool_calls = getattr(getattr(choice, "message", None), "tool_calls", None) + if raw_tool_calls: + message["tool_calls"] = [ + call.model_dump(exclude_none=True) + if hasattr(call, "model_dump") + else call.dict(exclude_none=True) + if hasattr(call, "dict") + else call + for call in raw_tool_calls + ] + finish_reason: Final = getattr(choice, "finish_reason", None) + if finish_reason: + message["finish_reason"] = finish_reason + serialized_choices.append({"index": len(serialized_choices), "message": message}) - tool_calls: list[Any] | None = None - finish_reason: str | None = None - if hasattr(response, "choices") and response.choices: - choice = response.choices[0] - finish_reason = getattr(choice, "finish_reason", None) - message = getattr(choice, "message", None) - raw_tool_calls = getattr(message, "tool_calls", None) - if raw_tool_calls: - tool_calls = [ - call.model_dump(exclude_none=True) - if hasattr(call, "model_dump") - else call.dict(exclude_none=True) - if hasattr(call, "dict") - else call - for call in raw_tool_calls - ] - - # Build A2A message - a2a_message: Final = { - "kind": "message", - "role": "agent", - "parts": [{"kind": "text", "text": content}], - "messageId": uuid4().hex, - } - if tool_calls: - a2a_message["tool_calls"] = tool_calls - if finish_reason: - a2a_message["finish_reason"] = finish_reason + a2a_message: Final = ( + serialized_choices[0]["message"] + if serialized_choices + else { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": ""}], + "messageId": uuid4().hex, + } + ) usage: Final = getattr(response, "usage", None) @@ -212,8 +215,10 @@ class A2ACompletionBridgeTransformation: } if usage is not None: a2a_response["usage"] = usage.model_dump(exclude_none=True) if hasattr(usage, "model_dump") else usage + if len(serialized_choices) > 1: + a2a_response["choices"] = serialized_choices - verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(content)) + verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(a2a_message["parts"][0]["text"])) return a2a_response diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 9da95139c4e..58b3b396a0e 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -5,6 +5,7 @@ A2A Streaming Response Iterator from typing import Final from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.types.utils import GenericStreamingChunk, ModelResponseStream from ..common_utils import A2AError, extract_text_from_a2a_response @@ -118,14 +119,16 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return None - def _get_tool_calls(self, chunk: dict) -> list[dict] | None: + def _get_tool_calls(self, chunk: dict) -> ChatCompletionToolCallChunk | None: result: Final = chunk.get("result", {}) if not isinstance(result, dict): return None tool_calls = result.get("tool_calls") - if isinstance(tool_calls, list): - return tool_calls + if isinstance(tool_calls, list) and tool_calls: + first_tool_call: Final = tool_calls[0] + return first_tool_call if isinstance(first_tool_call, dict) else None message = result.get("message") - if isinstance(message, dict) and isinstance(message.get("tool_calls"), list): - return message["tool_calls"] + if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) and message["tool_calls"]: + first_tool_call = message["tool_calls"][0] + return first_tool_call if isinstance(first_tool_call, dict) else None return None diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 69c3aebf4f6..5826442865d 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -213,14 +213,53 @@ async def _route_registered_provider( result_dict: Final = result if isinstance(result, Mapping) else {} nested_message: Final = result_dict.get("message") response_message: Final = nested_message if isinstance(nested_message, Mapping) else result_dict - tool_calls: Final = response_message.get("tool_calls") - normalized_tool_calls: Final = tool_calls if isinstance(tool_calls, list) else None - finish_reason: Final = response_message.get("finish_reason") - text: Final = extract_text_from_a2a_response(response) - model_response: Final = ModelResponse( - id=str(response.get("id") or request_id), - model=model_name, - choices=[ # mutable-ok: ModelResponse requires a choices list + response_choices: Final = response.get("choices") + choice_payloads: Final = ( + response_choices + if isinstance(response_choices, list) + else result_dict.get("choices") + ) + if isinstance(choice_payloads, list) and choice_payloads: + model_choices = [ + Choices( + finish_reason=( + choice.get("finish_reason") + if isinstance(choice, Mapping) and isinstance(choice.get("finish_reason"), str) + else choice.get("message", {}).get("finish_reason") + if isinstance(choice, Mapping) + and isinstance(choice.get("message"), Mapping) + and isinstance(choice.get("message", {}).get("finish_reason"), str) + else "stop" + ), + index=choice.get("index", choice_index) + if isinstance(choice, Mapping) and isinstance(choice.get("index", choice_index), int) + else choice_index, + message=Message( + content=extract_text_from_a2a_response( + {"result": choice.get("message", choice)} + if isinstance(choice, Mapping) + else {"result": {}} + ), + role="assistant", + tool_calls=( + choice.get("message", {}).get("tool_calls") + if isinstance(choice, Mapping) + and isinstance(choice.get("message"), Mapping) + and isinstance(choice.get("message", {}).get("tool_calls"), list) + else choice.get("tool_calls") + if isinstance(choice, Mapping) and isinstance(choice.get("tool_calls"), list) + else None + ), + ), + ) + for choice_index, choice in enumerate(choice_payloads) + ] + else: + tool_calls: Final = response_message.get("tool_calls") + normalized_tool_calls: Final = tool_calls if isinstance(tool_calls, list) else None + finish_reason: Final = response_message.get("finish_reason") + text: Final = extract_text_from_a2a_response(response) + model_choices = [ Choices( finish_reason=( finish_reason @@ -232,12 +271,16 @@ async def _route_registered_provider( index=0, message=Message(content=text, role="assistant", tool_calls=normalized_tool_calls), ) - ], + ] + model_response: Final = ModelResponse( + id=str(response.get("id") or request_id), + model=model_name, + choices=model_choices, ) raw_usage: Final = response.get("usage") usage: Final = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage if usage is not None: - setattr(model_response, "usage", usage) + model_response.usage = usage if isinstance(logging_obj, Logging): logging_obj.model_call_details["usage"] = usage @@ -477,7 +520,8 @@ async def route_a2a_agent_request( cardless_provider: Final = registered_provider == "watsonx_orchestrate" or ( registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model ) - if (not isinstance(agent_url, str) or not agent_url) and not cardless_provider: + has_configured_api_base: Final = isinstance(configured_api_base, str) and bool(configured_api_base) + if (not isinstance(agent_url, str) or not agent_url) and not has_configured_api_base and not cardless_provider: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type) raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False) diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 1b3e5f86020..23495d06629 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -222,8 +222,8 @@ async def test_handle_streaming_emits_proper_events(): ): events.append(event) - # Should have 4 events: task, working, artifact, completed - assert len(events) == 4 + # Should have 5 events: task, working, two artifacts, completed + assert len(events) == 5 # Event 1: task submitted assert events[0]["result"]["kind"] == "task" @@ -234,14 +234,18 @@ async def test_handle_streaming_emits_proper_events(): assert events[1]["result"]["status"]["state"] == "working" assert events[1]["result"]["final"] is False - # Event 3: artifact update with accumulated content + # Event 3: first artifact update assert events[2]["result"]["kind"] == "artifact-update" - assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello world" + assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello" - # Event 4: status completed - assert events[3]["result"]["kind"] == "status-update" - assert events[3]["result"]["status"]["state"] == "completed" - assert events[3]["result"]["final"] is True + # Event 4: second artifact update + assert events[3]["result"]["kind"] == "artifact-update" + assert events[3]["result"]["artifact"]["parts"][0]["text"] == " world" + + # Event 5: status completed + assert events[4]["result"]["kind"] == "status-update" + assert events[4]["result"]["status"]["state"] == "completed" + assert events[4]["result"]["final"] is True @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index f436fe27a57..dafa2839ace 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -44,7 +44,7 @@ async def test_async_iterator_preserves_tool_calls(): chunk = await iterator.__aiter__().__anext__() - assert chunk["tool_use"] == tool_calls + assert chunk["tool_use"] == tool_calls[0] assert chunk["finish_reason"] == "tool_calls" diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 8cbeaf6f6d1..766745425c3 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -194,6 +194,92 @@ async def test_route_a2a_cardless_bedrock_agentcore_uses_registered_model(): assert bridge.await_args.kwargs["api_base"] is None +@pytest.mark.asyncio +async def test_route_a2a_registered_provider_uses_configured_api_base_without_card_url(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={}, + litellm_params={ + "custom_llm_provider": "langflow", + "model": "flow", + "api_base": "https://flow.example.com", + }, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "result": {"kind": "message", "parts": [{"kind": "text", "text": "Hello back"}]}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ) as bridge, + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + await call + + assert bridge.await_args.kwargs["api_base"] == "https://flow.example.com" + + +@pytest.mark.asyncio +async def test_registered_provider_response_preserves_multiple_choices(): + from litellm.types.agents import AgentResponse + + agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, + ) + bridge_response = { + "jsonrpc": "2.0", + "id": "request-id", + "choices": [ + {"index": 0, "message": {"parts": [{"kind": "text", "text": "first"}]}}, + {"index": 1, "message": {"parts": [{"kind": "text", "text": "second"}]}}, + ], + "result": {}, + } + + with ( + patch( + "litellm.proxy.common_utils.registry_read_through.get_agent_with_read_through", + AsyncMock(return_value=agent), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + AsyncMock(return_value=True), + ), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=bridge_response), + ), + ): + call = await route_a2a_agent_request( + {"model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}]}, + "acompletion", + ) + response = await call + + assert [choice.message.content for choice in response.choices] == ["first", "second"] + + @pytest.mark.asyncio async def test_route_a2a_cardless_watsonx_orchestrate_uses_registered_model(): from litellm.types.agents import AgentResponse From 33a8945ca5d67e67220b5fcf4a66a0b6ba24c54b Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:32:53 +0900 Subject: [PATCH 07/20] fix: close second A2A review round --- .../litellm_completion_bridge/handler.py | 78 ++++++++---- .../transformation.py | 59 ++++++++- litellm/llms/a2a/chat/streaming_iterator.py | 58 ++++++++- litellm/proxy/agent_endpoints/a2a_routing.py | 117 ++++++++++-------- .../test_completion_bridge_streaming.py | 73 +++++++++-- .../a2a/chat/test_a2a_streaming_iterator.py | 50 ++++++++ .../proxy/test_route_a2a_models.py | 6 + 7 files changed, 353 insertions(+), 88 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index a5c4463da8d..71993135037 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -154,7 +154,7 @@ class A2ACompletionBridgeHandler: if a2a_provider_config is not None: verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider) - provider_params: Final = {key: value for key, value in params.items() if key != "messages"} + provider_params: Final = dict(params) provider_kwargs: Final[dict[str, Any]] = { "request_id": request_id, "params": provider_params, @@ -227,7 +227,7 @@ class A2ACompletionBridgeHandler: if a2a_provider_config is not None: verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider) - provider_params: Final = {key: value for key, value in params.items() if key != "messages"} + provider_params: Final = dict(params) provider_kwargs: Final[dict[str, Any]] = { "request_id": request_id, "params": provider_params, @@ -237,8 +237,14 @@ class A2ACompletionBridgeHandler: } if litellm_params.get("timeout") is not None: provider_kwargs["timeout"] = litellm_params["timeout"] - async for chunk in a2a_provider_config.handle_streaming(**provider_kwargs): - yield chunk + provider_stream: Final = a2a_provider_config.handle_streaming(**provider_kwargs) + try: + async for chunk in provider_stream: + yield chunk + finally: + close_provider_stream = getattr(provider_stream, "aclose", None) + if close_provider_stream is not None: + await close_provider_stream() return @@ -274,26 +280,52 @@ class A2ACompletionBridgeHandler: # 3. Forward content as artifact updates accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas + stream_usage: object | None = None + stream_finish_reason: str | None = None chunk_count = 0 - async for chunk in response: - chunk_count += 1 + try: + async for chunk in response: + chunk_count += 1 - # Extract delta content - content = "" - if chunk is not None and hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - if hasattr(choice, "delta") and choice.delta: - content = choice.delta.content or "" - tool_calls = getattr(choice.delta, "tool_calls", None) - if isinstance(tool_calls, (list, tuple)): - accumulated_tool_calls.extend(tool_calls) + raw_usage = getattr(chunk, "usage", None) + if isinstance(raw_usage, Mapping): + stream_usage = raw_usage + else: + dump_usage = getattr(raw_usage, "model_dump", None) + if callable(dump_usage): + dumped_usage = dump_usage(exclude_none=True) + if isinstance(dumped_usage, Mapping): + stream_usage = dumped_usage + else: + dict_usage = getattr(raw_usage, "dict", None) + if callable(dict_usage): + dumped_usage = dict_usage(exclude_none=True) + if isinstance(dumped_usage, Mapping): + stream_usage = dumped_usage - if content: - artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=content, - ) - yield artifact_event + # Extract delta content + content = "" + if chunk is not None and hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + raw_finish_reason = getattr(choice, "finish_reason", None) + if isinstance(raw_finish_reason, str) and raw_finish_reason: + stream_finish_reason = raw_finish_reason + if hasattr(choice, "delta") and choice.delta: + content = choice.delta.content or "" + tool_calls = getattr(choice.delta, "tool_calls", None) + if isinstance(tool_calls, (list, tuple)): + accumulated_tool_calls.extend(tool_calls) + + if content: + artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=content, + ) + yield artifact_event + finally: + close_response = getattr(response, "aclose", None) + if close_response is not None: + await close_response() # 4. Emit final status update (kind: "status-update", status: "completed", final: true) completed_event: Final = A2ACompletionBridgeTransformation.create_status_update_event( @@ -303,6 +335,10 @@ class A2ACompletionBridgeHandler: ) if accumulated_tool_calls: completed_event["result"]["tool_calls"] = accumulated_tool_calls + if stream_finish_reason: + completed_event["result"]["finish_reason"] = stream_finish_reason + if stream_usage is not None: + completed_event["usage"] = stream_usage yield completed_event verbose_logger.info( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index c24af243c83..02006e98fa0 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -151,6 +151,22 @@ class A2ACompletionBridgeTransformation: return [openai_message] + @staticmethod + def _model_dump(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + dump = getattr(value, "model_dump", None) + if callable(dump): + dumped = dump(exclude_none=True) + if isinstance(dumped, dict): + return dumped + dump = getattr(value, "dict", None) + if callable(dump): + dumped = dump(exclude_none=True) + if isinstance(dumped, dict): + return dumped + return {} + @staticmethod def openai_response_to_a2a_response( response: Any, @@ -170,16 +186,19 @@ class A2ACompletionBridgeTransformation: raw_choices: Final = getattr(response, "choices", None) if raw_choices: for choice in raw_choices: - content: Final = ( - getattr(getattr(choice, "message", None), "content", None) or "" - ) + raw_message = getattr(choice, "message", None) + message_fields: Final = A2ACompletionBridgeTransformation._model_dump(raw_message) + raw_content = message_fields.get("content") + if raw_content is None: + raw_content = getattr(raw_message, "content", None) + content: Final = raw_content if isinstance(raw_content, str) else "" message: Final = { "kind": "message", "role": "agent", "parts": [{"kind": "text", "text": content}], "messageId": uuid4().hex, } - raw_tool_calls = getattr(getattr(choice, "message", None), "tool_calls", None) + raw_tool_calls = message_fields.get("tool_calls") if raw_tool_calls: message["tool_calls"] = [ call.model_dump(exclude_none=True) @@ -189,10 +208,38 @@ class A2ACompletionBridgeTransformation: else call for call in raw_tool_calls ] - finish_reason: Final = getattr(choice, "finish_reason", None) + for field in ( + "annotations", + "audio", + "function_call", + "images", + "provider_specific_fields", + "reasoning_content", + "reasoning_items", + "thinking_blocks", + ): + value = message_fields.get(field) + if value is not None: + message[field] = value + choice_fields: Final = A2ACompletionBridgeTransformation._model_dump(choice) + finish_reason: Final = choice_fields.get("finish_reason") + if finish_reason is None: + raw_finish_reason = getattr(choice, "finish_reason", None) + finish_reason = raw_finish_reason if isinstance(raw_finish_reason, str) else None if finish_reason: message["finish_reason"] = finish_reason - serialized_choices.append({"index": len(serialized_choices), "message": message}) + choice_payload: Final[dict[str, Any]] = { + "index": len(serialized_choices), + "message": message, + } + logprobs = choice_fields.get("logprobs") + if logprobs is None: + raw_logprobs = getattr(choice, "logprobs", None) + logprobs = raw_logprobs if isinstance(raw_logprobs, dict) else None + if logprobs is not None: + choice_payload["logprobs"] = logprobs + message["logprobs"] = logprobs + serialized_choices.append(choice_payload) a2a_message: Final = ( serialized_choices[0]["message"] diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 58b3b396a0e..636aef0e724 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -2,8 +2,10 @@ A2A Streaming Response Iterator """ +from collections.abc import Mapping from typing import Final +import litellm from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -73,13 +75,14 @@ class A2AModelResponseIterator(BaseModelResponseIterator): # Determine finish reason finish_reason: Final = self._get_finish_reason(chunk) tool_calls: Final = self._get_tool_calls(chunk) + usage: Final = self._get_usage(chunk) # Return generic streaming chunk return GenericStreamingChunk( text=text, is_finished=bool(finish_reason or tool_calls), finish_reason=finish_reason or ("tool_calls" if tool_calls else ""), - usage=None, + usage=usage, index=0, tool_use=tool_calls, ) @@ -105,6 +108,14 @@ class A2AModelResponseIterator(BaseModelResponseIterator): # Check for task completion if isinstance(result, dict): + explicit_finish_reason: Final = result.get("finish_reason") + if isinstance(explicit_finish_reason, str) and explicit_finish_reason: + return explicit_finish_reason + message: Final = result.get("message") + if isinstance(message, dict): + message_finish_reason: Final = message.get("finish_reason") + if isinstance(message_finish_reason, str) and message_finish_reason: + return message_finish_reason status: Final = result.get("status", {}) if isinstance(status, dict): state: Final = status.get("state") @@ -119,16 +130,53 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return None + def _get_usage(self, chunk: dict) -> object | None: + raw_usage: object | None = chunk.get("usage") + result: Final = chunk.get("result", {}) + if raw_usage is None and isinstance(result, dict): + raw_usage = result.get("usage") + if raw_usage is None: + return None + if isinstance(raw_usage, Mapping): + try: + return litellm.Usage(**raw_usage) + except Exception: + return raw_usage + if hasattr(raw_usage, "model_dump"): + try: + return litellm.Usage(**raw_usage.model_dump(exclude_none=True)) + except Exception: + return raw_usage + return raw_usage + def _get_tool_calls(self, chunk: dict) -> ChatCompletionToolCallChunk | None: result: Final = chunk.get("result", {}) if not isinstance(result, dict): return None tool_calls = result.get("tool_calls") if isinstance(tool_calls, list) and tool_calls: - first_tool_call: Final = tool_calls[0] - return first_tool_call if isinstance(first_tool_call, dict) else None + return self._serialize_tool_call(tool_calls[0]) message = result.get("message") if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) and message["tool_calls"]: - first_tool_call = message["tool_calls"][0] - return first_tool_call if isinstance(first_tool_call, dict) else None + return self._serialize_tool_call(message["tool_calls"][0]) return None + + @staticmethod + def _serialize_tool_call(tool_call: object) -> ChatCompletionToolCallChunk | None: + if isinstance(tool_call, dict): + return tool_call + if hasattr(tool_call, "model_dump"): + return tool_call.model_dump(exclude_none=True) + if hasattr(tool_call, "dict"): + return tool_call.dict(exclude_none=True) + return None + + async def aclose(self) -> None: + streaming_response = self.streaming_response + self.streaming_response = None + try: + await super().aclose() + finally: + close_stream = getattr(streaming_response, "aclose", None) + if close_stream is not None: + await close_stream() diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 5826442865d..34e945f1415 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -52,6 +52,7 @@ _FORWARDED_REQUEST_PARAMS: Final = frozenset( "response_format", "seed", "service_tier", + "safety_identifier", "stop", "store", "temperature", @@ -64,6 +65,8 @@ _FORWARDED_REQUEST_PARAMS: Final = frozenset( "user", "verbosity", "web_search_options", + "output_config", + "prompt_cache_key", } ) _A2A_PRICING_PARAMS: Final = frozenset({"cost_per_query", "response_cost"}) | frozenset( @@ -159,6 +162,7 @@ async def _route_registered_provider( logging_obj: Final = data.get("litellm_logging_obj") if isinstance(logging_obj, Logging): + provider_params["no-log"] = True pricing_params = { key: litellm_params[key] for key in _A2A_PRICING_PARAMS @@ -168,7 +172,6 @@ async def _route_registered_provider( logging_obj.litellm_params.update(pricing_params) logging_obj.model_call_details["litellm_params"].update(pricing_params) logging_obj.custom_pricing = True - provider_params["no-log"] = True if stream: streaming_response: Final = A2ACompletionBridgeHandler.handle_streaming( @@ -214,64 +217,80 @@ async def _route_registered_provider( nested_message: Final = result_dict.get("message") response_message: Final = nested_message if isinstance(nested_message, Mapping) else result_dict response_choices: Final = response.get("choices") - choice_payloads: Final = ( - response_choices - if isinstance(response_choices, list) - else result_dict.get("choices") - ) + choice_payloads: Final = response_choices if isinstance(response_choices, list) else result_dict.get("choices") + + def _serialize_value(value: object) -> object: + if hasattr(value, "model_dump"): + return value.model_dump(exclude_none=True) + if hasattr(value, "dict"): + return value.dict(exclude_none=True) + return value + + def _build_message(message_payload: Mapping[str, object], content: str) -> Message: + message_kwargs: dict[str, object] = { + "content": content, + "role": "assistant", + } + raw_tool_calls = message_payload.get("tool_calls") + if isinstance(raw_tool_calls, list): + message_kwargs["tool_calls"] = raw_tool_calls + for field in ( + "audio", + "annotations", + "function_call", + "images", + "provider_specific_fields", + "reasoning_content", + "reasoning_items", + "thinking_blocks", + ): + value = message_payload.get(field) + if value is not None: + message_kwargs[field] = _serialize_value(value) + return Message(**message_kwargs) + if isinstance(choice_payloads, list) and choice_payloads: - model_choices = [ - Choices( - finish_reason=( - choice.get("finish_reason") - if isinstance(choice, Mapping) and isinstance(choice.get("finish_reason"), str) - else choice.get("message", {}).get("finish_reason") - if isinstance(choice, Mapping) - and isinstance(choice.get("message"), Mapping) - and isinstance(choice.get("message", {}).get("finish_reason"), str) + model_choices = [] + for choice_index, choice in enumerate(choice_payloads): + choice_mapping: Mapping[str, object] = choice if isinstance(choice, Mapping) else {} + raw_message = choice_mapping.get("message") + message_payload: Mapping[str, object] = raw_message if isinstance(raw_message, Mapping) else choice_mapping + choice_kwargs: dict[str, object] = { + "finish_reason": ( + choice_mapping.get("finish_reason") + if isinstance(choice_mapping.get("finish_reason"), str) + else message_payload.get("finish_reason") + if isinstance(message_payload.get("finish_reason"), str) else "stop" ), - index=choice.get("index", choice_index) - if isinstance(choice, Mapping) and isinstance(choice.get("index", choice_index), int) + "index": choice_mapping.get("index", choice_index) + if isinstance(choice_mapping.get("index", choice_index), int) else choice_index, - message=Message( - content=extract_text_from_a2a_response( - {"result": choice.get("message", choice)} - if isinstance(choice, Mapping) - else {"result": {}} - ), - role="assistant", - tool_calls=( - choice.get("message", {}).get("tool_calls") - if isinstance(choice, Mapping) - and isinstance(choice.get("message"), Mapping) - and isinstance(choice.get("message", {}).get("tool_calls"), list) - else choice.get("tool_calls") - if isinstance(choice, Mapping) and isinstance(choice.get("tool_calls"), list) - else None - ), + "message": _build_message( + message_payload, + extract_text_from_a2a_response({"result": message_payload}), ), - ) - for choice_index, choice in enumerate(choice_payloads) - ] + } + raw_logprobs = choice_mapping.get("logprobs", message_payload.get("logprobs")) + if raw_logprobs is not None: + choice_kwargs["logprobs"] = _serialize_value(raw_logprobs) + model_choices.append(Choices(**choice_kwargs)) else: tool_calls: Final = response_message.get("tool_calls") normalized_tool_calls: Final = tool_calls if isinstance(tool_calls, list) else None finish_reason: Final = response_message.get("finish_reason") text: Final = extract_text_from_a2a_response(response) - model_choices = [ - Choices( - finish_reason=( - finish_reason - if isinstance(finish_reason, str) - else "tool_calls" - if normalized_tool_calls - else "stop" - ), - index=0, - message=Message(content=text, role="assistant", tool_calls=normalized_tool_calls), - ) - ] + choice_kwargs = { + "finish_reason": ( + finish_reason if isinstance(finish_reason, str) else "tool_calls" if normalized_tool_calls else "stop" + ), + "index": 0, + "message": _build_message(response_message, text), + } + raw_logprobs = response_message.get("logprobs") + if raw_logprobs is not None: + choice_kwargs["logprobs"] = _serialize_value(raw_logprobs) + model_choices = [Choices(**choice_kwargs)] model_response: Final = ModelResponse( id=str(response.get("id") or request_id), model=model_name, diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 23495d06629..d2b74d416a0 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.types.utils import Choices, Message, ModelResponse + class TestA2AStreamingTransformation: """Test the A2A streaming transformation creates proper events.""" @@ -26,9 +28,7 @@ class TestA2AStreamingTransformation: "parts": [{"text": "Reply to ticket #4823"}], "metadata": {"skillId": "draft_reply"}, } - openai_messages = ( - A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) - ) + openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) # Metadata is forwarded on the run payload only, not duplicated on messages. assert "metadata" not in openai_messages[0] @@ -174,10 +174,7 @@ class TestA2AStreamingTransformation: assert "artifactId" in event["result"]["artifact"] assert event["result"]["artifact"]["name"] == "response" assert event["result"]["artifact"]["parts"][0]["kind"] == "text" - assert ( - event["result"]["artifact"]["parts"][0]["text"] - == "Hello, I am an AI assistant." - ) + assert event["result"]["artifact"]["parts"][0]["text"] == "Hello, I am an AI assistant." @pytest.mark.asyncio @@ -197,6 +194,8 @@ async def test_handle_streaming_emits_proper_events(): mock_chunk2.choices = [MagicMock()] mock_chunk2.choices[0].delta = MagicMock() mock_chunk2.choices[0].delta.content = " world" + mock_chunk2.choices[0].finish_reason = "length" + mock_chunk2.usage = {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5} async def mock_streaming_response(): yield mock_chunk1 @@ -246,6 +245,66 @@ async def test_handle_streaming_emits_proper_events(): assert events[4]["result"]["kind"] == "status-update" assert events[4]["result"]["status"]["state"] == "completed" assert events[4]["result"]["final"] is True + assert events[4]["result"]["finish_reason"] == "length" + assert events[4]["usage"]["total_tokens"] == 5 + + +@pytest.mark.asyncio +async def test_provider_config_receives_full_message_history(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + provider_config = MagicMock() + provider_config.handle_non_streaming = AsyncMock(return_value={"result": {}}) + messages = [ + {"role": "system", "content": "Be concise"}, + {"role": "user", "content": "Hello"}, + ] + params = { + "message": {"role": "user", "parts": []}, + "messages": messages, + } + + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config", + return_value=provider_config, + ): + await A2ACompletionBridgeHandler.handle_non_streaming( + request_id="req-1", + params=params, + litellm_params={"custom_llm_provider": "langflow", "model": "flow"}, + ) + + assert provider_config.handle_non_streaming.await_args.kwargs["params"]["messages"] == messages + + +def test_response_transform_preserves_audio_and_logprobs(): + from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( + A2ACompletionBridgeTransformation, + ) + + response = ModelResponse( + id="resp-1", + model="test-model", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="hello", + role="assistant", + audio={"data": "abc", "expires_at": 1, "transcript": "hello"}, + ), + logprobs={"content": []}, + ) + ], + ) + + transformed = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(response) + + assert transformed["result"]["audio"]["data"] == "abc" + assert transformed["result"]["logprobs"] == {"content": []} @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index dafa2839ace..db303719864 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -4,6 +4,7 @@ import pytest from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator from litellm.llms.a2a.common_utils import A2AError +from litellm.types.utils import Delta @pytest.mark.asyncio @@ -48,6 +49,55 @@ async def test_async_iterator_preserves_tool_calls(): assert chunk["finish_reason"] == "tool_calls" +@pytest.mark.asyncio +async def test_async_iterator_serializes_delta_tool_calls_and_usage(): + delta = Delta( + tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + ) + + async def _events(): + yield { + "jsonrpc": "2.0", + "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}, + "result": { + "tool_calls": [delta.tool_calls[0]], + "finish_reason": "length", + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + chunk = await iterator.__aiter__().__anext__() + + assert chunk["tool_use"]["id"] == "call-1" + assert chunk["finish_reason"] == "length" + assert chunk["usage"].total_tokens == 5 + + +@pytest.mark.asyncio +async def test_async_iterator_closes_nested_stream(): + closed = False + + async def _events(): + nonlocal closed + try: + yield {"jsonrpc": "2.0", "result": {"kind": "artifact-update"}} + raise AssertionError("stream should be closed before a second event") + finally: + closed = True + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + await iterator.__aiter__().__anext__() + await iterator.aclose() + + assert closed is True + + @pytest.mark.asyncio async def test_async_iterator_propagates_jsonrpc_errors(): async def _events(): diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 766745425c3..ca8a7ff5b2a 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -98,6 +98,9 @@ async def test_route_a2a_model_uses_registered_provider(): "temperature": 0.2, "timeout": 12.0, "tools": [{"type": "function", "function": {"name": "lookup"}}], + "output_config": {"format": "json"}, + "prompt_cache_key": "cache-key", + "safety_identifier": "safety-id", "proxy_server_request": { "headers": { "x-tenant": "tenant-1", @@ -144,6 +147,9 @@ async def test_route_a2a_model_uses_registered_provider(): assert bridge_kwargs["litellm_params"]["temperature"] == 0.2 assert bridge_kwargs["litellm_params"]["timeout"] == 12.0 assert bridge_kwargs["litellm_params"]["tools"] == data["tools"] + assert bridge_kwargs["litellm_params"]["output_config"] == data["output_config"] + assert bridge_kwargs["litellm_params"]["prompt_cache_key"] == data["prompt_cache_key"] + assert bridge_kwargs["litellm_params"]["safety_identifier"] == data["safety_identifier"] assert bridge_kwargs["litellm_params"]["guardrails"] == ["request-guardrail", "agent-guardrail"] assert bridge_kwargs["litellm_params"]["extra_headers"] == { "X-Tenant": "tenant-1", From c1add95b62645c6ebcbbd03487c6669204376823 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:57:17 +0900 Subject: [PATCH 08/20] fix: preserve A2A pricing and stream data --- .../litellm_completion_bridge/handler.py | 69 ++++++++++++++----- .../litellm_core_utils/streaming_handler.py | 7 +- litellm/llms/a2a/chat/streaming_iterator.py | 21 +++++- litellm/proxy/agent_endpoints/a2a_routing.py | 7 ++ litellm/proxy/common_request_processing.py | 10 +-- litellm/types/utils.py | 2 +- .../test_completion_bridge_streaming.py | 38 ++++++++++ .../a2a/chat/test_a2a_streaming_iterator.py | 25 +++++++ .../proxy/test_route_a2a_models.py | 34 +++++++++ 9 files changed, 185 insertions(+), 28 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 71993135037..8d84c65ab4e 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -280,6 +280,9 @@ class A2ACompletionBridgeHandler: # 3. Forward content as artifact updates accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas + choice_texts: dict[int, str] = {} + choice_tool_calls: dict[int, list[object]] = {} + choice_finish_reasons: dict[int, str] = {} stream_usage: object | None = None stream_finish_reason: str | None = None chunk_count = 0 @@ -304,24 +307,34 @@ class A2ACompletionBridgeHandler: stream_usage = dumped_usage # Extract delta content - content = "" - if chunk is not None and hasattr(chunk, "choices") and chunk.choices: - choice = chunk.choices[0] - raw_finish_reason = getattr(choice, "finish_reason", None) - if isinstance(raw_finish_reason, str) and raw_finish_reason: - stream_finish_reason = raw_finish_reason - if hasattr(choice, "delta") and choice.delta: - content = choice.delta.content or "" - tool_calls = getattr(choice.delta, "tool_calls", None) - if isinstance(tool_calls, (list, tuple)): - accumulated_tool_calls.extend(tool_calls) + choices = getattr(chunk, "choices", None) if chunk is not None else None + if isinstance(choices, (list, tuple)): + for choice_position, choice in enumerate(choices): + raw_index = getattr(choice, "index", choice_position) + choice_index = raw_index if isinstance(raw_index, int) else choice_position + choice_texts.setdefault(choice_index, "") + raw_finish_reason = getattr(choice, "finish_reason", None) + if isinstance(raw_finish_reason, str) and raw_finish_reason: + choice_finish_reasons[choice_index] = raw_finish_reason + if choice_index == 0 or stream_finish_reason is None: + stream_finish_reason = raw_finish_reason + content = "" + delta = getattr(choice, "delta", None) + if delta: + raw_content = getattr(delta, "content", None) + content = raw_content if isinstance(raw_content, str) else "" + choice_texts[choice_index] += content + tool_calls = getattr(delta, "tool_calls", None) + if isinstance(tool_calls, (list, tuple)): + accumulated_tool_calls.extend(tool_calls) + choice_tool_calls.setdefault(choice_index, []).extend(tool_calls) - if content: - artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( - ctx=ctx, - text=content, - ) - yield artifact_event + if content: + artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( + ctx=ctx, + text=content, + ) + yield artifact_event finally: close_response = getattr(response, "aclose", None) if close_response is not None: @@ -339,6 +352,28 @@ class A2ACompletionBridgeHandler: completed_event["result"]["finish_reason"] = stream_finish_reason if stream_usage is not None: completed_event["usage"] = stream_usage + if len(choice_texts) > 1: + completed_event["result"]["choices"] = [ + { + "index": choice_index, + "message": { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": choice_texts[choice_index]}], + **( + {"tool_calls": choice_tool_calls[choice_index]} + if choice_tool_calls.get(choice_index) + else {} + ), + }, + **( + {"finish_reason": choice_finish_reasons[choice_index]} + if choice_index in choice_finish_reasons + else {} + ), + } + for choice_index in sorted(choice_texts) + ] yield completed_event verbose_logger.info( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f6340426c1b..1e416a21e84 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1233,7 +1233,8 @@ class CustomStreamWrapper: ) if "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None: - completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] + tool_use = anthropic_response_obj["tool_use"] + completion_obj["tool_calls"] = tool_use if isinstance(tool_use, list) else [tool_use] if ( "provider_specific_fields" in anthropic_response_obj @@ -2559,6 +2560,8 @@ def convert_generic_chunk_to_model_response_stream( ) -> ModelResponseStream: from litellm.types.utils import Delta + tool_use = chunk.get("tool_use", None) + tool_calls = tool_use if isinstance(tool_use, list) else [tool_use] if tool_use is not None else None model_response_stream: Final = ModelResponseStream( id=str(uuid.uuid4()), model="", @@ -2567,7 +2570,7 @@ def convert_generic_chunk_to_model_response_stream( index=chunk.get("index", 0), delta=Delta( content=chunk["text"], - tool_calls=chunk.get("tool_use", None), + tool_calls=tool_calls, ), ) ], diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 636aef0e724..32f2d0babc2 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -149,16 +149,18 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return raw_usage return raw_usage - def _get_tool_calls(self, chunk: dict) -> ChatCompletionToolCallChunk | None: + def _get_tool_calls( + self, chunk: dict + ) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None: result: Final = chunk.get("result", {}) if not isinstance(result, dict): return None tool_calls = result.get("tool_calls") if isinstance(tool_calls, list) and tool_calls: - return self._serialize_tool_call(tool_calls[0]) + return self._serialize_tool_calls(tool_calls) message = result.get("message") if isinstance(message, dict) and isinstance(message.get("tool_calls"), list) and message["tool_calls"]: - return self._serialize_tool_call(message["tool_calls"][0]) + return self._serialize_tool_calls(message["tool_calls"]) return None @staticmethod @@ -171,6 +173,19 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return tool_call.dict(exclude_none=True) return None + @classmethod + def _serialize_tool_calls( + cls, tool_calls: list[object] + ) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None: + serialized: Final = [ + tool_call_value + for tool_call in tool_calls + if (tool_call_value := cls._serialize_tool_call(tool_call)) is not None + ] + if len(serialized) == 1: + return serialized[0] + return serialized or None + async def aclose(self) -> None: streaming_response = self.streaming_response self.streaming_response = None diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 34e945f1415..d6dcd58c02b 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -163,6 +163,13 @@ async def _route_registered_provider( logging_obj: Final = data.get("litellm_logging_obj") if isinstance(logging_obj, Logging): provider_params["no-log"] = True + provider_model: Final = litellm_params.get("model") + if isinstance(provider_model, str): + logging_obj.model_call_details["model"] = provider_model + logging_obj.model_call_details.setdefault("litellm_params", {})["model"] = provider_model + provider_name: Final = litellm_params.get("custom_llm_provider") + if isinstance(provider_name, str): + logging_obj.model_call_details["custom_llm_provider"] = provider_name pricing_params = { key: litellm_params[key] for key in _A2A_PRICING_PARAMS diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 16a10a16e70..5927214a202 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1841,11 +1841,6 @@ class ProxyBaseLLMRequestProcessing: merge_a2a_agent_guardrails_before_hooks, ) - self.data = await authorize_a2a_agent_before_hooks( - data=self.data, - user_api_key_dict=user_api_key_dict, - ) - logging_obj, self.data = litellm.utils.function_setup( original_function=route_type, rules_obj=litellm.utils.Rules(), @@ -1855,6 +1850,11 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_logging_obj"] = logging_obj + self.data = await authorize_a2a_agent_before_hooks( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) + self.data = await merge_a2a_agent_guardrails_before_hooks(self.data) # Merge model-level guardrails before pre_call_hook so DB/UI-configured diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 67eae2b4f21..61d7aca0430 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -317,7 +317,7 @@ class ModelInfo(ModelInfoBase, total=False): class GenericStreamingChunk(TypedDict, total=False): text: Required[str] - tool_use: ChatCompletionToolCallChunk | None + tool_use: ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None is_finished: Required[bool] finish_reason: Required[str] usage: Required[ChatCompletionUsageBlock | None] diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index d2b74d416a0..6279bc78dab 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -249,6 +249,44 @@ async def test_handle_streaming_emits_proper_events(): assert events[4]["usage"]["total_tokens"] == 5 +@pytest.mark.asyncio +async def test_handle_streaming_preserves_multiple_choices(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + mock_chunk = MagicMock() + first_choice = MagicMock() + first_choice.index = 0 + first_choice.finish_reason = None + first_choice.delta.content = "first" + second_choice = MagicMock() + second_choice.index = 1 + second_choice.finish_reason = "length" + second_choice.delta.content = "second" + mock_chunk.choices = [first_choice, second_choice] + + async def mock_streaming_response(): + yield mock_chunk + + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + mock_acompletion.return_value = mock_streaming_response() + events = [ + event + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-choices", + params={"message": {"role": "user", "parts": []}}, + litellm_params={"custom_llm_provider": "langgraph", "model": "agent", "n": 2}, + ) + ] + + choices = events[-1]["result"]["choices"] + assert [choice["index"] for choice in choices] == [0, 1] + assert choices[0]["message"]["parts"][0]["text"] == "first" + assert choices[1]["message"]["parts"][0]["text"] == "second" + assert choices[1]["finish_reason"] == "length" + + @pytest.mark.asyncio async def test_provider_config_receives_full_message_history(): from litellm.a2a_protocol.litellm_completion_bridge.handler import ( diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index db303719864..443254e69fd 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -49,6 +49,31 @@ async def test_async_iterator_preserves_tool_calls(): assert chunk["finish_reason"] == "tool_calls" +@pytest.mark.asyncio +async def test_async_iterator_preserves_parallel_tool_calls(): + tool_calls = [ + { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + }, + { + "id": "call-2", + "type": "function", + "function": {"name": "write", "arguments": "{}"}, + }, + ] + + async def _events(): + yield {"jsonrpc": "2.0", "result": {"tool_calls": tool_calls}} + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["tool_use"] == tool_calls + + @pytest.mark.asyncio async def test_async_iterator_serializes_delta_tool_calls_and_usage(): delta = Delta( diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index ca8a7ff5b2a..2a34f4404bc 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -11,6 +11,7 @@ from fastapi import HTTPException from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.agent_endpoints.a2a_routing import ( + _route_registered_provider, merge_a2a_agent_guardrails_before_hooks, route_a2a_agent_request, ) @@ -642,6 +643,39 @@ async def test_route_a2a_stream_uses_registered_provider(): assert response is wrapper +@pytest.mark.asyncio +async def test_registered_provider_logging_uses_provider_model_for_builtin_pricing(): + class FakeLogging: + def __init__(self) -> None: + self.model_call_details = {"litellm_params": {}} + self.litellm_params = self.model_call_details["litellm_params"] + self.custom_pricing = False + + logging_obj = FakeLogging() + response = {"result": {"message": {"parts": [{"kind": "text", "text": "hello"}]}}} + with ( + patch("litellm.litellm_core_utils.litellm_logging.Logging", FakeLogging), + patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=response), + ), + ): + await _route_registered_provider( + data={ + "messages": [{"role": "user", "content": "hello"}], + "litellm_logging_obj": logging_obj, + }, + model_name="a2a/agent", + api_base="https://provider.example", + litellm_params={"model": "gpt-4o", "custom_llm_provider": "openai"}, + static_headers=None, + ) + + assert logging_obj.model_call_details["model"] == "gpt-4o" + assert logging_obj.model_call_details["custom_llm_provider"] == "openai" + assert logging_obj.model_call_details["litellm_params"]["model"] == "gpt-4o" + + @pytest.mark.asyncio async def test_route_non_a2a_model_raises_error_if_not_in_router(): """Test that non-a2a models that aren't in router raise an error""" From a52bf355c63a73e47619e8c98a8b4d5f2ce14ef9 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:05:09 +0900 Subject: [PATCH 09/20] fix: preserve A2A stream metadata --- .../litellm_completion_bridge/handler.py | 28 ++++++++++ litellm/llms/a2a/chat/streaming_iterator.py | 56 +++++++++++++++++-- litellm/proxy/agent_endpoints/a2a_routing.py | 1 + .../test_completion_bridge_streaming.py | 43 ++++++++++++++ .../a2a/chat/test_a2a_streaming_iterator.py | 46 +++++++++++++++ 5 files changed, 170 insertions(+), 4 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 8d84c65ab4e..f99a91729d4 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -282,6 +282,8 @@ class A2ACompletionBridgeHandler: accumulated_tool_calls: Final[list[object]] = [] # mutable-ok: collect streaming tool-call deltas choice_texts: dict[int, str] = {} choice_tool_calls: dict[int, list[object]] = {} + choice_delta_fields: dict[int, dict[str, object]] = {} + choice_logprobs: dict[int, object] = {} choice_finish_reasons: dict[int, str] = {} stream_usage: object | None = None stream_finish_reason: str | None = None @@ -328,6 +330,26 @@ class A2ACompletionBridgeHandler: if isinstance(tool_calls, (list, tuple)): accumulated_tool_calls.extend(tool_calls) choice_tool_calls.setdefault(choice_index, []).extend(tool_calls) + delta_fields = A2ACompletionBridgeTransformation._model_dump(delta) + if delta_fields: + choice_fields = choice_delta_fields.setdefault(choice_index, {}) + for field, value in delta_fields.items(): + if field in {"content", "role", "tool_calls"} or value is None: + continue + previous = choice_fields.get(field) + if (isinstance(previous, str) and isinstance(value, str)) or ( + isinstance(previous, list) and isinstance(value, list) + ): + choice_fields[field] = previous + value + elif isinstance(previous, Mapping) and isinstance(value, Mapping): + choice_fields[field] = {**previous, **value} + else: + choice_fields[field] = value + + raw_logprobs = getattr(choice, "logprobs", None) + serialized_logprobs = A2ACompletionBridgeTransformation._model_dump(raw_logprobs) + if serialized_logprobs: + choice_logprobs[choice_index] = serialized_logprobs if content: artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( @@ -352,6 +374,10 @@ class A2ACompletionBridgeHandler: completed_event["result"]["finish_reason"] = stream_finish_reason if stream_usage is not None: completed_event["usage"] = stream_usage + if choice_delta_fields.get(0): + completed_event["result"].update(choice_delta_fields[0]) + if 0 in choice_logprobs: + completed_event["result"]["logprobs"] = choice_logprobs[0] if len(choice_texts) > 1: completed_event["result"]["choices"] = [ { @@ -365,12 +391,14 @@ class A2ACompletionBridgeHandler: if choice_tool_calls.get(choice_index) else {} ), + **choice_delta_fields.get(choice_index, {}), }, **( {"finish_reason": choice_finish_reasons[choice_index]} if choice_index in choice_finish_reasons else {} ), + **({"logprobs": choice_logprobs[choice_index]} if choice_index in choice_logprobs else {}), } for choice_index in sorted(choice_texts) ] diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 32f2d0babc2..1a598bb7a50 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -70,7 +70,56 @@ class A2AModelResponseIterator(BaseModelResponseIterator): try: # Extract text from A2A response - text: Final = extract_text_from_a2a_response(chunk) + result: Final = chunk.get("result", {}) + status: Final = result.get("status", {}) if isinstance(result, Mapping) else {} + is_working_status: Final = ( + isinstance(result, Mapping) + and result.get("kind") == "status-update" + and isinstance(status, Mapping) + and status.get("state") == "working" + ) + text: Final = "" if is_working_status else extract_text_from_a2a_response(chunk) + provider_fields: dict[str, object] = {} + if isinstance(result, Mapping) and not is_working_status: + control_fields = { + "artifacts", + "choices", + "contextId", + "final", + "finish_reason", + "history", + "id", + "kind", + "message", + "parts", + "status", + "taskId", + "tool_calls", + "usage", + } + provider_fields.update( + {key: value for key, value in result.items() if key not in control_fields and value is not None} + ) + choices = result.get("choices") + if isinstance(choices, list) and choices: + first_choice = choices[0] + if isinstance(first_choice, Mapping): + provider_fields.update( + { + key: value + for key, value in first_choice.items() + if key not in {"index", "message", "finish_reason"} and value is not None + } + ) + first_message = first_choice.get("message") + if isinstance(first_message, Mapping): + provider_fields.update( + { + key: value + for key, value in first_message.items() + if key not in {"kind", "role", "parts", "tool_calls"} and value is not None + } + ) # Determine finish reason finish_reason: Final = self._get_finish_reason(chunk) @@ -85,6 +134,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): usage=usage, index=0, tool_use=tool_calls, + provider_specific_fields=provider_fields or None, ) except Exception: # Return empty chunk on parse error @@ -149,9 +199,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): return raw_usage return raw_usage - def _get_tool_calls( - self, chunk: dict - ) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None: + def _get_tool_calls(self, chunk: dict) -> ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None: result: Final = chunk.get("result", {}) if not isinstance(result, dict): return None diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index d6dcd58c02b..38cd5e9f4f2 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -55,6 +55,7 @@ _FORWARDED_REQUEST_PARAMS: Final = frozenset( "safety_identifier", "stop", "store", + "stream_options", "temperature", "thinking", "timeout", diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 6279bc78dab..11f2e75107d 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -287,6 +287,49 @@ async def test_handle_streaming_preserves_multiple_choices(): assert choices[1]["finish_reason"] == "length" +@pytest.mark.asyncio +async def test_handle_streaming_preserves_non_text_delta_fields(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + delta = MagicMock() + delta.content = "" + delta.tool_calls = None + delta.model_dump.return_value = { + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "provider_specific_fields": {"trace_id": "trace-1"}, + } + choice = MagicMock() + choice.index = 0 + choice.finish_reason = "stop" + choice.delta = delta + choice.logprobs = {"content": []} + chunk = MagicMock() + chunk.choices = [choice] + + async def mock_streaming_response(): + yield chunk + + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + mock_acompletion.return_value = mock_streaming_response() + events = [ + event + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-fields", + params={"message": {"role": "user", "parts": []}}, + litellm_params={"custom_llm_provider": "langgraph", "model": "agent"}, + ) + ] + + result = events[-1]["result"] + assert result["audio"] == {"data": "abc"} + assert result["reasoning_content"] == "thinking" + assert result["provider_specific_fields"] == {"trace_id": "trace-1"} + assert result["logprobs"] == {"content": []} + + @pytest.mark.asyncio async def test_provider_config_receives_full_message_history(): from litellm.a2a_protocol.litellm_completion_bridge.handler import ( diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index 443254e69fd..9c4212ee135 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -28,6 +28,52 @@ async def test_async_iterator_accepts_decoded_a2a_events(): assert chunk["text"] == "Hello" +@pytest.mark.asyncio +async def test_async_iterator_ignores_status_message_text(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "status-update", + "status": { + "state": "working", + "message": {"parts": [{"kind": "text", "text": "Processing request..."}]}, + }, + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["text"] == "" + + +@pytest.mark.asyncio +async def test_async_iterator_preserves_non_text_fields(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "status-update", + "status": {"state": "completed"}, + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "logprobs": {"content": []}, + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + + chunk = await iterator.__aiter__().__anext__() + + assert chunk["provider_specific_fields"] == { + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "logprobs": {"content": []}, + } + + @pytest.mark.asyncio async def test_async_iterator_preserves_tool_calls(): tool_calls = [ From eba758223d33fc777b9a5a206738135f255b4b78 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:27:34 +0900 Subject: [PATCH 10/20] fix: forward static WXO headers --- .../litellm_completion_bridge/handler.py | 6 +++ .../providers/watsonx_orchestrate/config.py | 2 + .../providers/watsonx_orchestrate/handler.py | 52 +++++++++++++++---- litellm/proxy/agent_endpoints/a2a_routing.py | 2 + ...test_watsonx_orchestrate_transformation.py | 42 ++++++++++++++- 5 files changed, 92 insertions(+), 12 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index f99a91729d4..d4708e02786 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -127,6 +127,7 @@ class A2ACompletionBridgeHandler: litellm_params: dict[str, Any], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, + agent_static_headers: Mapping[str, str] | None = None, *, _skip_a2a_provider_routing: bool = False, ) -> dict[str, object]: @@ -140,6 +141,7 @@ class A2ACompletionBridgeHandler: api_base: API base URL from agent_card_params agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and admin extra_headers) to forward on the upstream HTTP call. + agent_static_headers: Configured headers for provider-specific routing. Returns: A2A SendMessageResponse dict @@ -161,6 +163,7 @@ class A2ACompletionBridgeHandler: "api_base": api_base, "litellm_params": litellm_params, "agent_extra_headers": agent_extra_headers, + "agent_static_headers": agent_static_headers, } if litellm_params.get("timeout") is not None: provider_kwargs["timeout"] = litellm_params["timeout"] @@ -194,6 +197,7 @@ class A2ACompletionBridgeHandler: litellm_params: dict[str, Any], api_base: str | None = None, agent_extra_headers: dict[str, str] | None = None, + agent_static_headers: Mapping[str, str] | None = None, *, _skip_a2a_provider_routing: bool = False, ) -> AsyncIterator[dict[str, object]]: @@ -213,6 +217,7 @@ class A2ACompletionBridgeHandler: api_base: API base URL from agent_card_params agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and admin extra_headers) to forward on the upstream HTTP call. + agent_static_headers: Configured headers for provider-specific routing. Yields: A2A streaming response events @@ -234,6 +239,7 @@ class A2ACompletionBridgeHandler: "api_base": api_base, "litellm_params": litellm_params, "agent_extra_headers": agent_extra_headers, + "agent_static_headers": agent_static_headers, } if litellm_params.get("timeout") is not None: provider_kwargs["timeout"] = litellm_params["timeout"] diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index ca84d3e07b4..1b8024ae075 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -32,6 +32,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): request_id=request_id, params=params, litellm_params=litellm_params, + static_headers=kwargs.get("agent_static_headers"), ) async def handle_streaming( @@ -52,5 +53,6 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): request_id=request_id, params=params, litellm_params=litellm_params, + static_headers=kwargs.get("agent_static_headers"), ): yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index c66b07c321c..f648cca40cd 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -6,7 +6,7 @@ import asyncio import hashlib import json import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from typing import Any, Final, NamedTuple, Protocol import httpx @@ -26,9 +26,36 @@ _IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token" _POLL_INTERVAL_S: Final = 2.0 _MAX_POLL_ATTEMPTS: Final = 90 _TOKEN_CACHE_TTL_BUFFER_S: Final = 60 +_WXO_RESERVED_HEADERS: Final = frozenset({"accept", "authorization", "content-type"}) _token_cache: Final[dict[str, tuple[str, float]]] = {} +def _build_wxo_headers( + token: str, + accept: str, + static_headers: Mapping[str, str] | None = None, +) -> dict[str, str]: + headers: dict[str, str] = ( + { + key: value + for key, value in static_headers.items() + if isinstance(key, str) + and isinstance(value, str) + and key.lower() not in _WXO_RESERVED_HEADERS + } + if isinstance(static_headers, Mapping) + else {} + ) + headers.update( + { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": accept, + } + ) + return headers + + class WXORequestParams(NamedTuple): cp4d_host: str instance_id: str @@ -277,6 +304,7 @@ class WatsonxOrchestrateHandler: request_id: str, params: dict[str, object], litellm_params: WXOLitellmParams, + static_headers: Mapping[str, str] | None = None, ) -> dict[str, object]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) @@ -289,11 +317,11 @@ class WatsonxOrchestrateHandler: client=client, ) base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) - auth_headers: Final = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "Accept": "application/json", - } + auth_headers: Final = _build_wxo_headers( + token=token, + accept="application/json", + static_headers=static_headers, + ) text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body( @@ -326,6 +354,7 @@ class WatsonxOrchestrateHandler: litellm_params: WXOLitellmParams, chunk_size: int = 50, delay_ms: int = 10, + static_headers: Mapping[str, str] | None = None, ) -> AsyncIterator[dict[str, object]]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) @@ -338,11 +367,11 @@ class WatsonxOrchestrateHandler: client=client, ) base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) - auth_headers: Final = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "Accept": "text/event-stream, application/json", - } + auth_headers: Final = _build_wxo_headers( + token=token, + accept="text/event-stream, application/json", + static_headers=static_headers, + ) text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body( wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id @@ -366,6 +395,7 @@ class WatsonxOrchestrateHandler: request_id=request_id, params=params, litellm_params=litellm_params, + static_headers=static_headers, ) response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 38cd5e9f4f2..ff8381c6531 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -188,6 +188,7 @@ async def _route_registered_provider( litellm_params=provider_params, api_base=api_base, agent_extra_headers=agent_extra_headers, + agent_static_headers=static_headers, ) completion_stream: Final = A2AModelResponseIterator( streaming_response=streaming_response, @@ -210,6 +211,7 @@ async def _route_registered_provider( litellm_params=provider_params, api_base=api_base, agent_extra_headers=agent_extra_headers, + agent_static_headers=static_headers, ) error_value: Final = response.get("error") if isinstance(error_value, dict): diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 43dfdaba02d..869769fe630 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -6,9 +6,11 @@ from pathlib import Path import httpx import pytest - from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager from litellm.a2a_protocol.providers.watsonx_orchestrate import handler as wxo_handler +from litellm.a2a_protocol.providers.watsonx_orchestrate.config import ( + WatsonxOrchestrateA2AConfig, +) from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( WatsonxOrchestrateHandler, ) @@ -343,6 +345,44 @@ async def test_poll_run_raises_asyncio_timeout_when_never_terminal(): assert client.get_calls == 2 +def test_build_wxo_headers_preserves_auth_headers(): + headers = wxo_handler._build_wxo_headers( + token="token", + accept="application/json", + static_headers={ + "x-tenant-id": "tenant-1", + "Authorization": "caller-token", + "content-type": "text/plain", + }, + ) + + assert headers["x-tenant-id"] == "tenant-1" + assert headers["Authorization"] == "Bearer token" + assert headers["Content-Type"] == "application/json" + assert headers["Accept"] == "application/json" + assert "content-type" not in headers + + +@pytest.mark.asyncio +async def test_wxo_config_forwards_static_headers(monkeypatch): + captured = {} + + async def fake_handle_non_streaming(**kwargs): + captured.update(kwargs) + return {"result": {}} + + monkeypatch.setattr(WatsonxOrchestrateHandler, "handle_non_streaming", fake_handle_non_streaming) + + await WatsonxOrchestrateA2AConfig().handle_non_streaming( + request_id="req-1", + params={}, + litellm_params={"model": "agent"}, + agent_static_headers={"x-tenant-id": "tenant-1"}, + ) + + assert captured["static_headers"] == {"x-tenant-id": "tenant-1"} + + @pytest.mark.asyncio async def test_handle_streaming_polls_non_sse_json_until_complete(monkeypatch): client = _JsonStreamClient({"status": "running", "run_id": "run-1"}) From a5997bb90824e4e8c35885f059fc649b4eb41c66 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:29:13 +0900 Subject: [PATCH 11/20] fix: address A2A review feedback --- .../litellm_completion_bridge/handler.py | 43 +++++++++++++------ .../transformation.py | 18 +++++--- litellm/llms/a2a/chat/streaming_iterator.py | 12 +++++- .../proxy/agent_endpoints/a2a_endpoints.py | 2 +- litellm/proxy/agent_endpoints/a2a_routing.py | 11 +++-- litellm/proxy/common_request_processing.py | 4 +- litellm/utils.py | 4 +- .../test_bedrock_agentcore_a2a.py | 4 +- .../test_completion_bridge_streaming.py | 11 +++-- .../proxy/test_route_a2a_models.py | 5 ++- 10 files changed, 81 insertions(+), 33 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index d4708e02786..04f013e86c6 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -89,11 +89,12 @@ class A2ACompletionBridgeHandler: "api_base": api_base, "stream": stream, } + configured_headers: Final[object] = litellm_params.get("extra_headers") or litellm_params.get("headers") # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) litellm_params_to_add: Final = { k: v for k, v in litellm_params.items() - if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS + if k not in ("model", "custom_llm_provider", "extra_headers", "headers") and k not in _AGENT_ONLY_PARAMS } completion_params.update(litellm_params_to_add) # Apply forward metadata AFTER the litellm_params merge so the helper @@ -105,10 +106,10 @@ class A2ACompletionBridgeHandler: params=params, ) - if agent_extra_headers: + if agent_extra_headers or configured_headers: completion_params["extra_headers"] = merge_agent_headers( dynamic_headers=agent_extra_headers, - static_headers=completion_params.get("extra_headers"), + static_headers=configured_headers if isinstance(configured_headers, Mapping) else None, ) return completion_params @@ -361,6 +362,7 @@ class A2ACompletionBridgeHandler: artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( ctx=ctx, text=content, + index=choice_index, ) yield artifact_event finally: @@ -380,13 +382,10 @@ class A2ACompletionBridgeHandler: completed_event["result"]["finish_reason"] = stream_finish_reason if stream_usage is not None: completed_event["usage"] = stream_usage - if choice_delta_fields.get(0): - completed_event["result"].update(choice_delta_fields[0]) - if 0 in choice_logprobs: - completed_event["result"]["logprobs"] = choice_logprobs[0] if len(choice_texts) > 1: - completed_event["result"]["choices"] = [ - { + choice_payloads: list[dict[str, object]] = [] + for choice_index in sorted(choice_texts): + choice_payload: dict[str, object] = { "index": choice_index, "message": { "kind": "message", @@ -397,7 +396,6 @@ class A2ACompletionBridgeHandler: if choice_tool_calls.get(choice_index) else {} ), - **choice_delta_fields.get(choice_index, {}), }, **( {"finish_reason": choice_finish_reasons[choice_index]} @@ -406,8 +404,29 @@ class A2ACompletionBridgeHandler: ), **({"logprobs": choice_logprobs[choice_index]} if choice_index in choice_logprobs else {}), } - for choice_index in sorted(choice_texts) - ] + if choice_delta_fields.get(choice_index): + choice_payload["delta"] = choice_delta_fields[choice_index] + choice_payloads.append(choice_payload) + completed_event["result"]["choices"] = choice_payloads + else: + metadata_indices = sorted(set(choice_delta_fields) | set(choice_logprobs)) + if metadata_indices: + completed_event["result"]["choices"] = [ + { + "index": choice_index, + **( + {"delta": choice_delta_fields[choice_index]} + if choice_delta_fields.get(choice_index) + else {} + ), + **( + {"logprobs": choice_logprobs[choice_index]} + if choice_index in choice_logprobs + else {} + ), + } + for choice_index in metadata_indices + ] yield completed_event verbose_logger.info( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 02006e98fa0..6e11acb71a6 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -262,6 +262,10 @@ class A2ACompletionBridgeTransformation: } if usage is not None: a2a_response["usage"] = usage.model_dump(exclude_none=True) if hasattr(usage, "model_dump") else usage + for field in ("system_fingerprint", "service_tier"): + value = getattr(response, field, None) + if value is not None: + a2a_response[field] = value if len(serialized_choices) > 1: a2a_response["choices"] = serialized_choices @@ -354,6 +358,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, + index: int | None = None, ) -> dict[str, Any]: """ Create an artifact update event with content. @@ -362,15 +367,18 @@ class A2ACompletionBridgeTransformation: ctx: Streaming context text: The text content for the artifact """ + artifact: Final[dict[str, Any]] = { + "artifactId": str(uuid4()), + "name": "response", + "parts": [{"kind": "text", "text": text}], + } + if index is not None: + artifact["index"] = index return { "id": ctx.request_id, "jsonrpc": "2.0", "result": { - "artifact": { - "artifactId": str(uuid4()), - "name": "response", - "parts": [{"kind": "text", "text": text}], - }, + "artifact": artifact, "contextId": ctx.context_id, "kind": "artifact-update", "taskId": ctx.task_id, diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 1a598bb7a50..605c8cbd375 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -71,6 +71,16 @@ class A2AModelResponseIterator(BaseModelResponseIterator): try: # Extract text from A2A response result: Final = chunk.get("result", {}) + chunk_index = 0 + if isinstance(result, Mapping): + artifact = result.get("artifact") + if isinstance(artifact, Mapping) and isinstance(artifact.get("index"), int): + chunk_index = artifact["index"] + choices = result.get("choices") + if isinstance(choices, list) and choices and isinstance(choices[0], Mapping): + raw_index = choices[0].get("index") + if isinstance(raw_index, int): + chunk_index = raw_index status: Final = result.get("status", {}) if isinstance(result, Mapping) else {} is_working_status: Final = ( isinstance(result, Mapping) @@ -132,7 +142,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): is_finished=bool(finish_reason or tool_calls), finish_reason=finish_reason or ("tool_calls" if tool_calls else ""), usage=usage, - index=0, + index=chunk_index, tool_use=tool_calls, provider_specific_fields=provider_fields or None, ) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index bd02cfdf907..ef36296fff5 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -860,7 +860,7 @@ async def invoke_agent_a2a( _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is not None: logging_obj._enqueue_deferred_logging = None - _enqueue_fn() + _enqueue_fn(response) response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index ff8381c6531..b1af6121839 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -37,6 +37,7 @@ _FORWARDED_REQUEST_PARAMS: Final = frozenset( "frequency_penalty", "functions", "function_call", + "guided_json", "include_server_side_tool_invocations", "logit_bias", "logprobs", @@ -305,6 +306,10 @@ async def _route_registered_provider( id=str(response.get("id") or request_id), model=model_name, choices=model_choices, + system_fingerprint=response.get("system_fingerprint") + if isinstance(response.get("system_fingerprint"), str) + else None, + service_tier=response.get("service_tier") if isinstance(response.get("service_tier"), str) else None, ) raw_usage: Final = response.get("usage") usage: Final = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage @@ -315,10 +320,10 @@ async def _route_registered_provider( if isinstance(logging_obj, Logging): - def _enqueue_logging() -> None: + def _enqueue_logging(final_response: ModelResponse | None = None) -> None: asyncio.create_task( logging_obj.dispatch_success_handlers( - model_response, + final_response if final_response is not None else model_response, cache_hit=False, prefer_async_handlers=True, ) @@ -544,7 +549,7 @@ async def route_a2a_agent_request( else registered_params_value or {} ) configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None - api_base: Final = configured_api_base if isinstance(configured_api_base, str) else agent_url + api_base: Final = configured_api_base if isinstance(configured_api_base, str) and configured_api_base else agent_url registered_model: Final = registered_params_value.get("model") if registered_params_value else None cardless_provider: Final = registered_provider == "watsonx_orchestrate" or ( registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5927214a202..04432b2ca21 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2544,6 +2544,7 @@ class ProxyBaseLLMRequestProcessing: ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( logging_obj=logging_obj, exception_raised=_exception_raised, + response=response, ) # Streaming cleanup: if an exception occurred AND the deferred @@ -3027,6 +3028,7 @@ class ProxyBaseLLMRequestProcessing: def _flush_deferred_async_logging( logging_obj: Any, exception_raised: bool, + response: Any | None = None, ) -> None: """ Fire the deferred async-success closure stored by wrapper_async, then @@ -3057,7 +3059,7 @@ class ProxyBaseLLMRequestProcessing: if exception_raised: return try: - _enqueue_fn() + _enqueue_fn(response) if response is not None else _enqueue_fn() except Exception as e: verbose_proxy_logger.exception("Error firing deferred logging: %s", e) diff --git a/litellm/utils.py b/litellm/utils.py index e5ce7157e77..fb636b094f1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1818,11 +1818,11 @@ def client(original_function): if not _is_litellm_internal_call: if getattr(logging_obj, "_defer_async_logging", False): - def _enqueue_deferred_logging() -> None: + def _enqueue_deferred_logging(final_response=None) -> None: asyncio.create_task( _client_async_logging_helper( logging_obj=logging_obj, - result=result, + result=final_response if final_response is not None else result, start_time=start_time, end_time=end_time, is_completion_with_fallbacks=is_completion_with_fallbacks, diff --git a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 5503a5668bf..a959386aae9 100644 --- a/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/test_litellm/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -10,10 +10,9 @@ Verifies that: """ import json - -import pytest from unittest.mock import AsyncMock, MagicMock, patch +import pytest SAMPLE_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789:runtime/my_agent" SAMPLE_MODEL = f"bedrock/agentcore/{SAMPLE_ARN}" @@ -482,6 +481,7 @@ class TestHandlerIntegration: api_base=None, litellm_params=SAMPLE_LITELLM_PARAMS, agent_extra_headers=None, + agent_static_headers=None, ) @pytest.mark.asyncio diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 11f2e75107d..29d4e6ffe36 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -324,10 +324,13 @@ async def test_handle_streaming_preserves_non_text_delta_fields(): ] result = events[-1]["result"] - assert result["audio"] == {"data": "abc"} - assert result["reasoning_content"] == "thinking" - assert result["provider_specific_fields"] == {"trace_id": "trace-1"} - assert result["logprobs"] == {"content": []} + choice_result = result["choices"][0] + assert choice_result["delta"] == { + "audio": {"data": "abc"}, + "reasoning_content": "thinking", + "provider_specific_fields": {"trace_id": "trace-1"}, + } + assert choice_result["logprobs"] == {"content": []} @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 2a34f4404bc..05bc118907c 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -62,7 +62,7 @@ async def test_route_a2a_model_bypasses_router(): "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", mock_registry, ): - result = await route_request( + await route_request( data=data, llm_router=mock_router, user_model=None, @@ -593,6 +593,7 @@ async def test_route_a2a_stream_uses_registered_provider(): litellm_params={"custom_llm_provider": "pydantic_ai_agents"}, ) logging_obj = Mock(spec=Logging) + logging_obj.model_call_details = {} data = { "model": "a2a/test-agent", "messages": [{"role": "user", "content": "Hello"}], @@ -745,7 +746,7 @@ def _router_without_models(): @pytest.mark.asyncio async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry agent_name = "a2a-sibling-replica-agent" From d2bfb5ed1bbc17c8679e4a2fe2ae2bc2f3cca344 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:05:32 +0900 Subject: [PATCH 12/20] fix: address A2A streaming review feedback --- .../litellm_completion_bridge/handler.py | 68 +++++++++++------- .../transformation.py | 3 +- .../providers/watsonx_orchestrate/config.py | 15 +++- .../providers/watsonx_orchestrate/handler.py | 11 ++- litellm/llms/a2a/chat/streaming_iterator.py | 53 +++++++++++++- litellm/proxy/common_request_processing.py | 13 ++++ ...test_watsonx_orchestrate_transformation.py | 10 ++- .../test_completion_bridge_streaming.py | 70 +++++++++++++++++++ .../a2a/chat/test_a2a_streaming_iterator.py | 31 ++++++++ 9 files changed, 241 insertions(+), 33 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 04f013e86c6..ffc94a15890 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -46,6 +46,21 @@ class A2ACompletionBridgeHandler: Static methods for handling A2A requests via LiteLLM completion. """ + @staticmethod + def _merge_stream_values(previous: object, current: object) -> object: + if isinstance(previous, Mapping) and isinstance(current, Mapping): + merged = dict(previous) + for key, value in current.items(): + merged[key] = ( + A2ACompletionBridgeHandler._merge_stream_values(merged[key], value) + if key in merged + else value + ) + return merged + if isinstance(previous, list) and isinstance(current, list): + return [*previous, *current] + return current + @staticmethod def _build_completion_params( params: dict[str, Any], @@ -94,7 +109,8 @@ class A2ACompletionBridgeHandler: litellm_params_to_add: Final = { k: v for k, v in litellm_params.items() - if k not in ("model", "custom_llm_provider", "extra_headers", "headers") and k not in _AGENT_ONLY_PARAMS + if k not in ("model", "custom_llm_provider", "extra_headers", "headers", "api_base", "stream") + and k not in _AGENT_ONLY_PARAMS } completion_params.update(litellm_params_to_add) # Apply forward metadata AFTER the litellm_params merge so the helper @@ -290,8 +306,9 @@ class A2ACompletionBridgeHandler: choice_texts: dict[int, str] = {} choice_tool_calls: dict[int, list[object]] = {} choice_delta_fields: dict[int, dict[str, object]] = {} - choice_logprobs: dict[int, object] = {} + choice_logprobs: dict[int, dict[str, object]] = {} choice_finish_reasons: dict[int, str] = {} + stream_metadata: dict[str, str] = {} stream_usage: object | None = None stream_finish_reason: str | None = None chunk_count = 0 @@ -315,6 +332,14 @@ class A2ACompletionBridgeHandler: if isinstance(dumped_usage, Mapping): stream_usage = dumped_usage + for metadata_name in ("system_fingerprint", "service_tier"): + metadata_value = getattr(chunk, metadata_name, None) + if not isinstance(metadata_value, str): + chunk_fields = A2ACompletionBridgeTransformation._model_dump(chunk) + metadata_value = chunk_fields.get(metadata_name) + if isinstance(metadata_value, str) and metadata_value: + stream_metadata[metadata_name] = metadata_value + # Extract delta content choices = getattr(chunk, "choices", None) if chunk is not None else None if isinstance(choices, (list, tuple)): @@ -356,7 +381,12 @@ class A2ACompletionBridgeHandler: raw_logprobs = getattr(choice, "logprobs", None) serialized_logprobs = A2ACompletionBridgeTransformation._model_dump(raw_logprobs) if serialized_logprobs: - choice_logprobs[choice_index] = serialized_logprobs + previous_logprobs = choice_logprobs.get(choice_index, {}) + merged_logprobs = A2ACompletionBridgeHandler._merge_stream_values( + previous_logprobs, serialized_logprobs + ) + if isinstance(merged_logprobs, dict): + choice_logprobs[choice_index] = merged_logprobs if content: artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( @@ -382,9 +412,18 @@ class A2ACompletionBridgeHandler: completed_event["result"]["finish_reason"] = stream_finish_reason if stream_usage is not None: completed_event["usage"] = stream_usage - if len(choice_texts) > 1: + for metadata_name, metadata_value in stream_metadata.items(): + completed_event[metadata_name] = metadata_value + choice_indices = sorted( + set(choice_texts) + | set(choice_tool_calls) + | set(choice_delta_fields) + | set(choice_logprobs) + | set(choice_finish_reasons) + ) + if choice_indices: choice_payloads: list[dict[str, object]] = [] - for choice_index in sorted(choice_texts): + for choice_index in choice_indices: choice_payload: dict[str, object] = { "index": choice_index, "message": { @@ -408,25 +447,6 @@ class A2ACompletionBridgeHandler: choice_payload["delta"] = choice_delta_fields[choice_index] choice_payloads.append(choice_payload) completed_event["result"]["choices"] = choice_payloads - else: - metadata_indices = sorted(set(choice_delta_fields) | set(choice_logprobs)) - if metadata_indices: - completed_event["result"]["choices"] = [ - { - "index": choice_index, - **( - {"delta": choice_delta_fields[choice_index]} - if choice_delta_fields.get(choice_index) - else {} - ), - **( - {"logprobs": choice_logprobs[choice_index]} - if choice_index in choice_logprobs - else {} - ), - } - for choice_index in metadata_indices - ] yield completed_event verbose_logger.info( diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 6e11acb71a6..644ea9ab75f 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -34,6 +34,7 @@ class A2AStreamingContext: self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) + self.artifact_id = str(uuid4()) self.input_message = input_message self.accumulated_text = "" self.has_emitted_task = False @@ -368,7 +369,7 @@ class A2ACompletionBridgeTransformation: text: The text content for the artifact """ artifact: Final[dict[str, Any]] = { - "artifactId": str(uuid4()), + "artifactId": ctx.artifact_id, "name": "response", "parts": [{"kind": "text", "text": text}], } diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index 1b8024ae075..59d6f6003b0 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -9,6 +9,7 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( WatsonxOrchestrateHandler, ) +from litellm.interactions.agents.utils import merge_agent_headers class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): @@ -28,11 +29,16 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): "litellm_params is required for WatsonxOrchestrateA2AConfig " "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" ) + forwarded_headers: Final = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) return await WatsonxOrchestrateHandler.handle_non_streaming( request_id=request_id, params=params, litellm_params=litellm_params, - static_headers=kwargs.get("agent_static_headers"), + static_headers=forwarded_headers, + timeout=kwargs.get("timeout"), ) async def handle_streaming( @@ -49,10 +55,15 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): "litellm_params is required for WatsonxOrchestrateA2AConfig " "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" ) + forwarded_headers: Final = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) async for chunk in WatsonxOrchestrateHandler.handle_streaming( request_id=request_id, params=params, litellm_params=litellm_params, - static_headers=kwargs.get("agent_static_headers"), + static_headers=forwarded_headers, + timeout=kwargs.get("timeout"), ): yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index f648cca40cd..2d3e5cb1569 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -305,10 +305,13 @@ class WatsonxOrchestrateHandler: params: dict[str, object], litellm_params: WXOLitellmParams, static_headers: Mapping[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, object]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0) + client: Final = WatsonxOrchestrateHandler._http_client( + timeout=timeout if timeout is not None else 90.0 + ) token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, @@ -355,10 +358,13 @@ class WatsonxOrchestrateHandler: chunk_size: int = 50, delay_ms: int = 10, static_headers: Mapping[str, str] | None = None, + timeout: float | None = None, ) -> AsyncIterator[dict[str, object]]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0) + client: Final = WatsonxOrchestrateHandler._http_client( + timeout=timeout if timeout is not None else 120.0 + ) token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, @@ -396,6 +402,7 @@ class WatsonxOrchestrateHandler: params=params, litellm_params=litellm_params, static_headers=static_headers, + timeout=timeout, ) response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 605c8cbd375..099a9f67047 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -3,12 +3,12 @@ A2A Streaming Response Iterator """ from collections.abc import Mapping -from typing import Final +from typing import Any, Final import litellm from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import ChatCompletionToolCallChunk -from litellm.types.utils import GenericStreamingChunk, ModelResponseStream +from litellm.types.utils import Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices from ..common_utils import A2AError, extract_text_from_a2a_response @@ -90,6 +90,13 @@ class A2AModelResponseIterator(BaseModelResponseIterator): ) text: Final = "" if is_working_status else extract_text_from_a2a_response(chunk) provider_fields: dict[str, object] = {} + provider_fields.update( + { + key: value + for key, value in chunk.items() + if key in {"system_fingerprint", "service_tier"} and value is not None + } + ) if isinstance(result, Mapping) and not is_working_status: control_fields = { "artifacts", @@ -136,6 +143,48 @@ class A2AModelResponseIterator(BaseModelResponseIterator): tool_calls: Final = self._get_tool_calls(chunk) usage: Final = self._get_usage(chunk) + if isinstance(result, Mapping): + choices = result.get("choices") + if isinstance(choices, list) and choices: + streaming_choices: list[StreamingChoices] = [] + for choice_position, raw_choice in enumerate(choices): + if not isinstance(raw_choice, Mapping): + continue + raw_index = raw_choice.get("index", choice_position) + choice_index = raw_index if isinstance(raw_index, int) else choice_position + delta_fields: dict[str, Any] = {} + raw_delta = raw_choice.get("delta") + if isinstance(raw_delta, Mapping): + delta_fields.update(raw_delta) + raw_message = raw_choice.get("message") + if isinstance(raw_message, Mapping): + message_text = extract_text_from_a2a_response({"result": {"message": raw_message}}) + if message_text and "content" not in delta_fields: + delta_fields["content"] = message_text + message_tool_calls = raw_message.get("tool_calls") + if message_tool_calls and "tool_calls" not in delta_fields: + delta_fields["tool_calls"] = message_tool_calls + raw_finish_reason = raw_choice.get("finish_reason") + choice_finish_reason = ( + raw_finish_reason + if isinstance(raw_finish_reason, str) and raw_finish_reason + else finish_reason + ) + streaming_choices.append( + StreamingChoices( + index=choice_index, + delta=Delta(**delta_fields), + finish_reason=choice_finish_reason, + logprobs=raw_choice.get("logprobs"), + ) + ) + if streaming_choices: + return ModelResponseStream( + choices=streaming_choices, + usage=usage, + provider_specific_fields=provider_fields or None, + ) + # Return generic streaming chunk return GenericStreamingChunk( text=text, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 04432b2ca21..035bc770aed 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1872,12 +1872,25 @@ class ProxyBaseLLMRequestProcessing: trust_client_model_info=False, ) + authorized_model = self.data.get("model") self.data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=self.data, call_type=route_type, ) + if self.data.get("model") != authorized_model: + self.data = await authorize_a2a_agent_before_hooks( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) + self.data = await merge_a2a_agent_guardrails_before_hooks(self.data) + self.data = _check_and_merge_model_level_guardrails( + data=self.data, + llm_router=llm_router, + trust_client_model_info=False, + ) + # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in # add_litellm_data_to_request predates that mutation. diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index 869769fe630..fa8b11ea182 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -364,7 +364,7 @@ def test_build_wxo_headers_preserves_auth_headers(): @pytest.mark.asyncio -async def test_wxo_config_forwards_static_headers(monkeypatch): +async def test_wxo_config_forwards_headers_and_timeout(monkeypatch): captured = {} async def fake_handle_non_streaming(**kwargs): @@ -377,10 +377,16 @@ async def test_wxo_config_forwards_static_headers(monkeypatch): request_id="req-1", params={}, litellm_params={"model": "agent"}, + agent_extra_headers={"x-request-id": "request-1"}, agent_static_headers={"x-tenant-id": "tenant-1"}, + timeout=12, ) - assert captured["static_headers"] == {"x-tenant-id": "tenant-1"} + assert captured["static_headers"] == { + "x-request-id": "request-1", + "x-tenant-id": "tenant-1", + } + assert captured["timeout"] == 12 @pytest.mark.asyncio diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 29d4e6ffe36..acd584b36db 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -240,6 +240,10 @@ async def test_handle_streaming_emits_proper_events(): # Event 4: second artifact update assert events[3]["result"]["kind"] == "artifact-update" assert events[3]["result"]["artifact"]["parts"][0]["text"] == " world" + assert ( + events[2]["result"]["artifact"]["artifactId"] + == events[3]["result"]["artifact"]["artifactId"] + ) # Event 5: status completed assert events[4]["result"]["kind"] == "status-update" @@ -249,6 +253,72 @@ async def test_handle_streaming_emits_proper_events(): assert events[4]["usage"]["total_tokens"] == 5 +def test_build_completion_params_keeps_bridge_routing_fields(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + params = A2ACompletionBridgeHandler._build_completion_params( + params={"message": {"role": "user", "parts": []}}, + litellm_params={ + "custom_llm_provider": "openai", + "model": "agent", + "api_base": "https://untrusted.example", + "stream": False, + }, + api_base="https://configured.example", + agent_extra_headers=None, + stream=True, + ) + + assert params["api_base"] == "https://configured.example" + assert params["stream"] is True + + +@pytest.mark.asyncio +async def test_handle_streaming_accumulates_logprobs_and_provider_metadata(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + chunks = [] + for token in ("a", "b"): + choice = MagicMock() + choice.index = 0 + choice.finish_reason = None + choice.delta.content = token + choice.logprobs = {"content": [{"token": token}]} + chunk = MagicMock() + chunk.choices = [choice] + chunk.system_fingerprint = "fp-1" + chunk.service_tier = "scale" + chunks.append(chunk) + chunks[-1].choices[0].finish_reason = "stop" + + async def mock_streaming_response(): + for chunk in chunks: + yield chunk + + with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion: + mock_acompletion.return_value = mock_streaming_response() + events = [ + event + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id="req-metadata", + params={"message": {"role": "user", "parts": []}}, + litellm_params={"custom_llm_provider": "openai", "model": "agent"}, + ) + ] + + result = events[-1] + assert result["system_fingerprint"] == "fp-1" + assert result["service_tier"] == "scale" + assert result["result"]["choices"][0]["logprobs"]["content"] == [ + {"token": "a"}, + {"token": "b"}, + ] + + @pytest.mark.asyncio async def test_handle_streaming_preserves_multiple_choices(): from litellm.a2a_protocol.litellm_completion_bridge.handler import ( diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index 9c4212ee135..8518cd3d2f7 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -120,6 +120,37 @@ async def test_async_iterator_preserves_parallel_tool_calls(): assert chunk["tool_use"] == tool_calls +@pytest.mark.asyncio +async def test_async_iterator_preserves_every_terminal_choice(): + async def _events(): + yield { + "jsonrpc": "2.0", + "result": { + "kind": "status-update", + "status": {"state": "completed"}, + "choices": [ + { + "index": 0, + "message": {"parts": [{"kind": "text", "text": "first"}]}, + "finish_reason": "stop", + }, + { + "index": 1, + "message": {"parts": [{"kind": "text", "text": "second"}]}, + "finish_reason": "length", + }, + ], + }, + } + + iterator = A2AModelResponseIterator(streaming_response=_events(), sync_stream=False) + chunk = await iterator.__aiter__().__anext__() + + assert [choice.index for choice in chunk.choices] == [0, 1] + assert [choice.delta.content for choice in chunk.choices] == ["first", "second"] + assert [choice.finish_reason for choice in chunk.choices] == ["stop", "length"] + + @pytest.mark.asyncio async def test_async_iterator_serializes_delta_tool_calls_and_usage(): delta = Delta( From 9af6eacf8410544f4e1f17975f78f6287d0fd31a Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:32:12 +0900 Subject: [PATCH 13/20] fix: harden A2A provider routing --- .../litellm_completion_bridge/handler.py | 14 ++++++++- litellm/proxy/agent_endpoints/a2a_routing.py | 16 +++++++--- .../test_completion_bridge_streaming.py | 31 +++++++++++++++++-- 3 files changed, 53 insertions(+), 8 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index ffc94a15890..3546ae0891d 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -174,6 +174,12 @@ class A2ACompletionBridgeHandler: verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider) provider_params: Final = dict(params) + if custom_llm_provider == "pydantic_ai_agents" or ( + custom_llm_provider == "bedrock" + and isinstance(litellm_params.get("model"), str) + and "agentcore" in litellm_params["model"] + ): + provider_params.pop("messages", None) provider_kwargs: Final[dict[str, Any]] = { "request_id": request_id, "params": provider_params, @@ -250,6 +256,12 @@ class A2ACompletionBridgeHandler: verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider) provider_params: Final = dict(params) + if custom_llm_provider == "pydantic_ai_agents" or ( + custom_llm_provider == "bedrock" + and isinstance(litellm_params.get("model"), str) + and "agentcore" in litellm_params["model"] + ): + provider_params.pop("messages", None) provider_kwargs: Final[dict[str, Any]] = { "request_id": request_id, "params": provider_params, @@ -429,7 +441,7 @@ class A2ACompletionBridgeHandler: "message": { "kind": "message", "role": "agent", - "parts": [{"kind": "text", "text": choice_texts[choice_index]}], + "parts": [{"kind": "text", "text": ""}], **( {"tool_calls": choice_tool_calls[choice_index]} if choice_tool_calls.get(choice_index) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index b1af6121839..5f5b5fc7705 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -147,7 +147,16 @@ async def _route_registered_provider( **_OBJECT_DICT_ADAPTER.validate_python(litellm_params), **{key: data[key] for key in _FORWARDED_REQUEST_PARAMS if key in data and data[key] is not None}, } - bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python(params) + registered_provider: Final = litellm_params.get("custom_llm_provider") + registered_model: Final = litellm_params.get("model") + native_provider: Final = registered_provider == "pydantic_ai_agents" or ( + registered_provider == "bedrock" + and isinstance(registered_model, str) + and "agentcore" in registered_model + ) + bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python( + {"message": params["message"]} if native_provider else params + ) configured_headers: Final = litellm_params.get("extra_headers") or litellm_params.get("headers") configured_headers_dict: Final = ( _HEADERS_ADAPTER.validate_python(configured_headers) if isinstance(configured_headers, dict) else None @@ -550,10 +559,7 @@ async def route_a2a_agent_request( ) configured_api_base: Final = registered_params_value.get("api_base") if registered_params_value else None api_base: Final = configured_api_base if isinstance(configured_api_base, str) and configured_api_base else agent_url - registered_model: Final = registered_params_value.get("model") if registered_params_value else None - cardless_provider: Final = registered_provider == "watsonx_orchestrate" or ( - registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model - ) + cardless_provider: Final = registered_provider is not None and registered_provider != "a2a" has_configured_api_base: Final = isinstance(configured_api_base, str) and bool(configured_api_base) if (not isinstance(agent_url, str) or not agent_url) and not has_configured_api_base and not cardless_provider: verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name) diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index acd584b36db..7fe0a6ced75 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -352,8 +352,7 @@ async def test_handle_streaming_preserves_multiple_choices(): choices = events[-1]["result"]["choices"] assert [choice["index"] for choice in choices] == [0, 1] - assert choices[0]["message"]["parts"][0]["text"] == "first" - assert choices[1]["message"]["parts"][0]["text"] == "second" + assert [choice["message"]["parts"][0]["text"] for choice in choices] == ["", ""] assert choices[1]["finish_reason"] == "length" @@ -433,6 +432,34 @@ async def test_provider_config_receives_full_message_history(): assert provider_config.handle_non_streaming.await_args.kwargs["params"]["messages"] == messages +@pytest.mark.asyncio +async def test_native_provider_config_drops_internal_message_history(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + provider_config = MagicMock() + provider_config.handle_non_streaming = AsyncMock(return_value={"result": {}}) + params = { + "message": {"role": "user", "parts": []}, + "messages": [{"role": "user", "content": "Hello"}], + } + + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2AProviderConfigManager.get_provider_config", + return_value=provider_config, + ): + await A2ACompletionBridgeHandler.handle_non_streaming( + request_id="req-native", + params=params, + litellm_params={"custom_llm_provider": "pydantic_ai_agents", "model": "agent"}, + ) + + assert provider_config.handle_non_streaming.await_args.kwargs["params"] == { + "message": params["message"] + } + + def test_response_transform_preserves_audio_and_logprobs(): from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( A2ACompletionBridgeTransformation, From 97f01b692855da7b295e57d7ee19126474b8fd95 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 12:38:59 +0900 Subject: [PATCH 14/20] chore: format A2A routing files --- .../litellm_completion_bridge/handler.py | 4 +--- .../providers/watsonx_orchestrate/handler.py | 12 +++--------- litellm/proxy/agent_endpoints/a2a_routing.py | 4 +--- 3 files changed, 5 insertions(+), 15 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 3546ae0891d..3cd2e9329dc 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -52,9 +52,7 @@ class A2ACompletionBridgeHandler: merged = dict(previous) for key, value in current.items(): merged[key] = ( - A2ACompletionBridgeHandler._merge_stream_values(merged[key], value) - if key in merged - else value + A2ACompletionBridgeHandler._merge_stream_values(merged[key], value) if key in merged else value ) return merged if isinstance(previous, list) and isinstance(current, list): diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 2d3e5cb1569..fe9d1ac3c15 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -39,9 +39,7 @@ def _build_wxo_headers( { key: value for key, value in static_headers.items() - if isinstance(key, str) - and isinstance(value, str) - and key.lower() not in _WXO_RESERVED_HEADERS + if isinstance(key, str) and isinstance(value, str) and key.lower() not in _WXO_RESERVED_HEADERS } if isinstance(static_headers, Mapping) else {} @@ -309,9 +307,7 @@ class WatsonxOrchestrateHandler: ) -> dict[str, object]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client: Final = WatsonxOrchestrateHandler._http_client( - timeout=timeout if timeout is not None else 90.0 - ) + client: Final = WatsonxOrchestrateHandler._http_client(timeout=timeout if timeout is not None else 90.0) token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, @@ -362,9 +358,7 @@ class WatsonxOrchestrateHandler: ) -> AsyncIterator[dict[str, object]]: wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client: Final = WatsonxOrchestrateHandler._http_client( - timeout=timeout if timeout is not None else 120.0 - ) + client: Final = WatsonxOrchestrateHandler._http_client(timeout=timeout if timeout is not None else 120.0) token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 5f5b5fc7705..36557d8bfeb 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -150,9 +150,7 @@ async def _route_registered_provider( registered_provider: Final = litellm_params.get("custom_llm_provider") registered_model: Final = litellm_params.get("model") native_provider: Final = registered_provider == "pydantic_ai_agents" or ( - registered_provider == "bedrock" - and isinstance(registered_model, str) - and "agentcore" in registered_model + registered_provider == "bedrock" and isinstance(registered_model, str) and "agentcore" in registered_model ) bridge_params: Final = _OBJECT_DICT_ADAPTER.validate_python( {"message": params["message"]} if native_provider else params From 4556ffc86df6e7e296f7a35072fffff8b4dad9f4 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:19:41 +0900 Subject: [PATCH 15/20] fix: close A2A provider review gaps --- .../providers/bedrock_agentcore/config.py | 2 ++ .../providers/bedrock_agentcore/handler.py | 4 ++++ .../a2a_protocol/providers/langflow/config.py | 11 ++++++++++ .../providers/watsonx_orchestrate/handler.py | 21 +++++++++++++++++-- .../litellm_core_utils/streaming_handler.py | 12 +++++++++++ .../proxy/agent_endpoints/a2a_endpoints.py | 5 ++++- litellm/proxy/agent_endpoints/a2a_routing.py | 15 +++++++------ litellm/proxy/common_request_processing.py | 7 +++++++ 8 files changed, 68 insertions(+), 9 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py index 2b37c0c4906..04541d8863e 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py @@ -38,6 +38,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): params=params, litellm_params=litellm_params, agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=kwargs.get("timeout"), ) async def handle_streaming( @@ -58,5 +59,6 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): params=params, litellm_params=litellm_params, agent_extra_headers=kwargs.get("agent_extra_headers"), + timeout=kwargs.get("timeout"), ): yield chunk diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index db57072ca38..64979848110 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -31,6 +31,7 @@ class BedrockAgentCoreA2AHandler: params: dict[str, Any], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, Any]: """ Handle non-streaming A2A request to AgentCore. @@ -62,6 +63,7 @@ class BedrockAgentCoreA2AHandler: url, headers=headers, data=body, + timeout=timeout, ) response.raise_for_status() response_data: Final = response.json() @@ -77,6 +79,7 @@ class BedrockAgentCoreA2AHandler: params: dict[str, Any], litellm_params: dict[str, Any], agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming A2A request to AgentCore. @@ -110,6 +113,7 @@ class BedrockAgentCoreA2AHandler: headers=headers, data=body, stream=True, + timeout=timeout, ) response.raise_for_status() diff --git a/litellm/a2a_protocol/providers/langflow/config.py b/litellm/a2a_protocol/providers/langflow/config.py index 54d403f88c0..eec32d4e274 100644 --- a/litellm/a2a_protocol/providers/langflow/config.py +++ b/litellm/a2a_protocol/providers/langflow/config.py @@ -6,6 +6,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, ) from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.interactions.agents.utils import merge_agent_headers from litellm.llms.langflow.a2a import merge_a2a_session_into_litellm_params @@ -28,11 +29,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) ) + forwarded_headers = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) return await A2ACompletionBridgeHandler.handle_non_streaming( request_id=request_id, params=params, litellm_params=litellm_params, api_base=api_base, + agent_extra_headers=forwarded_headers, _skip_a2a_provider_routing=True, ) @@ -51,11 +57,16 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): litellm_params = merge_a2a_session_into_litellm_params( litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) ) + forwarded_headers = merge_agent_headers( + dynamic_headers=kwargs.get("agent_extra_headers"), + static_headers=kwargs.get("agent_static_headers"), + ) async for chunk in A2ACompletionBridgeHandler.handle_streaming( request_id=request_id, params=params, litellm_params=litellm_params, api_base=api_base, + agent_extra_headers=forwarded_headers, _skip_a2a_provider_routing=True, ): yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index fe9d1ac3c15..0af22578ccc 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -210,12 +210,25 @@ class WatsonxOrchestrateHandler: client: AsyncHTTPHandler, max_attempts: int = _MAX_POLL_ATTEMPTS, interval_s: float = _POLL_INTERVAL_S, + timeout: float | None = None, ) -> _WXORun: url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}" + deadline: float | None = time.monotonic() + max(timeout, 0) if timeout is not None else None for attempt in range(max_attempts): - await asyncio.sleep(interval_s) - response = await client.get(url, headers=auth_headers) + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s") + if interval_s > 0: + await asyncio.sleep(min(interval_s, remaining)) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise asyncio.TimeoutError(f"WXO run '{run_id}' exceeded timeout of {timeout}s") + response = await asyncio.wait_for(client.get(url, headers=auth_headers), timeout=remaining) + else: + await asyncio.sleep(interval_s) + response = await client.get(url, headers=auth_headers) response.raise_for_status() result = WatsonxOrchestrateHandler._run_body(response) status = result.get("status", "") @@ -233,6 +246,7 @@ class WatsonxOrchestrateHandler: base_url: str, auth_headers: dict[str, str], client: AsyncHTTPHandler, + timeout: float | None = None, ) -> _WXORun: status = run_data.get("status", "") if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: @@ -244,6 +258,7 @@ class WatsonxOrchestrateHandler: run_id=run_id, auth_headers=auth_headers, client=client, + timeout=timeout, ) status = run_data.get("status", "") @@ -341,6 +356,7 @@ class WatsonxOrchestrateHandler: base_url=base_url, auth_headers=auth_headers, client=client, + timeout=timeout, ) response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data) @@ -418,6 +434,7 @@ class WatsonxOrchestrateHandler: base_url=base_url, auth_headers=auth_headers, client=client, + timeout=timeout, ) accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) else: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1e416a21e84..7cfaddc0b0e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1180,6 +1180,15 @@ class CustomStreamWrapper: completion_obj: dict[str, Any], ) -> _ProviderChunkResult: response_obj: dict[str, Any] = {} + if isinstance(chunk, ModelResponseStream) and self.custom_llm_provider == "a2a": + model_response = chunk + model_response.model = self.model + finish_reasons = [getattr(choice, "finish_reason", None) for choice in chunk.choices] + if finish_reasons and all(isinstance(reason, str) and reason for reason in finish_reasons): + self.received_finish_reason = finish_reasons[0] + self.sent_last_chunk = True + return _ProviderChunkEarlyReturn(model_response) + if ( isinstance(chunk, ModelResponseStream) and self.custom_llm_provider is not None @@ -1219,6 +1228,9 @@ class CustomStreamWrapper: raise StopIteration anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] + chunk_index = anthropic_response_obj.get("index") + if isinstance(chunk_index, int): + model_response.choices[0].index = chunk_index if anthropic_response_obj["is_finished"]: self.received_finish_reason = anthropic_response_obj["finish_reason"] diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index ef36296fff5..c2b3750bada 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -850,17 +850,20 @@ async def invoke_agent_a2a( agent_extra_headers=agent_extra_headers, ) + post_call_succeeded = False try: response = await proxy_logging_obj.post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, response=response, ) + post_call_succeeded = True finally: _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is not None: logging_obj._enqueue_deferred_logging = None - _enqueue_fn(response) + if post_call_succeeded: + _enqueue_fn(response) response_dict: Final[dict[str, Any]] = ( response.model_dump(mode="json", exclude_none=True) diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 36557d8bfeb..3e14d31f008 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -459,14 +459,17 @@ def _get_agent_dynamic_headers( return dynamic_headers -def _get_agent_identity_headers(user_api_key_dict: UserAPIKeyAuth | None) -> dict[str, str]: - if user_api_key_dict is None: - return {} +def _get_agent_identity_headers( + user_api_key_dict: UserAPIKeyAuth | None, + trace_id: object | None = None, +) -> dict[str, str]: headers: dict[str, str] = {} - if user_api_key_dict.user_id: + if user_api_key_dict is not None and user_api_key_dict.user_id: headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id - if user_api_key_dict.team_id: + if user_api_key_dict is not None and user_api_key_dict.team_id: headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id + if trace_id: + headers["X-LiteLLM-Trace-Id"] = str(trace_id) return headers @@ -585,7 +588,7 @@ async def route_a2a_agent_request( ) registered_static_headers = merge_agent_headers( dynamic_headers=registered_static_headers, - static_headers=_get_agent_identity_headers(user_api_key_dict), + static_headers=_get_agent_identity_headers(user_api_key_dict, data.get("litellm_trace_id")), ) if ( registered_provider diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 035bc770aed..c6f5cfbe5db 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1890,6 +1890,13 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, trust_client_model_info=False, ) + if isinstance(self.data.get("model"), str) and self.data["model"].startswith("a2a/"): + self.data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=self.data, + call_type=route_type, + guardrails_only=True, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in From 22952e4debfaae74127be497f0dd18ed935829df Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:48:06 +0900 Subject: [PATCH 16/20] fix: close final A2A review gaps --- .../litellm_completion_bridge/handler.py | 1 + .../proxy/agent_endpoints/a2a_endpoints.py | 10 ++++++++++ .../test_completion_bridge_streaming.py | 20 +++++++++++++++++++ 3 files changed, 31 insertions(+) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index 3cd2e9329dc..58e4366ad80 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -37,6 +37,7 @@ _AGENT_ONLY_PARAMS: Final = frozenset( "agent_id", "agent_card_params", A2A_USER_API_KEY_HASH_PARAM, + "databricks_oauth", } ) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c2b3750bada..e113be3ce7f 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -858,6 +858,16 @@ async def invoke_agent_a2a( response=response, ) post_call_succeeded = True + except HTTPException as e: + try: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=data, + ) + except Exception: + pass + raise finally: _enqueue_fn: Final = getattr(logging_obj, "_enqueue_deferred_logging", None) if _enqueue_fn is not None: diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 7fe0a6ced75..0495c0e4913 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -275,6 +275,26 @@ def test_build_completion_params_keeps_bridge_routing_fields(): assert params["stream"] is True +def test_build_completion_params_drops_proxy_only_databricks_oauth(): + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + + params = A2ACompletionBridgeHandler._build_completion_params( + params={"message": {"role": "user", "parts": []}}, + litellm_params={ + "custom_llm_provider": "databricks", + "model": "agent", + "databricks_oauth": {"client_id": "id"}, + }, + api_base="https://configured.example", + agent_extra_headers=None, + stream=False, + ) + + assert "databricks_oauth" not in params + + @pytest.mark.asyncio async def test_handle_streaming_accumulates_logprobs_and_provider_metadata(): from litellm.a2a_protocol.litellm_completion_bridge.handler import ( From f6c302ba5c6b61ac1609cdd10f78ef8225150b5e Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:31:49 +0900 Subject: [PATCH 17/20] fix: close A2A review gaps --- .../transformation.py | 1 + .../pydantic_ai_agents/transformation.py | 28 ++++++++- litellm/proxy/common_request_processing.py | 7 --- .../test_pydantic_ai_agent_transformation.py | 30 ++++++++- .../test_completion_bridge_streaming.py | 20 ++++++ .../proxy/test_common_request_processing.py | 62 +++++++++++++++++++ 6 files changed, 138 insertions(+), 10 deletions(-) diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 644ea9ab75f..b74cc0fa11c 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -217,6 +217,7 @@ class A2ACompletionBridgeTransformation: "provider_specific_fields", "reasoning_content", "reasoning_items", + "refusal", "thinking_blocks", ): value = message_fields.get(field) diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 024e8c179c2..5d14ed22b94 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,6 +6,7 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio +import time from collections.abc import AsyncIterator, Mapping, Sequence from typing import Any, Final, Protocol, cast, runtime_checkable from uuid import uuid4 @@ -101,6 +102,7 @@ class PydanticAITransformation: max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, + timeout: float | None = None, ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,11 +114,20 @@ class PydanticAITransformation: request_id: JSON-RPC request ID max_attempts: Maximum polling attempts poll_interval: Seconds between poll attempts + timeout: Total polling timeout in seconds Returns: Completed task response """ + deadline: Final[float | None] = ( + time.monotonic() + max(timeout, 0.0) if timeout is not None else None + ) for attempt in range(max_attempts): + remaining: float | None = ( + deadline - time.monotonic() if deadline is not None else None + ) + if remaining is not None and remaining <= 0: + break poll_request = { "jsonrpc": "2.0", "id": f"{request_id}-poll-{attempt}", @@ -131,6 +142,7 @@ class PydanticAITransformation: **(agent_extra_headers or {}), "Content-Type": "application/json", }, + timeout=remaining, ) response.raise_for_status() poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) @@ -146,9 +158,16 @@ class PydanticAITransformation: elif state in ("failed", "canceled"): raise Exception(f"Task {task_id} ended with state: {state}") - await asyncio.sleep(poll_interval) + if deadline is None: + await asyncio.sleep(poll_interval) + continue + remaining = deadline - time.monotonic() + if remaining <= 0: + break + await asyncio.sleep(min(poll_interval, remaining)) - raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds") + timeout_description = timeout if timeout is not None else max_attempts * poll_interval + raise TimeoutError(f"Task {task_id} did not complete within {timeout_description} seconds") @staticmethod async def _send_and_poll_raw( @@ -203,6 +222,9 @@ class PydanticAITransformation: llm_provider=cast(Any, "pydantic_ai_agent"), params={"timeout": timeout}, ) + deadline: Final[float | None] = ( + time.monotonic() + max(timeout, 0.0) if timeout is not None else None + ) response: Final = await client.post( endpoint, json=a2a_request, @@ -210,6 +232,7 @@ class PydanticAITransformation: **(agent_extra_headers or {}), "Content-Type": "application/json", }, + timeout=timeout, ) response.raise_for_status() response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) @@ -230,6 +253,7 @@ class PydanticAITransformation: task_id=task_id, request_id=request_id, agent_extra_headers=agent_extra_headers, + timeout=(max(deadline - time.monotonic(), 0.0) if deadline is not None else None), ) verbose_logger.info("Pydantic AI: Received completed response for request_id=%s", request_id) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c6f5cfbe5db..035bc770aed 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1890,13 +1890,6 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, trust_client_model_info=False, ) - if isinstance(self.data.get("model"), str) and self.data["model"].startswith("a2a/"): - self.data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, - data=self.data, - call_type=route_type, - guardrails_only=True, - ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py index c5626afa954..22679d02f06 100644 --- a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py @@ -4,10 +4,10 @@ Tests for Pydantic AI agents transformation. Tests the helper functions and response transformation without making real API calls. """ +from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( PydanticAITransformation, ) @@ -94,3 +94,31 @@ class TestPydanticAITransformation: assert result["result"]["kind"] == "message" assert result["result"]["role"] == "agent" assert result["result"]["parts"][0]["text"] == "The answer is 4." + + +@pytest.mark.asyncio +async def test_poll_for_completion_honors_request_timeout(): + client = MagicMock() + response = MagicMock() + response.json.return_value = { + "result": {"status": {"state": "working"}, "id": "task-1"} + } + response.raise_for_status.return_value = None + client.post = AsyncMock(return_value=response) + + with patch( + "litellm.a2a_protocol.providers.pydantic_ai_agents.transformation.time.monotonic", + side_effect=[100.0, 100.0, 100.02], + ): + with pytest.raises(TimeoutError, match="0.01 seconds"): + await PydanticAITransformation._poll_for_completion( + client=client, + endpoint="http://example.test", + task_id="task-1", + request_id="req-1", + poll_interval=1.0, + timeout=0.01, + ) + + client.post.assert_awaited_once() + assert client.post.await_args.kwargs["timeout"] == pytest.approx(0.01) diff --git a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py index 0495c0e4913..7314ec43b46 100644 --- a/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py +++ b/tests/test_litellm/a2a_protocol/test_completion_bridge_streaming.py @@ -508,6 +508,26 @@ def test_response_transform_preserves_audio_and_logprobs(): assert transformed["result"]["logprobs"] == {"content": []} +def test_response_transform_preserves_refusal(): + from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( + A2ACompletionBridgeTransformation, + ) + + message = MagicMock() + message.model_dump.return_value = { + "content": None, + "refusal": "I cannot help with that request.", + } + choice = MagicMock(message=message) + choice.model_dump.return_value = {"finish_reason": "stop"} + response = MagicMock(choices=[choice], usage=None) + + transformed = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(response) + + assert transformed["result"]["refusal"] == "I cannot help with that request." + assert transformed["result"]["parts"] == [{"kind": "text", "text": ""}] + + @pytest.mark.asyncio async def test_handle_streaming_forwards_api_key(): """Test that handle_streaming forwards api_key from litellm_params to acompletion.""" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58714a5e319..5d552d0a673 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6251,6 +6251,68 @@ class TestPerRequestModelGroupAlias: assert merged_for == ["group-b"] + @pytest.mark.asyncio + async def test_a2a_reroute_does_not_repeat_pre_call_hook(self, monkeypatch): + processing_obj = ProxyBaseLLMRequestProcessing( + data={"model": "source", "messages": [{"role": "user", "content": "hello"}]} + ) + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + async def mock_add_litellm_data_to_request(*args, **kwargs): + return kwargs.get("data", {}) + + hook_modes = [] + + async def reroute_once(user_api_key_dict, data, call_type, guardrails_only=False): + hook_modes.append(guardrails_only) + if len(hook_modes) == 1: + data["model"] = "a2a/agent" + return data + + async def passthrough(data, user_api_key_dict): + return data + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=reroute_once) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + mock_add_litellm_data_to_request, + ) + monkeypatch.setattr( + litellm.utils, + "function_setup", + lambda original_function, rules_obj, start_time, **kwargs: (MagicMock(), kwargs), + ) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "_check_and_merge_model_level_guardrails", + lambda data, llm_router, trust_client_model_info=True: data, + ) + monkeypatch.setattr( + "litellm.proxy.agent_endpoints.a2a_routing.authorize_a2a_agent_before_hooks", + passthrough, + ) + monkeypatch.setattr( + "litellm.proxy.agent_endpoints.a2a_routing.merge_a2a_agent_guardrails_before_hooks", + passthrough, + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="hash"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=None, + route_type="acompletion", + llm_router=None, + ) + + assert returned_data["model"] == "a2a/agent" + assert hook_modes == [False] + + class TestInjectCostIntoUsageDict: @staticmethod def _expected_cost(model, prompt_tokens, completion_tokens): From 3d6df166f4ea34244ce8dbfeae37c144937ce7a6 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:18:50 +0900 Subject: [PATCH 18/20] fix: close final A2A review gaps --- .../providers/watsonx_orchestrate/handler.py | 41 +++++++++++-------- .../litellm_core_utils/streaming_handler.py | 3 ++ litellm/llms/a2a/chat/streaming_iterator.py | 17 ++++++++ litellm/proxy/agent_endpoints/a2a_routing.py | 20 ++++++++- litellm/proxy/common_request_processing.py | 7 ++++ litellm/types/utils.py | 1 + ...test_watsonx_orchestrate_transformation.py | 12 ++++++ .../a2a/chat/test_a2a_streaming_iterator.py | 12 ++++++ .../proxy/test_common_request_processing.py | 4 +- .../proxy/test_route_a2a_models.py | 25 +++++++++++ 10 files changed, 121 insertions(+), 21 deletions(-) diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 0af22578ccc..f4eb65715fa 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -268,23 +268,28 @@ class WatsonxOrchestrateHandler: return run_data @staticmethod - async def _accumulate_wxo_sse_text(response: Any) -> str: - source: Final[_WXOView] = {"sse_source": response} - accumulated_text = "" - async for line in source["sse_source"].aiter_lines(): - if not line.startswith("data:"): - continue - data_str = line[5:].strip() - if not data_str or data_str == "[DONE]": - continue - try: - event = WatsonxOrchestrateHandler._decode_run_event(data_str) - except json.JSONDecodeError: - continue - chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event) - if chunk_text: - accumulated_text += chunk_text - return accumulated_text + async def _accumulate_wxo_sse_text(response: Any, timeout: float | None = None) -> str: + async def _collect() -> str: + source: Final[_WXOView] = {"sse_source": response} + accumulated_text = "" + async for line in source["sse_source"].aiter_lines(): + if not line.startswith("data:"): + continue + data_str = line[5:].strip() + if not data_str or data_str == "[DONE]": + continue + try: + event = WatsonxOrchestrateHandler._decode_run_event(data_str) + except json.JSONDecodeError: + continue + chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event) + if chunk_text: + accumulated_text += chunk_text + return accumulated_text + + if timeout is None: + return await _collect() + return await asyncio.wait_for(_collect(), timeout=max(timeout, 0)) @staticmethod def _extract_litellm_params(litellm_params: WXOLitellmParams) -> WXORequestParams: @@ -438,7 +443,7 @@ class WatsonxOrchestrateHandler: ) accumulated_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) else: - accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) + accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response, timeout=timeout) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( text=accumulated_text, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7cfaddc0b0e..bbb3856208c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1227,6 +1227,9 @@ class CustomStreamWrapper: if not _chunk_has_content and (not isinstance(chunk, dict) or "provider_specific_fields" not in chunk): raise StopIteration anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk) + chunk_id = anthropic_response_obj.get("id") + if isinstance(chunk_id, str) and chunk_id.strip(): + model_response = self.set_model_id(chunk_id, model_response) completion_obj["content"] = anthropic_response_obj["text"] chunk_index = anthropic_response_obj.get("index") if isinstance(chunk_index, int): diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 099a9f67047..c4f8d1b2edd 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -4,6 +4,7 @@ A2A Streaming Response Iterator from collections.abc import Mapping from typing import Any, Final +from uuid import uuid4 import litellm from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -33,6 +34,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): json_mode=json_mode, ) self.model = model + self.response_id: str | None = None def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream: """ @@ -69,6 +71,16 @@ class A2AModelResponseIterator(BaseModelResponseIterator): raise A2AError(status_code=500, message=f"A2A error: {error_message}") try: + if self.response_id is None: + raw_response_id = chunk.get("id") + raw_result = chunk.get("result") + if not isinstance(raw_response_id, str) and isinstance(raw_result, Mapping): + raw_response_id = raw_result.get("id") + self.response_id = ( + raw_response_id + if isinstance(raw_response_id, str) and raw_response_id.strip() + else f"chatcmpl-{uuid4().hex}" + ) # Extract text from A2A response result: Final = chunk.get("result", {}) chunk_index = 0 @@ -181,6 +193,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): if streaming_choices: return ModelResponseStream( choices=streaming_choices, + id=self.response_id, usage=usage, provider_specific_fields=provider_fields or None, ) @@ -188,6 +201,7 @@ class A2AModelResponseIterator(BaseModelResponseIterator): # Return generic streaming chunk return GenericStreamingChunk( text=text, + id=self.response_id, is_finished=bool(finish_reason or tool_calls), finish_reason=finish_reason or ("tool_calls" if tool_calls else ""), usage=usage, @@ -197,8 +211,11 @@ class A2AModelResponseIterator(BaseModelResponseIterator): ) except Exception: # Return empty chunk on parse error + if self.response_id is None: + self.response_id = f"chatcmpl-{uuid4().hex}" return GenericStreamingChunk( text="", + id=self.response_id, is_finished=False, finish_reason="", usage=None, diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index 3e14d31f008..7a5e06330da 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -260,6 +260,7 @@ async def _route_registered_provider( "provider_specific_fields", "reasoning_content", "reasoning_items", + "refusal", "thinking_blocks", ): value = message_payload.get(field) @@ -319,7 +320,24 @@ async def _route_registered_provider( service_tier=response.get("service_tier") if isinstance(response.get("service_tier"), str) else None, ) raw_usage: Final = response.get("usage") - usage: Final = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage + usage = litellm.Usage(**raw_usage) if isinstance(raw_usage, Mapping) else raw_usage + if usage is None and native_provider: + try: + from litellm.utils import token_counter + + prompt_tokens: Final = token_counter(model="gpt-3.5-turbo", messages=messages) + completion_tokens: Final = token_counter( + model="gpt-3.5-turbo", + text=extract_text_from_a2a_response(response), + count_response_tokens=True, + ) + usage = litellm.Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + except Exception: # noqa: BLE001 - token estimation must not fail the response + pass if usage is not None: model_response.usage = usage if isinstance(logging_obj, Logging): diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 035bc770aed..c6f5cfbe5db 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1890,6 +1890,13 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, trust_client_model_info=False, ) + if isinstance(self.data.get("model"), str) and self.data["model"].startswith("a2a/"): + self.data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=self.data, + call_type=route_type, + guardrails_only=True, + ) # Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may # have mutated `self.data` in place, and the audit-trail snapshot taken in diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 61d7aca0430..e3914511574 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -317,6 +317,7 @@ class ModelInfo(ModelInfoBase, total=False): class GenericStreamingChunk(TypedDict, total=False): text: Required[str] + id: str tool_use: ChatCompletionToolCallChunk | list[ChatCompletionToolCallChunk] | None is_finished: Required[bool] finish_reason: Required[str] diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index fa8b11ea182..7dadfb9c063 100644 --- a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -48,6 +48,12 @@ class _SSELines: yield line +class _HangingSSELines: + async def aiter_lines(self): + await asyncio.Event().wait() + yield "" + + class _InvalidJsonStreamResponse: headers = {"content-type": "application/json"} @@ -233,6 +239,12 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): ) +@pytest.mark.asyncio +async def test_accumulate_wxo_sse_text_respects_timeout(): + with pytest.raises(asyncio.TimeoutError): + await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(_HangingSSELines(), timeout=0.001) + + @pytest.mark.asyncio async def test_short_lived_tokens_are_not_served_from_cache(): client = _ShortTtlTokenClient() diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py index 8518cd3d2f7..1643604ad84 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_streaming_iterator.py @@ -28,6 +28,18 @@ async def test_async_iterator_accepts_decoded_a2a_events(): assert chunk["text"] == "Hello" +def test_chunk_parser_reuses_response_id_for_idless_artifacts(): + iterator = A2AModelResponseIterator(streaming_response=[], sync_stream=False) + first = iterator.chunk_parser( + {"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "one"}]}}} + ) + second = iterator.chunk_parser( + {"result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "two"}]}}} + ) + + assert first["id"] == second["id"] + + @pytest.mark.asyncio async def test_async_iterator_ignores_status_message_text(): async def _events(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 5d552d0a673..b755255fb48 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6252,7 +6252,7 @@ class TestPerRequestModelGroupAlias: @pytest.mark.asyncio - async def test_a2a_reroute_does_not_repeat_pre_call_hook(self, monkeypatch): + async def test_a2a_reroute_runs_target_guardrails(self, monkeypatch): processing_obj = ProxyBaseLLMRequestProcessing( data={"model": "source", "messages": [{"role": "user", "content": "hello"}]} ) @@ -6310,7 +6310,7 @@ class TestPerRequestModelGroupAlias: ) assert returned_data["model"] == "a2a/agent" - assert hook_modes == [False] + assert hook_modes == [False, True] class TestInjectCostIntoUsageDict: diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 05bc118907c..f0d0ff084d6 100644 --- a/tests/test_litellm/proxy/test_route_a2a_models.py +++ b/tests/test_litellm/proxy/test_route_a2a_models.py @@ -116,6 +116,7 @@ async def test_route_a2a_model_uses_registered_provider(): "kind": "message", "role": "agent", "parts": [{"kind": "text", "text": "Hello back"}], + "refusal": "I cannot complete that request.", "messageId": "message-id", }, } @@ -143,6 +144,7 @@ async def test_route_a2a_model_uses_registered_provider(): bridge.assert_awaited_once() generic_completion.assert_not_called() assert response.choices[0].message.content == "Hello back" + assert response.choices[0].message.refusal == "I cannot complete that request." bridge_kwargs = bridge.await_args.kwargs assert bridge_kwargs["litellm_params"]["max_tokens"] == 32 assert bridge_kwargs["litellm_params"]["temperature"] == 0.2 @@ -677,6 +679,29 @@ async def test_registered_provider_logging_uses_provider_model_for_builtin_prici assert logging_obj.model_call_details["litellm_params"]["model"] == "gpt-4o" +@pytest.mark.asyncio +async def test_native_registered_provider_estimates_usage_when_missing(monkeypatch): + response = {"result": {"message": {"parts": [{"kind": "text", "text": "hello"}]}}} + counter = Mock(side_effect=[3, 2]) + monkeypatch.setattr("litellm.utils.token_counter", counter) + + with patch( + "litellm.a2a_protocol.litellm_completion_bridge.handler.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value=response), + ): + result = await _route_registered_provider( + data={"messages": [{"role": "user", "content": "hello"}]}, + model_name="a2a/agent", + api_base="https://provider.example", + litellm_params={"model": "agent", "custom_llm_provider": "pydantic_ai_agents"}, + static_headers=None, + ) + + assert result.usage.prompt_tokens == 3 + assert result.usage.completion_tokens == 2 + assert result.usage.total_tokens == 5 + + @pytest.mark.asyncio async def test_route_non_a2a_model_raises_error_if_not_in_router(): """Test that non-a2a models that aren't in router raise an error""" From 711d74af3b0399b8e426f33f3758dabc064703d1 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:33:59 +0900 Subject: [PATCH 19/20] style: format pydantic_ai_agents transformation --- .../providers/pydantic_ai_agents/transformation.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 5d14ed22b94..86d9b4d2eae 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -119,13 +119,9 @@ class PydanticAITransformation: Returns: Completed task response """ - deadline: Final[float | None] = ( - time.monotonic() + max(timeout, 0.0) if timeout is not None else None - ) + deadline: Final[float | None] = time.monotonic() + max(timeout, 0.0) if timeout is not None else None for attempt in range(max_attempts): - remaining: float | None = ( - deadline - time.monotonic() if deadline is not None else None - ) + remaining: float | None = deadline - time.monotonic() if deadline is not None else None if remaining is not None and remaining <= 0: break poll_request = { @@ -222,9 +218,7 @@ class PydanticAITransformation: llm_provider=cast(Any, "pydantic_ai_agent"), params={"timeout": timeout}, ) - deadline: Final[float | None] = ( - time.monotonic() + max(timeout, 0.0) if timeout is not None else None - ) + deadline: Final[float | None] = time.monotonic() + max(timeout, 0.0) if timeout is not None else None response: Final = await client.post( endpoint, json=a2a_request, From 0cbf8802c149642c5e8487aebda7fb5e25638e80 Mon Sep 17 00:00:00 2001 From: aiedwardyi <41576951+aiedwardyi@users.noreply.github.com> Date: Fri, 28 Aug 2026 13:47:18 +0900 Subject: [PATCH 20/20] test: fix a2a guardrail stub arity and an unescaped raises match merge_a2a_agent_guardrails_before_hooks takes only data, so the shared passthrough stub failed the reroute test; the timeout assertion needed a raw pattern for RUF043. --- .../test_pydantic_ai_agent_transformation.py | 2 +- tests/test_litellm/proxy/test_common_request_processing.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py b/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py index 22679d02f06..8b40b64987f 100644 --- a/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py +++ b/tests/test_litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py @@ -110,7 +110,7 @@ async def test_poll_for_completion_honors_request_timeout(): "litellm.a2a_protocol.providers.pydantic_ai_agents.transformation.time.monotonic", side_effect=[100.0, 100.0, 100.02], ): - with pytest.raises(TimeoutError, match="0.01 seconds"): + with pytest.raises(TimeoutError, match=r"0\.01 seconds"): await PydanticAITransformation._poll_for_completion( client=client, endpoint="http://example.test", diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index b755255fb48..3b2beac0947 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -6273,6 +6273,9 @@ class TestPerRequestModelGroupAlias: async def passthrough(data, user_api_key_dict): return data + async def passthrough_data_only(data): + return data + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=reroute_once) monkeypatch.setattr( @@ -6296,7 +6299,7 @@ class TestPerRequestModelGroupAlias: ) monkeypatch.setattr( "litellm.proxy.agent_endpoints.a2a_routing.merge_a2a_agent_guardrails_before_hooks", - passthrough, + passthrough_data_only, ) returned_data, _ = await processing_obj.common_processing_pre_call_logic(