relationalai 0.11.4__py3-none-any.whl → 0.12.1__py3-none-any.whl

relationalai/tools/cli_controls.py

@@ -9,6 +9,7 @@ import shutil
 import sys
 import threading
 import time
+import importlib
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Sequence, TextIO, cast
@@ -116,6 +117,19 @@ def rich_str(string:str, style:str|None = None) -> str:
     console.print(string, style=style)
     return output.getvalue()
 
+def _load_ipython_display() -> tuple[Any, Callable[..., Any]]:
+    """Load IPython display helpers, raising if unavailable."""
+    try:
+        module = importlib.import_module("IPython.display")
+    except ImportError as exc:  # pragma: no cover - only triggered without IPython
+        raise RuntimeError(
+            "NotebookTaskProgress requires IPython when running in a notebook environment."
+        ) from exc
+
+    html_factory = getattr(module, "HTML")
+    display_fn = getattr(module, "display")
+    return html_factory, cast(Callable[..., Any], display_fn)
+
 def nat_path(path: Path, base: Path):
     resolved_path = path.resolve()
     resolved_base = base.resolve()
@@ -763,34 +777,22 @@ class _TimerManager:
 
     def _process_operation(self, task_id: str, op_type: str):
         """Process a completed delayed operation."""
+        progress = self._progress
         if op_type == "remove_highlighting":
-            if hasattr(self._progress, '_highlighted_tasks') and task_id in self._progress._highlighted_tasks:
-                del self._progress._highlighted_tasks[task_id]
-                # For TaskProgress, invalidate cache and update display
-                if hasattr(self._progress, '_invalidate_cache'):
-                    self._progress._invalidate_cache()
-                    self._progress._update_display()
-            elif hasattr(self._progress, 'highlighted_tasks') and task_id in self._progress.highlighted_tasks:
-                del self._progress.highlighted_tasks[task_id]
-                # For NotebookTaskProgress, no special update needed
+            if hasattr(progress, "_highlighted_tasks") and task_id in progress._highlighted_tasks:
+                del progress._highlighted_tasks[task_id]
+                if hasattr(progress, "_after_task_update"):
+                    progress._after_task_update()
         elif op_type == "delayed_removal":
-            if hasattr(self._progress, '_tasks') and task_id in self._progress._tasks:
-                del self._progress._tasks[task_id]
-                # For TaskProgress, invalidate cache and update display
-                if hasattr(self._progress, '_invalidate_cache'):
-                    self._progress._invalidate_cache()
-                    self._progress._update_display()
-            elif hasattr(self._progress, 'sub_tasks') and task_id in self._progress.sub_tasks:
-                del self._progress.sub_tasks[task_id]
-                # For NotebookTaskProgress, no special update needed
+            if hasattr(progress, "_tasks") and task_id in progress._tasks:
+                del progress._tasks[task_id]
+                if hasattr(progress, "_after_task_update"):
+                    progress._after_task_update()
         elif op_type == "delayed_hiding":
-            if hasattr(self._progress, '_tasks') and task_id in self._progress._tasks:
-                # Mark task as hidden but keep it in the data structure
-                self._progress._tasks[task_id].hidden = True
-                # For TaskProgress, invalidate cache and update display
-                if hasattr(self._progress, '_invalidate_cache'):
-                    self._progress._invalidate_cache()
-                    self._progress._update_display()
+            if hasattr(progress, "_tasks") and task_id in progress._tasks:
+                progress._tasks[task_id].hidden = True
+                if hasattr(progress, "_after_task_update"):
+                    progress._after_task_update()
 
     def stop(self):
         """Stop the timer manager."""
@@ -798,7 +800,176 @@ class _TimerManager:
         self._operations.clear()
 
 
