litellm/tests/unit/a2a_protocol/test_completion_bridge_streaming.py

"""
Test A2A completion bridge streaming transformation to proper A2A format.

Tests that the completion bridge emits proper A2A streaming events:
1. Task event (kind: "task") - Initial task with status "submitted"
2. Status update (kind: "status-update") - Status "working"
3. Artifact update (kind: "artifact-update") - Content delivery
4. Status update (kind: "status-update") - Final "completed" status
"""

from unittest.mock import AsyncMock, MagicMock, patch

import pytest


class TestA2AStreamingTransformation:
    """Test the A2A streaming transformation creates proper events."""

    def test_a2a_metadata_forwarded_to_completion_params(self):
        from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
            A2ACompletionBridgeTransformation,
        )

        message = {
            "role": "user",
            "parts": [{"text": "Reply to ticket #4823"}],
            "metadata": {"skillId": "draft_reply"},
        }
        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]

        completion_params: dict = {
            "model": "langgraph/agent",
            "messages": openai_messages,
        }
        A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
            completion_params=completion_params,
            a2a_message=message,
            params={"metadata": {"trace": "abc"}},
        )
        assert completion_params["extra_body"]["metadata"] == {
            "trace": "abc",
            "skillId": "draft_reply",
        }

    def test_configured_metadata_wins_over_forwarded_a2a_metadata(self):
        from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
            A2ACompletionBridgeTransformation,
        )

        # Agent-owner-configured run metadata in ``extra_body``.
        completion_params: dict = {
            "model": "langgraph/agent",
            "messages": [],
            "extra_body": {
                "metadata": {"owner_tag": "prod", "trace": "server-set"},
                "other": "keep",
            },
        }
        # Client tries to overwrite ``trace`` and inject a new key.
        message = {
            "role": "user",
            "parts": [{"text": "hi"}],
            "metadata": {"trace": "client-spoof", "skillId": "draft_reply"},
        }
        A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
            completion_params=completion_params,
            a2a_message=message,
            params={"metadata": {"trace": "client-spoof-2"}},
        )
        assert completion_params["extra_body"]["other"] == "keep"
        assert completion_params["extra_body"]["metadata"] == {
            "owner_tag": "prod",
            "trace": "server-set",
            "skillId": "draft_reply",
        }

    def test_langgraph_transform_preserves_message_metadata(self):
        from litellm.llms.langgraph.chat.transformation import LangGraphConfig

        config = LangGraphConfig()
        request = config.transform_request(
            model="langgraph/agent",
            messages=[
                {
                    "role": "user",
                    "content": "Reply to ticket #4823",
                    "metadata": {"skillId": "draft_reply"},
                }
            ],
            optional_params={},
            litellm_params={"stream": False},
            headers={},
        )
        assert request["input"]["messages"][-1]["metadata"] == {
            "skillId": "draft_reply",
        }

    def test_create_task_event(self):
        """Test that create_task_event produces proper A2A task event structure."""
        from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
            A2ACompletionBridgeTransformation,
            A2AStreamingContext,
        )

        input_message = {
            "role": "user",
            "parts": [{"kind": "text", "text": "Hello"}],
            "messageId": "msg-123",
        }
        ctx = A2AStreamingContext(request_id="req-456", input_message=input_message)

        event = A2ACompletionBridgeTransformation.create_task_event(ctx)

        # Validate structure
        assert event["jsonrpc"] == "2.0"
        assert event["id"] == "req-456"
        assert event["result"]["kind"] == "task"
        assert event["result"]["status"]["state"] == "submitted"
        assert "contextId" in event["result"]
        assert "id" in event["result"]  # task id
        assert "history" in event["result"]
        assert len(event["result"]["history"]) == 1
        assert event["result"]["history"][0]["role"] == "user"

    def test_create_status_update_working(self):
        """Test that create_status_update_event produces proper working status."""
        from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
            A2ACompletionBridgeTransformation,
            A2AStreamingContext,
        )

        ctx = A2AStreamingContext(
            request_id="req-456",
            input_message={"role": "user", "parts": []},
        )

        event = A2ACompletionBridgeTransformation.create_status_update_event(
            ctx=ctx,
            state="working",
            final=False,
            message_text="Processing...",
        )

        assert event["result"]["kind"] == "status-update"
        assert event["result"]["status"]["state"] == "working"
        assert event["result"]["final"] is False
        assert "taskId" in event["result"]
        assert "contextId" in event["result"]
        assert "timestamp" in event["result"]["status"]

    def test_create_artifact_update(self):
        """Test that create_artifact_update_event produces proper artifact event."""
        from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
            A2ACompletionBridgeTransformation,
            A2AStreamingContext,
        )

        ctx = A2AStreamingContext(
            request_id="req-456",
            input_message={"role": "user", "parts": []},
        )

        event = A2ACompletionBridgeTransformation.create_artifact_update_event(
            ctx=ctx,
            text="Hello, I am an AI assistant.",
        )

        assert event["result"]["kind"] == "artifact-update"
        assert "artifact" in event["result"]
        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."


