litellm/tests/test_litellm/a2a_protocol/test_main.py

"""Tests for litellm/a2a_protocol/main.py non-streaming send behavior."""

import httpx
import pytest

pytest.importorskip("a2a.compat.v0_3.conversions")

from a2a.compat.v0_3 import conversions as _conv
from a2a.compat.v0_3.types import (
    MessageSendParams,
    SendMessageRequest,
    SendStreamingMessageRequest,
)

import litellm
from litellm.a2a_protocol.main import _send_message, _stream_messages, create_a2a_client
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
from litellm.llms.custom_httpx.http_handler import (
    AsyncHTTPHandler,
    get_async_httpx_client,
    httpxSpecialProvider,
)


def _request() -> SendMessageRequest:
    params = MessageSendParams(
        message={
            "messageId": "m1",
            "role": "user",
            "parts": [{"kind": "text", "text": "hi"}],
        }
    )
    return SendMessageRequest(id="r1", params=params)


def _message_stream_response():
    sr = _conv.pb2_v10.StreamResponse()
    sr.message.message_id = "reply-1"
    sr.message.role = _conv.pb2_v10.Role.ROLE_AGENT
    sr.message.parts.add().text = "hello back"
    return sr


def _status_update_stream_response():
    sr = _conv.pb2_v10.StreamResponse()
    sr.status_update.task_id = "t1"
    sr.status_update.context_id = "c1"
    return sr


class _FakeClient:
    def __init__(self, *events):
        self._events = events

    async def send_message(self, _pb_request, context=None):
        self.context = context
        for event in self._events:
            yield event


@pytest.mark.asyncio
async def test_send_message_returns_message_result():
    response = await _send_message(_FakeClient(_message_stream_response()), _request())
    result = response.root.result
    assert type(result).__name__ == "Message"
    assert response.root.id == "r1"


@pytest.mark.asyncio
async def test_send_message_rejects_update_event_final_with_runtime_error():
    with pytest.raises(RuntimeError, match="Message or Task"):
        await _send_message(_FakeClient(_status_update_stream_response()), _request())


@pytest.mark.asyncio
async def test_streaming_trace_id_prefers_logging_trace_id():
    """The streaming X-LiteLLM-Trace-Id must use the logging object's trace id (same
    as the non-streaming path), not the JSON-RPC request id, so traces correlate."""
    from unittest.mock import AsyncMock, MagicMock, patch

    from a2a.compat.v0_3.types import (
        MessageSendParams,
        SendStreamingMessageRequest,
    )

    from litellm.a2a_protocol import main as a2a_main
    from litellm.litellm_core_utils.litellm_logging import Logging

    request = SendStreamingMessageRequest(
        id="rpc-1",
        params=MessageSendParams(
            message={
                "messageId": "m1",
                "role": "user",
                "parts": [{"kind": "text", "text": "hi"}],
            }
        ),
    )
    logging_obj = MagicMock(spec=Logging)
    logging_obj.litellm_trace_id = "trace-from-logging"

    captured: dict = {}

    async def _capture(*, base_url, extra_headers=None, streaming=False, **_):
        captured["extra_headers"] = extra_headers
        raise RuntimeError("stop")

    with patch.object(
        a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture)
    ):
        with pytest.raises(RuntimeError, match="stop"):
            async for _ in a2a_main.asend_message_streaming(
                request=request,
                api_base="http://upstream.local",
                litellm_logging_obj=logging_obj,
            ):
                pass

    assert captured["extra_headers"]["X-LiteLLM-Trace-Id"] == "trace-from-logging"