-class TaskProgress:
+class _TaskStateMixin:
+    """Shared task management helpers for notebook and terminal progress displays."""
+
+    enable_highlighting: bool = True
+
+    def _init_task_state(self, *, hide_on_completion: bool = False, show_duration_summary: bool = True) -> None:
+        self.hide_on_completion = hide_on_completion
+        self.show_duration_summary = show_duration_summary
+        self._tasks: dict[str, TaskInfo] = {}
+        self._next_task_id: int = 1
+        self._highlighted_tasks: dict[str, float] = {}
+        self._process_start_time: float | None = None
+        self._process_end_time: float | None = None
+        self.main_completed: bool = False
+        self.main_failed: bool = False
+        self._timer_manager = _TimerManager(self)
+
+    def _after_task_update(self) -> None:
+        """Hook for subclasses to react when task state changes."""
+        # Implemented by subclasses when they need to update the display immediately.
+        return None
+
+    def _generate_task_id(self) -> str:
+        """Generate a unique task ID."""
+        task_id = f"task_{self._next_task_id}"
+        self._next_task_id += 1
+        return task_id
+
+    def add_sub_task(self, description: str, task_id: str | None = None, category: str = "general") -> str:
+        """Add a new sub-task and return its unique ID."""
+        if task_id is None:
+            task_id = self._generate_task_id()
+
+        if task_id not in self._tasks:
+            self._tasks[task_id] = TaskInfo(description=description, category=category)
+            self._after_task_update()
+
+        return task_id
+
+    def update_sub_task(self, task_id: str, description: str) -> None:
+        """Update an existing sub-task description."""
+        if task_id in self._tasks:
+            task_info = self._tasks[task_id]
+            if self.enable_highlighting and task_info.description != description:
+                self._highlighted_tasks[task_id] = time.time() + HIGHLIGHT_DURATION
+                self._timer_manager.schedule_highlight_removal(task_id)
+
+            task_info.description = description
+            self._after_task_update()
+
+    def complete_sub_task(self, task_id: str, record_time: bool = True) -> None:
+        """Complete a sub-task by marking it as done."""
+        if task_id in self._tasks:
+            if task_id in self._highlighted_tasks:
+                del self._highlighted_tasks[task_id]
+
+            if not self._tasks[task_id].completed and record_time:
+                self._tasks[task_id].completed_time = time.time()
+            self._tasks[task_id].completed = True
+
+            self._after_task_update()
+            self._timer_manager.schedule_task_hiding(task_id)
+
+    def remove_sub_task(self, task_id: str, animate: bool = True) -> None:
+        """Remove a sub-task by ID with optional completion animation."""
+        if task_id in self._tasks:
+            if task_id in self._highlighted_tasks:
+                del self._highlighted_tasks[task_id]
+
+            if animate:
+                self.complete_sub_task(task_id)
+            else:
+                del self._tasks[task_id]
+                self._after_task_update()
+
+    def update_sub_status(self, sub_status: str) -> None:
+        """Legacy method for backward compatibility - creates/updates a default sub-task."""
+        self.add_sub_task(sub_status, "default")
+        self.update_sub_task("default", sub_status)
+
+    def update_main_status(self, message: str) -> None:
+        """Update the main status line with custom information."""
+        if getattr(self, "description", "") != message:
+            self.description = message
+            self._after_task_update()
+
+    def update_messages(self, updater: dict[str, str]) -> None:
+        """Update both main message and sub-status if provided."""
+        if "message" in updater:
+            self.description = updater["message"]
+            self._after_task_update()
+        if "sub_status" in updater:
+            self.update_sub_status(updater["sub_status"])
+        if "success_message" in updater:
+            self.success_message = updater["success_message"]
+        if "failure_message" in updater:
+            self.failure_message = updater["failure_message"]
+
+    def get_sub_task_count(self) -> int:
+        """Get the current number of active sub-tasks."""
+        return len(self._tasks)
+
+    def list_sub_tasks(self) -> list[str]:
+        """Get a list of all active sub-task IDs."""
+        return list(self._tasks.keys())
+
+    def get_task_status(self) -> str:
+        """Get a human-readable status of current task count vs limit."""
+        return f"› Active tasks: {len(self._tasks)}"
+
+    def get_task_duration(self, task_id: str) -> float:
+        """Get the duration of a specific task in seconds."""
+        if task_id in self._tasks:
+            return self._tasks[task_id].get_duration()
+        return 0.0
+
+    def get_completed_tasks(self) -> dict[str, TaskInfo]:
+        """Get all completed tasks with their timing information."""
+        return {task_id: task_info for task_id, task_info in self._tasks.items() if task_info.completed}
+
+    def get_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
+        """Get all tasks (completed or active) for a specific category."""
+        return {
+            task_id: task_info
+            for task_id, task_info in self._tasks.items()
+            if task_info.category == category
+        }
+
+    def get_completed_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
+        """Get all completed tasks for a specific category."""
+        return {
+            task_id: task_info
+            for task_id, task_info in self._tasks.items()
+            if task_info.category == category and task_info.completed
+        }
+
+    def _clear_all_tasks(self) -> None:
+        """Clear all tasks and related data."""
+        self._tasks.clear()
+        self._highlighted_tasks.clear()
+
+    def set_process_start_time(self) -> None:
+        """Set the overall process start time."""
+        self._process_start_time = time.time()
+
+    def set_process_end_time(self) -> None:
+        """Set the overall process end time."""
+        self._process_end_time = time.time()
+
+    def get_total_duration(self) -> float:
+        """Get the total duration from first task added to last task completed."""
+        if not self._tasks:
+            return 0.0
+
+        completed_tasks = self.get_completed_tasks()
+        if not completed_tasks:
+            return 0.0
+
+        start_times = [task.added_time for task in self._tasks.values()]
+        completion_times = [task.completed_time for task in completed_tasks.values() if task.completed_time > 0]
+
+        if not start_times or not completion_times:
+            return 0.0
+
+        earliest_start = min(start_times)
+        latest_completion = max(completion_times)
+        return latest_completion - earliest_start
+
+
+class TaskProgress(_TaskStateMixin):
     """A progress component that uses Rich's Live system to provide proper two-line display.
 
     This class provides:
@@ -828,40 +999,36 @@ class TaskProgress:
         self.leading_newline = leading_newline
         self.trailing_newline = trailing_newline
         self.transient = transient
-        self.hide_on_completion = hide_on_completion
-        self.show_duration_summary = show_duration_summary
+        self._init_task_state(
+            hide_on_completion=hide_on_completion,
+            show_duration_summary=show_duration_summary,
+        )
+        self.enable_highlighting = True
 
         # Detect CI environment to avoid cursor control issues
         from ..environments import CIEnvironment
         self.is_ci = isinstance(runtime_env, CIEnvironment)
+        self.is_jupyter = isinstance(runtime_env, JupyterEnvironment)
 
         # Core components
-        # In CI, don't force terminal to avoid ANSI escape sequences that cause multiple lines
-        self.console = Console(force_terminal=not self.is_ci)
+        # In CI or Jupyter, avoid forcing terminal rendering to prevent duplicate outputs
+        force_terminal = not self.is_ci and not self.is_jupyter
+        force_jupyter = True if self.is_jupyter else None
+        self.console = Console(
+            force_terminal=force_terminal,
+            force_jupyter=force_jupyter,
+        )
         self.live = None
         self.main_completed = False
         self.main_failed = False
 
-        # Task management - unified data structure
-        self._tasks = {}  # task_id -> TaskInfo
-        self._next_task_id = 1
-
-        # Overall process timing
-        self._process_start_time = None
-        self._process_end_time = None
-
         # Animation state
         self.spinner_index = 0
 
-        # Highlighting system
-        self._highlighted_tasks = {}  # task_id -> highlight_until_time
-
         # Performance optimizations
         self._render_cache = None
         self._last_state_hash = None
 
-        # Threading
-        self._timer_manager = _TimerManager(self)
         self._spinner_thread = None
 
     def _generate_task_id(self) -> str:
@@ -979,192 +1146,31 @@ class TaskProgress:
         if self.live:
             self.live.update(self._render_display())
 
-    def add_sub_task(self, description: str, task_id: str | None = None, category: str = "general") -> str:
-        """Add a new sub-task and return its unique ID.
-
-        Args:
-            description: Description of the subtask
-            task_id: Optional custom task ID, if not provided one will be generated
-            category: Category for this task (e.g., "indexing", "provisioning", "change_tracking")
-
-        Returns:
-            str: The task ID for this subtask
-        """
-        if task_id is None:
-            task_id = self._generate_task_id()
-
-        if task_id not in self._tasks:
-            self._tasks[task_id] = TaskInfo(description=description, category=category)
-            self._invalidate_cache()
-            self._update_display()
-
-        return task_id
-
-    def update_sub_task(self, task_id: str, description: str) -> None:
-        """Update an existing sub-task description.
-
-        When the description text changes from the previous value, the subtask
-        will be highlighted in yellow for 2 seconds to make the change visible.
-        """
-        if task_id in self._tasks:
-            task_info = self._tasks[task_id]
-
-            # Check if text has changed
-            if task_info.description != description:
-                # Text has changed - set up highlighting
-                self._highlighted_tasks[task_id] = time.time() + HIGHLIGHT_DURATION
-                self._timer_manager.schedule_highlight_removal(task_id)
+    def _after_task_update(self) -> None:
+        """Refresh the live display when task state changes."""
+        self._invalidate_cache()
+        self._update_display()
 
-            task_info.description = description
-            self._invalidate_cache()
-            self._update_display()
-
-    def complete_sub_task(self, task_id: str, record_time: bool = True) -> None:
-        """Complete a sub-task by marking it as done."""
-        if task_id in self._tasks:
-            # Remove any highlighting when completing
-            if task_id in self._highlighted_tasks:
-                del self._highlighted_tasks[task_id]
-
-            # Record completion time (only if not already completed and record_time is True)
-            if not self._tasks[task_id].completed and record_time:
-                self._tasks[task_id].completed_time = time.time()
-            self._tasks[task_id].completed = True
-
-            self._invalidate_cache()
-            self._update_display()
-
-            # Schedule hiding the task from display after a short delay
-            # but keep it in the data structure for summary generation
-            self._timer_manager.schedule_task_hiding(task_id)
-
-    def remove_sub_task(self, task_id: str, animate: bool = True) -> None:
-        """Remove a sub-task by ID with optional completion animation."""
-        if task_id in self._tasks:
-            # Remove any highlighting when removing
-            if task_id in self._highlighted_tasks:
-                del self._highlighted_tasks[task_id]
-
-            if animate:
-                self.complete_sub_task(task_id)
-            else:
-                del self._tasks[task_id]
-                self._invalidate_cache()
-                self._update_display()
-
-    def update_sub_status(self, sub_status: str):
-        """Legacy method for backward compatibility - creates/updates a default sub-task."""
-        self.add_sub_task(sub_status, "default")
-        self.update_sub_task("default", sub_status)
-
-    def update_main_status(self, message: str):
-        """Update the main status line with custom information."""
-        if self.description != message:  # Only update if changed
-            self.description = message
-            self._invalidate_cache()
-            self._update_display()
-
-    def update_messages(self, updater: dict[str, str]):
-        """Update both main message and sub-status if provided."""
-        if "message" in updater:
-            self.description = updater["message"]
-            self._invalidate_cache()
-            self._update_display()
-        if "sub_status" in updater:
-            self.update_sub_status(updater["sub_status"])
-        if "success_message" in updater:
-            self.success_message = updater["success_message"]
-        if "failure_message" in updater:
-            self.failure_message = updater["failure_message"]
-
-    def get_sub_task_count(self) -> int:
-        """Get the current number of active sub-tasks."""
-        return len(self._tasks)
-
-    def list_sub_tasks(self) -> list[str]:
-        """Get a list of all active sub-task IDs."""
-        return list(self._tasks.keys())
-
-    def get_task_status(self) -> str:
-        """Get a human-readable status of current task count vs limit."""
-        current_count = len(self._tasks)
-        return f"› Active tasks: {current_count}"
-
-    def get_task_duration(self, task_id: str) -> float:
-        """Get the duration of a specific task in seconds."""
-        if task_id in self._tasks:
-            return self._tasks[task_id].get_duration()
-        return 0.0
-
-
-    def get_completed_tasks(self) -> dict[str, TaskInfo]:
-        """Get all completed tasks with their timing information."""
-        return {task_id: task_info for task_id, task_info in self._tasks.items() if task_info.completed}
-
-    def get_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
-        """Get all tasks (completed and active) for a specific category."""
-        return {task_id: task_info for task_id, task_info in self._tasks.items() if task_info.category == category}
-
-    def get_completed_tasks_by_category(self, category: str) -> dict[str, TaskInfo]:
-        """Get all completed tasks for a specific category."""
-        return {task_id: task_info for task_id, task_info in self._tasks.items()
-                if task_info.category == category and task_info.completed}
-
-    def set_process_start_time(self) -> None:
-        """Set the overall process start time."""
-        self._process_start_time = time.time()
-
-    def set_process_end_time(self) -> None:
-        """Set the overall process end time."""
-        self._process_end_time = time.time()
-
-    def get_total_duration(self) -> float:
-        """Get the total duration from first task added to last task completed."""
-        if not self._tasks:
-            return 0.0
-
-        completed_tasks = self.get_completed_tasks()
-        if not completed_tasks:
-            return 0.0
-
-        # Find earliest start time and latest completion time
-        start_times = [task.added_time for task in self._tasks.values()]
-        completion_times = [task.completed_time for task in completed_tasks.values() if task.completed_time > 0]
-
-        if not start_times or not completion_times:
-            return 0.0
-
-        earliest_start = min(start_times)
-        latest_completion = max(completion_times)
-
-        return latest_completion - earliest_start
+    def _clear_all_tasks(self) -> None:
+        """Clear tasks and refresh display."""
+        super()._clear_all_tasks()
+        self._after_task_update()
 
     def generate_summary(self, categories: dict[str, str] | None = None) -> str:
-        """Generate a summary of completed tasks by category.
-
-        Args:
-            categories: Optional dict mapping category names to display names.
-                       Defaults to standard UseIndexPoller categories.
-
-        Returns:
-            Formatted summary string, or empty string if no meaningful tasks.
-        """
+        """Generate a summary of completed tasks by category."""
         if categories is None:
             categories = DEFAULT_SUMMARY_CATEGORIES
 
-        # Get completed tasks by category and calculate durations
-        category_durations = {}
+        category_durations: dict[str, float] = {}
         for category_name in categories:
             tasks = self.get_completed_tasks_by_category(category_name)
             category_durations[category_name] = _calculate_category_duration(category_name, tasks)
 
-        # If there's nothing meaningful to show, return empty string
         if not any(category_durations.values()):
             return ""
 
         total_duration = self.get_total_duration()
 
-        # Build Rich table directly from data (not from formatted strings)
         try:
             from rich.console import Console
             from rich.table import Table
@@ -1173,14 +1179,12 @@ class TaskProgress:
             table.add_column("Operation", style="white")
             table.add_column("Duration", style="green", justify="right")
 
-            # Add total duration row
             if total_duration > 0:
                 table.add_row(
                     INITIALIZATION_COMPLETED_TEXT,
                     format_duration(total_duration)
                 )
 
-            # Add category rows
             for category_name, display_name in categories.items():
                 duration = category_durations[category_name]
                 if duration > MIN_CATEGORY_DURATION_SECONDS:
@@ -1189,7 +1193,6 @@ class TaskProgress:
                         format_duration(duration)
                     )
 
-            # Add blank row for spacing
             table.add_row("", "")
 
             console = Console()
@@ -1198,8 +1201,7 @@ class TaskProgress:
             return capture.get()
 
         except ImportError:
-            # Fallback to simple text format
-            lines = []
+            lines: list[str] = []
             if total_duration > 0:
                 lines.append(f"{INITIALIZATION_COMPLETED_TEXT} {format_duration(total_duration)}")
 
@@ -1305,11 +1307,6 @@ class TaskProgress:
             print()
         self._cleanup()
 
-    def _clear_all_tasks(self):
-        """Clear all tasks and related data."""
-        self._tasks.clear()
-        self._highlighted_tasks.clear()
-
     def _cleanup(self):
         """Clean up resources."""
         if self.live:
@@ -1347,9 +1344,9 @@ leading_newline: bool = False, trailing_newline: bool = False, show_duration_sum
     Automatically detects if we're in a notebook environment (Snowflake, Jupyter, etc.)
     and returns the appropriate progress class.
     """