@pytest.mark.asyncio
async def test_handle_streaming_emits_proper_events():
    """Test that handle_streaming emits events in correct order with proper structure."""
    from litellm.a2a_protocol.litellm_completion_bridge.handler import (
        A2ACompletionBridgeHandler,
    )

    # Mock litellm.acompletion to return a streaming response
    mock_chunk1 = MagicMock()
    mock_chunk1.choices = [MagicMock()]
    mock_chunk1.choices[0].delta = MagicMock()
    mock_chunk1.choices[0].delta.content = "Hello"

    mock_chunk2 = MagicMock()
    mock_chunk2.choices = [MagicMock()]
    mock_chunk2.choices[0].delta = MagicMock()
    mock_chunk2.choices[0].delta.content = " world"

    async def mock_streaming_response():
        yield mock_chunk1
        yield mock_chunk2

    with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
        mock_acompletion.return_value = mock_streaming_response()

        params = {
            "message": {
                "role": "user",
                "parts": [{"kind": "text", "text": "Hi"}],
                "messageId": "msg-123",
            }
        }

        events = []
        async for event in A2ACompletionBridgeHandler.handle_streaming(
            request_id="req-456",
            params=params,
            litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
            api_base="http://localhost:2024",
        ):
            events.append(event)

        # Should have 4 events: task, working, artifact, completed
        assert len(events) == 4

        # Event 1: task submitted
        assert events[0]["result"]["kind"] == "task"
        assert events[0]["result"]["status"]["state"] == "submitted"

        # Event 2: status working
        assert events[1]["result"]["kind"] == "status-update"
        assert events[1]["result"]["status"]["state"] == "working"
        assert events[1]["result"]["final"] is False

        # Event 3: artifact update with accumulated content
        assert events[2]["result"]["kind"] == "artifact-update"
        assert events[2]["result"]["artifact"]["parts"][0]["text"] == "Hello world"

        # 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


@pytest.mark.asyncio
async def test_handle_streaming_forwards_api_key():
    """Test that handle_streaming forwards api_key from litellm_params to acompletion."""
    from litellm.a2a_protocol.litellm_completion_bridge.handler import (
        A2ACompletionBridgeHandler,
    )

    mock_chunk = MagicMock()
    mock_chunk.choices = [MagicMock()]
    mock_chunk.choices[0].delta = MagicMock()
    mock_chunk.choices[0].delta.content = "Response"

    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()

        params = {
            "message": {
                "role": "user",
                "parts": [{"kind": "text", "text": "Hi"}],
                "messageId": "msg-123",
            }
        }

        events = []
        async for event in A2ACompletionBridgeHandler.handle_streaming(
            request_id="req-456",
            params=params,
            litellm_params={
                "custom_llm_provider": "azure_ai",
                "model": "agents/asst_123",
                "api_key": "test-api-key-12345",
            },
            api_base="https://example.azure.com/",
        ):
            events.append(event)

        # Verify acompletion was called with api_key
        mock_acompletion.assert_called_once()
        call_kwargs = mock_acompletion.call_args.kwargs
        assert call_kwargs["api_key"] == "test-api-key-12345"
        assert call_kwargs["api_base"] == "https://example.azure.com/"
        assert call_kwargs["model"] == "azure_ai/agents/asst_123"


