feat(ws): expose websocket proxy/connect options on Client

lark_oapi/ws/client.py

@@ -5,8 +5,8 @@
 import json
 import random
 import time
-from typing import Callable, Dict, Mapping, Optional
-from urllib.parse import urlparse, parse_qs
+from typing import Any, Callable, Dict, Mapping, Optional
+from urllib.parse import urlparse, parse_qs, urlunparse, urlencode
 
 import requests
 import websockets
@@ -28,12 +28,6 @@
 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)
-
 
 def _get_by_key(headers: RepeatedCompositeFieldContainer, key: str) -> str:
     for header in headers:
@@ -43,6 +37,31 @@ def _get_by_key(headers: RepeatedCompositeFieldContainer, key: str) -> str:
     raise HeaderNotFoundException(key)
 
 
+# Query parameters on the WS endpoint URL whose values are credentials
+# (access_key, ticket) and must never appear in logs.
+_SENSITIVE_QUERY_KEYS = ("access_key", "ticket")
+
+
+def _redact_conn_url(url: Optional[str]) -> Optional[str]:
+    """Return a copy of ``url`` safe for logging.
+
+    The WS endpoint URL carries ``access_key`` and ``ticket`` credentials as
+    query parameters; those values are masked while the rest of the URL
+    (scheme, host, path, other params) is preserved. The returned URL must
+    only be used for logging - never for connecting.
+    """
+    if not url:
+        return url
+    u = urlparse(url)
+    if not u.query:
+        return url
+    q = parse_qs(u.query, keep_blank_values=True)
+    for key in _SENSITIVE_QUERY_KEYS:
+        if key in q:
+            q[key] = ["***"]
+    return urlunparse(u._replace(query=urlencode(q, doseq=True)))
+
+
 def _new_ping_frame(service_id: int) -> Frame:
     frame = Frame()
     header = frame.headers.add()
@@ -125,7 +144,8 @@ def __init__(self,
                  source: Optional[str] = None,
                  extra_ua_tags: Optional[list] = None,
                  headers: Optional[Mapping[str, str]] = None,
-                 client_assertion_provider=None) -> None:
+                 client_assertion_provider=None,
+                 ws_connect_kwargs: Optional[Mapping[str, Any]] = None) -> None:
         self._app_id: str = app_id
         self._app_secret: str = app_secret
         self._log_level: LogLevel = log_level
@@ -134,6 +154,11 @@ def __init__(self,
         self._domain: str = domain
         self._client_assertion_provider = client_assertion_provider
         self._headers: Dict[str, str] = dict(headers or {})
+        # Extra kwargs forwarded to websockets.connect(). None keeps the
+        # historical direct-connect behavior (see _resolved_ws_connect_kwargs).
+        self._ws_connect_kwargs: Optional[Dict[str, Any]] = (
+            dict(ws_connect_kwargs) if ws_connect_kwargs is not None else None
+        )
         # UA used on the endpoint-discovery POST (and any future HTTP/WS
         # handshakes from this client). ``extra_ua_tags`` is internal — sub-
         # modules (e.g. FeishuChannel) pass ``["channel"]`` here.
@@ -151,7 +176,13 @@ def __init__(self,
         self._reconnect_interval: int = 120
         self._ping_interval: int = 120
         self._cache: ExpiringCache = ExpiringCache(clear_interval=30)
-        self._lock = asyncio.Lock()
+        # Event loop owned by this client instance (created lazily on first
+        # use). Previously a module-level global was shared by every client,
+        # which broke multi-bot setups (issues #119 / #133).
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        # asyncio primitives bind to an event loop; created lazily on the
+        # client's own loop so it never binds to a foreign loop.
+        self._lock: Optional[asyncio.Lock] = None
         # Observer hooks for higher-level wrappers (e.g. FeishuChannel) to
         # react to reconnect lifecycle. ``on_reconnecting`` fires when the
         # client decides a connection was lost and starts retrying;
@@ -161,7 +192,48 @@ def __init__(self,
         self.on_reconnected: Callable[[], None] = lambda: None
         logger.setLevel(log_level.value)
 
+    def _get_loop(self) -> asyncio.AbstractEventLoop:
+        """Return the event loop this client runs on.
+
+        Every client owns a dedicated loop, created lazily on first use, so
+        multiple clients (multi-bot, one thread per bot) never share a loop,
+        and a client constructed inside an already-running loop (e.g. inside
+        ``asyncio.run``) is not tied to it. See issues #119 / #133.
+        """
+        if self._loop is None or self._loop.is_closed():
+            self._loop = asyncio.new_event_loop()
+        return self._loop
+
+    def _resolved_ws_connect_kwargs(self) -> Dict[str, Any]:
+        """Return the kwargs passed to ``websockets.connect()``.
+
+        Default (``ws_connect_kwargs=None``) preserves the SDK's historical
+        direct connection: on websockets >= 15, where the library enables
+        environment proxy discovery by default, it is explicitly disabled.
+
+        When the caller supplies ``ws_connect_kwargs``, the forced
+        direct-connect flag is dropped so explicit proxies (or, on
+        websockets >= 15, environment discovery via HTTP_PROXY /
+        HTTPS_PROXY / ALL_PROXY / NO_PROXY) work; caller values win over
+        SDK defaults (issue #143).
+        """
+        base = _ws_connect_kwargs()
+        if self._ws_connect_kwargs is None:
+            return base
+        base.pop("proxy", None)
+        base.update(self._ws_connect_kwargs)
+        return base
+
+    def _get_lock(self) -> asyncio.Lock:
+        # asyncio primitives bind to the running loop when created; create the
+        # lock lazily (from a coroutine running on the client's own loop) so
+        # it never binds to a foreign loop.
+        if self._lock is None:
+            self._lock = asyncio.Lock()
+        return self._lock
+
     def start(self) -> None:
+        loop = self._get_loop()
         try:
             loop.run_until_complete(self._connect())
         except ClientException as e:
@@ -191,7 +263,8 @@ async def _ping_loop(self):
                 await asyncio.sleep(self._ping_interval)
 
     async def _connect(self) -> None:
-        await self._lock.acquire()
+        lock = self._get_lock()
+        await lock.acquire()
         if self._conn is not None:
             return
         try:
@@ -201,26 +274,26 @@ async def _connect(self) -> None:
             conn_id = q[DEVICE_ID][0]
             service_id = q[SERVICE_ID][0]
 
-            conn = await websockets.connect(conn_url, **_ws_connect_kwargs())
+            conn = await websockets.connect(conn_url, **self._resolved_ws_connect_kwargs())
             self._conn = conn
             self._conn_url = conn_url
             self._conn_id = conn_id
             self._service_id = service_id
 
-            logger.info(self._fmt_log("connected to {}", conn_url))
-            loop.create_task(self._receive_message_loop())
+            logger.info(self._fmt_log("connected to {}", _redact_conn_url(conn_url)))
+            asyncio.get_running_loop().create_task(self._receive_message_loop())
         except InvalidHandshake as e:
             _parse_ws_conn_exception(e)
         finally:
-            self._lock.release()
+            lock.release()
 
     async def _receive_message_loop(self):
         try:
             while True:
                 if self._conn is None:
                     raise ConnectionClosedException("connection is closed")
                 msg = await self._conn.recv()
-                loop.create_task(self._handle_message(msg))
+                asyncio.get_running_loop().create_task(self._handle_message(msg))
         except Exception as e:
             logger.error(self._fmt_log("receive message loop exit, err: {}", e))
             await self._disconnect()
@@ -413,7 +486,7 @@ async def _disconnect(self):
             if self._conn is None:
                 return
             await self._conn.close()
-            logger.info(self._fmt_log("disconnected to {}", self._conn_url))
+            logger.info(self._fmt_log("disconnected to {}", _redact_conn_url(self._conn_url)))
         finally:
             self._conn = None
             self._conn_url = ""

lark_oapi/ws/tests/test_ws_connect_kwargs.py

@@ -0,0 +1,50 @@
+import inspect
+
+import websockets
+
+from lark_oapi.ws.client import Client, _ws_connect_kwargs
+
+
+def make_client(**kwargs):
+    return Client("app_id", "app_secret", **kwargs)
+
+
+def test_default_preserves_historical_direct_connect():
+    c = make_client()
+    assert c._resolved_ws_connect_kwargs() == _ws_connect_kwargs()
+
+
+def test_explicit_proxy_kwarg_wins_over_forced_direct_connect():
+    c = make_client(ws_connect_kwargs={"proxy": "http://proxy.example:8080"})
+    resolved = c._resolved_ws_connect_kwargs()
+    assert resolved["proxy"] == "http://proxy.example:8080"
+
+
+def test_env_discovery_enabled_when_caller_opt_in_without_proxy():
+    # On websockets >= 15 the SDK forces proxy=None by default; an explicit
+    # (even empty) ws_connect_kwargs drops that flag so HTTP_PROXY etc. work.
+    c = make_client(ws_connect_kwargs={})
+    resolved = c._resolved_ws_connect_kwargs()
+    assert "proxy" not in resolved
+    # Other legacy kwargs (if any on this websockets version) are preserved.
+    base = _ws_connect_kwargs()
+    base.pop("proxy", None)
+    assert resolved == base
+
+
+def test_caller_values_merge_with_sdk_defaults():
+    c = make_client(ws_connect_kwargs={"max_size": 2 ** 21})
+    resolved = c._resolved_ws_connect_kwargs()
+    assert resolved["max_size"] == 2 ** 21
+    if "proxy" in inspect.signature(websockets.connect).parameters:
+        # proxy: None was dropped, not merged.
+        assert "proxy" not in resolved
+
+
+def test_connect_signature_detection_still_works():
+    params = inspect.signature(websockets.connect).parameters
+    base = _ws_connect_kwargs()
+    if "proxy" in params:
+        assert base == {"proxy": None}
+    else:
+        assert base == {}