diff --git a/litellm/llms/a2a/chat/streaming_iterator.py b/litellm/llms/a2a/chat/streaming_iterator.py index 1983c18a6b3..4cf7237afde 100644 --- a/litellm/llms/a2a/chat/streaming_iterator.py +++ b/litellm/llms/a2a/chat/streaming_iterator.py @@ -2,6 +2,7 @@ A2A Streaming Response Iterator """ +from itertools import accumulate from typing import Final from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -10,6 +11,39 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream from ..common_utils import extract_text_from_a2a_response +def _ignoring_whitespace(text: str) -> str: + """ + Comparison key for snapshot detection. + + A2A joins a multi-part message's text parts with spaces, so the same content arrives + with different whitespace depending on how the server chunked it: two delta events + ("Hello", "world") accumulate to "Helloworld" while a single two-part snapshot of the + same content renders "Hello world". Comparing without whitespace makes them equal. + """ + return "".join(text.split()) + + +def _index_after(text: str, non_space_count: int) -> int: + """Index in `text` just past its first `non_space_count` non-whitespace characters.""" + if non_space_count <= 0: + return 0 + running: Final = accumulate(0 if char.isspace() else 1 for char in text) + return next( + (index + 1 for index, total in enumerate(running) if total >= non_space_count), + len(text), + ) + + +def _trailing_whitespace(text: str) -> int: + """How many whitespace characters `text` ends with.""" + return len(text) - len(text.rstrip()) + + +def _leading_whitespace(text: str) -> int: + """How many whitespace characters `text` begins with.""" + return len(text) - len(text.lstrip()) + + class A2AModelResponseIterator(BaseModelResponseIterator): """ Iterator for parsing A2A streaming responses. @@ -30,6 +64,8 @@ class A2AModelResponseIterator(BaseModelResponseIterator): json_mode=json_mode, ) self.model = model + # Text already emitted downstream, used to collapse cumulative snapshots. + self._emitted_text: str = "" def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream: """ @@ -57,8 +93,8 @@ class A2AModelResponseIterator(BaseModelResponseIterator): } """ try: - # Extract text from A2A response - text: Final = extract_text_from_a2a_response(chunk) + # Extract text from A2A response, then reduce it to what is actually new. + text: Final = self._to_incremental_text(extract_text_from_a2a_response(chunk)) # Determine finish reason finish_reason: Final = self._get_finish_reason(chunk) @@ -83,6 +119,60 @@ class A2AModelResponseIterator(BaseModelResponseIterator): tool_use=None, ) + def _to_incremental_text(self, text: str) -> str: + """ + Reduce an A2A event's text to the portion not yet emitted. + + A2A servers interleave true deltas with cumulative snapshots of the whole reply: a + terminal non-partial ``status-update`` and an ``artifact-update`` carrying + ``append: false`` both repeat everything produced so far. Forwarding those verbatim + makes the client render the reply two or three times over, so emit only new text. + + Handles both streaming styles: servers that send deltas ("O", "K") and servers that + send growing snapshots ("O", "OK") collapse to the same output. + + A snapshot need not repeat the whole stream. An agent that delegates reports + sub-agent progress into the same task before answering, so its closing events carry + only its own answer, extending nothing. Text that merely repeats the tail of what has + been delivered is therefore read as a snapshot as well. + + A2A marks no event as delta-or-snapshot, so an event whose text equals everything + emitted so far is necessarily ambiguous. It is read as a snapshot, because servers + repeating the whole reply at the end of a stream are common while a delta that + exactly reproduces the accumulated text is not, and duplicating a reply is far + worse for a reader than dropping one repeated fragment. + """ + if not text: + return "" + + emitted_key: Final = _ignoring_whitespace(self._emitted_text) + text_key: Final = _ignoring_whitespace(text) + + if emitted_key and text_key.startswith(emitted_key): + # A cumulative snapshot: emit only its tail, which is empty when the snapshot + # just repeats everything sent so far. The offset lands just past the last + # matched non-whitespace character, so any whitespace already delivered at the + # end of the emitted text would otherwise be sent a second time. Drop only that + # overlap, so a snapshot introducing further whitespace (a paragraph break, say) + # keeps it. + tail: Final = text[_index_after(text, len(emitted_key)) :] + already_sent: Final = min(_trailing_whitespace(self._emitted_text), _leading_whitespace(tail)) + suffix: Final = tail[already_sent:] + self._emitted_text += suffix + return suffix + + if emitted_key and text_key and emitted_key.endswith(text_key): + # A terminal snapshot that repeats only the agent's own answer instead of the + # whole stream. A delegating agent reports sub-agent progress into the same task + # before answering, so its closing ``status-update`` and ``artifact-update`` carry + # just the answer: that extends nothing already emitted and so escapes the + # prefix test above, and would be forwarded twice more. Everything such a + # snapshot carries has already been delivered, so emit nothing. + return "" + + self._emitted_text += text + return text + def _get_finish_reason(self, chunk: dict) -> str | None: """Extract finish reason from A2A chunk""" result: Final = chunk.get("result", {}) diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 57eadfe36d2..3dd56ee3aa1 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -62,6 +62,17 @@ def convert_messages_to_prompt(messages: list[AllMessageValues]) -> str: return "\n".join(conversation_parts) +def _is_user_authored(message: object) -> bool: + """ + Whether an A2A message was authored by the caller rather than the agent. + + A2A ``Message.role`` is either ``"user"`` or ``"agent"``. Agents commonly echo the + inbound message back on the first ``status-update`` (``state: "submitted"``), and that + text must never surface as assistant output in a chat completion. + """ + return isinstance(message, dict) and message.get("role") == "user" + + def extract_text_from_a2a_message(message: dict[str, Any], depth: int = 0, max_depth: int = 10) -> str: """ Extract text content from A2A message parts. @@ -114,13 +125,15 @@ def extract_text_from_a2a_response(response_dict: Mapping[str, object], max_dept # 4. Task with status message: {"result": {"kind": "task", "status": {"message": {"parts": [...]}}}} # 5. Streaming artifact-update: {"result": {"kind": "artifact-update", "artifact": {"parts": [...]}}} - # Check if result itself has parts (direct message) - if "parts" in result: + # Check if result itself has parts (direct message). A caller-authored message is + # skipped rather than returned, so extraction falls through to the agent's own + # output further down (artifacts, typically) instead of yielding nothing. + if "parts" in result and not _is_user_authored(result): return extract_text_from_a2a_message(result, depth=0, max_depth=max_depth) # Check for nested message message: Final = result.get("message") - if message: + if message and not _is_user_authored(message): return extract_text_from_a2a_message(message, depth=0, max_depth=max_depth) # Check for streaming artifact-update (singular artifact) @@ -132,7 +145,7 @@ def extract_text_from_a2a_response(response_dict: Mapping[str, object], max_dept status: Final = result.get("status", {}) if isinstance(status, dict): status_message: Final = status.get("message") - if status_message: + if status_message and not _is_user_authored(status_message): return extract_text_from_a2a_message(status_message, depth=0, max_depth=max_depth) # Handle task result with artifacts (plural, array) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 20b4708c193..5a80292ffe8 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -538,6 +538,21 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr return getattr(litellm, f"{route_type}")(**data) elif llm_router is not None: _raise_if_model_fully_blocked(llm_router=llm_router, model_name=data.get("model"), team_id=team_id) + + # A2A agents are addressed by an unambiguous "a2a/" model prefix and are never backed + # by a router deployment, so they must be resolved before any router-based branch. + # Otherwise the branches below swallow the request and it fails with + # "no healthy deployments": `map_team_model` claims it for team-scoped keys, and the + # wildcard/default-deployment fallback claims it whenever a pattern model group exists. + if route_type == "acompletion" and _is_a2a_agent_model(data.get("model", "")): + from litellm.proxy.agent_endpoints.a2a_routing import ( + route_a2a_agent_request, + ) + + a2a_result: Final = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict) + if a2a_result is not None: + return a2a_result + # Evals API: always route to litellm directly (not through router) # But extract model credentials if a model is provided if route_type in [ @@ -698,15 +713,6 @@ async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited pr except Exception: # If router fails (e.g., model not found in router), fall back to direct call return getattr(litellm, f"{route_type}")(**data) - elif _is_a2a_agent_model(data.get("model", "")): - from litellm.proxy.agent_endpoints.a2a_routing import ( - route_a2a_agent_request, - ) - - result: Final = await route_a2a_agent_request(data, route_type, user_api_key_dict=user_api_key_dict) - if result is not None: - return result - # Fall through to raise exception below if result is None elif user_model is not None or route_type == "allm_passthrough_route": return getattr(litellm, f"{route_type}")(**data) diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py index 2e11c68244c..9df586b45b1 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py @@ -1,7 +1,11 @@ -"""Tests for litellm/llms/a2a/chat/transformation.py response transform.""" +"""Tests for the A2A chat bridge: litellm/llms/a2a/chat/ response transform and streaming.""" from unittest.mock import MagicMock +import pytest + +from litellm.llms.a2a.chat.streaming_iterator import A2AModelResponseIterator +from litellm.llms.a2a.common_utils import extract_text_from_a2a_response from litellm.llms.a2a.chat.transformation import A2AConfig from litellm.types.utils import ModelResponse @@ -40,3 +44,225 @@ def test_transform_response_sets_usage(): assert result.usage.prompt_tokens > 0 assert result.usage.completion_tokens > 0 assert result.usage.total_tokens == (result.usage.prompt_tokens + result.usage.completion_tokens) + + +def _iterator() -> A2AModelResponseIterator: + return A2AModelResponseIterator(streaming_response=iter(()), sync_stream=True) + + +def _status_update( + *, + text: str | None = None, + role: str = "agent", + state: str = "working", + final: bool = False, +) -> dict: + status: dict = {"state": state} + if text is not None: + status["message"] = { + "kind": "message", + "role": role, + "parts": [{"kind": "text", "text": text}], + } + return { + "jsonrpc": "2.0", + "id": "1", + "result": {"kind": "status-update", "final": final, "status": status}, + } + + +def _artifact_update(*texts: str, append: bool = False) -> dict: + return { + "jsonrpc": "2.0", + "id": "1", + "result": { + "kind": "artifact-update", + "append": append, + "lastChunk": True, + "artifact": {"parts": [{"kind": "text", "text": text} for text in texts]}, + }, + } + + +# Event sequence captured from a real kagent A2A agent replying "OK": the submitted +# status-update echoes the caller's own message, then two true deltas are followed by two +# cumulative snapshots of the whole reply. +KAGENT_OK_STREAM = [ + _status_update(text="Reply with exactly: OK", role="user", state="submitted"), + _status_update(), + _status_update(text="O"), + _status_update(text="K"), + _status_update(text="OK"), + _artifact_update("OK"), + _status_update(state="completed", final=True), +] + + +def test_kagent_stream_yields_reply_exactly_once(): + """ + Regression: the caller's echoed message must not be emitted as assistant output, and + cumulative snapshots must not repeat the reply. + + Previously this stream rendered as "user: Reply with exactly: OKOKOKOK". + """ + iterator = _iterator() + assert "".join(iterator.chunk_parser(e)["text"] for e in KAGENT_OK_STREAM) == "OK" + + +def test_kagent_stream_finishes_on_completed_state(): + iterator = _iterator() + chunks = [iterator.chunk_parser(e) for e in KAGENT_OK_STREAM] + assert [c["finish_reason"] for c in chunks if c["is_finished"]] == ["stop"] + + +@pytest.mark.parametrize( + "texts, expected", + [ + pytest.param(["Hello", " world"], "Hello world", id="incremental_deltas"), + pytest.param(["O", "OK", "OKAY"], "OKAY", id="cumulative_snapshots"), + pytest.param(["O", "K", "OK"], "OK", id="deltas_then_final_snapshot"), + pytest.param(["O", "K", "OK", "OK"], "OK", id="deltas_then_repeated_snapshots"), + pytest.param(["OK", "OK"], "OK", id="one_delta_then_equal_snapshot"), + pytest.param(["Hello ", "Hello world"], "Hello world", id="snapshot_keeps_emitted_space_once"), + pytest.param(["OK\n", "OK\n"], "OK\n", id="snapshot_keeps_emitted_newline_once"), + pytest.param( + ["Hello ", "Hello \n\nworld"], + "Hello \n\nworld", + id="snapshot_keeps_its_own_new_whitespace", + ), + # Known limitation: A2A marks no event as delta-or-snapshot, so a delta that + # exactly reproduces the accumulated text is indistinguishable from a snapshot + # and collapses. Duplicating a whole reply is the worse failure of the two. + pytest.param(["a", "a", "a"], "a", id="identical_deltas_collapse"), + pytest.param(["Hello", "world", "Hello world"], "Helloworld", id="multipart_snapshot_respaced"), + pytest.param( + ["Hello", "world", "Hello world again"], + "Helloworld again", + id="multipart_snapshot_extends", + ), + pytest.param(["", "OK", ""], "OK", id="empty_events_ignored"), + # A whitespace-only delta carries no non-whitespace text, so it must not be + # mistaken for a snapshot repeating the tail of the stream. + pytest.param(["Hello", " "], "Hello ", id="whitespace_only_delta_kept"), + pytest.param(["chatter", "answer", "answer"], "chatteranswer", id="snapshot_repeats_only_tail"), + ], +) +def test_incremental_text_reduction(texts, expected): + """Delta-style and snapshot-style servers must collapse to the same output.""" + iterator = _iterator() + assert "".join(iterator._to_incremental_text(t) for t in texts) == expected + + +# A server that chunks its reply into separate delta events but sends the final artifact +# as one multi-part message: A2A joins those parts with a space, so the snapshot reads +# "Hello world" while the deltas accumulated to "Helloworld". +MULTIPART_SNAPSHOT_STREAM = [ + _status_update(text="Hello"), + _status_update(text="world"), + _artifact_update("Hello", "world"), + _status_update(state="completed", final=True), +] + + +def test_multipart_snapshot_is_not_re_emitted(): + """Regression: whitespace introduced by part-joining must not defeat snapshot detection.""" + iterator = _iterator() + rendered = "".join(iterator.chunk_parser(e)["text"] for e in MULTIPART_SNAPSHOT_STREAM) + assert rendered == "Helloworld" + + +# A one-token reply: a single delta followed by the terminal cumulative snapshot. This is +# the ordinary shape of a short A2A answer, so the snapshot must not be forwarded again. +SINGLE_DELTA_STREAM = [ + _status_update(text="Reply with exactly: OK", role="user", state="submitted"), + _status_update(text="OK"), + _artifact_update("OK"), + _status_update(state="completed", final=True), +] + + +def test_single_delta_then_snapshot_is_not_duplicated(): + """Regression: snapshot suppression must not depend on how many deltas preceded it.""" + iterator = _iterator() + assert "".join(iterator.chunk_parser(e)["text"] for e in SINGLE_DELTA_STREAM) == "OK" + + +def _task(*, status_role: str | None = None, status_text: str = "", artifact_text: str | None = None) -> dict: + result: dict = {"kind": "task"} + if status_role is not None: + result["status"] = { + "state": "completed", + "message": {"role": status_role, "parts": [{"kind": "text", "text": status_text}]}, + } + if artifact_text is not None: + result["artifacts"] = [{"parts": [{"kind": "text", "text": artifact_text}]}] + return {"jsonrpc": "2.0", "id": "1", "result": result} + + +@pytest.mark.parametrize( + "response, expected", + [ + # Regression: a task echoing the caller in status.message while carrying the real + # answer in artifacts must not come back empty. Skipping a caller-authored message + # has to fall through to the agent's own output, not short-circuit extraction. + pytest.param( + _task(status_role="user", status_text="check active alarms", artifact_text="THE ANSWER"), + "THE ANSWER", + id="user_echo_falls_through_to_artifacts", + ), + pytest.param( + _task(status_role="agent", status_text="agent status", artifact_text="THE ANSWER"), + "agent status", + id="agent_status_message_preferred", + ), + pytest.param(_task(artifact_text="THE ANSWER"), "THE ANSWER", id="artifacts_only"), + pytest.param(_task(status_role="user", status_text="echo"), "", id="user_echo_alone_is_empty"), + pytest.param( + {"result": {"kind": "message", "role": "user", "parts": [{"kind": "text", "text": "echo"}]}}, + "", + id="direct_user_message_is_empty", + ), + pytest.param( + {"result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": "hi"}]}}, + "hi", + id="direct_agent_message", + ), + ], +) +def test_extract_text_skips_caller_authored_messages(response, expected): + """Caller-authored parts are skipped, never returned, and never hide the agent's reply.""" + assert extract_text_from_a2a_response(response) == expected + + +# A delegating agent reports sub-agent progress into the same task before producing its own +# answer. Modelled on kagent's ADK executor, which enqueues a status-update for every ADK event +# and then re-sends the aggregated answer as BOTH a terminal status-update and a final artifact +# (kagent-adk/src/kagent/adk/_agent_executor.py: run loop, then task result publication). +ANSWER = "Diagnosis: optic degraded on emm001a-jnx-01." + +DELEGATING_AGENT_STREAM = [ + _status_update(text="Diagnose the optical alarm on emm001a-jnx-01.", role="user", state="submitted"), + _status_update(text="Calling telemetry-agent..."), + _status_update(text="telemetry-agent: no anomalies found."), + _status_update(text="Diagnosis: optic degraded"), + _status_update(text=" on emm001a-jnx-01."), + _status_update(text=ANSWER), + _artifact_update(ANSWER), + _status_update(state="completed", final=True), +] + + +def test_delegating_agent_answer_is_not_repeated_after_tool_chatter(): + """ + Regression: a terminal snapshot repeats only the agent's answer, not the progress text + streamed ahead of it, so snapshot suppression cannot depend on the snapshot extending + everything emitted so far. + + Without this, a delegating agent renders its answer three times: once from the deltas, + once from the terminal status-update and once from the final artifact. + """ + iterator = _iterator() + rendered = "".join(iterator.chunk_parser(e)["text"] for e in DELEGATING_AGENT_STREAM) + + assert rendered.count(ANSWER) == 1 + assert rendered == f"Calling telemetry-agent...telemetry-agent: no anomalies found.{ANSWER}" diff --git a/tests/test_litellm/proxy/test_route_a2a_models.py b/tests/test_litellm/proxy/test_route_a2a_models.py index 35308474949..d1e85d56817 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 @@ -180,3 +178,89 @@ async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_re assert call_kwargs["model"] == f"a2a/{agent_name}" assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com" prisma_client.db.litellm_agentstable.find_unique.assert_awaited() + + +def _router_without_a2a_deployment( + *, + team_model: str | None = None, + patterns: tuple[str, ...] = (), + default_deployment: dict | None = None, +) -> Mock: + """A router that serves no A2A deployment, optionally tripping one shadowing branch.""" + router = Mock() + router.model_names = ["gpt-4", "gpt-3.5-turbo"] + router.deployment_names = [] + router.has_model_id = Mock(return_value=False) + router.is_recognized_model = Mock(return_value=False) + router.get_routing_group = Mock(return_value=None) + router.model_group_alias = None + router.router_general_settings = Mock(pass_through_all_models=False) + router.default_deployment = default_deployment + router.pattern_router = Mock(patterns=list(patterns)) + router.map_team_model = Mock(return_value=team_model) + return router + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "router_kwargs, extra_data", + [ + pytest.param({}, {}, id="no_shadowing_branch"), + pytest.param( + {"team_model": "a2a/test-agent"}, + {"metadata": {"user_api_key_team_id": "team-123"}}, + id="team_scoped_key", + ), + pytest.param({"patterns": ("openrouter/*",)}, {}, id="wildcard_model_group"), + pytest.param({"default_deployment": {"model_name": "*"}}, {}, id="default_deployment"), + ], +) +async def test_a2a_model_resolves_before_router_branches(router_kwargs, extra_data): + """ + Regression: an `a2a/` model must reach A2A routing even when a router branch would + otherwise claim it. + + `map_team_model` claims the request for any team-scoped key, and the + wildcard/default-deployment fallback claims it whenever a pattern model group exists. + Both previously shadowed the A2A branch, so the call failed with + "no healthy deployments" instead of reaching the agent. + """ + from litellm.types.agents import AgentResponse + + data = { + "model": "a2a/test-agent", + "messages": [{"role": "user", "content": "Hello"}], + **extra_data, + } + + mock_agent = AgentResponse( + agent_id="test-agent-id", + agent_name="test-agent", + agent_card_params={"url": "http://agent.example.com"}, + litellm_params=None, + ) + mock_registry = Mock() + mock_registry.get_agent_by_id = Mock(return_value=None) + mock_registry.get_agent_by_name = Mock(return_value=mock_agent) + + mock_acompletion = AsyncMock(return_value={"id": "test-response"}) + + # The bug is that the request never leaves the router, so there is no HTTP boundary to + # fake: which collaborator gets called *is* the behaviour under test. The registry is + # Prisma-backed and has no injection seam available to a unit test. + with patch("litellm.acompletion", mock_acompletion): # test-quality-ok: the dispatch target is the assertion + with patch( # test-quality-ok: Prisma-backed registry has no unit-test injection seam + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + mock_registry, + ): + await route_request( + data=data, + llm_router=_router_without_a2a_deployment(**router_kwargs), + user_model=None, + route_type="acompletion", + ) + + mock_acompletion.assert_called_once() + call_kwargs = mock_acompletion.call_args.kwargs + assert call_kwargs["model"] == "a2a/test-agent" + assert call_kwargs["api_base"] == "http://agent.example.com"