mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 16:21:50 +03:00
fix(dingtalk): stop late inbound task creation
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user