Bug fix: drain in-flight Feishu messages during shutdown

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)

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

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(