fix(channels): await real cross-thread tasks on shutdown (#4816)

backend/app/channels/base.py

@@ -9,6 +9,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Awaitable, Callable, Coroutine
 from concurrent.futures import CancelledError as FutureCancelledError
 from concurrent.futures import Future
+from dataclasses import dataclass
 from typing import Any, TypeVar
 
 from app.channels.commands import extract_connect_code
@@ -29,6 +30,18 @@ logger = logging.getLogger(__name__)
 T = TypeVar("T")
 
 
+@dataclass(eq=False, slots=True)
+class _ThreadsafeSubmission:
+    coroutine: Coroutine[Any, Any, Any]
+    loop: asyncio.AbstractEventLoop
+    name: str
+    msg_id: Any
+    reservation: InboundReservation | None
+    completion: Future[Any]
+    task: asyncio.Task[Any] | None = None
+    cancel_requested: bool = False
+
+
 class Channel(ABC):
     """Base class for all IM channel implementations.
 
@@ -48,9 +61,9 @@ class Channel(ABC):
         # Provider SDK callbacks often run on a dedicated thread and submit
         # preparation work to the Gateway loop. Submission and shutdown share
         # this lock so stop() cannot miss a future created concurrently.
-        self._threadsafe_futures: set[Future[Any]] = set()
-        self._threadsafe_futures_lock = threading.Lock()
-        self._threadsafe_future_intake_open = True
+        self._threadsafe_submissions: set[_ThreadsafeSubmission] = set()
+        self._threadsafe_submissions_lock = threading.Lock()
+        self._threadsafe_submission_intake_open = True
 
     @property
     def is_running(self) -> bool:
@@ -139,12 +152,10 @@ class Channel(ABC):
 
     def _open_threadsafe_future_intake(self) -> None:
         """Allow a newly started provider to submit work to its main loop."""
-        with self._threadsafe_futures_lock:
-            pending = [future for future in self._threadsafe_futures if not future.done()]
-            if pending:
+        with self._threadsafe_submissions_lock:
+            if self._threadsafe_submissions:
                 raise RuntimeError(f"cannot restart {self.name} while cross-thread work is still running")
-            self._threadsafe_futures.clear()
-            self._threadsafe_future_intake_open = True
+            self._threadsafe_submission_intake_open = True
 
     def _submit_threadsafe_coroutine(
         self,
@@ -154,60 +165,111 @@ class Channel(ABC):
         name: str,
         msg_id: Any,
         reservation: InboundReservation | None = None,
-    ) -> Future[T] | None:
-        """Submit and retain provider-thread work until completion or shutdown."""
+    ) -> bool:
+        """Submit provider-thread work while retaining its real asyncio Task."""
 
-        with self._threadsafe_futures_lock:
-            if not self._threadsafe_future_intake_open or loop is None or not loop.is_running():
+        with self._threadsafe_submissions_lock:
+            if not self._threadsafe_submission_intake_open or loop is None or not loop.is_running():
                 coroutine.close()
                 if reservation is not None:
                     reservation.release()
-                return None
-            try:
-                future = asyncio.run_coroutine_threadsafe(coroutine, loop)
-            except RuntimeError:
-                coroutine.close()
-                if reservation is not None:
-                    reservation.release()
-                return None
-            self._threadsafe_futures.add(future)
+                return False
 
-        future.add_done_callback(
-            lambda completed: self._finalize_threadsafe_future(
-                completed,
+            submission = _ThreadsafeSubmission(
+                coroutine=coroutine,
+                loop=loop,
                 name=name,
                 msg_id=msg_id,
                 reservation=reservation,
+                completion=Future(),
             )
-        )
-        return future
+            self._threadsafe_submissions.add(submission)
+            try:
+                loop.call_soon_threadsafe(self._start_threadsafe_submission, submission)
+            except RuntimeError:
+                self._threadsafe_submissions.discard(submission)
+                coroutine.close()
+                if reservation is not None:
+                    reservation.release()
+                return False
+        return True
 
-    def _finalize_threadsafe_future(
+    def _start_threadsafe_submission(self, submission: _ThreadsafeSubmission) -> None:
+        """Create the owned Task on its event loop or finish a pre-start cancel."""
+        task: asyncio.Task[Any] | None = None
+        startup_error: BaseException | None = None
+        with self._threadsafe_submissions_lock:
+            if submission.cancel_requested:
+                self._threadsafe_submissions.discard(submission)
+                cancelled_before_start = True
+            else:
+                cancelled_before_start = False
+                try:
+                    task = submission.loop.create_task(submission.coroutine)
+                except BaseException as exc:
+                    self._threadsafe_submissions.discard(submission)
+                    startup_error = exc
+                else:
+                    submission.task = task
+
+        if cancelled_before_start:
+            submission.coroutine.close()
+            if submission.reservation is not None:
+                submission.reservation.release()
+            submission.completion.cancel()
+            return
+
+        if startup_error is not None:
+            submission.coroutine.close()
+            if submission.reservation is not None:
+                submission.reservation.release()
+            submission.completion.set_exception(startup_error)
+            self._log_future_error(submission.completion, submission.name, submission.msg_id)
+            return
+
+        assert task is not None
+        task.add_done_callback(lambda completed: self._finalize_threadsafe_submission(submission, completed))
+
+    def _finalize_threadsafe_submission(
         self,
-        future: Future[Any],
-        *,
-        name: str,
-        msg_id: Any,
-        reservation: InboundReservation | None,
+        submission: _ThreadsafeSubmission,
+        task: asyncio.Task[Any],
     ) -> None:
-        with self._threadsafe_futures_lock:
-            self._threadsafe_futures.discard(future)
-        if reservation is not None:
-            reservation.release()
-        self._log_future_error(future, name, msg_id)
+        with self._threadsafe_submissions_lock:
+            self._threadsafe_submissions.discard(submission)
+        if submission.reservation is not None:
+            submission.reservation.release()
+
+        if task.cancelled():
+            submission.completion.cancel()
+        else:
+            try:
+                submission.completion.set_result(task.result())
+            except BaseException as exc:
+                submission.completion.set_exception(exc)
+        self._log_future_error(submission.completion, submission.name, submission.msg_id)
 
     async def _close_and_drain_threadsafe_futures(self) -> None:
-        """Close cross-thread submission, then cancel and await owned futures."""
-        with self._threadsafe_futures_lock:
-            self._threadsafe_future_intake_open = False
-            futures = tuple(self._threadsafe_futures)
+        """Close submission, then cancel and await the owned asyncio Tasks."""
+        with self._threadsafe_submissions_lock:
+            self._threadsafe_submission_intake_open = False
+            submissions = tuple(self._threadsafe_submissions)
+            tasks: list[tuple[asyncio.AbstractEventLoop, asyncio.Task[Any]]] = []
+            for submission in submissions:
+                submission.cancel_requested = True
+                if submission.task is not None:
+                    tasks.append((submission.loop, submission.task))
 
-        for future in futures:
-            future.cancel()
-        if futures:
-            await asyncio.gather(*(asyncio.wrap_future(future) for future in futures), return_exceptions=True)
-            with self._threadsafe_futures_lock:
-                self._threadsafe_futures.difference_update(future for future in futures if future.done())
+        for loop, task in tasks:
+            try:
+                loop.call_soon_threadsafe(task.cancel)
+            except RuntimeError:
+                logger.warning("[%s] event loop closed before cross-thread task cancellation", self.name)
+        if submissions:
+            await asyncio.gather(
+                *(asyncio.shield(asyncio.wrap_future(submission.completion)) for submission in submissions),
+                return_exceptions=True,
+            )
 
     def _pending_connect_code(self, text: str) -> str | None:
         """Return the one-time bind code if *text* is a ``/connect <code>`` command

backend/app/channels/dingtalk.py

@@ -398,7 +398,7 @@ class DingTalkChannel(Channel):
             connect_code = self._pending_connect_code(text)
             if connect_code:
                 if self._main_loop and self._main_loop.is_running():
-                    future = self._submit_threadsafe_coroutine(
+                    scheduled = self._submit_threadsafe_coroutine(
                         self._bind_connection_from_connect_code(
                             conversation_type=conversation_type,
                             sender_staff_id=sender_staff_id,
@@ -410,7 +410,7 @@ class DingTalkChannel(Channel):
                         name="bind_connection",
                         msg_id=msg_id,
                     )
-                    if future is None:
+                    if not scheduled:
                         logger.info("[DingTalk] main loop stopped before channel connection bind could be scheduled")
                 else:
                     logger.warning("[DingTalk] main loop not running, cannot bind channel connection")
@@ -492,14 +492,14 @@ class DingTalkChannel(Channel):
                     with self._incoming_messages_lock:
                         self._incoming_messages[source_key] = message
                 logger.info("[DingTalk] publishing inbound message to bus (type=%s, msg_id=%s)", msg_type.value, msg_id)
-                future = self._submit_threadsafe_coroutine(
+                scheduled = self._submit_threadsafe_coroutine(
                     self._prepare_inbound(chat_id, inbound, reservation=reservation),
                     self._main_loop,
                     name="prepare_inbound",
                     msg_id=msg_id,
                     reservation=reservation,
                 )
-                if future is None:
+                if not scheduled:
                     logger.info("[DingTalk] main loop stopped before reserved inbound could be scheduled")
             else:
                 logger.warning("[DingTalk] main loop not running, cannot publish inbound message")

backend/app/channels/feishu.py

@@ -844,7 +844,7 @@ class FeishuChannel(Channel):
             if reservation is None:
                 return
             logger.info("[Feishu] publishing inbound message to bus (type=%s, msg_id=%s)", inbound.msg_type.value, msg_id)
-            future = self._submit_threadsafe_coroutine(
+            scheduled = self._submit_threadsafe_coroutine(
                 self._prepare_inbound(
                     msg_id,
                     inbound,
@@ -856,20 +856,20 @@ class FeishuChannel(Channel):
                 msg_id=msg_id,
                 reservation=reservation,
             )
-            if future is None:
+            if not scheduled:
                 logger.info("[Feishu] main loop stopped before reserved inbound could be scheduled")
         else:
             logger.warning("[Feishu] main loop not running, cannot publish inbound message")
 
     def _schedule_batch_flush(self, key: tuple[str, str], source_message_id: str) -> None:
         if self._main_loop and self._main_loop.is_running():
-            future = self._submit_threadsafe_coroutine(
+            scheduled = self._submit_threadsafe_coroutine(
                 self._flush_pending_inbound_batch_after(key, source_message_id),
                 self._main_loop,
                 name="flush_inbound_batch",
                 msg_id=source_message_id,
             )
-            if future is None:
+            if not scheduled:
                 logger.info("[Feishu] main loop stopped before inbound batch flush could be scheduled")
         else:
             logger.warning("[Feishu] main loop not running, cannot flush inbound batch")
@@ -1115,7 +1115,7 @@ class FeishuChannel(Channel):
             connect_code = self._pending_connect_code(text)
             if connect_code:
                 if self._main_loop and self._main_loop.is_running():
-                    future = self._submit_threadsafe_coroutine(
+                    scheduled = self._submit_threadsafe_coroutine(
                         self._bind_connection_from_connect_code(
                             message_id=msg_id,
                             chat_id=chat_id,
@@ -1126,7 +1126,7 @@ class FeishuChannel(Channel):
                         name="bind_connection",
                         msg_id=msg_id,
                     )
-                    if future is None:
+                    if not scheduled:
                         logger.info("[Feishu] main loop stopped before channel connection bind could be scheduled")
                 else:
                     logger.warning("[Feishu] main loop not running, cannot bind channel connection")

backend/app/channels/slack.py

@@ -338,7 +338,7 @@ class SlackChannel(Channel):
         connect_code = self._pending_connect_code(text)
         if connect_code:
             if self._loop and self._loop.is_running():
-                future = self._submit_threadsafe_coroutine(
+                scheduled = self._submit_threadsafe_coroutine(
                     self._bind_connection_from_connect_code(
                         event=event,
                         team_id=str(team_id or ""),
@@ -348,7 +348,7 @@ class SlackChannel(Channel):
                     name="bind_connection",
                     msg_id=event.get("ts"),
                 )
-                if future is None:
+                if not scheduled:
                     logger.info("[Slack] main loop stopped before channel connection bind could be scheduled")
             return
 
@@ -398,14 +398,14 @@ class SlackChannel(Channel):
                     # thread; no coroutine/Future waits for queue capacity.
                     self._loop.call_soon_threadsafe(self._commit_reserved_inbound, reservation, inbound)
                 else:
-                    future = self._submit_threadsafe_coroutine(
+                    scheduled = self._submit_threadsafe_coroutine(
                         self._publish_inbound_with_connection(inbound, reservation=reservation, team_id=team_id),
                         self._loop,
                         name="publish_inbound",
                         msg_id=event.get("ts", thread_ts),
                         reservation=reservation,
                     )
-                    if future is None:
+                    if not scheduled:
                         logger.info("[Slack] main loop stopped before reserved inbound could be scheduled")
             except RuntimeError:
                 reservation.release()

backend/app/channels/telegram.py

@@ -857,7 +857,7 @@ class TelegramChannel(Channel):
                 return
             try:
                 inbound = await self._attach_connection_identity(inbound)
-                future = self._submit_threadsafe_coroutine(
+                scheduled = self._submit_threadsafe_coroutine(
                     self._process_incoming_with_reply(
                         chat_id,
                         update.message.message_id,
@@ -869,7 +869,7 @@ class TelegramChannel(Channel):
                     msg_id=update.message.message_id,
                     reservation=reservation,
                 )
-                if future is None:
+                if not scheduled:
                     logger.info("[Telegram] main loop stopped before reserved command could be scheduled")
             except Exception:
                 reservation.release()
@@ -928,7 +928,7 @@ class TelegramChannel(Channel):
                 return
             try:
                 inbound = await self._attach_connection_identity(inbound)
-                future = self._submit_threadsafe_coroutine(
+                scheduled = self._submit_threadsafe_coroutine(
                     self._process_incoming_with_reply(
                         chat_id,
                         update.message.message_id,
@@ -940,7 +940,7 @@ class TelegramChannel(Channel):
                     msg_id=update.message.message_id,
                     reservation=reservation,
                 )
-                if future is None:
+                if not scheduled:
                     logger.info("[Telegram] main loop stopped before reserved inbound could be scheduled")
             except Exception:
                 reservation.release()

backend/tests/test_channel_intake_backpressure.py

@@ -375,6 +375,7 @@ async def test_provider_stop_drains_cross_thread_preparation_futures() -> None:
     for channel in providers:
         started = asyncio.Event()
         cancelled = asyncio.Event()
+        release_after_cancel = asyncio.Event()
         finished = asyncio.Event()
 
         async def preparation() -> None:
@@ -383,29 +384,108 @@ async def test_provider_stop_drains_cross_thread_preparation_futures() -> None:
                 await asyncio.Event().wait()
             except asyncio.CancelledError:
                 cancelled.set()
-                raise
+                await release_after_cancel.wait()
             finally:
                 finished.set()
 
         channel._running = True
         channel._main_loop = loop
         channel._open_threadsafe_future_intake()
-        future = await asyncio.to_thread(
+        scheduled = await asyncio.to_thread(
             channel._submit_threadsafe_coroutine,
             preparation(),
             loop,
             name="test_preparation",
             msg_id="message-1",
         )
-        assert future is not None
+        assert scheduled is True
         await asyncio.wait_for(started.wait(), timeout=1)
 
-        await asyncio.wait_for(channel.stop(), timeout=1)
+        stop_task = asyncio.create_task(channel.stop())
+        await asyncio.wait_for(cancelled.wait(), timeout=1)
+        await asyncio.sleep(0)
+
+        assert not stop_task.done()
+        assert not finished.is_set()
+
+        release_after_cancel.set()
+        await asyncio.wait_for(stop_task, timeout=1)
 
-        assert cancelled.is_set()
         assert finished.is_set()
-        assert future.done()
-        assert channel._threadsafe_futures == set()
+        assert channel._threadsafe_submissions == set()
+
+
+@pytest.mark.asyncio
+async def test_shutdown_closes_cross_thread_submission_before_task_start() -> None:
+    channel = SlackChannel(MessageBus(), config={})
+    channel._open_threadsafe_future_intake()
+    coroutine_started = False
+
+    async def preparation() -> None:
+        nonlocal coroutine_started
+        coroutine_started = True
+
+    scheduled = channel._submit_threadsafe_coroutine(
+        preparation(),
+        asyncio.get_running_loop(),
+        name="test_preparation",
+        msg_id="message-1",
+    )
+    assert scheduled is True
+
+    await asyncio.wait_for(channel._close_and_drain_threadsafe_futures(), timeout=1)
+
+    assert coroutine_started is False
+    assert channel._threadsafe_submissions == set()
+
+
+@pytest.mark.asyncio
+async def test_cancelled_cross_thread_drain_remains_retryable() -> None:
+    channel = SlackChannel(MessageBus(), config={})
+    channel._open_threadsafe_future_intake()
+    started = asyncio.Event()
+    cancellation_seen = asyncio.Event()
+    release_after_cancel = asyncio.Event()
+    finished = asyncio.Event()
+
+    async def preparation() -> None:
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            cancellation_seen.set()
+            await release_after_cancel.wait()
+        finally:
+            finished.set()
+
+    assert channel._submit_threadsafe_coroutine(
+        preparation(),
+        asyncio.get_running_loop(),
+        name="test_preparation",
+        msg_id="message-1",
+    )
+    await asyncio.wait_for(started.wait(), timeout=1)
+    submission = next(iter(channel._threadsafe_submissions))
+
+    first_drain = asyncio.create_task(channel._close_and_drain_threadsafe_futures())
+    await asyncio.wait_for(cancellation_seen.wait(), timeout=1)
+    first_drain.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await first_drain
+
+    assert not finished.is_set()
+    assert channel._threadsafe_submissions
+    assert not submission.completion.done()
+
+    second_drain = asyncio.create_task(channel._close_and_drain_threadsafe_futures())
+    await asyncio.sleep(0)
+    assert not second_drain.done()
+
+    release_after_cancel.set()
+    await asyncio.wait_for(second_drain, timeout=1)
+
+    assert finished.is_set()
+    assert channel._threadsafe_submissions == set()
 
 
 def test_channel_service_threads_intake_limits_into_bus_and_worker_pool() -> None:

backend/tests/test_channels.py

@@ -7844,9 +7844,9 @@ class TestSlackSendRetry:
 
 class TestSlackAllowedUsers:
     @staticmethod
-    def _submit_coro(coro, loop):
+    def _submit_coro(coro, loop, **_kwargs):
         coro.close()
-        return MagicMock()
+        return True
 
     @staticmethod
     def _immediate_loop():
@@ -7946,8 +7946,9 @@ class TestSlackAllowedUsers:
             "ts": "1710000000.000100",
         }
 
-        with patch(
-            "app.channels.slack.asyncio.run_coroutine_threadsafe",
+        with patch.object(
+            channel,
+            "_submit_threadsafe_coroutine",
             side_effect=self._submit_coro,
         ) as submit:
             channel._handle_message_event(event)