@pytest.mark.asyncio
async def test_handle_non_streaming_forwards_api_key():
    """Test that handle_non_streaming forwards api_key from litellm_params to acompletion."""
    from litellm.a2a_protocol.litellm_completion_bridge.handler import (
        A2ACompletionBridgeHandler,
    )

    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message = MagicMock()
    mock_response.choices[0].message.content = "Hello!"
    mock_response.id = "resp-123"

    with patch("litellm.acompletion", new_callable=AsyncMock) as mock_acompletion:
        mock_acompletion.return_value = mock_response

        params = {
            "message": {
                "role": "user",
                "parts": [{"kind": "text", "text": "Hi"}],
                "messageId": "msg-123",
            }
        }

        await A2ACompletionBridgeHandler.handle_non_streaming(
            request_id="req-456",
            params=params,
            litellm_params={
                "custom_llm_provider": "azure_ai",
                "model": "agents/asst_456",
                "api_key": "my-secret-api-key",
            },
            api_base="https://my-azure.com/",
        )

        # Verify acompletion was called with api_key
        mock_acompletion.assert_called_once()
        call_kwargs = mock_acompletion.call_args.kwargs
        assert call_kwargs["api_key"] == "my-secret-api-key"
        assert call_kwargs["api_base"] == "https://my-azure.com/"
        assert call_kwargs["model"] == "azure_ai/agents/asst_456"


@pytest.mark.asyncio
async def test_handle_streaming_keeps_agent_card_path_out_of_the_completion_call():
    """agent_card_path describes where an A2A agent serves its card; a completion-bridge agent carrying
    it must not pass it to litellm.acompletion, where an unknown kwarg breaks the provider call."""
    from litellm.a2a_protocol.litellm_completion_bridge.handler import (
        A2ACompletionBridgeHandler,
    )

    async def mock_streaming_response():
        chunk = MagicMock()
        chunk.choices = [MagicMock()]
        chunk.choices[0].delta = MagicMock()
        chunk.choices[0].delta.content = "Hello"
        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-card-path",
                params={"message": {"role": "user", "parts": [{"kind": "text", "text": "Hi"}], "messageId": "m1"}},
                litellm_params={
                    "custom_llm_provider": "langgraph",
                    "model": "agent",
                    "agent_card_path": "agentCard/v1.0",
                },
                api_base="http://localhost:2024",
            )
        ]

    assert len(events) == 4
    assert "agent_card_path" not in mock_acompletion.call_args.kwargs


AGENTCORE_ARN = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/unit-agent-abc123"
AGENTCORE_INVOCATIONS_URL = (
    "https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/"
    + AGENTCORE_ARN.replace(":", "%3A").replace("/", "%2F")
    + "/invocations"
)


@pytest.mark.asyncio
async def test_a2a_completion_bridge_bedrock_agentcore(respx_mock, monkeypatch):
    """
    A2A streaming via the completion bridge against a scripted Bedrock AgentCore
    upstream: asserts the outbound invocation request and the emitted A2A event
    sequence.
    """
    import litellm

    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
    import httpx
    from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest
    from uuid import uuid4

    from litellm.a2a_protocol import asend_message_streaming

    sse_body = (
        'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"Machine "}}]}\n\n'
        'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"learning."}}]}\n\n'
        'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n'
        "data: [DONE]\n\n"
    )
    invoke_route = respx_mock.post(
        url__regex=r"https://bedrock-agentcore\.us-west-2\.amazonaws\.com/runtimes/.*/invocations.*"
    ).mock(
        return_value=httpx.Response(200, content=sse_body.encode(), headers={"content-type": "text/event-stream"})
    )

    send_message_payload = {
        "message": {
            "role": "user",
            "parts": [{"kind": "text", "text": "Explain machine learning in simple terms"}],
            "messageId": uuid4().hex,
        }
    }
    request = SendStreamingMessageRequest(
        id=str(uuid4()),
        params=MessageSendParams(**send_message_payload),  # type: ignore
    )

    chunks = []
    async for chunk in asend_message_streaming(
        request=request,
        api_base=None,
        litellm_params={
            "custom_llm_provider": "bedrock",
            "model": f"bedrock/agentcore/{AGENTCORE_ARN}",
            "api_key": "unit-jwt-token",
        },
    ):
        chunks.append(chunk)

    assert invoke_route.call_count == 1
    invoke_request = invoke_route.calls.last.request
    encoded_arn = AGENTCORE_ARN.replace(":", "%3A").replace("/", "%2F")
    assert encoded_arn in str(invoke_request.url)
    assert invoke_request.headers["authorization"] == "Bearer unit-jwt-token"

    assert len(chunks) == 3
    content = "".join(
        chunk["choices"][0]["delta"]["content"]
        for chunk in chunks
        if chunk["choices"][0].get("delta", {}).get("content")
    )
    assert content == "Machine learning."
    assert chunks[-1]["choices"][0].get("finish_reason") == "stop"