diff --git a/Source/exehead/Ui.c b/Source/exehead/Ui.c
index 4590bea9..fb100512 100644
--- a/Source/exehead/Ui.c
+++ b/Source/exehead/Ui.c
@@ -577,7 +577,7 @@ skipPage:
     m_curwnd = (HWND)wParam;
     goto skipPage;
   }
-  if (uMsg == WM_CLOSE && m_page == g_blocks[NB_PAGES].num - 1)
+  if (uMsg == WM_QUERYENDSESSION || (uMsg == WM_CLOSE && m_page == g_blocks[NB_PAGES].num - 1))
   {
     if (!IsWindowEnabled(m_hwndCancel))
     {
@@ -632,11 +632,6 @@ skipPage:
       SendMessage(m_curwnd, WM_COMMAND, wParam, lParam);
     }
   }
-  if (uMsg == WM_ENDSESSION && wParam)
-  {
-    // the session can end any time after we process this message so we better clean up now
-    CleanUp();
-  }
   return HandleStaticBkColor();
 }