diff --git a/lark_oapi/ws/client.py b/lark_oapi/ws/client.py index 8ee99183..fe80faf9 100644 --- a/lark_oapi/ws/client.py +++ b/lark_oapi/ws/client.py @@ -28,11 +28,11 @@ from lark_oapi.ws.pb.google.protobuf.internal.containers import RepeatedCompositeFieldContainer from lark_oapi.ws.pb.pbbp2_pb2 import Frame -try: - loop = asyncio.get_event_loop() -except RuntimeError: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) +# Backward-compatible module-level reference. Deprecated — new code should +# rely on the per-instance ``Client._loop`` instead. Retained so that any +# third-party code importing ``from lark_oapi.ws.client import loop`` does +# not break immediately; it will point to the *last* loop set by start(). +loop: Optional[asyncio.AbstractEventLoop] = None def _get_by_key(headers: RepeatedCompositeFieldContainer, key: str) -> str: @@ -65,11 +65,6 @@ def _ordinal(n: int): return str(n) + suffix -async def _select(): - while True: - await asyncio.sleep(3600) - - def _ws_connect_kwargs(): params = inspect.signature(websockets.connect).parameters if "proxy" in params: @@ -142,6 +137,8 @@ def __init__(self, self._conn_url: str = "" self._service_id: str = "" self._conn_id: str = "" + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._stop_event: Optional[asyncio.Event] = None # Local defaults; the Feishu WS endpoint authoritatively replaces these # via _configure() on every handshake (and may push updates mid-session # via CONTROL frames). Matches node-sdk parent SDK — user-facing @@ -162,24 +159,78 @@ def __init__(self, logger.setLevel(log_level.value) def start(self) -> None: + """Start the WebSocket client (blocking). + + Creates a new event loop for this client instance so that multiple + Client objects can run concurrently in separate threads without + interfering with each other. + """ + global loop + + self._loop = asyncio.new_event_loop() + asyncio.set_event_loop(self._loop) + # Keep module-level reference updated for backward compatibility. + loop = self._loop + self._stop_event = asyncio.Event() + + try: + self._loop.run_until_complete(self._run()) + except ClientException: + raise + finally: + self._loop.close() + self._loop = None + self._stop_event = None + + async def start_async(self) -> None: + """Start the WebSocket client within an already-running event loop. + + Use this when integrating with async frameworks (FastAPI, aiohttp, etc.) + where an event loop is already running:: + + async def main(): + client = Client(app_id, app_secret, event_handler=handler) + await client.start_async() + """ + self._loop = asyncio.get_running_loop() + self._stop_event = asyncio.Event() + try: + await self._run() + finally: + self._loop = None + self._stop_event = None + + def stop(self) -> None: + """Gracefully stop the client (thread-safe). + + Can be called from any thread. The client will disconnect and + ``start()`` / ``start_async()`` will return. + """ + if self._loop and self._stop_event: + self._loop.call_soon_threadsafe(self._stop_event.set) + + async def _run(self) -> None: + """Core run loop: connect, dispatch, wait for stop signal.""" try: - loop.run_until_complete(self._connect()) + await self._connect() except ClientException as e: logger.error(self._fmt_log("connect failed, err: {}", e)) raise e except Exception as e: logger.error(self._fmt_log("connect failed, err: {}", e)) - loop.run_until_complete(self._disconnect()) + await self._disconnect() if self._auto_reconnect: - loop.run_until_complete(self._reconnect()) + await self._reconnect() else: raise e - loop.create_task(self._ping_loop()) - loop.run_until_complete(_select()) + self._loop.create_task(self._ping_loop()) + # Block until stop() is called or the event loop is closed. + await self._stop_event.wait() + await self._disconnect() async def _ping_loop(self): - while True: + while not self._stop_event.is_set(): try: if self._conn is not None: frame = _new_ping_frame(int(self._service_id)) @@ -187,8 +238,13 @@ async def _ping_loop(self): logger.debug(self._fmt_log("ping success")) except Exception as e: logger.warn(self._fmt_log("ping failed, err: {}", e)) - finally: - await asyncio.sleep(self._ping_interval) + # Use wait_for so that stop() can interrupt the sleep. + try: + await asyncio.wait_for(self._stop_event.wait(), + timeout=self._ping_interval) + return # stop_event was set + except asyncio.TimeoutError: + pass async def _connect(self) -> None: await self._lock.acquire() @@ -208,7 +264,7 @@ async def _connect(self) -> None: self._service_id = service_id logger.info(self._fmt_log("connected to {}", conn_url)) - loop.create_task(self._receive_message_loop()) + self._loop.create_task(self._receive_message_loop()) except InvalidHandshake as e: _parse_ws_conn_exception(e) finally: @@ -219,14 +275,23 @@ async def _receive_message_loop(self): while True: if self._conn is None: raise ConnectionClosedException("connection is closed") - msg = await self._conn.recv() - loop.create_task(self._handle_message(msg)) + if self._stop_event.is_set(): + return + msg = await asyncio.wait_for(self._conn.recv(), timeout=300) + self._loop.create_task(self._handle_message(msg)) except Exception as e: + if self._stop_event.is_set(): + return logger.error(self._fmt_log("receive message loop exit, err: {}", e)) await self._disconnect() if self._auto_reconnect: - await self._reconnect() + try: + await self._reconnect() + except Exception as re: + logger.error(self._fmt_log("reconnect failed permanently, err: {}", re)) + self._stop_event.set() else: + self._stop_event.set() raise e def _get_conn_url(self) -> str: @@ -369,24 +434,40 @@ async def _reconnect(self): # 首次重连随机抖动 if self._reconnect_nonce > 0: nonce = random.random() * self._reconnect_nonce - await asyncio.sleep(nonce) + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=nonce) + return # stop() was called during nonce delay + except asyncio.TimeoutError: + pass # 重连 if self._reconnect_count >= 0: for i in range(self._reconnect_count): + if self._stop_event.is_set(): + return if await self._try_connect(i): self._fire_on_reconnected() return - await asyncio.sleep(self._reconnect_interval) + try: + await asyncio.wait_for(self._stop_event.wait(), + timeout=self._reconnect_interval) + return # stop() was called during interval + except asyncio.TimeoutError: + pass raise ServerUnreachableException( f"unable to connect to the server after trying {self._reconnect_count} times") else: i = 0 - while True: + while not self._stop_event.is_set(): if await self._try_connect(i): self._fire_on_reconnected() return - await asyncio.sleep(self._reconnect_interval) + try: + await asyncio.wait_for(self._stop_event.wait(), + timeout=self._reconnect_interval) + return # stop() was called during interval + except asyncio.TimeoutError: + pass i += 1 def _fire_on_reconnected(self) -> None: diff --git a/lark_oapi/ws/tests/test_multi_instance.py b/lark_oapi/ws/tests/test_multi_instance.py new file mode 100644 index 00000000..e5bbcfb8 --- /dev/null +++ b/lark_oapi/ws/tests/test_multi_instance.py @@ -0,0 +1,192 @@ +"""Tests for per-instance event loop isolation and graceful stop. + +Verifies that multiple Client instances can run concurrently in separate +threads without interfering with each other (the root cause of the +"Future attached to a different loop" error in multi-bot deployments). +""" +import asyncio +import threading +import time +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from lark_oapi.ws import client as ws_client +from lark_oapi.ws.client import Client + + +class _FakeConn: + """Minimal fake WebSocket connection.""" + + def __init__(self): + self._closed = False + self._recv_event = asyncio.Event() + + async def recv(self): + await self._recv_event.wait() + raise Exception("connection closed") + + async def send(self, data): + pass + + async def close(self): + self._closed = True + self._recv_event.set() + + +def _make_client(): + """Create a Client with no real network.""" + return Client("app_id", "app_secret", auto_reconnect=False) + + +def _patch_connect(client, fake_conn): + """Patch _get_conn_url and websockets.connect to avoid network.""" + client._get_conn_url = lambda: "ws://fake/callback?device_id=d1&service_id=1" + original_connect = client._connect + + async def patched_connect(): + await client._lock.acquire() + if client._conn is not None: + client._lock.release() + return + try: + client._conn = fake_conn + client._conn_url = "ws://fake" + client._conn_id = "d1" + client._service_id = "1" + client._loop.create_task(client._receive_message_loop()) + finally: + client._lock.release() + + client._connect = patched_connect + + +class TestPerInstanceLoop: + """Each Client.start() creates its own event loop.""" + + def test_two_clients_have_independent_loops(self): + """Two clients started in separate threads use different loops.""" + loops_seen = [] + barrier = threading.Barrier(2, timeout=5) + + def run_client(idx): + client = _make_client() + conn = _FakeConn() + _patch_connect(client, conn) + + # Intercept _run to capture the loop then stop + original_run = client._run + + async def capture_and_stop(): + loops_seen.append(asyncio.get_running_loop()) + barrier.wait() # sync both threads + client._stop_event.set() + + client._run = capture_and_stop + client.start() + + t1 = threading.Thread(target=run_client, args=(1,)) + t2 = threading.Thread(target=run_client, args=(2,)) + t1.start() + t2.start() + t1.join(timeout=5) + t2.join(timeout=5) + + assert len(loops_seen) == 2 + assert loops_seen[0] is not loops_seen[1] + + def test_stop_terminates_start(self): + """Calling stop() from another thread causes start() to return.""" + client = _make_client() + conn = _FakeConn() + _patch_connect(client, conn) + + started = threading.Event() + + original_run = client._run + + async def run_with_signal(): + started.set() + await original_run() + + client._run = run_with_signal + + t = threading.Thread(target=client.start) + t.start() + + assert started.wait(timeout=3), "start() did not begin" + time.sleep(0.1) # let it settle + client.stop() + t.join(timeout=3) + assert not t.is_alive(), "start() did not return after stop()" + + def test_stop_interrupts_reconnect(self): + """stop() interrupts the reconnect sleep rather than waiting 120s.""" + client = Client("app_id", "app_secret", auto_reconnect=True) + + connect_attempts = [] + + async def failing_connect(): + connect_attempts.append(time.monotonic()) + raise Exception("simulated network failure") + + client._connect = failing_connect + + t0 = time.monotonic() + + def run(): + client.start() + + t = threading.Thread(target=run) + t.start() + time.sleep(0.5) # let it start reconnect loop + client.stop() + t.join(timeout=5) + + elapsed = time.monotonic() - t0 + # Should stop within a few seconds, not wait the full 120s interval + assert elapsed < 10, f"Took {elapsed:.1f}s, expected < 10s" + + +@pytest.mark.asyncio +class TestStartAsync: + """start_async() reuses the caller's event loop.""" + + async def test_uses_running_loop(self): + """start_async() uses asyncio.get_running_loop(), not a new one.""" + client = _make_client() + conn = _FakeConn() + _patch_connect(client, conn) + + current_loop = asyncio.get_running_loop() + + async def stop_after_start(): + await asyncio.sleep(0.1) + client.stop() + + current_loop.create_task(stop_after_start()) + await client.start_async() + + # After start_async returns, instance loop should be cleaned up + assert client._loop is None + assert client._stop_event is None + + +class TestBackwardCompatibility: + """Module-level `loop` variable is still updated for old code.""" + + def test_module_loop_updated_after_start(self): + """After start(), ws_client.loop points to the last used loop.""" + client = _make_client() + + async def immediate_stop(): + client._stop_event.set() + + client._run = immediate_stop + client.start() + + # loop was set during start() but is now closed + # The important thing is it was set (not None) + # It will be the closed loop from the last start() call + assert ws_client.loop is not None diff --git a/lark_oapi/ws/tests/test_websockets_compat.py b/lark_oapi/ws/tests/test_websockets_compat.py index 538d1ee9..ff7041d7 100644 --- a/lark_oapi/ws/tests/test_websockets_compat.py +++ b/lark_oapi/ws/tests/test_websockets_compat.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +import asyncio import pytest from lark_oapi.ws import client as ws_client @@ -17,6 +18,10 @@ async def close(self): pass +async def _noop(): + pass + + def test_parse_ws_connection_exception_reads_new_invalid_status_response_headers(): exc = RuntimeError("handshake failed") exc.response = SimpleNamespace( @@ -92,17 +97,15 @@ async def fake_connect(uri, *, proxy=True): return _FakeConn() client = ws_client.Client("app_id", "app_secret") + client._loop = asyncio.get_running_loop() + client._stop_event = asyncio.Event() monkeypatch.setattr( client, "_get_conn_url", lambda: "ws://example.test/callback?device_id=device&service_id=42", ) monkeypatch.setattr(ws_client.websockets, "connect", fake_connect) - monkeypatch.setattr( - ws_client.loop, - "create_task", - lambda coro: coro.close() if hasattr(coro, "close") else None, - ) + monkeypatch.setattr(client, "_receive_message_loop", _noop) await client._connect() await client._disconnect() @@ -122,17 +125,15 @@ async def fake_connect(uri): return _FakeConn() client = ws_client.Client("app_id", "app_secret") + client._loop = asyncio.get_running_loop() + client._stop_event = asyncio.Event() monkeypatch.setattr( client, "_get_conn_url", lambda: "ws://example.test/callback?device_id=device&service_id=42", ) monkeypatch.setattr(ws_client.websockets, "connect", fake_connect) - monkeypatch.setattr( - ws_client.loop, - "create_task", - lambda coro: coro.close() if hasattr(coro, "close") else None, - ) + monkeypatch.setattr(client, "_receive_message_loop", _noop) await client._connect() await client._disconnect()