First commit

tests/test_a2a_http_streaming.py

@@ -0,0 +1,81 @@
+"""Verify SSE flush and disconnect/reconnect over a real local HTTP connection."""
+import asyncio
+import json
+import socket
+import threading
+import time
+from types import SimpleNamespace
+
+import httpx
+import pytest
+import uvicorn
+from fastapi import FastAPI
+pytest.importorskip("a2a", reason="Install agent-framework[a2a] to run A2A tests")
+
+from agent_framework.a2a.server import install_a2a
+from agent_framework.global_supervisor.models import BackendDefinition
+
+
+@pytest.mark.parametrize('legacy', [False, True])
+def test_sse_flush_disconnect_and_resubscribe(legacy):
+    release = threading.Event()
+    finished = threading.Event()
+    async def handler(request):
+        while not release.is_set():
+            await asyncio.sleep(.005)
+        finished.set()
+        return {'text': 'stream complete'}
+    app = FastAPI()
+    registry = SimpleNamespace(list=lambda: [BackendDefinition(backend_id='x', url='http://native')])
+    install_a2a(app, handler, registry, 'http://testserver')
+    sock = socket.socket()
+    sock.bind(('127.0.0.1', 0))
+    server = uvicorn.Server(uvicorn.Config(app, log_level='error', lifespan='off'))
+    thread = threading.Thread(target=lambda: server.run(sockets=[sock]), daemon=True)
+    thread.start()
+    try:
+        deadline = time.monotonic() + 3
+        while not server.started and time.monotonic() < deadline:
+            time.sleep(.005)
+        assert server.started
+        base_url = f'http://127.0.0.1:{sock.getsockname()[1]}'
+        with httpx.Client(base_url=base_url, headers={'A2A-Version': '0.3' if legacy else '1.0'}, timeout=2, trust_env=False) as client:
+            message = {'messageId': 'm', 'role': 'user' if legacy else 'ROLE_USER',
+                       'parts': [{'kind': 'text', 'text': 'run'}] if legacy else [{'text': 'run'}]}
+            if legacy: message['kind'] = 'message'
+            def envelope(method, params):
+                return {'jsonrpc': '2.0', 'id': 1, 'method': method, 'params': params}
+            with client.stream('POST', '/a2a', json=envelope('message/stream' if legacy else 'SendStreamingMessage', {'message': message})) as response:
+                assert response.status_code == 200
+                for line in response.iter_lines():
+                    if line.startswith('data:'):
+                        event = json.loads(line[5:].strip())
+                        first = event['result'] if legacy else event['result']['task']
+                        task_id = first['id']
+                        break
+                else:
+                    pytest.fail('SSE did not deliver an initial Task')
+                assert not release.is_set() and not finished.is_set()
+            # Closing the first SSE connection must not cancel the agent task.
+            current = client.post('/a2a', json=envelope('tasks/get' if legacy else 'GetTask', {'id': task_id})).json()
+            assert current['result']['status']['state'] in ({'submitted', 'working'} if legacy else {'TASK_STATE_SUBMITTED', 'TASK_STATE_WORKING'})
+            with client.stream('POST', '/a2a', json=envelope('tasks/resubscribe' if legacy else 'SubscribeToTask', {'id': task_id})) as response:
+                events = []
+                for line in response.iter_lines():
+                    if not line.startswith('data:'):
+                        continue
+                    event = json.loads(line[5:].strip())
+                    assert 'error' not in event, event
+                    events.append(event)
+                    if len(events) == 1:
+                        release.set()
+                encoded = json.dumps(events)
+                assert 'stream complete' in encoded
+                assert ('completed' if legacy else 'TASK_STATE_COMPLETED') in encoded
+                assert finished.is_set()
+    finally:
+        release.set()
+        server.should_exit = True
+        thread.join(timeout=3)
+        sock.close()
+        assert not thread.is_alive()