diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 24727aacd75..cc66e21b9d0 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -62,7 +62,7 @@ async def _handle_stream_message( if not A2A_SDK_AVAILABLE: # Return a streaming response that yields an error async def _error_stream(): - yield json.dumps( + yield "data: " + json.dumps( { "jsonrpc": "2.0", "id": request_id, @@ -71,9 +71,9 @@ async def _handle_stream_message( "message": "Server error: 'a2a' package not installed", }, } - ) + "\n" + ) + "\n\n" - return StreamingResponse(_error_stream(), media_type="application/x-ndjson") + return StreamingResponse(_error_stream(), media_type="text/event-stream") from a2a.types import ( MessageSendParams, @@ -96,22 +96,27 @@ async def _handle_stream_message( ): # Chunk may be dict or object depending on bridge vs standard path if hasattr(chunk, "model_dump"): - yield json.dumps( + yield "data: " + json.dumps( chunk.model_dump(mode="json", exclude_none=True) - ) + "\n" + ) + "\n\n" else: - yield json.dumps(chunk) + "\n" + yield "data: " + json.dumps(chunk) + "\n\n" except Exception as e: - verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") - yield json.dumps( + # Log full exception details server-side for debugging + verbose_proxy_logger.exception("Error streaming A2A response") + # Return a generic error message to the client without exposing internal details + yield "data: " + json.dumps( { "jsonrpc": "2.0", "id": request_id, - "error": {"code": -32603, "message": f"Streaming error: {str(e)}"}, + "error": { + "code": -32603, + "message": "Streaming error", + }, } - ) + "\n" + ) + "\n\n" - return StreamingResponse(stream_response(), media_type="application/x-ndjson") + return StreamingResponse(stream_response(), media_type="text/event-stream") @router.get( diff --git a/tests/a2a_manual/test_a2a_sse_manual.py b/tests/a2a_manual/test_a2a_sse_manual.py new file mode 100644 index 00000000000..8205e5189ce --- /dev/null +++ b/tests/a2a_manual/test_a2a_sse_manual.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +""" +Manual test: Verify A2A streaming returns text/event-stream (SSE format) + +This script: +1. Starts a simple A2A agent on port 10001 +2. Starts LiteLLM proxy on port 4000 with the agent registered +3. Makes a streaming request to the proxy's A2A gateway +4. Validates Content-Type header and SSE body format + +Run: python test_a2a_sse_manual.py + +Expected output: + āœ… Content-Type: text/event-stream + āœ… SSE framing: data: {...}\n\n +""" + +import asyncio +import json +import os +import signal +import subprocess +import sys +import tempfile +import time + +import httpx + +AGENT_PORT = 10001 +PROXY_PORT = 4000 +TIMEOUT = 60 + + +def start_agent(): + """Start a simple A2A agent using LiteLLM's a2a_protocol.""" + agent_code = ''' +import asyncio +from litellm.a2a_protocol import A2AServer, BaseA2AAgent + +class TestAgent(BaseA2AAgent): + @property + def name(self) -> str: + return "test-agent" + + @property + def description(self) -> str: + return "Test agent for SSE validation" + + @property + def streaming(self) -> bool: + return True + + async def invoke(self, query, session_id=None): + return {"role": "agent", "parts": [{"kind": "text", "text": "Hello from test agent!"}]} + + async def invoke_streaming(self, query, session_id=None): + for word in ["Hello", "from", "streaming", "agent!"]: + yield {"role": "agent", "parts": [{"kind": "text", "text": word + " "}]} + await asyncio.sleep(0.1) + +async def main(): + server = A2AServer(agent=TestAgent(), host="0.0.0.0", port=10001) + await server.start() + +if __name__ == "__main__": + asyncio.run(main()) +''' + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f: + f.write(agent_code) + agent_file = f.name + + proc = subprocess.Popen( + [sys.executable, agent_file], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + return proc, agent_file + + +def start_proxy(): + """Start LiteLLM proxy with A2A agent registered.""" + config = { + "model_list": [ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "gpt-4o-mini", + } + } + ], + "a2a_config": { + "agents": [ + { + "agent_id": "test-agent", + "api_base": f"http://localhost:{AGENT_PORT}", + } + ] + }, + "general_settings": { + "master_key": "sk-test-key" + } + } + + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + import yaml + yaml.dump(config, f) + config_file = f.name + + proc = subprocess.Popen( + [sys.executable, "-m", "litellm", "--config", config_file, "--port", str(PROXY_PORT)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + return proc, config_file + + +async def wait_for_server(url: str, timeout: int = 30): + """Wait for a server to be ready.""" + start = time.time() + async with httpx.AsyncClient() as client: + while time.time() - start < timeout: + try: + resp = await client.get(url, timeout=2) + if resp.status_code < 500: + return True + except Exception: + pass + await asyncio.sleep(0.5) + return False + + +async def test_streaming(): + """Test the A2A streaming endpoint.""" + print("\n" + "=" * 60) + print("Testing A2A Streaming SSE Format") + print("=" * 60) + + url = f"http://localhost:{PROXY_PORT}/a2a/test-agent" + headers = { + "Content-Type": "application/json", + "Authorization": "Bearer sk-test-key", + } + payload = { + "jsonrpc": "2.0", + "id": "test-123", + "method": "message/stream", + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-001", + } + } + } + + print(f"\nšŸ“” POST {url}") + print(f" Method: message/stream") + + async with httpx.AsyncClient() as client: + async with client.stream("POST", url, json=payload, headers=headers, timeout=30) as resp: + # Check Content-Type header + content_type = resp.headers.get("content-type", "") + print(f"\nšŸ“‹ Response Headers:") + print(f" Content-Type: {content_type}") + + if "text/event-stream" in content_type: + print(" āœ… Correct! Uses text/event-stream (SSE)") + elif "application/x-ndjson" in content_type: + print(" āŒ WRONG! Still using application/x-ndjson") + return False + else: + print(f" āš ļø Unexpected Content-Type") + + # Check body format + print(f"\nšŸ“¦ Response Body (SSE events):") + events = [] + async for line in resp.aiter_lines(): + if line.strip(): + events.append(line) + print(f" {line[:100]}{'...' if len(line) > 100 else ''}") + + # Validate SSE framing + print(f"\nšŸ” SSE Format Validation:") + all_valid = True + for i, event in enumerate(events): + if event.startswith("data: "): + try: + payload_str = event[6:] # Remove "data: " prefix + json.loads(payload_str) + print(f" Event {i+1}: āœ… Valid SSE (data: )") + except json.JSONDecodeError as e: + print(f" Event {i+1}: āŒ Invalid JSON: {e}") + all_valid = False + else: + print(f" Event {i+1}: āŒ Missing 'data: ' prefix") + all_valid = False + + if all_valid and events: + print(f"\nāœ… ALL TESTS PASSED - A2A streaming uses correct SSE format") + return True + else: + print(f"\nāŒ TESTS FAILED") + return False + + +async def main(): + agent_proc = None + proxy_proc = None + agent_file = None + config_file = None + + try: + print("šŸš€ Starting test A2A agent on port", AGENT_PORT) + agent_proc, agent_file = start_agent() + + print("ā³ Waiting for agent to be ready...") + if not await wait_for_server(f"http://localhost:{AGENT_PORT}/.well-known/agent.json"): + print("āŒ Agent failed to start") + # Print agent stderr for debugging + if agent_proc.stderr: + print("Agent stderr:", agent_proc.stderr.read().decode()) + return 1 + print("āœ… Agent ready") + + print("\nšŸš€ Starting LiteLLM proxy on port", PROXY_PORT) + proxy_proc, config_file = start_proxy() + + print("ā³ Waiting for proxy to be ready...") + if not await wait_for_server(f"http://localhost:{PROXY_PORT}/health"): + print("āŒ Proxy failed to start") + if proxy_proc.stderr: + print("Proxy stderr:", proxy_proc.stderr.read().decode()) + return 1 + print("āœ… Proxy ready") + + # Run the test + success = await test_streaming() + return 0 if success else 1 + + except Exception as e: + print(f"āŒ Error: {e}") + import traceback + traceback.print_exc() + return 1 + + finally: + # Cleanup + if agent_proc: + agent_proc.terminate() + agent_proc.wait() + if proxy_proc: + proxy_proc.terminate() + proxy_proc.wait() + if agent_file and os.path.exists(agent_file): + os.unlink(agent_file) + if config_file and os.path.exists(config_file): + os.unlink(config_file) + + +if __name__ == "__main__": + sys.exit(asyncio.run(main())) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 9588c3b55c3..135e24953a9 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -4,6 +4,7 @@ Mock tests for A2A endpoints. Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request. """ +import json import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -168,3 +169,128 @@ async def test_invoke_agent_a2a_adds_litellm_data(): # Verify proxy_server_request was added assert "proxy_server_request" in captured_data assert captured_data["proxy_server_request"]["method"] == "POST" + + +@pytest.mark.asyncio +async def test_handle_stream_message_returns_sse_content_type(): + """ + Test that _handle_stream_message returns Content-Type: text/event-stream + with SSE-framed body (data: ...\\n\\n), not application/x-ndjson. + + The A2A protocol spec requires text/event-stream for streaming responses. + The official a2a-sdk client validates this header. + + Ref: https://github.com/BerriAI/litellm/issues/20278 + """ + # Mock chunk with model_dump + mock_chunk = MagicMock() + mock_chunk.model_dump.return_value = { + "jsonrpc": "2.0", + "id": "test-id", + "result": {"kind": "status-update"}, + } + + async def mock_streaming(*args, **kwargs): + yield mock_chunk + + # Try to use real a2a.types if available + try: + from a2a.types import ( + MessageSendParams, + SendStreamingMessageRequest, + ) + except ImportError: + + class MessageSendParams: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + class SendStreamingMessageRequest: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = MessageSendParams + mock_a2a_types.SendStreamingMessageRequest = SendStreamingMessageRequest + + with patch.dict( + sys.modules, + {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, + ), patch( + "litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", + True, + ), patch( + "litellm.a2a_protocol.asend_message_streaming", + side_effect=mock_streaming, + ): + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _handle_stream_message, + ) + + response = await _handle_stream_message( + api_base="http://backend:10001", + request_id="test-id", + params={ + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } + }, + ) + + # Verify Content-Type is text/event-stream (required by A2A spec) + assert response.media_type == "text/event-stream" + + # Collect streamed body and verify SSE framing + body_parts = [] + async for chunk in response.body_iterator: + body_parts.append(chunk) + + assert len(body_parts) > 0 + for part in body_parts: + # Each SSE event must start with "data: " and end with "\n\n" + assert part.startswith("data: "), ( + f"SSE event must start with 'data: ', got: {part!r}" + ) + assert part.endswith("\n\n"), ( + f"SSE event must end with '\\n\\n', got: {part!r}" + ) + # The payload between "data: " and "\n\n" must be valid JSON + payload = part[len("data: "):-2] + parsed = json.loads(payload) + assert isinstance(parsed, dict) + + +@pytest.mark.asyncio +async def test_handle_stream_message_error_uses_sse_format(): + """ + Test that when A2A SDK is not available, the error stream also uses + text/event-stream with SSE framing. + """ + with patch( + "litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", + False, + ): + from litellm.proxy.agent_endpoints.a2a_endpoints import ( + _handle_stream_message, + ) + + response = await _handle_stream_message( + api_base=None, + request_id="err-id", + params={}, + ) + + assert response.media_type == "text/event-stream" + + body_parts = [] + async for chunk in response.body_iterator: + body_parts.append(chunk) + + assert len(body_parts) == 1 + part = body_parts[0] + assert part.startswith("data: ") + assert part.endswith("\n\n") + payload = json.loads(part[len("data: "):-2]) + assert payload["error"]["code"] == -32603