diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index cc66e21b9d0..24727aacd75 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 "data: " + json.dumps( + yield 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="text/event-stream") + return StreamingResponse(_error_stream(), media_type="application/x-ndjson") from a2a.types import ( MessageSendParams, @@ -96,27 +96,22 @@ async def _handle_stream_message( ): # Chunk may be dict or object depending on bridge vs standard path if hasattr(chunk, "model_dump"): - yield "data: " + json.dumps( + yield json.dumps( chunk.model_dump(mode="json", exclude_none=True) - ) + "\n\n" + ) + "\n" else: - yield "data: " + json.dumps(chunk) + "\n\n" + yield json.dumps(chunk) + "\n" except Exception as e: - # 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( + verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") + yield json.dumps( { "jsonrpc": "2.0", "id": request_id, - "error": { - "code": -32603, - "message": "Streaming error", - }, + "error": {"code": -32603, "message": f"Streaming error: {str(e)}"}, } - ) + "\n\n" + ) + "\n" - return StreamingResponse(stream_response(), media_type="text/event-stream") + return StreamingResponse(stream_response(), media_type="application/x-ndjson") @router.get( diff --git a/tests/a2a_manual/test_a2a_sse_manual.py b/tests/a2a_manual/test_a2a_sse_manual.py deleted file mode 100644 index 8205e5189ce..00000000000 --- a/tests/a2a_manual/test_a2a_sse_manual.py +++ /dev/null @@ -1,261 +0,0 @@ -#!/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 135e24953a9..9588c3b55c3 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -4,7 +4,6 @@ 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 @@ -169,128 +168,3 @@ 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