mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
fix(channels): make Feishu and DingTalk transports stoppable in-process (#6798)
* fix(channels): make Feishu and DingTalk transports stoppable in-process
FeishuChannel.stop() and DingTalkChannel.stop() could not end their SDK
threads. Both SDKs run on a thread-private event loop that only returned
on a fatal error: lark-oapi's ws.Client.start() parks forever in
run_until_complete(_select()) and has no public stop, and the channel
never kept the client handle; dingtalk-stream's start_forever() is a
while-True loop that also catches CancelledError, and the disconnect()
the channel called does not exist on the SDK. So every stop joined for
5s, raised ChannelStopTimeout, and ChannelService retained the instance
while the platform connection stayed alive on the old process. That
burns the full budget on shutdown and on POST /api/channels/{name}/restart,
and it blocks the per-provider leader lease: a demoted leader could not
release the transport before a peer took over.
Both channels now publish the SDK thread's loop and client under a lock;
stop() schedules a shutdown on that loop and then joins. Feishu disarms
lark's _auto_reconnect before closing the socket and stopping the loop,
because start() otherwise swallows the interrupted run_until_complete in
its reconnect branch and dials again. DingTalk drives its own copy of the
SDK connect loop that re-checks the running flag before every dial, runs
the blocking ticket request off the loop, cancels keepalive with the
socket, and logs the endpoint rather than the ticket-bearing payload. A
shared shutdown_private_loop cancels the SDK's leftover tasks with a
bounded drain and closes the loop; a stop requested before the client
exists never dials. ChannelStopTimeout now only fires for a thread that
genuinely wedges inside an SDK call. Telegram already stopped this way.
Tests: a fake lark ws.Client with the SDK's start() shape (including the
reconnect branch that swallows a bare loop stop), a real local websockets
server that observes the DingTalk close, the stop-before-client race for
both, and a blocking-io anchor that stop() ends a live private loop
without blocking the Gateway loop. All eight were red on main.
* docs(changelog): link the stoppable-transports entry to #6798
* fix(channels): release the socket on a stalled close and bound the DingTalk ticket request
Review follow-ups on #6798 (Totoro-qaq):
Feishu: the bounded _disconnect() wait cancelled websockets' close before
its own close-timeout abort ran, and lark clears _conn in its finally, so
a peer that never answers the close frame left the socket to garbage
collection. Worse, a stop landing while _connect() still held lark's
_lock snapshotted _conn as None. _shutdown_ws now settles the in-flight
dial by taking that lock under one shared close deadline, snapshots the
connection, runs _disconnect() under the same deadline, and aborts the
transport itself (abort_connection_transport in base.py) when the
handshake did not finish, yielding once so connection_lost runs before
loop.stop(). DingTalk's websocket.close() path gets the same abort.
DingTalk: the SDK's open_connection() calls requests.post with no
timeout, and running it in an executor let a stalled ticket endpoint
block interpreter exit (executor workers are joined at shutdown). The
channel now sends the same request through httpx on the stream loop
under a 30s deadline, so stop() cancels it and no worker thread lingers.
Tests drive the real SDK clients against loopback peers: a WebSocket
peer that completes the handshake and never ACKs the close (both
channels assert the connection's socket reports fileno() == -1 right
after stop(); red with the abort disabled), a real websockets server for
the clean path, and a loopback ticket server (ok / 500 / never responds)
with a no-lingering-non-daemon-thread check and a request-shape check.
This commit is contained in:
parent
2cf7424a80
commit
5841c8e9b7
26
CHANGELOG.md
26
CHANGELOG.md
@ -1226,6 +1226,31 @@ This release closes that milestone with **621 merged pull requests**.
|
|||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
|
- **channels:** `FeishuChannel.stop()` and `DingTalkChannel.stop()` now end
|
||||||
|
their SDK threads instead of timing out. Both SDKs run on a thread-private
|
||||||
|
event loop that only returned on a fatal error, so every stop joined for
|
||||||
|
5s, raised `ChannelStopTimeout`, and left the platform connection alive on
|
||||||
|
the old process — a restart, a redeploy, and the upcoming per-provider
|
||||||
|
leader hand-off could not release the connection. Each channel now
|
||||||
|
publishes that loop under a lock; `stop()` schedules a shutdown on it, then
|
||||||
|
joins. Feishu disarms lark-oapi's `_auto_reconnect` before closing the
|
||||||
|
socket and stopping the loop, because the SDK's `start()` otherwise swallows
|
||||||
|
the interrupted loop in its reconnect branch and dials again; DingTalk
|
||||||
|
drives its own copy of dingtalk-stream's connect loop (the SDK's is
|
||||||
|
`while True` and catches `CancelledError`), re-checks the running flag
|
||||||
|
before every dial, sends the ticket request itself through `httpx` with a
|
||||||
|
deadline (the SDK's `requests.post` has none, so a stalled endpoint could
|
||||||
|
pin the thread or process exit), and logs only the endpoint rather than
|
||||||
|
the ticket-bearing connection payload. When a peer never answers the close
|
||||||
|
frame, both channels abort the transport themselves instead of leaving the
|
||||||
|
socket to garbage collection, because the bounded close wait cancels
|
||||||
|
websockets' own close-timeout abort. A shared `shutdown_private_loop`
|
||||||
|
cancels the SDK's leftover tasks with a bounded drain and closes the loop,
|
||||||
|
and a stop requested before the client exists never dials.
|
||||||
|
`ChannelStopTimeout` is now reserved for a thread that genuinely wedges
|
||||||
|
inside an SDK call. Telegram already stopped its loop this way and is
|
||||||
|
unchanged. ([#6798])
|
||||||
|
|
||||||
- **frontend:** The user-management settings page no longer renders a doomed
|
- **frontend:** The user-management settings page no longer renders a doomed
|
||||||
query on the static-website demo, and a non-admin `?settings=users` deep
|
query on the static-website demo, and a non-admin `?settings=users` deep
|
||||||
link no longer opens the admin-only panel. The page gated only on
|
link no longer opens the admin-only panel. The page gated only on
|
||||||
@ -11608,3 +11633,4 @@ with **180 merged pull requests** since the first 2.0 milestone tag.
|
|||||||
[#6702]: https://github.com/bytedance/deer-flow/pull/6702
|
[#6702]: https://github.com/bytedance/deer-flow/pull/6702
|
||||||
[#6706]: https://github.com/bytedance/deer-flow/pull/6706
|
[#6706]: https://github.com/bytedance/deer-flow/pull/6706
|
||||||
[#6746]: https://github.com/bytedance/deer-flow/pull/6746
|
[#6746]: https://github.com/bytedance/deer-flow/pull/6746
|
||||||
|
[#6798]: https://github.com/bytedance/deer-flow/pull/6798
|
||||||
|
|||||||
File diff suppressed because one or more lines are too long
@ -43,6 +43,50 @@ class ChannelStopTimeout(RuntimeError):
|
|||||||
"""A channel still owns a live provider worker after bounded teardown."""
|
"""A channel still owns a live provider worker after bounded teardown."""
|
||||||
|
|
||||||
|
|
||||||
|
PRIVATE_LOOP_DRAIN_TIMEOUT_SECONDS = 1.0
|
||||||
|
|
||||||
|
|
||||||
|
def shutdown_private_loop(loop: asyncio.AbstractEventLoop, *, drain_timeout: float = PRIVATE_LOOP_DRAIN_TIMEOUT_SECONDS) -> None:
|
||||||
|
"""Cancel what is still scheduled on a thread-private SDK loop, then close it.
|
||||||
|
|
||||||
|
Runs on the thread that owns ``loop`` once its ``run_until_complete`` or
|
||||||
|
``run_forever`` has returned. Provider SDKs leave their own tasks behind
|
||||||
|
(lark-oapi's ping and receive loops, dingtalk-stream's keepalive); cancelling
|
||||||
|
them with a bounded drain keeps the thread exit inside the channel's stop
|
||||||
|
budget, and a task that refuses to finish is abandoned with the loop rather
|
||||||
|
than joined.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
|
||||||
|
for task in pending:
|
||||||
|
task.cancel()
|
||||||
|
if pending:
|
||||||
|
loop.run_until_complete(asyncio.wait(pending, timeout=drain_timeout))
|
||||||
|
loop.run_until_complete(loop.shutdown_asyncgens())
|
||||||
|
except Exception:
|
||||||
|
logger.debug("Error draining a private channel loop before close", exc_info=True)
|
||||||
|
finally:
|
||||||
|
asyncio.set_event_loop(None)
|
||||||
|
loop.close()
|
||||||
|
|
||||||
|
|
||||||
|
def abort_connection_transport(connection: Any) -> bool:
|
||||||
|
"""Tear down a ``websockets`` connection's transport if its close handshake did not finish.
|
||||||
|
|
||||||
|
``connection.close()`` waits for the peer's close frame and only aborts the
|
||||||
|
transport after its own ``close_timeout``; a caller that bounds that wait more
|
||||||
|
tightly (channel stop budgets do) cancels before the abort, and the socket
|
||||||
|
then lives until garbage collection. Returns True when a transport was
|
||||||
|
aborted; the caller should yield to its loop once so ``connection_lost``
|
||||||
|
runs before the loop is stopped.
|
||||||
|
"""
|
||||||
|
transport = getattr(connection, "transport", None)
|
||||||
|
if transport is None or transport.is_closing():
|
||||||
|
return False
|
||||||
|
transport.abort()
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
@dataclass(eq=False, slots=True)
|
@dataclass(eq=False, slots=True)
|
||||||
class _ThreadsafeSubmission:
|
class _ThreadsafeSubmission:
|
||||||
coroutine: Coroutine[Any, Any, Any]
|
coroutine: Coroutine[Any, Any, Any]
|
||||||
|
|||||||
@ -5,16 +5,18 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
import platform
|
||||||
import re
|
import re
|
||||||
import stat
|
import stat
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
from urllib.parse import quote_plus
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
from app.channels.base import Channel, ChannelStopTimeout
|
from app.channels.base import Channel, ChannelStopTimeout, abort_connection_transport, shutdown_private_loop
|
||||||
from app.channels.commands import is_known_channel_command, strip_leading_mentions
|
from app.channels.commands import is_known_channel_command, strip_leading_mentions
|
||||||
from app.channels.connection_identity import attach_connection_identity
|
from app.channels.connection_identity import attach_connection_identity
|
||||||
from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment
|
from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment
|
||||||
@ -36,6 +38,16 @@ DINGTALK_API_BASE = "https://api.dingtalk.com"
|
|||||||
|
|
||||||
_TOKEN_REFRESH_MARGIN_SECONDS = 300
|
_TOKEN_REFRESH_MARGIN_SECONDS = 300
|
||||||
|
|
||||||
|
# stop() budget for the dingtalk-stream thread: ask its private loop to close
|
||||||
|
# the socket (bounded by the close timeout) and stop, then join the thread.
|
||||||
|
DINGTALK_STOP_TIMEOUT_SECONDS = 5
|
||||||
|
DINGTALK_STREAM_CLOSE_TIMEOUT_SECONDS = 2
|
||||||
|
# Deadline for the Stream ticket request (the SDK's own has none).
|
||||||
|
DINGTALK_STREAM_TICKET_TIMEOUT_SECONDS = 30
|
||||||
|
# Retry pacing copied from DingTalkStreamClient.start() (dingtalk-stream 0.24).
|
||||||
|
_DINGTALK_STREAM_RETRY_SECONDS = 10
|
||||||
|
_DINGTALK_STREAM_ERROR_RETRY_SECONDS = 3
|
||||||
|
|
||||||
_CONVERSATION_TYPE_P2P = "1"
|
_CONVERSATION_TYPE_P2P = "1"
|
||||||
_CONVERSATION_TYPE_GROUP = "2"
|
_CONVERSATION_TYPE_GROUP = "2"
|
||||||
|
|
||||||
@ -182,7 +194,11 @@ class DingTalkChannel(Channel):
|
|||||||
self._card_template_id: str = config.get("card_template_id", "")
|
self._card_template_id: str = config.get("card_template_id", "")
|
||||||
self._card_track_ids: dict[str, str] = {}
|
self._card_track_ids: dict[str, str] = {}
|
||||||
self._dingtalk_client: Any = None
|
self._dingtalk_client: Any = None
|
||||||
|
# The SDK thread's private loop and client, published under _stream_lock
|
||||||
|
# once constructed so stop() can shut them down from the Gateway loop.
|
||||||
self._stream_client: Any = None
|
self._stream_client: Any = None
|
||||||
|
self._stream_loop: asyncio.AbstractEventLoop | None = None
|
||||||
|
self._stream_lock = threading.Lock()
|
||||||
self._incoming_messages: dict[str, Any] = {}
|
self._incoming_messages: dict[str, Any] = {}
|
||||||
self._incoming_messages_lock = threading.Lock()
|
self._incoming_messages_lock = threading.Lock()
|
||||||
self._card_repliers: dict[str, Any] = {}
|
self._card_repliers: dict[str, Any] = {}
|
||||||
@ -233,33 +249,64 @@ class DingTalkChannel(Channel):
|
|||||||
async def stop(self) -> None:
|
async def stop(self) -> None:
|
||||||
self._running = False
|
self._running = False
|
||||||
self.bus.unsubscribe_outbound(self._on_outbound)
|
self.bus.unsubscribe_outbound(self._on_outbound)
|
||||||
|
# Signal the SDK thread first so its loop winds down while the
|
||||||
|
# Gateway-side drain below runs.
|
||||||
|
self._request_stream_shutdown()
|
||||||
await self._close_and_drain_threadsafe_futures()
|
await self._close_and_drain_threadsafe_futures()
|
||||||
|
|
||||||
stream_client = self._stream_client
|
|
||||||
if stream_client is not None:
|
|
||||||
try:
|
|
||||||
if hasattr(stream_client, "disconnect"):
|
|
||||||
stream_client.disconnect()
|
|
||||||
except Exception:
|
|
||||||
logger.debug("[DingTalk] error disconnecting stream client", exc_info=True)
|
|
||||||
|
|
||||||
self._dingtalk_client = None
|
self._dingtalk_client = None
|
||||||
self._stream_client = None
|
|
||||||
with self._incoming_messages_lock:
|
with self._incoming_messages_lock:
|
||||||
self._incoming_messages.clear()
|
self._incoming_messages.clear()
|
||||||
self._card_repliers.clear()
|
self._card_repliers.clear()
|
||||||
self._card_track_ids.clear()
|
self._card_track_ids.clear()
|
||||||
if self._thread:
|
if self._thread:
|
||||||
# The SDK thread only returns on a fatal error, so this join
|
# The thread exits once its loop has stopped; keep the join off the
|
||||||
# normally waits out its full timeout; keep it off the event loop.
|
# event loop because a wedged SDK call can still use up the budget.
|
||||||
thread = self._thread
|
thread = self._thread
|
||||||
await asyncio.to_thread(thread.join, timeout=5)
|
await asyncio.to_thread(thread.join, timeout=DINGTALK_STOP_TIMEOUT_SECONDS)
|
||||||
if thread.is_alive():
|
if thread.is_alive():
|
||||||
raise ChannelStopTimeout("DingTalk SDK thread is still running after stop timeout")
|
raise ChannelStopTimeout("DingTalk SDK thread is still running after stop timeout")
|
||||||
if self._thread is thread:
|
if self._thread is thread:
|
||||||
self._thread = None
|
self._thread = None
|
||||||
|
with self._stream_lock:
|
||||||
|
self._stream_loop = None
|
||||||
|
self._stream_client = None
|
||||||
logger.info("DingTalk channel stopped")
|
logger.info("DingTalk channel stopped")
|
||||||
|
|
||||||
|
def _request_stream_shutdown(self) -> None:
|
||||||
|
"""Ask the SDK thread to close its socket and stop its loop (non-blocking)."""
|
||||||
|
with self._stream_lock:
|
||||||
|
loop = self._stream_loop
|
||||||
|
client = self._stream_client
|
||||||
|
if loop is None:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
loop.call_soon_threadsafe(self._begin_stream_shutdown, loop, client)
|
||||||
|
except RuntimeError:
|
||||||
|
# The loop closed between the snapshot and the call: the thread is
|
||||||
|
# already on its way out and the join below observes that.
|
||||||
|
logger.debug("[DingTalk] stream loop already closed before shutdown could be scheduled")
|
||||||
|
|
||||||
|
def _begin_stream_shutdown(self, loop: asyncio.AbstractEventLoop, client: Any) -> None:
|
||||||
|
loop.create_task(self._shutdown_stream(loop, client))
|
||||||
|
|
||||||
|
async def _shutdown_stream(self, loop: asyncio.AbstractEventLoop, client: Any) -> None:
|
||||||
|
"""Runs on the SDK loop: close the live websocket, then stop the loop."""
|
||||||
|
try:
|
||||||
|
websocket = getattr(client, "websocket", None) if client is not None else None
|
||||||
|
if websocket is not None:
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(websocket.close(), timeout=DINGTALK_STREAM_CLOSE_TIMEOUT_SECONDS)
|
||||||
|
except Exception:
|
||||||
|
logger.debug("[DingTalk] stream close did not complete cleanly", exc_info=True)
|
||||||
|
if abort_connection_transport(websocket):
|
||||||
|
logger.warning("[DingTalk] stream close handshake did not finish within %ss; transport aborted", DINGTALK_STREAM_CLOSE_TIMEOUT_SECONDS)
|
||||||
|
await asyncio.sleep(0) # let connection_lost run before the loop stops
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
# The thread is already draining this loop; leave it to that.
|
||||||
|
raise
|
||||||
|
loop.stop()
|
||||||
|
|
||||||
def _resolve_routing(self, msg: OutboundMessage) -> tuple[str, str, str]:
|
def _resolve_routing(self, msg: OutboundMessage) -> tuple[str, str, str]:
|
||||||
"""Return (conversation_type, sender_staff_id, conversation_id).
|
"""Return (conversation_type, sender_staff_id, conversation_id).
|
||||||
|
|
||||||
@ -398,22 +445,126 @@ class DingTalkChannel(Channel):
|
|||||||
# -- stream client (runs in dedicated thread) --------------------------
|
# -- stream client (runs in dedicated thread) --------------------------
|
||||||
|
|
||||||
def _run_stream(self, client_id: str, client_secret: str) -> None:
|
def _run_stream(self, client_id: str, client_secret: str) -> None:
|
||||||
|
loop = asyncio.new_event_loop()
|
||||||
|
asyncio.set_event_loop(loop)
|
||||||
try:
|
try:
|
||||||
import dingtalk_stream
|
import dingtalk_stream
|
||||||
|
|
||||||
credential = dingtalk_stream.Credential(client_id, client_secret)
|
credential = dingtalk_stream.Credential(client_id, client_secret)
|
||||||
client = dingtalk_stream.DingTalkStreamClient(credential)
|
client = dingtalk_stream.DingTalkStreamClient(credential)
|
||||||
self._stream_client = client
|
|
||||||
client.register_callback_handler(
|
client.register_callback_handler(
|
||||||
dingtalk_stream.chatbot.ChatbotMessage.TOPIC,
|
dingtalk_stream.chatbot.ChatbotMessage.TOPIC,
|
||||||
_DingTalkMessageHandler(self),
|
_DingTalkMessageHandler(self),
|
||||||
)
|
)
|
||||||
client.start_forever()
|
with self._stream_lock:
|
||||||
|
if not self._running:
|
||||||
|
# stop() ran before the client existed: never dial.
|
||||||
|
return
|
||||||
|
self._stream_loop = loop
|
||||||
|
self._stream_client = client
|
||||||
|
loop.run_until_complete(self._serve_stream(client))
|
||||||
except Exception:
|
except Exception:
|
||||||
if self._running:
|
if self._running:
|
||||||
logger.exception("DingTalk Stream Push error")
|
logger.exception("DingTalk Stream Push error")
|
||||||
|
else:
|
||||||
|
logger.debug("[DingTalk] stream loop exited after stop", exc_info=True)
|
||||||
|
self._running = False
|
||||||
finally:
|
finally:
|
||||||
self._stream_client = None
|
with self._stream_lock:
|
||||||
|
if self._stream_loop is loop:
|
||||||
|
self._stream_loop = None
|
||||||
|
self._stream_client = None
|
||||||
|
shutdown_private_loop(loop)
|
||||||
|
|
||||||
|
async def _serve_stream(self, client: Any) -> None:
|
||||||
|
"""Run dingtalk-stream's connection loop so ``stop()`` can end it.
|
||||||
|
|
||||||
|
Mirrors ``DingTalkStreamClient.start()`` (dingtalk-stream 0.24): open a
|
||||||
|
connection ticket, hold the websocket, hand each frame to the SDK's
|
||||||
|
``background_task``, and retry after 10s on a lost connection or 3s on
|
||||||
|
anything else. The SDK's own loop is ``while True`` and also catches
|
||||||
|
``CancelledError``, so it can only be left by killing the thread; this
|
||||||
|
copy re-checks ``_running`` before every dial and fetches the ticket
|
||||||
|
asynchronously with a deadline so ``stop()`` can cancel it.
|
||||||
|
"""
|
||||||
|
import websockets
|
||||||
|
|
||||||
|
client.pre_start()
|
||||||
|
while self._running:
|
||||||
|
try:
|
||||||
|
connection = await self._open_stream_connection(client)
|
||||||
|
if not self._running:
|
||||||
|
break
|
||||||
|
if not connection:
|
||||||
|
logger.error("[DingTalk] open connection failed")
|
||||||
|
await asyncio.sleep(_DINGTALK_STREAM_RETRY_SECONDS)
|
||||||
|
continue
|
||||||
|
# The ticket in ``connection`` is a credential; log the endpoint only.
|
||||||
|
logger.info("[DingTalk] stream endpoint is %s", connection.get("endpoint"))
|
||||||
|
uri = f"{connection['endpoint']}?ticket={quote_plus(connection['ticket'])}"
|
||||||
|
async with websockets.connect(uri) as websocket:
|
||||||
|
client.websocket = websocket
|
||||||
|
keepalive = asyncio.create_task(client.keepalive(websocket))
|
||||||
|
try:
|
||||||
|
async for raw_message in websocket:
|
||||||
|
asyncio.create_task(client.background_task(json.loads(raw_message)))
|
||||||
|
finally:
|
||||||
|
keepalive.cancel()
|
||||||
|
except websockets.exceptions.ConnectionClosedError as exc:
|
||||||
|
if not self._running:
|
||||||
|
break
|
||||||
|
logger.error("[DingTalk] stream connection lost: %s", exc)
|
||||||
|
await asyncio.sleep(_DINGTALK_STREAM_RETRY_SECONDS)
|
||||||
|
except Exception:
|
||||||
|
if not self._running:
|
||||||
|
break
|
||||||
|
logger.exception("[DingTalk] stream error")
|
||||||
|
await asyncio.sleep(_DINGTALK_STREAM_ERROR_RETRY_SECONDS)
|
||||||
|
|
||||||
|
async def _open_stream_connection(self, client: Any) -> dict[str, Any] | None:
|
||||||
|
"""Fetch a Stream ticket: ``DingTalkStreamClient.open_connection()`` with a deadline.
|
||||||
|
|
||||||
|
The SDK's version calls ``requests.post`` with no timeout, so a stalled
|
||||||
|
ticket endpoint pins whichever thread runs it (and, from an executor,
|
||||||
|
Gateway process exit, which joins executor workers). This sends the same
|
||||||
|
request through ``httpx`` on the stream loop under
|
||||||
|
``DINGTALK_STREAM_TICKET_TIMEOUT_SECONDS``, so ``stop()`` cancels it like
|
||||||
|
any other task. Returns ``None`` on failure like the SDK does.
|
||||||
|
"""
|
||||||
|
from dingtalk_stream.version import VERSION_STRING
|
||||||
|
|
||||||
|
topics: list[dict[str, str]] = []
|
||||||
|
if getattr(client, "_is_event_required", False):
|
||||||
|
topics.append({"type": "EVENT", "topic": "*"})
|
||||||
|
for topic in client.callback_handler_map:
|
||||||
|
topics.append({"type": "CALLBACK", "topic": topic})
|
||||||
|
headers = {
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
"Accept": "application/json",
|
||||||
|
"User-Agent": f"DingTalkStream/1.0 SDK/{VERSION_STRING} Python/{platform.python_version()} (+https://github.com/open-dingtalk/dingtalk-stream-sdk-python)",
|
||||||
|
}
|
||||||
|
body = {
|
||||||
|
"clientId": client.credential.client_id,
|
||||||
|
"clientSecret": client.credential.client_secret,
|
||||||
|
"subscriptions": topics,
|
||||||
|
"ua": f"dingtalk-sdk-python/v{VERSION_STRING}-union",
|
||||||
|
"localIp": client.get_host_ip(),
|
||||||
|
}
|
||||||
|
try:
|
||||||
|
async with httpx.AsyncClient(timeout=DINGTALK_STREAM_TICKET_TIMEOUT_SECONDS) as http:
|
||||||
|
response = await http.post(client.OPEN_CONNECTION_API, headers=headers, json=body)
|
||||||
|
response.raise_for_status()
|
||||||
|
payload = response.json()
|
||||||
|
except httpx.HTTPStatusError as exc:
|
||||||
|
logger.error("[DingTalk] open connection failed: HTTP %s", exc.response.status_code)
|
||||||
|
return None
|
||||||
|
except (httpx.HTTPError, ValueError) as exc:
|
||||||
|
logger.error("[DingTalk] open connection failed: %s", type(exc).__name__)
|
||||||
|
return None
|
||||||
|
if not isinstance(payload, dict) or not payload.get("endpoint") or not payload.get("ticket"):
|
||||||
|
logger.error("[DingTalk] open connection returned no endpoint/ticket")
|
||||||
|
return None
|
||||||
|
return payload
|
||||||
|
|
||||||
def _on_chatbot_message(self, message: Any) -> None:
|
def _on_chatbot_message(self, message: Any) -> None:
|
||||||
if not self._running:
|
if not self._running:
|
||||||
|
|||||||
@ -16,7 +16,7 @@ from concurrent.futures import Future
|
|||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, Literal, TypeVar
|
from typing import Any, Literal, TypeVar
|
||||||
|
|
||||||
from app.channels.base import Channel, ChannelStopTimeout, ChannelUnavailable
|
from app.channels.base import Channel, ChannelStopTimeout, ChannelUnavailable, abort_connection_transport, shutdown_private_loop
|
||||||
from app.channels.commands import is_known_channel_command, strip_leading_mentions
|
from app.channels.commands import is_known_channel_command, strip_leading_mentions
|
||||||
from app.channels.connection_identity import attach_connection_identity, lookup_thread_id
|
from app.channels.connection_identity import attach_connection_identity, lookup_thread_id
|
||||||
from app.channels.feishu_card_interaction import (
|
from app.channels.feishu_card_interaction import (
|
||||||
@ -50,6 +50,10 @@ from deerflow.uploads.manager import (
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
PENDING_CLARIFICATION_TTL_SECONDS = 30 * 60
|
PENDING_CLARIFICATION_TTL_SECONDS = 30 * 60
|
||||||
|
# stop() budget for the lark-oapi thread: ask its private loop to close the
|
||||||
|
# socket (bounded by the close timeout) and stop, then join the thread.
|
||||||
|
FEISHU_STOP_TIMEOUT_SECONDS = 5
|
||||||
|
FEISHU_WS_CLOSE_TIMEOUT_SECONDS = 2
|
||||||
_PENDING_CARD_METADATA_KEY = "_feishu_pending_card_id"
|
_PENDING_CARD_METADATA_KEY = "_feishu_pending_card_id"
|
||||||
# Channel-store deadlines. ``_on_message`` runs on lark-oapi's private event loop,
|
# Channel-store deadlines. ``_on_message`` runs on lark-oapi's private event loop,
|
||||||
# which also carries the SDK's websocket reads, pongs and its 120s ping loop:
|
# which also carries the SDK's websocket reads, pongs and its 120s ping loop:
|
||||||
@ -168,6 +172,11 @@ class FeishuChannel(Channel):
|
|||||||
super().__init__(name="feishu", bus=bus, config=config)
|
super().__init__(name="feishu", bus=bus, config=config)
|
||||||
self._thread: threading.Thread | None = None
|
self._thread: threading.Thread | None = None
|
||||||
self._main_loop: asyncio.AbstractEventLoop | None = None
|
self._main_loop: asyncio.AbstractEventLoop | None = None
|
||||||
|
# The SDK thread's private loop and client, published under _ws_lock once
|
||||||
|
# constructed so stop() can shut them down from the Gateway loop.
|
||||||
|
self._ws_loop: asyncio.AbstractEventLoop | None = None
|
||||||
|
self._ws_client: Any = None
|
||||||
|
self._ws_lock = threading.Lock()
|
||||||
self._api_client = None
|
self._api_client = None
|
||||||
self._CreateMessageReactionRequest = None
|
self._CreateMessageReactionRequest = None
|
||||||
self._CreateMessageReactionRequestBody = None
|
self._CreateMessageReactionRequestBody = None
|
||||||
@ -363,6 +372,13 @@ class FeishuChannel(Channel):
|
|||||||
We work around this by creating a plain asyncio event loop for this
|
We work around this by creating a plain asyncio event loop for this
|
||||||
thread and patching the SDK's module-level reference before calling
|
thread and patching the SDK's module-level reference before calling
|
||||||
``start()``.
|
``start()``.
|
||||||
|
|
||||||
|
``stop()`` ends this thread through ``_request_ws_shutdown``: lark-oapi
|
||||||
|
1.7 has no public stop, so ``start()`` only returns once this loop is
|
||||||
|
stopped from a callback scheduled on it — and only after that callback
|
||||||
|
disarmed ``_auto_reconnect``, because otherwise ``start()`` swallows the
|
||||||
|
interrupted ``run_until_complete`` in its reconnect branch and dials
|
||||||
|
again.
|
||||||
"""
|
"""
|
||||||
loop = asyncio.new_event_loop()
|
loop = asyncio.new_event_loop()
|
||||||
asyncio.set_event_loop(loop)
|
asyncio.set_event_loop(loop)
|
||||||
@ -383,11 +399,93 @@ class FeishuChannel(Channel):
|
|||||||
log_level=lark.LogLevel.INFO,
|
log_level=lark.LogLevel.INFO,
|
||||||
domain=domain,
|
domain=domain,
|
||||||
)
|
)
|
||||||
|
with self._ws_lock:
|
||||||
|
if not self._running:
|
||||||
|
# stop() ran before the client existed: never dial.
|
||||||
|
return
|
||||||
|
self._ws_loop = loop
|
||||||
|
self._ws_client = ws_client
|
||||||
ws_client.start()
|
ws_client.start()
|
||||||
except Exception:
|
except Exception:
|
||||||
if self._running:
|
if self._running:
|
||||||
logger.exception("Feishu WebSocket error")
|
logger.exception("Feishu WebSocket error")
|
||||||
|
else:
|
||||||
|
logger.debug("[Feishu] SDK loop exited after stop", exc_info=True)
|
||||||
self._running = False
|
self._running = False
|
||||||
|
finally:
|
||||||
|
with self._ws_lock:
|
||||||
|
if self._ws_loop is loop:
|
||||||
|
self._ws_loop = None
|
||||||
|
self._ws_client = None
|
||||||
|
shutdown_private_loop(loop)
|
||||||
|
|
||||||
|
def _request_ws_shutdown(self) -> None:
|
||||||
|
"""Ask the SDK thread to close its socket and stop its loop (non-blocking)."""
|
||||||
|
with self._ws_lock:
|
||||||
|
loop = self._ws_loop
|
||||||
|
ws_client = self._ws_client
|
||||||
|
if loop is None:
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
loop.call_soon_threadsafe(self._begin_ws_shutdown, loop, ws_client)
|
||||||
|
except RuntimeError:
|
||||||
|
# The loop closed between the snapshot and the call: the thread is
|
||||||
|
# already on its way out and the join below observes that.
|
||||||
|
logger.debug("[Feishu] SDK loop already closed before shutdown could be scheduled")
|
||||||
|
|
||||||
|
def _begin_ws_shutdown(self, loop: asyncio.AbstractEventLoop, ws_client: Any) -> None:
|
||||||
|
loop.create_task(self._shutdown_ws(loop, ws_client))
|
||||||
|
|
||||||
|
async def _shutdown_ws(self, loop: asyncio.AbstractEventLoop, ws_client: Any) -> None:
|
||||||
|
"""Runs on the SDK loop: disarm reconnects, close the socket, stop the loop."""
|
||||||
|
try:
|
||||||
|
if ws_client is not None:
|
||||||
|
# Must precede any await: once the loop stops, start() lands in
|
||||||
|
# its ``except Exception`` branch and consults this flag.
|
||||||
|
ws_client._auto_reconnect = False
|
||||||
|
deadline = loop.time() + FEISHU_WS_CLOSE_TIMEOUT_SECONDS
|
||||||
|
# _disconnect() clears _conn in its finally (also when cancelled
|
||||||
|
# mid-acquire), so take the handle first — once any in-flight
|
||||||
|
# dial has settled — and release the socket here if the close
|
||||||
|
# handshake does not finish, rather than whenever the
|
||||||
|
# connection object happens to be garbage-collected.
|
||||||
|
connection = await self._settled_ws_connection(ws_client, deadline)
|
||||||
|
disconnect = getattr(ws_client, "_disconnect", None)
|
||||||
|
if disconnect is not None:
|
||||||
|
try:
|
||||||
|
async with asyncio.timeout_at(deadline):
|
||||||
|
await disconnect()
|
||||||
|
except Exception:
|
||||||
|
logger.debug("[Feishu] WebSocket close did not complete cleanly", exc_info=True)
|
||||||
|
if abort_connection_transport(connection):
|
||||||
|
logger.warning("[Feishu] WebSocket close handshake did not finish within %ss; transport aborted", FEISHU_WS_CLOSE_TIMEOUT_SECONDS)
|
||||||
|
await asyncio.sleep(0) # let connection_lost run before the loop stops
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
# The thread is already draining this loop; leave it to that.
|
||||||
|
raise
|
||||||
|
loop.stop()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
async def _settled_ws_connection(ws_client: Any, deadline: float) -> Any:
|
||||||
|
"""Return the SDK's live connection once any in-flight ``_connect()`` has settled.
|
||||||
|
|
||||||
|
lark holds ``_lock`` across the whole dial and assigns ``_conn`` under
|
||||||
|
it, so a stop that lands mid-handshake would otherwise snapshot
|
||||||
|
``None`` and leave the socket the dial produces to the SDK's cancelled
|
||||||
|
close. Waiting for the lock is bounded by the close deadline.
|
||||||
|
"""
|
||||||
|
lock = getattr(ws_client, "_lock", None)
|
||||||
|
if lock is None:
|
||||||
|
return getattr(ws_client, "_conn", None)
|
||||||
|
try:
|
||||||
|
async with asyncio.timeout_at(deadline):
|
||||||
|
await lock.acquire()
|
||||||
|
except TimeoutError:
|
||||||
|
return getattr(ws_client, "_conn", None)
|
||||||
|
try:
|
||||||
|
return getattr(ws_client, "_conn", None)
|
||||||
|
finally:
|
||||||
|
lock.release()
|
||||||
|
|
||||||
def _on_ignored_message_event(self, event) -> None:
|
def _on_ignored_message_event(self, event) -> None:
|
||||||
logger.debug("[Feishu] ignoring non-content message event: %s", type(event).__name__)
|
logger.debug("[Feishu] ignoring non-content message event: %s", type(event).__name__)
|
||||||
@ -395,6 +493,9 @@ class FeishuChannel(Channel):
|
|||||||
async def stop(self) -> None:
|
async def stop(self) -> None:
|
||||||
self._running = False
|
self._running = False
|
||||||
self.bus.unsubscribe_outbound(self._on_outbound)
|
self.bus.unsubscribe_outbound(self._on_outbound)
|
||||||
|
# Signal the SDK thread first so its loop winds down while the
|
||||||
|
# Gateway-side drain below runs.
|
||||||
|
self._request_ws_shutdown()
|
||||||
await self._close_and_drain_threadsafe_futures()
|
await self._close_and_drain_threadsafe_futures()
|
||||||
|
|
||||||
background_tasks = tuple(self._background_tasks)
|
background_tasks = tuple(self._background_tasks)
|
||||||
@ -408,14 +509,17 @@ class FeishuChannel(Channel):
|
|||||||
self._background_tasks.clear()
|
self._background_tasks.clear()
|
||||||
self._running_card_tasks.clear()
|
self._running_card_tasks.clear()
|
||||||
if self._thread:
|
if self._thread:
|
||||||
# The SDK thread only returns on a fatal error, so this join
|
# The thread exits once its loop has stopped; keep the join off the
|
||||||
# normally waits out its full timeout; keep it off the event loop.
|
# event loop because a wedged SDK call can still use up the budget.
|
||||||
thread = self._thread
|
thread = self._thread
|
||||||
await asyncio.to_thread(thread.join, timeout=5)
|
await asyncio.to_thread(thread.join, timeout=FEISHU_STOP_TIMEOUT_SECONDS)
|
||||||
if thread.is_alive():
|
if thread.is_alive():
|
||||||
raise ChannelStopTimeout("Feishu SDK thread is still running after stop timeout")
|
raise ChannelStopTimeout("Feishu SDK thread is still running after stop timeout")
|
||||||
if self._thread is thread:
|
if self._thread is thread:
|
||||||
self._thread = None
|
self._thread = None
|
||||||
|
with self._ws_lock:
|
||||||
|
self._ws_loop = None
|
||||||
|
self._ws_client = None
|
||||||
logger.info("Feishu channel stopped")
|
logger.info("Feishu channel stopped")
|
||||||
|
|
||||||
async def send(self, msg: OutboundMessage, *, _max_retries: int = 3) -> None:
|
async def send(self, msg: OutboundMessage, *, _max_retries: int = 3) -> None:
|
||||||
|
|||||||
@ -6,8 +6,8 @@ every run, stream and channel the Gateway serves:
|
|||||||
|
|
||||||
- Slack: ``SocketModeClient.close()`` joins the SDK's message-processor thread
|
- Slack: ``SocketModeClient.close()`` joins the SDK's message-processor thread
|
||||||
and waits for in-flight listeners.
|
and waits for in-flight listeners.
|
||||||
- Feishu / DingTalk: ``stop()`` joins the SDK thread with a 5s timeout, and
|
- Feishu / DingTalk: ``stop()`` asks the SDK thread's private event loop to
|
||||||
that thread only returns on a fatal error.
|
shut down, then joins that thread with a 5s timeout.
|
||||||
- Discord: ``stop()`` joins the client thread with a 10s timeout, which a
|
- Discord: ``stop()`` joins the client thread with a 10s timeout, which a
|
||||||
timed-out client close or a slow ``_run_client()`` drain can use up.
|
timed-out client close or a slow ``_run_client()`` drain can use up.
|
||||||
|
|
||||||
@ -102,3 +102,35 @@ async def test_stop_joins_sdk_thread_off_loop(channel_cls) -> None:
|
|||||||
|
|
||||||
assert not sdk_thread.is_alive()
|
assert not sdk_thread.is_alive()
|
||||||
assert channel._thread is None
|
assert channel._thread is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("channel_cls", "loop_attr"),
|
||||||
|
[(FeishuChannel, "_ws_loop"), (DingTalkChannel, "_stream_loop")],
|
||||||
|
)
|
||||||
|
async def test_stop_ends_live_sdk_loop_off_loop(channel_cls, loop_attr) -> None:
|
||||||
|
sdk_loop = asyncio.new_event_loop()
|
||||||
|
ready = threading.Event()
|
||||||
|
|
||||||
|
def run_sdk_loop() -> None:
|
||||||
|
asyncio.set_event_loop(sdk_loop)
|
||||||
|
sdk_loop.call_soon(ready.set)
|
||||||
|
try:
|
||||||
|
sdk_loop.run_forever()
|
||||||
|
finally:
|
||||||
|
sdk_loop.close()
|
||||||
|
|
||||||
|
sdk_thread = threading.Thread(target=run_sdk_loop, daemon=True)
|
||||||
|
sdk_thread.start()
|
||||||
|
assert await asyncio.to_thread(ready.wait, 5)
|
||||||
|
channel = channel_cls(MessageBus(), {})
|
||||||
|
channel._thread = sdk_thread
|
||||||
|
channel._running = True
|
||||||
|
setattr(channel, loop_attr, sdk_loop)
|
||||||
|
|
||||||
|
await channel.stop()
|
||||||
|
|
||||||
|
assert not sdk_thread.is_alive()
|
||||||
|
assert sdk_loop.is_closed()
|
||||||
|
assert channel._thread is None
|
||||||
|
assert getattr(channel, loop_attr) is None
|
||||||
|
|||||||
142
backend/tests/support/local_sockets.py
Normal file
142
backend/tests/support/local_sockets.py
Normal file
@ -0,0 +1,142 @@
|
|||||||
|
"""Loopback peers for exercising SDK transports at the socket boundary.
|
||||||
|
|
||||||
|
Both run on the calling test's event loop and only ever bind ``127.0.0.1:0``.
|
||||||
|
|
||||||
|
- ``StalledWebSocketServer`` completes the WebSocket opening handshake, then
|
||||||
|
reads and discards every frame without answering, so a client's close
|
||||||
|
handshake never gets its ACK. ``closed`` is set once the peer observes EOF,
|
||||||
|
i.e. once the client actually released the TCP connection.
|
||||||
|
- ``LocalHttpServer`` answers every request with a fixed status and JSON body,
|
||||||
|
or (``stall=True``) accepts the connection and never responds. Each request
|
||||||
|
body is recorded so tests can assert on what the SDK sent.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import base64
|
||||||
|
import hashlib
|
||||||
|
import json
|
||||||
|
from collections.abc import Callable
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
_WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
|
||||||
|
|
||||||
|
|
||||||
|
async def wait_until(predicate: Callable[[], Any], *, timeout: float = 5.0, interval: float = 0.01) -> Any:
|
||||||
|
"""Poll ``predicate`` (which may read another thread's state) until it is truthy."""
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
deadline = loop.time() + timeout
|
||||||
|
while True:
|
||||||
|
value = predicate()
|
||||||
|
if value:
|
||||||
|
return value
|
||||||
|
if loop.time() >= deadline:
|
||||||
|
raise TimeoutError("condition not met in time")
|
||||||
|
await asyncio.sleep(interval)
|
||||||
|
|
||||||
|
|
||||||
|
def transport_socket(connection: Any) -> Any:
|
||||||
|
"""The socket object under a ``websockets`` connection, for release checks after stop."""
|
||||||
|
return connection.transport.get_extra_info("socket")
|
||||||
|
|
||||||
|
|
||||||
|
def _header(request: bytes, name: str) -> str:
|
||||||
|
prefix = name.lower() + ":"
|
||||||
|
for line in request.decode("latin-1").split("\r\n"):
|
||||||
|
if line.lower().startswith(prefix):
|
||||||
|
return line.split(":", 1)[1].strip()
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
class StalledWebSocketServer:
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.accepted = asyncio.Event()
|
||||||
|
self.closed = asyncio.Event()
|
||||||
|
self.port = 0
|
||||||
|
self._server: asyncio.base_events.Server | None = None
|
||||||
|
|
||||||
|
async def __aenter__(self) -> StalledWebSocketServer:
|
||||||
|
self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0)
|
||||||
|
self.port = self._server.sockets[0].getsockname()[1]
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, *exc_info: object) -> None:
|
||||||
|
assert self._server is not None
|
||||||
|
self._server.close()
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(self._server.wait_closed(), timeout=5)
|
||||||
|
except TimeoutError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@property
|
||||||
|
def url(self) -> str:
|
||||||
|
return f"ws://127.0.0.1:{self.port}/"
|
||||||
|
|
||||||
|
async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||||||
|
try:
|
||||||
|
request = await reader.readuntil(b"\r\n\r\n")
|
||||||
|
key = _header(request, "Sec-WebSocket-Key")
|
||||||
|
accept = base64.b64encode(hashlib.sha1((key + _WS_GUID).encode("ascii")).digest()).decode("ascii")
|
||||||
|
writer.write(("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n" + f"Sec-WebSocket-Accept: {accept}\r\n\r\n").encode("ascii"))
|
||||||
|
await writer.drain()
|
||||||
|
self.accepted.set()
|
||||||
|
while await reader.read(65536):
|
||||||
|
pass # swallow pings, data and the close frame; never answer
|
||||||
|
except (asyncio.IncompleteReadError, ConnectionError):
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
self.closed.set()
|
||||||
|
writer.close()
|
||||||
|
|
||||||
|
|
||||||
|
class LocalHttpServer:
|
||||||
|
def __init__(self, *, status: int = 200, body: Any = None, stall: bool = False) -> None:
|
||||||
|
self.status = status
|
||||||
|
self.body = body if body is not None else {}
|
||||||
|
self.stall = stall
|
||||||
|
self.requests: list[dict[str, Any]] = []
|
||||||
|
self.accepted = asyncio.Event()
|
||||||
|
self.closed = asyncio.Event()
|
||||||
|
self.port = 0
|
||||||
|
self._server: asyncio.base_events.Server | None = None
|
||||||
|
|
||||||
|
async def __aenter__(self) -> LocalHttpServer:
|
||||||
|
self._server = await asyncio.start_server(self._handle, "127.0.0.1", 0)
|
||||||
|
self.port = self._server.sockets[0].getsockname()[1]
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, *exc_info: object) -> None:
|
||||||
|
assert self._server is not None
|
||||||
|
self._server.close()
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(self._server.wait_closed(), timeout=5)
|
||||||
|
except TimeoutError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
@property
|
||||||
|
def url(self) -> str:
|
||||||
|
return f"http://127.0.0.1:{self.port}/open"
|
||||||
|
|
||||||
|
async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||||||
|
try:
|
||||||
|
head = await reader.readuntil(b"\r\n\r\n")
|
||||||
|
length = int(_header(head, "Content-Length") or 0)
|
||||||
|
raw = await reader.readexactly(length) if length else b""
|
||||||
|
try:
|
||||||
|
self.requests.append(json.loads(raw) if raw else {})
|
||||||
|
except ValueError:
|
||||||
|
self.requests.append({"raw": raw.decode("latin-1")})
|
||||||
|
self.accepted.set()
|
||||||
|
if self.stall:
|
||||||
|
while await reader.read(65536):
|
||||||
|
pass # hold the request open until the client gives up
|
||||||
|
return
|
||||||
|
payload = json.dumps(self.body).encode("utf-8")
|
||||||
|
writer.write((f"HTTP/1.1 {self.status} X\r\nContent-Type: application/json\r\nContent-Length: {len(payload)}\r\nConnection: close\r\n\r\n").encode("ascii") + payload)
|
||||||
|
await writer.drain()
|
||||||
|
except (asyncio.IncompleteReadError, ConnectionError):
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
self.closed.set()
|
||||||
|
writer.close()
|
||||||
@ -5,6 +5,7 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import threading
|
||||||
|
import time
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
from unittest.mock import AsyncMock, MagicMock
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
@ -1808,6 +1809,183 @@ class TestStop:
|
|||||||
|
|
||||||
_run(go())
|
_run(go())
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _channel() -> DingTalkChannel:
|
||||||
|
return DingTalkChannel(MessageBus(), config={"client_id": "id", "client_secret": "secret"})
|
||||||
|
|
||||||
|
def test_stop_closes_live_stream_connection_and_exits_thread(self, monkeypatch):
|
||||||
|
pytest.importorskip("dingtalk_stream")
|
||||||
|
import dingtalk_stream
|
||||||
|
from support.local_sockets import LocalHttpServer
|
||||||
|
from websockets.asyncio.server import serve
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
connected = asyncio.Event()
|
||||||
|
closed = asyncio.Event()
|
||||||
|
|
||||||
|
async def handler(websocket):
|
||||||
|
connected.set()
|
||||||
|
try:
|
||||||
|
async for _ in websocket:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
closed.set()
|
||||||
|
|
||||||
|
async with serve(handler, "127.0.0.1", 0) as server:
|
||||||
|
port = server.sockets[0].getsockname()[1]
|
||||||
|
async with LocalHttpServer(body={"endpoint": f"ws://127.0.0.1:{port}", "ticket": "ticket"}) as tickets:
|
||||||
|
monkeypatch.setattr(dingtalk_stream.DingTalkStreamClient, "OPEN_CONNECTION_API", tickets.url)
|
||||||
|
channel = self._channel()
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
assert thread is not None and thread.is_alive()
|
||||||
|
await asyncio.wait_for(connected.wait(), 5)
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
await asyncio.wait_for(closed.wait(), 5)
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
assert channel._stream_client is None
|
||||||
|
assert channel.is_running is False
|
||||||
|
|
||||||
|
# The ticket request mirrors DingTalkStreamClient.open_connection().
|
||||||
|
request = tickets.requests[0]
|
||||||
|
assert request["clientId"] == "id"
|
||||||
|
assert request["clientSecret"] == "secret"
|
||||||
|
assert {"type": "CALLBACK", "topic": dingtalk_stream.chatbot.ChatbotMessage.TOPIC} in request["subscriptions"]
|
||||||
|
assert request["ua"].startswith("dingtalk-sdk-python/v")
|
||||||
|
assert "localIp" in request
|
||||||
|
|
||||||
|
_run(go())
|
||||||
|
|
||||||
|
def test_stop_aborts_transport_when_peer_never_acks_close(self, monkeypatch, caplog):
|
||||||
|
pytest.importorskip("dingtalk_stream")
|
||||||
|
import dingtalk_stream
|
||||||
|
from support.local_sockets import LocalHttpServer, StalledWebSocketServer, transport_socket, wait_until
|
||||||
|
|
||||||
|
from app.channels import dingtalk as dingtalk_module
|
||||||
|
|
||||||
|
caplog.set_level(logging.WARNING, logger="app.channels.dingtalk")
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
async with StalledWebSocketServer() as peer, LocalHttpServer(body={"endpoint": peer.url, "ticket": "ticket"}) as tickets:
|
||||||
|
monkeypatch.setattr(dingtalk_stream.DingTalkStreamClient, "OPEN_CONNECTION_API", tickets.url)
|
||||||
|
monkeypatch.setattr(dingtalk_module, "DINGTALK_STREAM_CLOSE_TIMEOUT_SECONDS", 0.5)
|
||||||
|
channel = self._channel()
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
await asyncio.wait_for(peer.accepted.wait(), 5)
|
||||||
|
websocket = await wait_until(lambda: channel._stream_client is not None and channel._stream_client.websocket)
|
||||||
|
sock = transport_socket(websocket)
|
||||||
|
assert sock.fileno() != -1
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
# The socket is released by stop() itself, not by a later GC pass.
|
||||||
|
assert sock.fileno() == -1
|
||||||
|
assert "transport aborted" in caplog.text
|
||||||
|
await asyncio.wait_for(peer.closed.wait(), 5)
|
||||||
|
|
||||||
|
_run(go())
|
||||||
|
|
||||||
|
def test_stop_interrupts_open_connection_retry_sleep(self, monkeypatch):
|
||||||
|
pytest.importorskip("dingtalk_stream")
|
||||||
|
import dingtalk_stream
|
||||||
|
from support.local_sockets import LocalHttpServer
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
async with LocalHttpServer(status=500, body={"error": "nope"}) as tickets:
|
||||||
|
monkeypatch.setattr(dingtalk_stream.DingTalkStreamClient, "OPEN_CONNECTION_API", tickets.url)
|
||||||
|
channel = self._channel()
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
await asyncio.wait_for(tickets.accepted.wait(), 5)
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
|
||||||
|
_run(go())
|
||||||
|
|
||||||
|
def test_stop_cancels_stalled_ticket_request_without_lingering_threads(self, monkeypatch):
|
||||||
|
pytest.importorskip("dingtalk_stream")
|
||||||
|
import dingtalk_stream
|
||||||
|
from support.local_sockets import LocalHttpServer
|
||||||
|
|
||||||
|
def forbidden_open_connection(self):
|
||||||
|
raise AssertionError("the SDK's timeout-less open_connection() must not run")
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
async with LocalHttpServer(stall=True) as tickets:
|
||||||
|
monkeypatch.setattr(dingtalk_stream.DingTalkStreamClient, "OPEN_CONNECTION_API", tickets.url)
|
||||||
|
monkeypatch.setattr(dingtalk_stream.DingTalkStreamClient, "open_connection", forbidden_open_connection)
|
||||||
|
await asyncio.to_thread(lambda: None) # warm the test loop's own executor before the baseline
|
||||||
|
baseline = {thread.ident for thread in threading.enumerate()}
|
||||||
|
channel = self._channel()
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
await asyncio.wait_for(tickets.accepted.wait(), 5)
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
assert not thread.is_alive()
|
||||||
|
# The cancelled request released its connection ...
|
||||||
|
await asyncio.wait_for(tickets.closed.wait(), 5)
|
||||||
|
# ... and nothing non-daemon is left to block interpreter exit.
|
||||||
|
lingering = [t.name for t in threading.enumerate() if t.ident not in baseline and t.is_alive() and not t.daemon]
|
||||||
|
assert lingering == []
|
||||||
|
|
||||||
|
_run(go())
|
||||||
|
|
||||||
|
def test_stop_before_stream_client_exists_never_connects(self, monkeypatch):
|
||||||
|
pytest.importorskip("dingtalk_stream")
|
||||||
|
import dingtalk_stream
|
||||||
|
from support.local_sockets import LocalHttpServer
|
||||||
|
|
||||||
|
gate = threading.Event()
|
||||||
|
|
||||||
|
class _GatedStreamClient(dingtalk_stream.DingTalkStreamClient):
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
gate.wait(timeout=5)
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr(dingtalk_stream, "DingTalkStreamClient", _GatedStreamClient)
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
async with LocalHttpServer(body={"endpoint": "ws://127.0.0.1:1", "ticket": "ticket"}) as tickets:
|
||||||
|
monkeypatch.setattr(dingtalk_stream.DingTalkStreamClient, "OPEN_CONNECTION_API", tickets.url)
|
||||||
|
channel = self._channel()
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
|
||||||
|
stop_task = asyncio.create_task(channel.stop())
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
assert not stop_task.done()
|
||||||
|
gate.set()
|
||||||
|
await stop_task
|
||||||
|
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
assert tickets.requests == []
|
||||||
|
|
||||||
|
_run(go())
|
||||||
|
|
||||||
def test_stop_joins_stream_thread_off_the_event_loop(self):
|
def test_stop_joins_stream_thread_off_the_event_loop(self):
|
||||||
async def go():
|
async def go():
|
||||||
channel = DingTalkChannel(MessageBus(), config={})
|
channel = DingTalkChannel(MessageBus(), config={})
|
||||||
|
|||||||
@ -1,5 +1,6 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
|
import logging
|
||||||
import tempfile
|
import tempfile
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
@ -209,6 +210,242 @@ def test_feishu_stop_retains_live_ws_thread_after_join_timeout():
|
|||||||
asyncio.run(go())
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
|
class _FakeLarkWsClient:
|
||||||
|
"""Stand-in for ``lark_oapi.ws.Client`` (1.7.x) with the same ``start()`` shape.
|
||||||
|
|
||||||
|
The SDK drives the module-level ``lark_oapi.ws.client.loop`` itself: it runs
|
||||||
|
``_connect()`` to completion, falls into a reconnect branch that swallows
|
||||||
|
*any* exception while ``_auto_reconnect`` is set, then parks forever in
|
||||||
|
``_select()``. A stop that only calls ``loop.stop()`` is therefore absorbed
|
||||||
|
by that reconnect branch; the channel has to disarm it first.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, *, connect_forever: bool = False, construct_gate: threading.Event | None = None, **_kwargs) -> None:
|
||||||
|
self._auto_reconnect = True
|
||||||
|
self.connect_forever = connect_forever
|
||||||
|
self.loop: asyncio.AbstractEventLoop | None = None
|
||||||
|
self.start_calls = 0
|
||||||
|
self.disconnect_calls = 0
|
||||||
|
self.reconnect_calls = 0
|
||||||
|
self.connecting = threading.Event()
|
||||||
|
self.selected = threading.Event()
|
||||||
|
if construct_gate is not None:
|
||||||
|
construct_gate.wait(timeout=5)
|
||||||
|
|
||||||
|
async def _connect(self) -> None:
|
||||||
|
self.connecting.set()
|
||||||
|
if self.connect_forever:
|
||||||
|
await asyncio.Event().wait()
|
||||||
|
|
||||||
|
async def _disconnect(self) -> None:
|
||||||
|
self.disconnect_calls += 1
|
||||||
|
|
||||||
|
async def _reconnect(self) -> None:
|
||||||
|
self.reconnect_calls += 1
|
||||||
|
await asyncio.Event().wait()
|
||||||
|
|
||||||
|
async def _select(self) -> None:
|
||||||
|
self.selected.set()
|
||||||
|
while True:
|
||||||
|
await asyncio.sleep(3600)
|
||||||
|
|
||||||
|
def start(self) -> None:
|
||||||
|
import lark_oapi.ws.client as ws_client_mod
|
||||||
|
|
||||||
|
loop = ws_client_mod.loop
|
||||||
|
self.loop = loop
|
||||||
|
self.start_calls += 1
|
||||||
|
try:
|
||||||
|
loop.run_until_complete(self._connect())
|
||||||
|
except Exception:
|
||||||
|
loop.run_until_complete(self._disconnect())
|
||||||
|
if self._auto_reconnect:
|
||||||
|
loop.run_until_complete(self._reconnect())
|
||||||
|
else:
|
||||||
|
raise
|
||||||
|
loop.run_until_complete(self._select())
|
||||||
|
|
||||||
|
|
||||||
|
def _feishu_channel_with_fake_ws(monkeypatch, **client_kwargs):
|
||||||
|
import lark_oapi
|
||||||
|
|
||||||
|
created: list[_FakeLarkWsClient] = []
|
||||||
|
constructed = threading.Event()
|
||||||
|
|
||||||
|
def factory(**_kwargs):
|
||||||
|
client = _FakeLarkWsClient(**client_kwargs)
|
||||||
|
created.append(client)
|
||||||
|
constructed.set()
|
||||||
|
return client
|
||||||
|
|
||||||
|
monkeypatch.setattr(lark_oapi.ws, "Client", factory)
|
||||||
|
channel = FeishuChannel(MessageBus(), {"app_id": "test", "app_secret": "test"})
|
||||||
|
return channel, created, constructed
|
||||||
|
|
||||||
|
|
||||||
|
def test_feishu_stop_disconnects_sdk_and_exits_its_thread(monkeypatch):
|
||||||
|
async def go():
|
||||||
|
channel, created, constructed = _feishu_channel_with_fake_ws(monkeypatch)
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
assert thread is not None and thread.is_alive()
|
||||||
|
assert await asyncio.to_thread(constructed.wait, 5)
|
||||||
|
client = created[0]
|
||||||
|
assert await asyncio.to_thread(client.selected.wait, 5)
|
||||||
|
assert channel.is_running is True
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
assert channel.is_running is False
|
||||||
|
assert client.start_calls == 1
|
||||||
|
assert client._auto_reconnect is False
|
||||||
|
assert client.disconnect_calls >= 1
|
||||||
|
assert client.reconnect_calls == 0
|
||||||
|
assert client.loop is not None and client.loop.is_closed()
|
||||||
|
|
||||||
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
|
def test_feishu_stop_interrupts_sdk_connect_without_reconnecting(monkeypatch):
|
||||||
|
async def go():
|
||||||
|
channel, created, constructed = _feishu_channel_with_fake_ws(monkeypatch, connect_forever=True)
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
assert await asyncio.to_thread(constructed.wait, 5)
|
||||||
|
client = created[0]
|
||||||
|
assert await asyncio.to_thread(client.connecting.wait, 5)
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
assert client._auto_reconnect is False
|
||||||
|
assert client.reconnect_calls == 0
|
||||||
|
assert client.loop is not None and client.loop.is_closed()
|
||||||
|
|
||||||
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
|
def test_feishu_stop_before_sdk_client_exists_never_starts_it(monkeypatch):
|
||||||
|
async def go():
|
||||||
|
gate = threading.Event()
|
||||||
|
channel, created, constructed = _feishu_channel_with_fake_ws(monkeypatch, construct_gate=gate)
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
|
||||||
|
stop_task = asyncio.create_task(channel.stop())
|
||||||
|
await asyncio.sleep(0.05)
|
||||||
|
assert not stop_task.done()
|
||||||
|
gate.set()
|
||||||
|
await stop_task
|
||||||
|
|
||||||
|
assert await asyncio.to_thread(constructed.wait, 5)
|
||||||
|
assert created[0].start_calls == 0
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
|
||||||
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
|
def _point_lark_at(monkeypatch, ws_url: str) -> None:
|
||||||
|
"""Skip lark's endpoint-discovery POST and dial ``ws_url`` with the real client."""
|
||||||
|
import lark_oapi.ws.client as ws_client_mod
|
||||||
|
|
||||||
|
monkeypatch.setattr(ws_client_mod.Client, "_get_conn_url", lambda self: f"{ws_url}?device_id=d1&service_id=1")
|
||||||
|
|
||||||
|
|
||||||
|
def test_feishu_stop_closes_real_sdk_connection(monkeypatch):
|
||||||
|
from websockets.asyncio.server import serve
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
connected = asyncio.Event()
|
||||||
|
closed = asyncio.Event()
|
||||||
|
|
||||||
|
async def handler(websocket):
|
||||||
|
connected.set()
|
||||||
|
try:
|
||||||
|
async for _ in websocket:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
closed.set()
|
||||||
|
|
||||||
|
async with serve(handler, "127.0.0.1", 0) as server:
|
||||||
|
port = server.sockets[0].getsockname()[1]
|
||||||
|
_point_lark_at(monkeypatch, f"ws://127.0.0.1:{port}/")
|
||||||
|
channel = FeishuChannel(MessageBus(), {"app_id": "test", "app_secret": "test"})
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
await asyncio.wait_for(connected.wait(), 5)
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
await asyncio.wait_for(closed.wait(), 5)
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
|
||||||
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
|
def test_feishu_stop_aborts_transport_when_peer_never_acks_close(monkeypatch, caplog):
|
||||||
|
"""Real lark-oapi client against a peer that completes the handshake but never answers the close frame.
|
||||||
|
|
||||||
|
The bounded ``_disconnect()`` wait is cancelled before websockets' own
|
||||||
|
``close_timeout`` abort runs, so the channel must release the socket itself;
|
||||||
|
``peer.closed`` only fires once the TCP connection is actually gone.
|
||||||
|
"""
|
||||||
|
from support.local_sockets import StalledWebSocketServer, transport_socket, wait_until
|
||||||
|
|
||||||
|
caplog.set_level(logging.WARNING, logger="app.channels.feishu")
|
||||||
|
|
||||||
|
async def go():
|
||||||
|
async with StalledWebSocketServer() as peer:
|
||||||
|
_point_lark_at(monkeypatch, peer.url)
|
||||||
|
monkeypatch.setattr(feishu_module, "FEISHU_WS_CLOSE_TIMEOUT_SECONDS", 0.5)
|
||||||
|
channel = FeishuChannel(MessageBus(), {"app_id": "test", "app_secret": "test"})
|
||||||
|
await channel.start()
|
||||||
|
thread = channel._thread
|
||||||
|
await asyncio.wait_for(peer.accepted.wait(), 5)
|
||||||
|
# Wait for lark to finish its dial (it assigns _conn under _lock).
|
||||||
|
connection = await wait_until(lambda: channel._ws_client is not None and not channel._ws_client._lock.locked() and channel._ws_client._conn)
|
||||||
|
sock = transport_socket(connection)
|
||||||
|
assert sock.fileno() != -1
|
||||||
|
|
||||||
|
started = time.monotonic()
|
||||||
|
await channel.stop()
|
||||||
|
elapsed = time.monotonic() - started
|
||||||
|
|
||||||
|
assert elapsed < 5, f"stop() took {elapsed:.2f}s"
|
||||||
|
assert not thread.is_alive()
|
||||||
|
assert channel._thread is None
|
||||||
|
# The socket is released by stop() itself, not by a later GC pass.
|
||||||
|
assert sock.fileno() == -1
|
||||||
|
assert "transport aborted" in caplog.text
|
||||||
|
await asyncio.wait_for(peer.closed.wait(), 5)
|
||||||
|
|
||||||
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
|
def test_feishu_stop_without_sdk_thread_is_a_no_op_for_the_loop_signal():
|
||||||
|
async def go():
|
||||||
|
channel = FeishuChannel(MessageBus(), {"app_id": "test", "app_secret": "test"})
|
||||||
|
channel._running = True
|
||||||
|
await channel.stop()
|
||||||
|
assert channel._thread is None
|
||||||
|
|
||||||
|
asyncio.run(go())
|
||||||
|
|
||||||
|
|
||||||
def test_feishu_event_handler_ignores_non_content_message_events():
|
def test_feishu_event_handler_ignores_non_content_message_events():
|
||||||
import lark_oapi as lark
|
import lark_oapi as lark
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user