def test_streaming_logging_obj_carries_call_type_into_model_call_details():
    """The streaming logging object is built by hand rather than through
    ``update_environment_variables``, which is the only place ``call_type`` normally
    reaches ``model_call_details``. Callbacks read the call type from there, so
    without this the streamed turn arrives at every logger with no call type at all
    and OTel's GenAI metrics label it ``chat`` instead of ``invoke_agent``."""
    from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest

    from litellm.a2a_protocol.main import _build_streaming_logging_obj

    request = SendStreamingMessageRequest(
        id="rpc-call-type",
        params=MessageSendParams(
            message={"messageId": "m1", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}
        ),
    )

    logging_obj = _build_streaming_logging_obj(
        request=request,
        agent_name="some-agent",
        agent_id=None,
        litellm_params=None,
        metadata=None,
        proxy_server_request=None,
    )

    assert logging_obj.model_call_details["call_type"] == "asend_message_streaming"


_AGENT_CARD = {
    "protocolVersion": "0.3.0",
    "name": "recording-agent",
    "url": "http://127.0.0.1:9/",
    "preferredTransport": "JSONRPC",
    "version": "1.0.0",
    "capabilities": {"streaming": True},
    "defaultInputModes": ["text/plain"],
    "defaultOutputModes": ["text/plain"],
    "skills": [],
}

_RPC_REPLY = {
    "jsonrpc": "2.0",
    "id": "reply",
    "result": {
        "messageId": "reply-1",
        "role": "agent",
        "parts": [{"kind": "text", "text": "pong"}],
        "kind": "message",
    },
}

_AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"}
_AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"}


class _RequestRecorder:
    """Records the headers httpx put on the wire, per outbound request."""

    def __init__(self):
        self.card_requests = []
        self.rpc_requests = []
        self.client = None

    def __call__(self, request: httpx.Request) -> httpx.Response:
        headers = {k.lower(): v for k, v in request.headers.items()}
        if request.method == "GET":
            self.card_requests.append(headers)
            return httpx.Response(200, json=_AGENT_CARD)
        self.rpc_requests.append(headers)
        return httpx.Response(200, json=_RPC_REPLY)


def _a2a_client_cache_key(timeout: float) -> str:
    return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider


async def _seed_shared_a2a_client() -> _RequestRecorder:
    """Put the one A2A client the cache will hand out behind a mock transport.

    Seeding has to happen on the test's own event loop, because the client cache keys on
    it. The injected client is a real httpx.AsyncClient, so the merge of per-request
    headers over client defaults, which is what these tests are about, stays real.
    """
    recorder = _RequestRecorder()
    handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT)
    owned_client = handler.client
    handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder))
    await owned_client.aclose()

    litellm.in_memory_llm_clients_cache.set_cache(key=_a2a_client_cache_key(DEFAULT_A2A_AGENT_TIMEOUT), value=handler)
    seeded = get_async_httpx_client(
        llm_provider=httpxSpecialProvider.A2AProvider,
        params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT},
    )
    assert seeded is handler, "cache key drifted from get_async_httpx_client; these tests would test nothing"

    recorder.client = handler.client
    return recorder


@pytest.fixture
def isolated_client_cache():
    previous = getattr(litellm, "in_memory_llm_clients_cache", None)
    litellm.in_memory_llm_clients_cache = LLMClientCache()
    yield litellm.in_memory_llm_clients_cache
    litellm.in_memory_llm_clients_cache = previous


def _send_request(request_id):
    return SendMessageRequest(
        id=request_id,
        params=MessageSendParams(
            message={"messageId": request_id, "role": "user", "parts": [{"kind": "text", "text": "hi"}]}
        ),
    )


@pytest.mark.asyncio
async def test_extra_headers_never_land_on_the_shared_cached_client(isolated_client_cache):
    """get_async_httpx_client hands back a process-wide shared client, so a caller's
    headers written onto it would outlive the request that supplied them."""
    recorder = await _seed_shared_a2a_client()

    await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)

    assert "x-agent-token" not in recorder.client.headers
    assert "x-tenant" not in recorder.client.headers


@pytest.mark.asyncio
async def test_callers_with_different_headers_reuse_one_pooled_client(isolated_client_cache):
    """Headers must not segregate the connection pool. Every A2A caller on one timeout
    shares one cached client, so header sets cannot multiply cached clients."""
    recorder = await _seed_shared_a2a_client()

    client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)
    client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS)
    client_none = await create_a2a_client(base_url="http://127.0.0.1:9")

    assert client_a._litellm_httpx_client is recorder.client
    assert client_b._litellm_httpx_client is recorder.client
    assert client_none._litellm_httpx_client is recorder.client

    cached = [key for key in isolated_client_cache.cache_dict if "a2a_provider" in key]
    assert len(cached) == 1, f"expected one pooled A2A client, cached: {cached}"


@pytest.mark.parametrize("order", [("a", "b", "none"), ("b", "none", "a"), ("none", "a", "b")])
@pytest.mark.asyncio
async def test_each_caller_sends_only_its_own_headers(order, isolated_client_cache):
    """Whatever order callers arrive in, each request carries that caller's headers and
    no other caller's, and a caller with no extra_headers sends none."""
    recorder = await _seed_shared_a2a_client()
    headers_by_caller = {"a": _AGENT_A_HEADERS, "b": _AGENT_B_HEADERS, "none": None}

    for caller in order:
        a2a_client = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=headers_by_caller[caller])
        await _send_message(a2a_client, _send_request(caller))

    received = dict(zip(order, recorder.rpc_requests, strict=True))
    assert received["a"]["x-agent-token"] == "token-for-a"
    assert received["a"]["x-tenant"] == "tenant-a"
    assert received["b"]["x-agent-token"] == "token-for-b"
    assert received["b"]["x-tenant"] == "tenant-b"
    assert "x-agent-token" not in received["none"]
    assert "x-tenant" not in received["none"]


@pytest.mark.asyncio
async def test_streaming_send_carries_only_its_own_caller_headers(isolated_client_cache):
    """The streaming path shares the same pooled client, so it needs the same guard."""
    recorder = await _seed_shared_a2a_client()

    client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS, streaming=True)
    await _send_message(client_a, _send_request("a"))

    client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS, streaming=True)
    streaming_request = SendStreamingMessageRequest(
        id="b",
        params=MessageSendParams(message={"messageId": "b", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}),
    )
    async for _ in _stream_messages(client_b, streaming_request):
        pass

    received = dict(zip(("a", "b"), recorder.rpc_requests, strict=True))
    assert received["a"]["x-agent-token"] == "token-for-a"
    assert received["b"]["x-agent-token"] == "token-for-b"
    assert received["b"]["x-tenant"] == "tenant-b"


@pytest.mark.asyncio
async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cache):
    """Agent cards can sit behind the same auth as the agent, so the card fetch must stay
    authenticated once the headers stop living on the client."""
    recorder = await _seed_shared_a2a_client()

    await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)

    assert recorder.card_requests, "no agent card request was made"
    assert recorder.card_requests[-1]["x-agent-token"] == "token-for-a"


@pytest.mark.asyncio
async def test_the_pooled_a2a_client_arrives_with_cookie_persistence_disabled(isolated_client_cache):
    """create_a2a_client takes its client from the shared builder rather than building one,
    and the builder is what refuses to persist cookies. This pins the join between those
    two facts, so the A2A path cannot quietly start acquiring a client that keeps a jar.

    test_callers_with_different_headers_reuse_one_pooled_client pins the other half, that
    create_a2a_client hands back exactly this cached client."""
    handler = get_async_httpx_client(
        llm_provider=httpxSpecialProvider.A2AProvider,
        params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT},
    )
    request = httpx.Request("GET", "https://agent-a.example.com/")
    handler.client.cookies.extract_cookies(
        httpx.Response(200, headers={"set-cookie": "SESSION=only-agent-a-may-hold-this"}, request=request)
    )

    assert dict(handler.client.cookies) == {}, "the pooled A2A client kept an upstream's cookie"
    await handler.close()