82 lines
3.8 KiB
Python
82 lines
3.8 KiB
Python
"""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()
|