From 091d7335a3e0042f19554d79c4ac7e44cc7a537d Mon Sep 17 00:00:00 2001
From: Andy <scene@scene.com>
Date: Thu, 31 Jul 2025 18:25:18 +0000
Subject: [PATCH] feat: Implement terminal cleanup on exit and signal handling
 in ComfyConsole

---
 unshackle/commands/dl.py  |  5 +++--
 unshackle/core/console.py | 40 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py
index ac18949..33ad6c4 100644
--- a/unshackle/commands/dl.py
+++ b/unshackle/commands/dl.py
@@ -765,7 +765,8 @@ class dl:
                 DOWNLOAD_LICENCE_ONLY.set()
 
             try:
-                with Live(Padding(download_table, (1, 5)), console=console, refresh_per_second=5):
+                # Use transient mode to prevent display remnants
+                with Live(Padding(download_table, (1, 5)), console=console, refresh_per_second=5, transient=True):
                     with ThreadPoolExecutor(downloads) as pool:
                         for download in futures.as_completed(
                             (
@@ -1009,7 +1010,7 @@ class dl:
 
                             multiplex_tasks.append((task_id, task_tracks))
 
-                    with Live(Padding(progress, (0, 5, 1, 5)), console=console):
+                    with Live(Padding(progress, (0, 5, 1, 5)), console=console, transient=True):
                         for task_id, task_tracks in multiplex_tasks:
                             progress.start_task(task_id)  # TODO: Needed?
                             muxed_path, return_code, errors = task_tracks.mux(
diff --git a/unshackle/core/console.py b/unshackle/core/console.py
index 30d3b46..81ff810 100644
--- a/unshackle/core/console.py
+++ b/unshackle/core/console.py
@@ -1,4 +1,7 @@
+import atexit
 import logging
+import signal
+import sys
 from datetime import datetime
 from types import ModuleType
 from typing import IO, Callable, Iterable, List, Literal, Mapping, Optional, Union
@@ -167,6 +170,8 @@ class ComfyConsole(Console):
             time.monotonic.
     """
 
+    _cleanup_registered = False
+
     def __init__(
         self,
         *,
@@ -233,6 +238,9 @@ class ComfyConsole(Console):
         if log_renderer:
             self._log_render = log_renderer
 
+        # Register terminal cleanup handlers
+        self._register_cleanup()
+
     def status(
         self,
         status: RenderableType,
@@ -283,6 +291,38 @@ class ComfyConsole(Console):
 
         return status_renderable
 
+    def _register_cleanup(self):
+        """Register terminal cleanup handlers."""
+        if not ComfyConsole._cleanup_registered:
+            ComfyConsole._cleanup_registered = True
+
+            # Register cleanup on normal exit
+            atexit.register(self._cleanup_terminal)
+
+            # Register cleanup on signals
+            signal.signal(signal.SIGINT, self._signal_handler)
+            signal.signal(signal.SIGTERM, self._signal_handler)
+
+    def _cleanup_terminal(self):
+        """Restore terminal to a clean state."""
+        try:
+            # Show cursor using ANSI escape codes
+            sys.stdout.write("\x1b[?25h")  # Show cursor
+            sys.stdout.write("\x1b[0m")  # Reset attributes
+            sys.stdout.flush()
+
+            # Also use Rich's method
+            self.show_cursor(True)
+        except Exception:
+            # Silently fail if cleanup fails
+            pass
+
+    def _signal_handler(self, signum, frame):
+        """Handle signals with cleanup."""
+        self._cleanup_terminal()
+        # Exit after cleanup
+        sys.exit(1)
+
 
 catppuccin_mocha = {
     # Colors based on "CatppuccinMocha" from Gogh themes