diff --git a/changelog.md b/changelog.md
index 4e111855..136c85c4 100644
--- a/changelog.md
+++ b/changelog.md
@@ -10,6 +10,7 @@ Features
 Bug Fixes
 ---------
 * Improve query cancellation on control-c.
+* Improve refresh of some format strings in the toolbar.
 
 
 Internal
diff --git a/mycli/main.py b/mycli/main.py
index fe8b5186..de6021c5 100755
--- a/mycli/main.py
+++ b/mycli/main.py
@@ -1423,9 +1423,11 @@ def get_output_margin(self, status: str | None = None) -> int:
         """Get the output margin (number of rows for the prompt, footer and
         timing message."""
         if not self.prompt_lines:
-            # self.prompt_app.app.render_counter failed in the test suite
-            app = get_app()
-            self.prompt_lines = self.get_prompt(self.prompt_format, app.render_counter).count('\n') + 1
+            if self.prompt_app and self.prompt_app.app:
+                render_counter = self.prompt_app.app.render_counter
+            else:
+                render_counter = 0
+            self.prompt_lines = self.get_prompt(self.prompt_format, render_counter).count('\n') + 1
         margin = self.get_reserved_space() + self.prompt_lines
         if special.is_timing_enabled():
             margin += 1
@@ -1623,10 +1625,13 @@ def set_external_multiplex_pane_title(self) -> None:
         sys.stderr.flush()
 
     def get_custom_toolbar(self, toolbar_format: str) -> ANSI:
-        if self.prompt_app and self.prompt_app.app.current_buffer.text:
+        if not self.prompt_app:
+            return ANSI('')
+        if not self.prompt_app.app:
+            return ANSI('')
+        if self.prompt_app.app.current_buffer.text:
             return self.last_custom_toolbar_message
-        app = get_app()
-        toolbar = self.get_prompt(toolbar_format, app.render_counter)
+        toolbar = self.get_prompt(toolbar_format, self.prompt_app.app.render_counter)
         toolbar = toolbar.replace("\\x1b", "\x1b")
         self.last_custom_toolbar_message = ANSI(toolbar)
         return self.last_custom_toolbar_message
diff --git a/test/test_main.py b/test/test_main.py
index fb486492..bcfce05e 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -650,6 +650,7 @@ def server_type(self):
 
     class PromptBuffer:
         output = TestOutput()
+        app = None
 
     m.prompt_app = PromptBuffer()
     m.sqlexecute = TestExecute()