diff --git a/sdk/python/agentfield/memory_events.py b/sdk/python/agentfield/memory_events.py index e58bcae1b..6c23e576a 100644 --- a/sdk/python/agentfield/memory_events.py +++ b/sdk/python/agentfield/memory_events.py @@ -3,7 +3,7 @@ import re from datetime import datetime from functools import wraps -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Coroutine, Dict, List, Optional, Set, Union from urllib.parse import urlencode import websockets @@ -103,6 +103,13 @@ def __init__(self, base_url: str, execution_context, api_key: Optional[str] = No self._max_reconnect_attempts = 5 self._reconnect_delay = 1.0 self._connect_timeout = 5.0 # Timeout for initial connection attempt + self._background_tasks: Set[asyncio.Task] = set() + self._connection_generation = 0 + + def _start_background_task(self, coroutine: Coroutine[Any, Any, None]) -> None: + task = asyncio.create_task(coroutine) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) def _is_connected(self) -> bool: """ @@ -139,10 +146,13 @@ async def connect( scope: Scope to filter events by scope_id: Scope ID to filter events by """ + generation = self._connection_generation if self._connect_lock is None: self._connect_lock = asyncio.Lock() async with self._connect_lock: + if generation != self._connection_generation: + return if self._is_connected(): return @@ -164,20 +174,25 @@ async def connect( if query_params: ws_url += "?" + urlencode(query_params) - self.websocket = await asyncio.wait_for( + websocket = await asyncio.wait_for( websockets.connect(ws_url, **{_HEADERS_KWARG: headers}), timeout=self._connect_timeout, ) + if generation != self._connection_generation: + await websocket.close() + return + self.websocket = websocket self.is_listening = True self._reconnect_attempts = 0 - asyncio.create_task(self._listen()) + self._start_background_task(self._listen()) except Exception as e: log_error(f"Failed to connect to memory events: {e}") # Background the retry loop so startup is not blocked by # repeated reconnection attempts when the server is # unreachable. - asyncio.create_task(self._handle_reconnect()) + if generation == self._connection_generation: + self._start_background_task(self._handle_reconnect()) async def _listen(self): """Listens for incoming messages and dispatches them to subscribers.""" @@ -269,7 +284,7 @@ def subscribe( # We rely on client-side pattern matching, so we don't need to send # pattern filters to the server. if not self._is_connected(): - asyncio.create_task(self.connect()) + self._start_background_task(self.connect()) return subscription @@ -445,12 +460,25 @@ def unsubscribe_all(self): async def close(self): """Closes the WebSocket connection and cleans up subscriptions.""" + self._connection_generation += 1 self.is_listening = False self.unsubscribe_all() - if self.websocket: - await self.websocket.close() - self.websocket = None + websocket = self.websocket + self.websocket = None + tasks = [ + task + for task in self._background_tasks + if task is not asyncio.current_task() + ] + for task in tasks: + task.cancel() + try: + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + finally: + if websocket: + await websocket.close() class ScopedMemoryEventClient: diff --git a/sdk/python/tests/test_memory_events.py b/sdk/python/tests/test_memory_events.py index ceecd4e57..b279481f7 100644 --- a/sdk/python/tests/test_memory_events.py +++ b/sdk/python/tests/test_memory_events.py @@ -73,7 +73,6 @@ async def test_memory_event_client_dispatches_literal_key_patterns( client = MemoryEventClient("http://agentfield", ctx) incoming = asyncio.Queue() handled = asyncio.Event() - listener_closed = asyncio.Event() received = [] class DummyWebSocket: @@ -82,7 +81,6 @@ class DummyWebSocket: async def recv(self): message = await incoming.get() if message is None: - listener_closed.set() return json.dumps({"scope": "global", "key": "closed", "action": "set"}) return message @@ -116,7 +114,7 @@ async def observe(event): assert received == [matching_key] finally: await client.close() - await asyncio.wait_for(listener_closed.wait(), timeout=1) + assert not websocket.open @pytest.mark.asyncio @@ -227,6 +225,282 @@ def test_websockets_version_detection(): assert _HEADERS_KWARG == "extra_headers" +@pytest.mark.asyncio +async def test_close_does_not_reconnect_after_listener_receives_close(monkeypatch): + client = MemoryEventClient( + "http://agentfield", SimpleNamespace(to_headers=lambda: {}) + ) + client._reconnect_delay = 0 + sockets = [] + tasks = [] + create_task = asyncio.create_task + + class Socket: + open = True + + def __init__(self): + self.receiving = asyncio.Event() + self.closed = asyncio.Event() + + async def recv(self): + self.receiving.set() + await self.closed.wait() + raise websockets.exceptions.ConnectionClosedOK(None, None) + + async def close(self): + self.open = False + self.closed.set() + + async def connect(url, **kwargs): + socket = Socket() + sockets.append(socket) + return socket + + def capture_task(coro, *args, **kwargs): + task = create_task(coro, *args, **kwargs) + tasks.append(task) + return task + + monkeypatch.setattr("agentfield.memory_events.websockets.connect", connect) + monkeypatch.setattr(asyncio, "create_task", capture_task) + try: + await client.connect() + await asyncio.wait_for(sockets[0].receiving.wait(), timeout=1) + listener = tasks[0] + + await client.close() + await asyncio.wait_for( + asyncio.gather(listener, return_exceptions=True), timeout=1 + ) + + assert len(sockets) == 1 + assert sockets[0].closed.is_set() + assert client.websocket is None + assert not client.is_listening + finally: + for socket in sockets: + await socket.close() + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_close_discards_connection_that_finishes_after_shutdown(monkeypatch): + client = MemoryEventClient( + "http://agentfield", SimpleNamespace(to_headers=lambda: {}) + ) + connecting = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + + class Socket: + open = True + + async def recv(self): + await asyncio.Event().wait() + + async def close(self): + self.open = False + closed.set() + + async def connect(url, **kwargs): + connecting.set() + await release.wait() + return Socket() + + monkeypatch.setattr("agentfield.memory_events.websockets.connect", connect) + pending = asyncio.create_task(client.connect()) + try: + await asyncio.wait_for(connecting.wait(), timeout=1) + await client.close() + release.set() + await asyncio.wait_for(pending, timeout=1) + + assert closed.is_set() + assert client.websocket is None + assert not client.is_listening + finally: + release.set() + await pending + await client.close() + + +@pytest.mark.asyncio +async def test_close_discards_connection_queued_behind_inflight_connect(monkeypatch): + client = MemoryEventClient( + "http://agentfield", SimpleNamespace(to_headers=lambda: {}) + ) + connecting = asyncio.Event() + queued = asyncio.Event() + release = asyncio.Event() + attempts = [] + + class Socket: + open = True + + async def recv(self): + await asyncio.Event().wait() + + async def close(self): + self.open = False + + async def connect(url, **kwargs): + socket = Socket() + attempts.append(socket) + connecting.set() + await release.wait() + return socket + + async def connect_again(): + queued.set() + await client.connect() + + monkeypatch.setattr("agentfield.memory_events.websockets.connect", connect) + pending = [asyncio.create_task(client.connect())] + try: + await asyncio.wait_for(connecting.wait(), timeout=1) + pending.append(asyncio.create_task(connect_again())) + await asyncio.wait_for(queued.wait(), timeout=1) + + await client.close() + release.set() + await asyncio.wait_for(asyncio.gather(*pending), timeout=1) + + assert len(attempts) == 1 + assert not attempts[0].open + assert client.websocket is None + assert not client.is_listening + finally: + release.set() + await asyncio.gather(*pending, return_exceptions=True) + await client.close() + + +@pytest.mark.asyncio +async def test_close_cancels_retry_backoff(monkeypatch): + client = MemoryEventClient( + "http://agentfield", SimpleNamespace(to_headers=lambda: {}) + ) + sleeping = asyncio.Event() + stopped = asyncio.Event() + release = asyncio.Event() + attempts = [] + + async def connect(url, **kwargs): + attempts.append(url) + raise ConnectionRefusedError("offline") + + async def sleep(delay): + sleeping.set() + try: + await release.wait() + finally: + stopped.set() + + monkeypatch.setattr("agentfield.memory_events.websockets.connect", connect) + monkeypatch.setattr("agentfield.memory_events.asyncio.sleep", sleep) + try: + await client.connect() + await asyncio.wait_for(sleeping.wait(), timeout=1) + + await client.close() + + assert stopped.is_set() + assert len(attempts) == 1 + assert client.websocket is None + finally: + release.set() + await client.close() + + +@pytest.mark.asyncio +async def test_cancelled_close_still_closes_transport(monkeypatch): + client = MemoryEventClient( + "http://agentfield", SimpleNamespace(to_headers=lambda: {}) + ) + receiving = asyncio.Event() + stopping = asyncio.Event() + release = asyncio.Event() + closed = asyncio.Event() + + class Socket: + open = True + + async def recv(self): + receiving.set() + try: + await asyncio.Event().wait() + finally: + stopping.set() + await release.wait() + + async def close(self): + self.open = False + closed.set() + + socket = Socket() + + async def connect(url, **kwargs): + return socket + + monkeypatch.setattr("agentfield.memory_events.websockets.connect", connect) + await client.connect() + await asyncio.wait_for(receiving.wait(), timeout=1) + closing = asyncio.create_task(client.close()) + try: + await asyncio.wait_for(stopping.wait(), timeout=1) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + + assert closed.is_set() + finally: + release.set() + if not closing.done(): + closing.cancel() + await asyncio.gather(closing, return_exceptions=True) + await client.close() + await socket.close() + + +@pytest.mark.asyncio +async def test_client_can_explicitly_reconnect_after_close(): + received = [] + events = asyncio.Queue() + connections = 0 + + async def handler(socket, path=None): + nonlocal connections + connections += 1 + await socket.send( + json.dumps( + {"scope": "global", "key": f"connection.{connections}", "action": "set"} + ) + ) + await socket.wait_closed() + + async def callback(event): + received.append(event.key) + events.put_nowait(event.key) + + async with websockets.serve(handler, "127.0.0.1", 0) as server: + port = server.sockets[0].getsockname()[1] + client = MemoryEventClient( + f"http://127.0.0.1:{port}", SimpleNamespace(to_headers=lambda: {}) + ) + try: + for key in ("connection.1", "connection.2"): + await client.connect() + client.subscribe("connection.*", callback) + assert await asyncio.wait_for(events.get(), timeout=1) == key + await client.close() + finally: + await client.close() + + assert received == ["connection.1", "connection.2"] + + @pytest.mark.asyncio async def test_connect_passes_correct_headers_kwarg(monkeypatch): """Verify connect() uses the version-appropriate header parameter.""" diff --git a/sdk/python/tests/test_memory_events_additional.py b/sdk/python/tests/test_memory_events_additional.py index 55706b77d..4f64c636d 100644 --- a/sdk/python/tests/test_memory_events_additional.py +++ b/sdk/python/tests/test_memory_events_additional.py @@ -123,7 +123,7 @@ async def test_subscribe_on_change_close_and_scoped_client(monkeypatch, executio def fake_create_task(coro): scheduled.append(coro) coro.close() - return SimpleNamespace() + return asyncio.get_running_loop().create_future() monkeypatch.setattr("agentfield.memory_events.asyncio.create_task", fake_create_task)