Merge 1cf002fd2b into 0c98afa780

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", {})

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)

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)

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}"

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"