Repository: textualize/textual
Issue: Removing a widget from an action, when it initiated the action, causes the app to hang

The following code may be a little convoluted in that it's using `@click` from inside a renderable, this started out as an investigation as to why dependency links in [PISpy](https://github.com/davep/pispy) were hanging the application after upgrading Textual to 0.10.1 -- when I get a moment I'll see if I can reduce it the the smallest recreation possible.

Anyway, given this:

```python
from textual.app        import App, ComposeResult
from textual.containers import Vertical
from textual.widgets    import Header, Footer, Label

class Output( Vertical ):

    async def clear( self ) -> None:
        await self.query( "*" ).remove()

    async def populate( self ) -> None:
        await self.clear()
        await self.mount( *[
            Label( f"[@click=app.populate]Repopulate this widget[/] {n}")
            for n in range( 20 )
        ] )

    async def on_mount( self ) -> None:
        await self.populate()

class AsyncClickApp( App[ None ] ):

    def compose( self ) -> ComposeResult:
        yield Header()
        yield Output()
        yield Footer()

    async def action_populate( self ) -> None:
        await self.query_one( Output ).populate()

if __name__ == "__main__":
    AsyncClickApp().run()
```

Clicking on any of the links causes the application to hang with 0.10.1. If I run this with 0.9.1 it works fine, clearing the `Vertical` of its content and repopulating it.

Changed files:
src/textual/app.py
tests/test_widget.py
src/textual/await_remove.py
src/textual/message_pump.py
src/textual/_context.py
src/textual/widget.py

diff --git a/src/textual/_context.py b/src/textual/_context.py
index 04b264d33..625152a95 100644
--- a/src/textual/_context.py
+++ b/src/textual/_context.py
@@ -4,6 +4,7 @@ from contextvars import ContextVar
 
 if TYPE_CHECKING:
     from .app import App
+    from .message_pump import MessagePump
 
 
 class NoActiveAppError(RuntimeError):
@@ -11,3 +12,4 @@ class NoActiveAppError(RuntimeError):
 
 
 active_app: ContextVar["App"] = ContextVar("active_app")
+active_message_pump: ContextVar["MessagePump"] = ContextVar("active_message_pump")
diff --git a/src/textual/app.py b/src/textual/app.py
index 068b94e0b..8712353e1 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -46,7 +46,7 @@ from ._animator import DEFAULT_EASING, Animatable, Animator, EasingFunction
 from ._ansi_sequences import SYNC_END, SYNC_START
 from ._asyncio import create_task
 from ._callback import invoke
-from ._context import active_app
+from ._context import active_app, active_message_pump
 from ._event_broker import NoHandler, extract_handler_actions
 from ._filter import LineFilter, Monochrome
 from ._path import _make_path_object_relative
@@ -1095,8 +1095,12 @@ class App(Generic[ReturnType], DOMNode):
 
         Args:
             *widgets: The widget(s) to mount.
-            before: Optional location to mount before.
-            after: Optional location to mount after.
+            before: Optional location to mount before. An `int` is the index
+                of the child to mount before, a `str` is a `query_one` query to
+                find the widget to mount before.
+            after: Optional location to mount after. An `int` is the index
+                of the child to mount after, a `str` is a `query_one` query to
+                find the widget to mount after.
 
         Returns:
             An awaitable object that waits for widgets to be mounted.
@@ -1113,6 +1117,7 @@ class App(Generic[ReturnType], DOMNode):
     def mount_all(
         self,
         widgets: Iterable[Widget],
+        *,
         before: int | str | Widget | None = None,
         after: int | str | Widget | None = None,
     ) -> AwaitMount:
@@ -1120,8 +1125,12 @@ class App(Generic[ReturnType], DOMNode):
 
         Args:
             widgets: An iterable of widgets.
-            before: Optional location to mount before.
-            after: Optional location to mount after.
+            before: Optional location to mount before. An `int` is the index
+                of the child to mount before, a `str` is a `query_one` query to
+                find the widget to mount before.
+            after: Optional location to mount after. An `int` is the index
+                of the child to mount after, a `str` is a `query_one` query to
+                find the widget to mount after.
 
         Returns:
             An awaitable object that waits for widgets to be mounted.
@@ -2103,7 +2112,7 @@ class App(Generic[ReturnType], DOMNode):
         """Remove nodes from DOM, and return an awaitable that awaits cleanup.
 
         Args:
-            widgets: List of nodes to remvoe.
+            widgets: List of nodes to remove.
 
         Returns:
             Awaitable that returns when the nodes have been fully removed.
@@ -2127,17 +2136,19 @@ class App(Generic[ReturnType], DOMNode):
         removed_widgets = self._detach_from_dom(widgets)
 
         finished_event = asyncio.Event()
-        create_task(
+        remove_task = create_task(
             prune_widgets_task(removed_widgets, finished_event), name="prune nodes"
         )
 
-        return AwaitRemove(finished_event)
+        await_remove = AwaitRemove(finished_event, remove_task)
+        self.call_next(await_remove)
+        return await_remove
 
     async def _prune_nodes(self, widgets: list[Widget]) -> None:
         """Remove nodes and children.
 
         Args:
