From eb803b2ca18c625a45f0153fd807cbd9e1b5ecea Mon Sep 17 00:00:00 2001 From: Xuxchloris <7482714452@qq.com> Date: Fri, 14 Aug 2026 19:51:14 +0000 Subject: [PATCH] feat(ws): expose websocket proxy/connect options on Client lark.ws.Client forced a direct connection (proxy=None on websockets >= 15) with no supported way to honor HTTP_PROXY/HTTPS_PROXY/ALL_PROXY/NO_PROXY or pass explicit websockets.connect() options. Add an optional ws_connect_kwargs parameter to Client.__init__: when provided, the forced direct-connect flag is dropped (so proxies work) and caller values merge over SDK defaults; the default stays exactly as before (issue #143). 5 new unit tests cover the default path, explicit proxy, env-discovery opt-in, and kwarg merging. --- lark_oapi/ws/client.py | 107 ++++++++++++++++--- lark_oapi/ws/tests/test_ws_connect_kwargs.py | 50 +++++++++ 2 files changed, 140 insertions(+), 17 deletions(-) create mode 100644 lark_oapi/ws/tests/test_ws_connect_kwargs.py diff --git a/lark_oapi/ws/client.py b/lark_oapi/ws/client.py index 8ee991838..8413da474 100644 --- a/lark_oapi/ws/client.py +++ b/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,18 +274,18 @@ 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: @@ -220,7 +293,7 @@ async def _receive_message_loop(self): 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 = "" diff --git a/lark_oapi/ws/tests/test_ws_connect_kwargs.py b/lark_oapi/ws/tests/test_ws_connect_kwargs.py new file mode 100644 index 000000000..c4e4ce6a7 --- /dev/null +++ b/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 == {}