-    from ..environments import runtime_env, SnowbookEnvironment, JupyterEnvironment
+    from ..environments import runtime_env, SnowbookEnvironment, NotebookRuntimeEnvironment
 
-    if isinstance(runtime_env, (SnowbookEnvironment, JupyterEnvironment)):
+    if isinstance(runtime_env, (SnowbookEnvironment, NotebookRuntimeEnvironment)):
         # Use NotebookTaskProgress for Snowflake and Jupyter notebooks
         return NotebookTaskProgress(
             description=description,
@@ -1393,7 +1390,7 @@ class SubTaskContext:
         return False  # Don't suppress exceptions
 
 
-class NotebookTaskProgress:
+class NotebookTaskProgress(_TaskStateMixin):
     """A progress component specifically designed for notebook environments like Snowflake.
 
     This class copies the EXACT working Spinner code and adapts it for notebook use.
@@ -1413,15 +1410,8 @@ class NotebookTaskProgress:
         self.failure_message = failure_message
         self.leading_newline = leading_newline
         self.trailing_newline = trailing_newline
-        self.show_duration_summary = show_duration_summary
-
-        # Task management - unified data structure
-        self._tasks = {}  # task_id -> TaskInfo
-        self._next_task_id = 1
-
-        # Overall process timing
-        self._process_start_time = None
-        self._process_end_time = None
+        self._init_task_state(show_duration_summary=show_duration_summary)
+        self.enable_highlighting = False
 
         self.spinner_generator = itertools.cycle(SPINNER_FRAMES)
 
@@ -1438,15 +1428,10 @@ class NotebookTaskProgress:
         self._update_lock = threading.Lock()
 
         # Add sub-task support for TaskProgress compatibility
-        # Note: _tasks and _next_task_id already initialized above (lines 1393-1394)
-        self.main_completed = False
         self.spinner_thread = None
         self._current_subtask = ""
         self.busy = False  # Initialize busy state
 
-        # Timer manager for delayed operations
-        self._timer_manager = _TimerManager(self)
-
 
     def _generate_task_id(self) -> str:
         """Generate a unique task ID."""
@@ -1500,31 +1485,78 @@ class NotebookTaskProgress:
         """Update the display - notebook environments only."""
         # Use lock to prevent race conditions between spinner thread and main thread
         with self._update_lock:
-            if message is None:
-                message = self.get_message(starting=starting)
             if self.is_jupyter:
-                # @NOTE: IPython isn't available in CI. This won't ever get invoked w/out IPython available though.
-                from IPython.display import HTML, display # pyright: ignore[reportMissingImports]
-                content = HTML(f"<span style='font-family: monospace;'>{message}</span>")
+                _, display_fn = _load_ipython_display()
+
+                if message is None:
+                    lines = self._build_jupyter_lines(starting=starting)
+                elif message == "":
+                    lines = []
+                else:
+                    lines = [message]
+
+                rendered = "\n".join(lines)
+                content = {"text/plain": rendered}
                 if self.display is not None:
-                    self.display.update(content)
+                    self.display.update(content, raw=True)
                 else:
-                    self.display = display(content, display_id=True)
-            else:
-                # Use the EXACT same approach as the working Spinner code
-                rich_string = rich_str(message)
-                def width(word):
-                    return sum(wcwidth(c) for c in word)
-                diff = width(self.last_message) - width(rich_string)
+                    self.display = display_fn(content, display_id=True, raw=True)
+                return
+
+            if message is None:
+                message = self.get_message(starting=starting)
 
-                sys.stdout.write("\r")           # Move to beginning
-                sys.stdout.write(" " * DEFAULT_TERMINAL_WIDTH)  # Clear with spaces
-                sys.stdout.write("\r")           # Move back to beginning
+            rich_string = rich_str(message)
 
-                sys.stdout.write(message + (" " * diff))  # Write text directly
-                if self.in_notebook:
-                    sys.stdout.flush()                    # Force output
-                self.last_message = rich_string
+            def width(word):
+                return sum(wcwidth(c) for c in word)
+
+            diff = width(self.last_message) - width(rich_string)
+
+            sys.stdout.write("\r")           # Move to beginning
+            sys.stdout.write(" " * DEFAULT_TERMINAL_WIDTH)  # Clear with spaces
+            sys.stdout.write("\r")           # Move back to beginning
+
+            sys.stdout.write(message + (" " * diff))  # Write text directly
+            if self.in_notebook:
+                sys.stdout.flush()                    # Force output
+            self.last_message = rich_string
+
+    def _build_jupyter_lines(self, starting: bool) -> list[str]:
+        """Compose the main status and subtasks for Jupyter display."""
+        if self.busy or starting:
+            spinner = SPINNER_FRAMES[0] if starting else next(self.spinner_generator)
+            main_line = f"{spinner} {self.description}"
+        else:
+            main_text = self.success_message or self.description
+            main_line = f"{SUCCESS_ICON} {main_text}"
+
+        visible_tasks = self._collect_visible_tasks()
+        lines = [main_line]
+
+        for marker, tasks in (
+            (ARROW, visible_tasks["incomplete"]),
+            (CHECK_MARK, visible_tasks["completed"]),
+        ):
+            for task_info in tasks:
+                lines.append(f"   {marker} {task_info.description}")
+
+        return lines
+
+    def _collect_visible_tasks(self) -> dict[str, list["TaskInfo"]]:
+        """Separate visible tasks into incomplete and completed lists."""
+        incomplete: list["TaskInfo"] = []
+        completed: list["TaskInfo"] = []
+
+        for task_info in self._tasks.values():
+            if task_info.hidden:
+                continue
+            if task_info.completed:
+                completed.append(task_info)
+            else:
+                incomplete.append(task_info)
+
+        return {"incomplete": incomplete, "completed": completed}
 
     def reset_cursor(self):
         """Reset cursor to beginning of line - notebook environments only."""
@@ -1553,114 +1585,54 @@ class NotebookTaskProgress:
             # The spinner will now show the subtask instead of main task
 
     def add_sub_task(self, description: str, task_id: str | None = None, category: str = "general") -> str:
-        """Add a new sub-task and return its unique ID.
-
-        Args:
-            description: Description of the subtask
-            task_id: Optional custom task ID, if not provided one will be generated
-            category: Category for this task (e.g., "indexing", "provisioning", "change_tracking")
-
-        Returns:
-            str: The task ID for this subtask
-        """
-        if task_id is None:
-            task_id = self._generate_task_id()
-
-        if task_id not in self._tasks:
-            self._tasks[task_id] = TaskInfo(description=description, category=category)
-
-            # Show the subtask by updating the main task text
-            self._update_subtask_display(description)
-
+        task_id = super().add_sub_task(description, task_id, category)
+        # Update spinner display with the active subtask
+        if task_id in self._tasks:
+            self._update_subtask_display(self._tasks[task_id].description)
         return task_id
 
     def update_sub_task(self, task_id: str, description: str) -> None:
-        """Update an existing sub-task description."""
+        super().update_sub_task(task_id, description)
         if task_id in self._tasks:
-            self._tasks[task_id].description = description
-            # Show the updated subtask by updating the main task text
             self._update_subtask_display(description)
 
     def complete_sub_task(self, task_id: str, record_time: bool = True) -> None:
-        """Complete a sub-task by marking it as done."""
-        if task_id in self._tasks:
-            # Record completion time (only if not already completed and record_time is True)
-            if not self._tasks[task_id].completed and record_time:
-                self._tasks[task_id].completed_time = time.time()
-            self._tasks[task_id].completed = True
-
-            # Clear the subtask display when completed
-            self._current_subtask = ""
-            self._current_display = ""
-            # The spinner will now show the main task again
-
-            # Schedule hiding the task from display after a short delay
-            # but keep it in the data structure for summary generation
-            self._timer_manager.schedule_task_hiding(task_id)
+        super().complete_sub_task(task_id, record_time=record_time)
+        # Clear the subtask display when completed
+        self._current_subtask = ""
+        self._current_display = ""
 
     def remove_sub_task(self, task_id: str, animate: bool = True) -> None:
         """Remove a sub-task by ID."""
+        task_description: str | None = None
         if task_id in self._tasks:
-            # Store description before deletion
             task_description = self._tasks[task_id].description
-            del self._tasks[task_id]
-            # Clear subtask display if this was the current one
-            if hasattr(self, '_current_subtask') and self._current_subtask == task_description:
-                self._current_subtask = ""
-                self._current_display = ""
-                # The spinner will now show the main task again
+
+        # Notebook display should drop the subtask immediately to clear the UI.
+        super().remove_sub_task(task_id, animate=False)
+
+        if task_description and getattr(self, "_current_subtask", "") == task_description:
+            self._current_subtask = ""
+            self._current_display = ""
+            # The spinner will now show the main task again
 
     def update_sub_status(self, sub_status: str):
         """Legacy method for backward compatibility - creates/updates a default sub-task."""
-        self.add_sub_task(sub_status, "default")
-        self.update_sub_task("default", sub_status)
+        super().update_sub_status(sub_status)
 
     def update_main_status(self, message: str):
         """Update the main status line with real-time updates."""
-        self.description = message
-        # Clear any existing subtask when main status changes
+        super().update_main_status(message)
         self._current_subtask = ""
         self._current_display = ""
         # The spinner will now show the updated main task
 
     def update_messages(self, updater: dict[str, str]):
         """Update both main message and sub-status if provided."""
-        if "message" in updater:
-            self.update_main_status(updater["message"])
-        if "sub_status" in updater:
-            self.update_sub_status(updater["sub_status"])
-        if "success_message" in updater:
-            self.success_message = updater["success_message"]
-        if "failure_message" in updater:
-            self.failure_message = updater["failure_message"]
+        super().update_messages(updater)
 
-    def get_sub_task_count(self) -> int:
-        """Get the current number of active sub-tasks."""
-        return len(self._tasks)
-
-    def list_sub_tasks(self) -> list[str]:
-        """Get a list of all active sub-task IDs."""
-        return list(self._tasks.keys())
-
-    def get_task_status(self) -> str:
-        """Get a human-readable status of current task count vs limit."""
-        current_count = len(self._tasks)
-        return f"› Active tasks: {current_count}"
-
-    def _invalidate_cache(self):
-        """Invalidate the render cache to force re-rendering."""
-        # NotebookTaskProgress doesn't use caching, but we need this for API compatibility
-        pass
-
-    def _update_display(self):
-        """Update the display if live."""
-        # NotebookTaskProgress updates display through the update() method
-        # This is here for API compatibility with TaskProgress
-        pass
-
-    def _clear_all_tasks(self):
-        """Clear all tasks and related data."""
-        self._tasks.clear()
+    def _clear_all_tasks(self) -> None:
+        super()._clear_all_tasks()
         self._current_subtask = ""
         self._current_display = ""
 
@@ -1700,14 +1672,20 @@ class NotebookTaskProgress:
         # Clear the spinner line completely
         self._clear_spinner_line()
 
-        if self.success_message != "":
+        final_message: str | None = None
+        if self.success_message:
             final_message = f"{SUCCESS_ICON} {self.success_message}"
-            # For Jupyter, use update() to properly handle IPython display
-            # For Snowflake, use print() to get a new line
+        elif summary:
+            final_message = f"{SUCCESS_ICON} Done"
+
+        if final_message:
             if self.is_jupyter:
-                self.update(final_message)
+                if self.display is not None:
+                    self.display.update({"text/plain": final_message}, raw=True)
+                else:
+                    _, display_fn = _load_ipython_display()
+                    self.display = display_fn({"text/plain": final_message}, display_id=True, raw=True)
             else:
-                # Print the success message on a clean line
                 print(final_message)
         elif self.success_message == "":
             # When there's no success message, clear the display for notebooks
@@ -1718,9 +1696,13 @@ class NotebookTaskProgress:
 
         # Print summary if there are completed tasks
         if summary:
-            # For all notebook environments: display was cleared above, now print summary
-            print()
-            print(summary.strip())  # Summary includes visual separator line
+            if self.is_jupyter:
+                # Use IPython display to avoid blank stdout lines in notebooks
+                _, display_fn = _load_ipython_display()
+                display_fn({"text/plain": summary.strip()}, raw=True)
+            else:
+                print()
+                print(summary.strip())  # Summary includes visual separator line
 
         # Skip trailing newline for Jupyter - it interferes with IPython display
         if self.trailing_newline and not self.is_jupyter: