mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
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:
parent
2f833669bd
commit
ca1cae6185
@ -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
|
||||
|
||||
57
backend/tests/test_discord_thread_mapping_cancellation.py
Normal file
57
backend/tests/test_discord_thread_mapping_cancellation.py
Normal 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()
|
||||
Loading…
x
Reference in New Issue
Block a user