From 604b04db2aef078fe2d0884210739c8436d95f1e Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Fri, 14 Jun 2024 10:52:53 +0100
Subject: [PATCH] hardening

---
 src/textual/app.py             | 11 ++++++++---
 src/textual/widget.py          |  2 ++
 src/textual/widgets/_header.py | 11 +++++++++--
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/textual/app.py b/src/textual/app.py
index 2b82abb246..a9237dc6ec 100644
--- a/src/textual/app.py
+++ b/src/textual/app.py
@@ -2617,9 +2617,14 @@ async def recompose(self) -> None:
 
         Recomposing will remove children and call `self.compose` again to remount.
         """
-        async with self.screen.batch():
-            await self.screen.query("*").exclude(".-textual-system").remove()
-            await self.screen.mount_all(compose(self))
+        if self._exit:
+            return
+        try:
+            async with self.screen.batch():
+                await self.screen.query("*").exclude(".-textual-system").remove()
+                await self.screen.mount_all(compose(self))
+        except ScreenStackError:
+            pass
 
     def _register_child(
         self, parent: DOMNode, child: Widget, before: int | None, after: int | None
diff --git a/src/textual/widget.py b/src/textual/widget.py
index b1321e378a..5c701695c3 100644
--- a/src/textual/widget.py
+++ b/src/textual/widget.py
@@ -1005,6 +1005,8 @@ def mount_all(
             Only one of ``before`` or ``after`` can be provided. If both are
             provided a ``MountError`` will be raised.
         """
+        if self.app._exit:
+            return AwaitMount(self, [])
         await_mount = self.mount(*widgets, before=before, after=after)
         return await_mount
 
diff --git a/src/textual/widgets/_header.py b/src/textual/widgets/_header.py
index a549881829..3505514813 100644
--- a/src/textual/widgets/_header.py
+++ b/src/textual/widgets/_header.py
@@ -7,6 +7,7 @@
 from rich.text import Text
 
 from ..app import RenderResult
+from ..dom import NoScreen
 from ..events import Click, Mount
 from ..reactive import Reactive
 from ..widget import Widget
@@ -213,10 +214,16 @@ def screen_sub_title(self) -> str:
 
     def _on_mount(self, _: Mount) -> None:
         async def set_title() -> None:
-            self.query_one(HeaderTitle).text = self.screen_title
+            try:
+                self.query_one(HeaderTitle).text = self.screen_title
+            except NoScreen:
+                pass
 
         async def set_sub_title() -> None:
-            self.query_one(HeaderTitle).sub_text = self.screen_sub_title
+            try:
+                self.query_one(HeaderTitle).sub_text = self.screen_sub_title
+            except NoScreen:
+                pass
 
         self.watch(self.app, "title", set_title)
         self.watch(self.app, "sub_title", set_sub_title)