From ca1cae6185552a4ca4cff28395819ea48a2f9d92 Mon Sep 17 00:00:00 2001
From: NanPan <111261006+poijygfdyy@users.noreply.github.com>
Date: Tue, 6 Oct 2026 15:49:27 +0800
Subject: [PATCH] fix(discord): drain thread mapping persistence across
 cancellation (#6356)

* fix(discord): drain thread mapping persistence across cancellation

* test(discord): retain mapping flush across cancellation
---
 backend/app/channels/discord.py               | 13 +++--
 ...est_discord_thread_mapping_cancellation.py | 57 +++++++++++++++++++
 2 files changed, 66 insertions(+), 4 deletions(-)
 create mode 100644 backend/tests/test_discord_thread_mapping_cancellation.py

diff --git a/backend/app/channels/discord.py b/backend/app/channels/discord.py
index 17c8a72f26..4aeeffa602 100644
--- a/backend/app/channels/discord.py
+++ b/backend/app/channels/discord.py
@@ -15,6 +15,7 @@ from app.channels.base import Channel
 from app.channels.commands import is_known_channel_command
 from app.channels.connection_identity import attach_connection_identity
 from app.channels.message_bus import InboundMessage, InboundMessageType, InboundReservation, MessageBus, OutboundMessage, ResolvedAttachment
+from deerflow.utils.file_io import await_drained
 
 logger = logging.getLogger(__name__)
 
@@ -253,6 +254,10 @@ class DiscordChannel(Channel):
             except Exception:
                 logger.exception("[Discord] failed to persist thread mappings")
 
+    async def _persist_thread_mappings_async(self) -> None:
+        """Persist mappings without letting cancellation detach the worker thread."""
+        await await_drained(asyncio.to_thread(self._persist_thread_mappings))
+
     @staticmethod
     def _read_attachment_bytes(path: str) -> bytes:
         """Read an attachment file synchronously (intended for ``asyncio.to_thread``)."""
@@ -274,7 +279,7 @@ class DiscordChannel(Channel):
         # with {} — the #2897 data loss this PR exists to prevent.
         if self._thread_store_loaded:
             try:
-                await asyncio.to_thread(self._persist_thread_mappings)
+                await self._persist_thread_mappings_async()
             except Exception:
                 logger.warning("[Discord] failed to flush thread mappings during shutdown")
 
@@ -685,7 +690,7 @@ class DiscordChannel(Channel):
                     if thread_obj is not None:
                         target_thread_id = str(thread_obj.id)
                         self._record_thread_mapping(channel_id, target_thread_id)
-                        await asyncio.to_thread(self._persist_thread_mappings)
+                        await self._persist_thread_mappings_async()
                         thread_id = target_thread_id
                         chat_id = channel_id
                         typing_target = thread_obj
@@ -712,7 +717,7 @@ class DiscordChannel(Channel):
                 if thread_obj is not None:
                     target_thread_id = str(thread_obj.id)
                     self._record_thread_mapping(channel_id, target_thread_id)
-                    await asyncio.to_thread(self._persist_thread_mappings)
+                    await self._persist_thread_mappings_async()
                     thread_id = target_thread_id
                     chat_id = channel_id
                     typing_target = thread_obj  # Type into the new thread
@@ -735,7 +740,7 @@ class DiscordChannel(Channel):
                 else:
                     target_thread_id = str(thread_obj.id)
                     self._record_thread_mapping(channel_id, target_thread_id)
-                    await asyncio.to_thread(self._persist_thread_mappings)
+                    await self._persist_thread_mappings_async()
                     thread_id = target_thread_id
                     chat_id = channel_id
                     typing_target = thread_obj  # Type into the new thread
diff --git a/backend/tests/test_discord_thread_mapping_cancellation.py b/backend/tests/test_discord_thread_mapping_cancellation.py
new file mode 100644
index 0000000000..a57c78e484
--- /dev/null
+++ b/backend/tests/test_discord_thread_mapping_cancellation.py
@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from unittest.mock import AsyncMock
+
+import pytest
+
+from app.channels.discord import DiscordChannel
+from app.channels.message_bus import MessageBus
+
+
+@pytest.mark.asyncio
+async def test_discord_stop_drains_thread_mapping_write_across_repeated_cancellation(tmp_path) -> None:
+    channel = DiscordChannel(bus=MessageBus(), config={"bot_token": "token"})
+    channel._thread_store_path = tmp_path / "discord_threads.json"
+    channel._thread_store_loaded = True
+    channel._discord_loop = None
+    channel._client = None
+    channel._thread = None
+    channel._cancel_ephemeral_tasks = AsyncMock()
+    channel._record_thread_mapping("chan", "thread")
+
+    started = threading.Event()
+    allow = threading.Event()
+    original = channel._persist_thread_mappings
+
+    def blocking_persist() -> None:
+        started.set()
+        assert allow.wait(timeout=2)
+        original()
+
+    channel._persist_thread_mappings = blocking_persist
+    task = asyncio.create_task(channel.stop())
+    try:
+        assert await asyncio.to_thread(started.wait, 2)
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done()
+
+        task.cancel()
+        for _ in range(5):
+            await asyncio.sleep(0)
+        assert not task.done()
+
+        allow.set()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    finally:
+        allow.set()
+        if not task.done():
+            task.cancel()
+            await asyncio.gather(task, return_exceptions=True)
+
+    assert channel._thread_store_path.exists()