diff --git a/src/dialog_handler.py b/src/dialog_handler.py
index edbca98e..7489714a 100644
--- a/src/dialog_handler.py
+++ b/src/dialog_handler.py
@@ -81,14 +81,21 @@ def close_thread(event: Event, box: CustomDialog, close_time: int):
             title = self.__choose_language("box_title")
         fill_string = "-" * 70
         fancy_message = f"{fill_string}\n{message}\n{fill_string}"
-        self.message_box = CustomDialog(fancy_message, title, use_ok)
         event = Event()
+        self.message_box = CustomDialog(fancy_message, title, use_ok, event.set)
         # If there is a close time, start auto close
         if close_time is not None:
-            auto_closer = Thread(target=close_thread, args=(event, self.message_box, close_time), daemon=True)
+            auto_closer = Thread(
+                target=close_thread,
+                args=(
+                    event,
+                    self.message_box,
+                    close_time,
+                ),
+                daemon=True,
+            )
             auto_closer.start()
         # Need to set event, in case thread is still waiting
-        event.set()
 
     def user_okay(self, text: str):
         """Prompts the user for the given message and asks for confirmation.
diff --git a/src/ui/setup_custom_dialog.py b/src/ui/setup_custom_dialog.py
index 9915a1ba..daa8d0e1 100644
--- a/src/ui/setup_custom_dialog.py
+++ b/src/ui/setup_custom_dialog.py
@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+from typing import Callable
+
 from PyQt5.QtWidgets import QMainWindow
 
 from src.dialog_handler import UI_LANGUAGE
@@ -8,13 +12,14 @@
 class CustomDialog(QMainWindow, Ui_CustomDialog):
     """Class for the Team selection Screen."""
 
-    def __init__(self, message: str, title: str = "Information", use_ok=False):
+    def __init__(self, message: str, title: str = "Information", use_ok=False, close_callback: Callable | None = None):
         super().__init__()
         self.setupUi(self)
         DP_CONTROLLER.initialize_window_object(self)
         self.informationLabel.setText(message)
         self.setWindowTitle(title)
         self.closeButton.clicked.connect(self.close_clicked)
+        self.close_callback = close_callback
 
         # self.closeButton.setText(close_text)
         UI_LANGUAGE.adjust_custom_dialog(self, use_ok)
@@ -22,4 +27,6 @@ def __init__(self, message: str, title: str = "Information", use_ok=False):
         DP_CONTROLLER.set_display_settings(self)
 
     def close_clicked(self):
+        if self.close_callback is not None:
+            self.close_callback()
         self.close()