-            widgets: _description_
+            widgets: Widgets to remove.
         """
         async with self._dom_lock:
             for widget in widgets:
diff --git a/src/textual/await_remove.py b/src/textual/await_remove.py
index cd794d8c6..f8d61e3ff 100644
--- a/src/textual/await_remove.py
+++ b/src/textual/await_remove.py
@@ -1,19 +1,24 @@
 """Provides the type of an awaitable remove."""
 
-from asyncio import Event
+from asyncio import Event, Task
 from typing import Generator
 
 
 class AwaitRemove:
     """An awaitable returned by App.remove and DOMQuery.remove."""
 
-    def __init__(self, finished_flag: Event) -> None:
+    def __init__(self, finished_flag: Event, task: Task) -> None:
         """Initialise the instance of ``AwaitRemove``.
 
         Args:
             finished_flag: The asyncio event to wait on.
+            task: The task which does the remove (required to keep a reference).
         """
         self.finished_flag = finished_flag
+        self._task = task
+
+    async def __call__(self) -> None:
+        return await self
 
     def __await__(self) -> Generator[None, None, None]:
         async def await_prune() -> None:
diff --git a/src/textual/message_pump.py b/src/textual/message_pump.py
index 7cbaf2875..c469b381d 100644
--- a/src/textual/message_pump.py
+++ b/src/textual/message_pump.py
@@ -17,7 +17,7 @@ from weakref import WeakSet
 from . import Logger, events, log, messages
 from ._asyncio import create_task
 from ._callback import invoke
-from ._context import NoActiveAppError, active_app
+from ._context import NoActiveAppError, active_app, active_message_pump
 from ._time import time
 from .case import camel_to_snake
 from .errors import DuplicateKeyHandlers
@@ -313,12 +313,18 @@ class MessagePump(metaclass=MessagePumpMeta):
         Reactive._reset_object(self)
         await self._message_queue.put(None)
         if wait and self._task is not None and asyncio.current_task() != self._task:
-            # Ensure everything is closed before returning
-            await self._task
+            try:
+                running_widget = active_message_pump.get()
+            except LookupError:
+                running_widget = None
+
+            if running_widget is None or running_widget is not self:
+                await self._task
 
     def _start_messages(self) -> None:
         """Start messages task."""
         if self.app._running:
+            active_message_pump.set(self)
             self._task = create_task(
                 self._process_messages(), name=f"message pump {self}"
             )
@@ -357,7 +363,6 @@ class MessagePump(metaclass=MessagePumpMeta):
     async def _process_messages_loop(self) -> None:
         """Process messages until the queue is closed."""
         _rich_traceback_guard = True
-
         while not self._closed:
             try:
                 message = await self._get_message()
diff --git a/src/textual/widget.py b/src/textual/widget.py
index e6c9c2626..46708f6cd 100644
--- a/src/textual/widget.py
+++ b/src/textual/widget.py
@@ -618,6 +618,37 @@ class Widget(DOMNode):
         self.call_next(await_mount)
         return await_mount
 
+    def mount_all(
+        self,
+        widgets: Iterable[Widget],
+        *,
+        before: int | str | Widget | None = None,
+        after: int | str | Widget | None = None,
+    ) -> AwaitMount:
+        """Mount widgets from an iterable.
+
+        Args:
+            widgets: An iterable of widgets.
+            before: Optional location to mount before. An `int` is the index
+                of the child to mount before, a `str` is a `query_one` query to
+                find the widget to mount before.
+            after: Optional location to mount after. An `int` is the index
+                of the child to mount after, a `str` is a `query_one` query to
+                find the widget to mount after.
+
+        Returns:
+            An awaitable object that waits for widgets to be mounted.
+
+        Raises:
+            MountError: If there is a problem with the mount request.
+
+        Note:
+            Only one of ``before`` or ``after`` can be provided. If both are
+            provided a ``MountError`` will be raised.
+        """
+        await_mount = self.mount(*widgets, before=before, after=after)
+        return await_mount
+
     def move_child(
         self,
         child: int | Widget,
diff --git a/tests/test_widget.py b/tests/test_widget.py
index e3c0a618c..a06cf7857 100644
--- a/tests/test_widget.py
+++ b/tests/test_widget.py
@@ -1,13 +1,12 @@
 import pytest
-import rich
 
 from textual._node_list import DuplicateIds
 from textual.app import App, ComposeResult
 from textual.css.errors import StyleValueError
 from textual.css.query import NoMatches
-from textual.dom import DOMNode
 from textual.geometry import Size
 from textual.widget import Widget, MountError
+from textual.widgets import Label
 
 
 @pytest.mark.parametrize(
@@ -157,3 +156,28 @@ def test_widget_mount_ids_must_be_unique_mounting_multiple_calls(parent):
     parent.mount(widget1)
     with pytest.raises(DuplicateIds):
         parent.mount(widget2)
+
+
+# Regression test for https://github.com/Textualize/textual/issues/1634
+async def test_remove():
+    class RemoveMeLabel(Label):
+        async def on_mount(self) -> None:
+            await self.action("app.remove_all")
+
+    class Container(Widget):
+        async def clear(self) -> None:
+            await self.query("*").remove()
+
+    class RemoveApp(App):
+        def compose(self) -> ComposeResult:
+            yield Container(RemoveMeLabel())
+
+        async def action_remove_all(self) -> None:
+            await self.query_one(Container).clear()
+            self.exit(123)
+
+    app = RemoveApp()
+    async with app.run_test() as pilot:
+        await pilot.press("r")
+        await pilot.pause()
+    assert app.return_value == 123