From 86d961035c3e8a21cd9a3a092a0fb5f4540ce84a Mon Sep 17 00:00:00 2001 From: fy17ohhh Date: Sat, 15 Aug 2026 01:12:03 +0800 Subject: [PATCH] Bug fix: drain in-flight Feishu messages during shutdown --- tests/integrations/agentkit/test_app.py | 30 +++-- tests/test_feishu_channel_extension.py | 114 ++++++++++++++++++ veadk/extensions/feishu_channel.py | 21 +++- veadk/integrations/agentkit/app.py | 152 +++--------------------- 4 files changed, 169 insertions(+), 148 deletions(-) diff --git a/tests/integrations/agentkit/test_app.py b/tests/integrations/agentkit/test_app.py index b356f1819..79f8a3f5b 100644 --- a/tests/integrations/agentkit/test_app.py +++ b/tests/integrations/agentkit/test_app.py @@ -14,6 +14,7 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace from typing import Any, cast @@ -397,29 +398,40 @@ def test_feishu_lifecycle_starts_and_stops_with_application( ) -> None: runners: list[dict[str, Any]] = [] events: list[str] = [] + loops: list[object] = [] class _FakeRunner: def __init__(self, **kwargs: Any) -> None: runners.append(kwargs) - async def fake_start(app: FastAPI, runner: object) -> None: - del app, runner - events.append("start") + class _FakeFeishuChannel: + def start(self) -> None: + events.append("start") + loops.append(asyncio.get_running_loop()) - async def fake_stop(app: FastAPI) -> None: - del app - events.append("stop") + async def shutdown(self) -> None: + events.append("shutdown") + loops.append(asyncio.get_running_loop()) monkeypatch.setattr(veadk, "Runner", _FakeRunner) - monkeypatch.setattr(agentkit_app, "_start_feishu_channel", fake_start) - monkeypatch.setattr(agentkit_app, "_stop_feishu_channel", fake_stop) + monkeypatch.setenv("FEISHU_APP_ID", "cli_test") + monkeypatch.setenv("FEISHU_APP_SECRET", "secret") + channel = _FakeFeishuChannel() + monkeypatch.setattr( + agentkit_app, + "_build_feishu_channel", + lambda runner, app_id, app_secret: channel, + ) root_agent = _root_agent() app = agentkit_app.create_agentkit_app(root_agent, enable_feishu=True) with TestClient(app): assert events == ["start"] + assert app.state.feishu_channel is channel - assert events == ["start", "stop"] + assert events == ["start", "shutdown"] + assert loops[0] is loops[1] + assert app.state.feishu_channel is None assert runners[0]["agent"] is root_agent assert runners[0]["app_name"] == "agent" assert isinstance(runners[0]["short_term_memory"], _FakeShortTermMemory) diff --git a/tests/test_feishu_channel_extension.py b/tests/test_feishu_channel_extension.py index 24c321cb9..053e3b68a 100644 --- a/tests/test_feishu_channel_extension.py +++ b/tests/test_feishu_channel_extension.py @@ -392,3 +392,117 @@ async def test_extension_prefers_sync_start_stop_over_async_connect(): assert channel.stop_called is True assert channel.start_loop_running is False assert channel.stop_loop_running is False + + +@pytest.mark.anyio +async def test_extension_rebinds_lark_ws_loop_before_sync_start(monkeypatch): + from lark_channel.ws import client as lark_ws_client + + asgi_loop = asyncio.get_running_loop() + observed = {} + + class LoopInspectingChannel(FakeChannel): + def start(self): + worker_loop = asyncio.get_event_loop() + observed["worker_loop"] = worker_loop + observed["sdk_loop"] = lark_ws_client.loop + observed["worker_loop_running"] = worker_loop.is_running() + lark_ws_client.loop.run_until_complete(asyncio.sleep(0)) + + LoopInspectingChannel.__module__ = "lark_channel.testing" + monkeypatch.setattr(lark_ws_client, "loop", asgi_loop) + extension = FeishuChannelExtension( + runner=FakeRunner(), + channel=LoopInspectingChannel(), + ) + + await extension.connect() + + assert observed["sdk_loop"] is observed["worker_loop"] + assert observed["sdk_loop"] is not asgi_loop + assert observed["worker_loop_running"] is False + + +@pytest.mark.anyio +async def test_extension_shutdown_disconnects_before_draining_inflight_messages(): + events = [] + message_started = asyncio.Event() + release_message = asyncio.Event() + + class DrainChannel(FakeChannel): + def stop(self): + events.append("channel.stopped") + + extension = FeishuChannelExtension( + runner=FakeRunner(), + channel=DrainChannel(), + ) + + async def process_message(): + events.append("message.started") + message_started.set() + await release_message.wait() + events.append("message.completed") + + task = asyncio.create_task(process_message()) + extension._inflight.add(task) + task.add_done_callback(extension._inflight.discard) + await message_started.wait() + + shutdown = asyncio.create_task(extension.shutdown(drain_timeout=1)) + await asyncio.sleep(0) + while "channel.stopped" not in events: + await asyncio.sleep(0) + + assert events == ["message.started", "channel.stopped"] + assert not shutdown.done() + + release_message.set() + await shutdown + + assert events == [ + "message.started", + "channel.stopped", + "message.completed", + ] + assert not extension._inflight + + +@pytest.mark.anyio +async def test_extension_start_and_shutdown_manage_blocking_channel(): + events = [] + + class BlockingStartChannel(FakeChannel): + def __init__(self): + super().__init__() + self.started = threading.Event() + self.stopped = threading.Event() + self.start_returned = threading.Event() + + def start(self): + events.append("channel.started") + self.started.set() + self.stopped.wait(timeout=1) + events.append("channel.start_returned") + self.start_returned.set() + + def stop(self): + events.append("channel.stopped") + self.stopped.set() + + channel = BlockingStartChannel() + extension = FeishuChannelExtension(runner=FakeRunner(), channel=channel) + + extension.start() + assert await asyncio.to_thread(channel.started.wait, 1) + + await extension.shutdown(drain_timeout=1) + assert await asyncio.to_thread(channel.start_returned.wait, 1) + + assert events == [ + "channel.started", + "channel.stopped", + "channel.start_returned", + ] + assert extension.is_draining is True + assert extension._retry_task is None diff --git a/veadk/extensions/feishu_channel.py b/veadk/extensions/feishu_channel.py index a108a12cc..51bcf42f4 100644 --- a/veadk/extensions/feishu_channel.py +++ b/veadk/extensions/feishu_channel.py @@ -71,10 +71,23 @@ def _read_attr(obj: Any, *path: str) -> Any: return current -def _call_in_fresh_event_loop(method: Callable[[], Any]) -> Any: +def _call_in_fresh_event_loop( + method: Callable[[], Any], + *, + bind_lark_ws_loop: bool = False, +) -> Any: loop = asyncio.new_event_loop() try: asyncio.set_event_loop(loop) + if bind_lark_ws_loop: + # lark_channel.ws.client captures a module-level event loop when it + # is imported. AgentKit imports the SDK from a running ASGI loop, + # while the SDK's synchronous start() later calls + # loop.run_until_complete(). Rebind that SDK-global loop to this + # dedicated worker loop before starting the WebSocket. + from lark_channel.ws import client as lark_ws_client + + lark_ws_client.loop = loop result = method() if inspect.isawaitable(result): return loop.run_until_complete(result) @@ -417,7 +430,11 @@ async def connect(self) -> Any: connect = getattr(self.channel, "start", None) or self.channel.connect if inspect.iscoroutinefunction(connect): return await connect() - return await asyncio.to_thread(_call_in_fresh_event_loop, connect) + return await asyncio.to_thread( + _call_in_fresh_event_loop, + connect, + bind_lark_ws_loop=type(self.channel).__module__.startswith("lark_channel."), + ) async def disconnect(self) -> Any: disconnect = getattr(self.channel, "stop", None) or getattr( diff --git a/veadk/integrations/agentkit/app.py b/veadk/integrations/agentkit/app.py index 5ad2a9b16..277ffb106 100644 --- a/veadk/integrations/agentkit/app.py +++ b/veadk/integrations/agentkit/app.py @@ -20,8 +20,6 @@ import inspect import json import os -import threading -import traceback from collections.abc import Callable, Mapping from contextlib import asynccontextmanager from pathlib import Path @@ -58,6 +56,7 @@ if TYPE_CHECKING: from agentkit.identity import RuntimeIdentity + from veadk.extensions.feishu_channel import FeishuChannelExtension from veadk.runner import Runner @@ -218,62 +217,9 @@ def _agent_node( } -def _get_feishu_channel_method( - channel: object, - names: tuple[str, ...], -) -> Callable[[], Any] | None: - raw_channel = getattr(channel, "channel", None) - for target in (raw_channel, channel): - if target is None: - continue - for name in names: - method = getattr(target, name, None) - if callable(method): - return method - return None - - -def _call_feishu_channel_method( - loop: asyncio.AbstractEventLoop, - method: Callable[[], Any], -) -> Any: - result = method() - if inspect.isawaitable(result): - return loop.run_until_complete(result) - return result - - -def _connect_feishu_channel( - loop: asyncio.AbstractEventLoop, - channel: object, -) -> Any: - connect = _get_feishu_channel_method(channel, ("start", "connect")) - if connect is None: - raise AttributeError("Feishu channel has no start/connect method") - return _call_feishu_channel_method(loop, connect) - - -def _disconnect_feishu_channel( - loop: asyncio.AbstractEventLoop, - channel: object, -) -> Any: - disconnect = _get_feishu_channel_method(channel, ("stop", "disconnect")) - if disconnect is None: - return None - return _call_feishu_channel_method(loop, disconnect) - - -def _stop_feishu_channel_from_lifespan(channel: object) -> None: - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - _disconnect_feishu_channel(loop, channel) - finally: - asyncio.set_event_loop(None) - loop.close() - - -def _build_feishu_channel(runner: Runner, app_id: str, app_secret: str) -> object: +def _build_feishu_channel( + runner: Runner, app_id: str, app_secret: str +) -> FeishuChannelExtension: from veadk.extensions import FeishuChannelExtension return FeishuChannelExtension( @@ -286,54 +232,6 @@ def _build_feishu_channel(runner: Runner, app_id: str, app_secret: str) -> objec ) -def _run_feishu_channel( - runner: Runner, - app_id: str, - app_secret: str, - stop_event: threading.Event, - state: dict[str, Any], -) -> None: - loop = asyncio.new_event_loop() - state["loop"] = loop - asyncio.set_event_loop(loop) - try: - while not stop_event.is_set(): - channel = None - try: - channel = _build_feishu_channel(runner, app_id, app_secret) - state["channel"] = channel - print("feishu channel connecting in dedicated thread", flush=True) - _connect_feishu_channel(loop, channel) - print("feishu channel disconnected; reconnecting in 5s", flush=True) - except Exception as exc: # The channel reconnects after transport errors. - stage = "initialization" if channel is None else "connect" - print( - f"feishu channel {stage} failed: " - f"{type(exc).__name__}: {exc}; reconnecting in 5s", - flush=True, - ) - if channel is None: - print(traceback.format_exc(), flush=True) - finally: - if channel is not None: - try: - _disconnect_feishu_channel(loop, channel) - except Exception as exc: # Cleanup must not stop reconnection. - print( - "feishu channel disconnect failed: " - f"{type(exc).__name__}: {exc}", - flush=True, - ) - finally: - if state.get("channel") is channel: - state["channel"] = None - stop_event.wait(5) - finally: - asyncio.set_event_loop(None) - state["loop"] = None - loop.close() - - async def _start_feishu_channel(app: FastAPI, runner: Runner) -> None: app_id = os.getenv("FEISHU_APP_ID") app_secret = os.getenv("FEISHU_APP_SECRET") @@ -344,40 +242,20 @@ async def _start_feishu_channel(app: FastAPI, runner: Runner) -> None: ) return - app.state.feishu_channel_state = {"channel": None, "loop": None} - app.state.feishu_channel_stop_event = threading.Event() - app.state.feishu_channel_thread = threading.Thread( - target=_run_feishu_channel, - args=( - runner, - app_id, - app_secret, - app.state.feishu_channel_stop_event, - app.state.feishu_channel_state, - ), - name="feishu-channel", - daemon=True, - ) - app.state.feishu_channel_thread.start() - print("feishu channel background thread started", flush=True) + channel = _build_feishu_channel(runner, app_id, app_secret) + app.state.feishu_channel = channel + channel.start() + print("feishu channel reconnect loop started", flush=True) async def _stop_feishu_channel(app: FastAPI) -> None: - stop_event = getattr(app.state, "feishu_channel_stop_event", None) - if stop_event is not None: - stop_event.set() - state = getattr(app.state, "feishu_channel_state", None) or {} - channel = state.get("channel") - if channel is not None: - await asyncio.to_thread(_stop_feishu_channel_from_lifespan, channel) - thread = getattr(app.state, "feishu_channel_thread", None) - if thread is not None: - await asyncio.to_thread(thread.join, 2) - if thread.is_alive(): - print( - "feishu channel background thread did not stop within 2s", - flush=True, - ) + channel = getattr(app.state, "feishu_channel", None) + if channel is None: + return + try: + await channel.shutdown() + finally: + app.state.feishu_channel = None def _configure_feishu_lifecycle(