From c579551bb19f131476132dcaa84af2130bb7d9f9 Mon Sep 17 00:00:00 2001
From: chengyongru <chengyongru.ai@gmail.com>
Date: Thu, 2 Jul 2026 14:51:03 +0800
Subject: [PATCH] fix(dingtalk): stop stream task on shutdown

---
 nanobot/channels/dingtalk.py            | 37 +++++++++++++++++
 tests/channels/test_dingtalk_channel.py | 55 +++++++++++++++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/nanobot/channels/dingtalk.py b/nanobot/channels/dingtalk.py
index 73136dd1..e3984aba 100644
--- a/nanobot/channels/dingtalk.py
+++ b/nanobot/channels/dingtalk.py
@@ -6,6 +6,8 @@ import mimetypes
 import os
 import time
 import zipfile
+from contextlib import suppress
+from inspect import isawaitable
 from io import BytesIO
 from pathlib import Path
 from typing import Any
@@ -204,6 +206,7 @@ class DingTalkChannel(BaseChannel):
         self.config: DingTalkConfig = config
         self._client: Any = None
         self._http: httpx.AsyncClient | None = None
+        self._start_task: asyncio.Task | None = None
 
         # Access Token management for sending messages
         self._access_token: str | None = None
@@ -214,6 +217,8 @@ class DingTalkChannel(BaseChannel):
 
     async def start(self) -> None:
         """Start the DingTalk bot with Stream Mode."""
+        current_task = asyncio.current_task()
+        self._start_task = current_task
         try:
             if not DINGTALK_AVAILABLE:
                 self.logger.error(
@@ -255,10 +260,25 @@ class DingTalkChannel(BaseChannel):
 
         except Exception:
             self.logger.exception("Failed to start channel")
+        finally:
+            self._running = False
+            if self._start_task is current_task:
+                self._start_task = None
 
     async def stop(self) -> None:
         """Stop the DingTalk bot."""
         self._running = False
+        await self._close_stream_client()
+        start_task = self._start_task
+        if start_task and start_task is not asyncio.current_task() and not start_task.done():
+            start_task.cancel()
+            await asyncio.sleep(0)
+            if not start_task.done():
+                start_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await start_task
+        self._client = None
+
         # Close the shared HTTP client
         if self._http:
             await self._http.aclose()
@@ -268,6 +288,23 @@ class DingTalkChannel(BaseChannel):
             task.cancel()
         self._background_tasks.clear()
 
+    async def _close_stream_client(self) -> None:
+        client = self._client
+        if client is None:
+            return
+        close = getattr(client, "close", None)
+        if close is None:
+            websocket = getattr(client, "websocket", None)
+            close = getattr(websocket, "close", None)
+        if close is None:
+            return
+        try:
+            result = close()
+            if isawaitable(result):
+                await result
+        except Exception:
+            self.logger.debug("DingTalk stream client close failed", exc_info=True)
+
     async def _get_access_token(self) -> str | None:
         """Get or refresh Access Token."""
         if self._access_token and time.time() < self._token_expiry:
diff --git a/tests/channels/test_dingtalk_channel.py b/tests/channels/test_dingtalk_channel.py
index 3a464e85..0d731f45 100644
--- a/tests/channels/test_dingtalk_channel.py
+++ b/tests/channels/test_dingtalk_channel.py
@@ -403,6 +403,61 @@ async def test_start_configures_http_timeout(monkeypatch) -> None:
     await channel.stop()
 
 
+@pytest.mark.asyncio
+async def test_stop_cancels_stream_client_after_sdk_swallows_first_cancel(monkeypatch) -> None:
+    channel = DingTalkChannel(
+        DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]),
+        MessageBus(),
+    )
+    created: dict[str, object] = {}
+
+    class _FakeWebsocket:
+        def __init__(self) -> None:
+            self.closed = False
+
+        async def close(self) -> None:
+            self.closed = True
+
+    class _CancelSwallowingStreamClient:
+        def __init__(self, _credential):
+            self.websocket = _FakeWebsocket()
+            self.started = asyncio.Event()
+            self.cancelled_once = asyncio.Event()
+            created["client"] = self
+
+        def register_callback_handler(self, _topic, _handler):
+            pass
+
+        async def start(self):
+            self.started.set()
+            while True:
+                try:
+                    await asyncio.Future()
+                except asyncio.CancelledError:
+                    self.cancelled_once.set()
+                    await asyncio.sleep(3600)
+
+    monkeypatch.setattr(dingtalk_module, "DINGTALK_AVAILABLE", True)
+    monkeypatch.setattr(dingtalk_module, "Credential", lambda *a, **k: object())
+    monkeypatch.setattr(dingtalk_module, "DingTalkStreamClient", _CancelSwallowingStreamClient)
+    monkeypatch.setattr(dingtalk_module, "ChatbotMessage", SimpleNamespace(TOPIC="topic"))
+
+    start_task = asyncio.create_task(channel.start())
+    while "client" not in created:
+        await asyncio.sleep(0)
+    client = created["client"]
+    await asyncio.wait_for(client.started.wait(), timeout=0.5)
+
+    start_task.cancel()
+    await asyncio.wait_for(client.cancelled_once.wait(), timeout=0.5)
+    assert not start_task.done()
+
+    await asyncio.wait_for(channel.stop(), timeout=0.5)
+
+    assert client.websocket.closed is True
+    assert start_task.cancelled()
+
+
 @pytest.mark.asyncio
 async def test_download_dingtalk_file(tmp_path, monkeypatch) -> None:
     """Test the two-step file download flow (get URL then download content)."""