fix(dingtalk): stop late inbound task creation

This commit is contained in:
Xubin Ren
2026-08-30 17:56:34 +08:00
parent e8385d9257
commit d019658501
2 changed files with 75 additions and 0 deletions
+9
View File
@@ -182,6 +182,12 @@ class NanobotDingTalkHandler(_CallbackHandlerBase):
)
)
if not self.channel._accepting_inbound_tasks:
self.channel.logger.debug(
"Skipping DingTalk inbound dispatch during channel shutdown"
)
return AckMessage.STATUS_OK, "OK"
self.channel.logger.info("Received message from {} ({}): {}", sender_name, sender_id, content)
# Forward to Nanobot via _on_message (non-blocking).
@@ -256,6 +262,7 @@ class DingTalkChannel(BaseChannel):
# Hold references to background tasks to prevent GC
self._background_tasks: set[asyncio.Task[None]] = set()
self._accepting_inbound_tasks = True
def _on_background_task_done(self, task: asyncio.Task[None]) -> None:
self._background_tasks.discard(task)
@@ -282,6 +289,7 @@ class DingTalkChannel(BaseChannel):
self.logger.error("client_id and client_secret not configured")
return
self._accepting_inbound_tasks = True
self._running = True
self._http = httpx.AsyncClient(
timeout=httpx.Timeout(10.0, connect=10.0, read=30.0, write=30.0, pool=10.0)
@@ -319,6 +327,7 @@ class DingTalkChannel(BaseChannel):
async def stop(self) -> None:
"""Stop the DingTalk bot."""
self._accepting_inbound_tasks = False
self._running = False
await self._close_stream_client()
start_task = self._start_task
@@ -506,6 +506,72 @@ async def test_handler_processes_file_message(monkeypatch) -> None:
assert "/tmp/nanobot_dingtalk/user1/report.xlsx" in msg.content
@pytest.mark.asyncio
async def test_handler_does_not_spawn_message_task_after_stop_during_download(
monkeypatch,
) -> None:
channel = DingTalkChannel(
DingTalkConfig(client_id="app", client_secret="secret", allow_from=["user1"]),
MessageBus(),
)
handler = NanobotDingTalkHandler(channel)
download_started = asyncio.Event()
release_download = asyncio.Event()
message_task_started = asyncio.Event()
class _FakeFileChatbotMessage:
text = None
extensions = {}
image_content = None
rich_text_content = None
sender_staff_id = "user1"
sender_id = "fallback-user"
sender_nick = "Alice"
message_type = "file"
@staticmethod
def from_dict(_data):
return _FakeFileChatbotMessage()
async def delayed_download(*_args):
download_started.set()
await release_download.wait()
return "/tmp/nanobot_dingtalk/user1/report.xlsx"
async def block_message(*_args) -> None:
message_task_started.set()
await asyncio.Future()
monkeypatch.setattr(dingtalk_module, "ChatbotMessage", _FakeFileChatbotMessage)
monkeypatch.setattr(dingtalk_module, "AckMessage", SimpleNamespace(STATUS_OK="OK"))
monkeypatch.setattr(channel, "_download_dingtalk_file", delayed_download)
monkeypatch.setattr(channel, "_on_message", block_message)
process_task = asyncio.create_task(handler.process(SimpleNamespace(data={
"conversationType": "1",
"content": {"downloadCode": "abc123", "fileName": "report.xlsx"},
"text": {"content": ""},
})))
await download_started.wait()
try:
await channel.stop()
release_download.set()
assert await process_task == ("OK", "OK")
await asyncio.sleep(0)
assert not message_task_started.is_set()
assert not channel._background_tasks
finally:
release_download.set()
if not process_task.done():
process_task.cancel()
pending = tuple(channel._background_tasks)
for task in pending:
task.cancel()
await asyncio.gather(process_task, *pending, return_exceptions=True)
def _rich_text_message(rich_text_list):
class _FakeRichTextChatbotMessage:
text = None