fix(discord): drain thread mapping persistence across cancellation (#6356)

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

* test(discord): retain mapping flush across cancellation
This commit is contained in:
NanPan 2026-10-06 15:49:27 +08:00 committed by GitHub
parent 2f833669bd
commit ca1cae6185
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 66 additions and 4 deletions

View File

@ -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

View File

@ -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()