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

relationalai/tools/cli_controls.py

@@ -11,7 +11,7 @@ import threading
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, List, Sequence, TextIO, cast
+from typing import Any, Callable, Dict, List, Sequence, TextIO, cast
 
 # Third-party imports
 import rich
@@ -26,6 +26,7 @@ from wcwidth import wcwidth
 
 # Local imports
 from relationalai import debugging
+from relationalai.util.format import format_duration
 from ..environments import (
     HexEnvironment,
     JupyterEnvironment,
@@ -38,18 +39,64 @@ from ..environments import (
 # Constants
 #--------------------------------------------------
 
+# Display symbols
+ARROW = "➜"
+CHECK_MARK = "✓"
+SUCCESS_ICON = "✅"
+FAIL_ICON = "❌"
+
+# Spinner animation frames
+SPINNER_FRAMES = ["▰▱▱▱", "▰▰▱▱", "▰▰▰▱", "▰▰▰▰", "▱▰▰▰", "▱▱▰▰", "▱▱▱▰", "▱▱▱▱"]
+
+# Terminal display constants
+DEFAULT_TERMINAL_WIDTH = 80
+SEPARATOR_WIDTH = 40
+
+# Task progress constants
+INITIALIZATION_COMPLETED_TEXT = "Parallel init finished in"
+MIN_CATEGORY_DURATION_SECONDS = 0.25  # Only show categories with duration > 250ms
+
+# Task category constants
+TASK_CATEGORY_INDEXING = "indexing"
+TASK_CATEGORY_PROVISIONING = "provisioning"
+TASK_CATEGORY_CHANGE_TRACKING = "change_tracking"
+TASK_CATEGORY_CACHE = "cache"
+TASK_CATEGORY_RELATIONS = "relations"
+TASK_CATEGORY_STATUS = "status"
+TASK_CATEGORY_VALIDATION = "validation"
+TASK_CATEGORY_OTHER = "other"
+
+# Default summary categories
+DEFAULT_SUMMARY_CATEGORIES = {
+    TASK_CATEGORY_INDEXING: "Indexing",
+    TASK_CATEGORY_PROVISIONING: "Provisioning",
+    TASK_CATEGORY_CHANGE_TRACKING: "Change tracking",
+    TASK_CATEGORY_RELATIONS: "Relations",
+    TASK_CATEGORY_STATUS: "Status",
+    TASK_CATEGORY_VALIDATION: "Validation",
+    TASK_CATEGORY_OTHER: "Other"
+}
+
+# Parallel task categories (for duration calculation)
+PARALLEL_TASK_CATEGORIES = {
+    TASK_CATEGORY_INDEXING,
+    TASK_CATEGORY_PROVISIONING,
+    TASK_CATEGORY_VALIDATION,
+    TASK_CATEGORY_CHANGE_TRACKING
+}
+
+# Prompt constants
 REFETCH = "[REFETCH LIST]"
 MANUAL_ENTRY = "[MANUAL ENTRY]"
 
-# TaskProgress timing constants
+# Timing constants
 HIGHLIGHT_DURATION = 2.0
 COMPLETION_DISPLAY_DURATION = 8.0
 TIMER_CHECK_INTERVAL = 0.1
 SPINNER_UPDATE_INTERVAL = 0.15
-
-# Display symbols
-SUCCESS_ICON = "✅"
-FAIL_ICON = "❌"
+INITIAL_DISPLAY_DELAY = 0.25
+BRIEF_PAUSE = 0.1
+LIVE_REFRESH_RATE = 10
 
 #--------------------------------------------------
 # Style
@@ -159,7 +206,7 @@ def _enumerate_choices(choices: inquirer_utils.InquirerPyListChoices) -> inquire
     else:
         return _enumerate_static_choices(choices)
 
-def _fuzzy(message:str, choices:inquirer_utils.InquirerPyListChoices, default:str|None = None, multiselect=False, show_index=False, **kwargs) -> str|list[str]:
+def _fuzzy(message:str, choices:inquirer_utils.InquirerPyListChoices, default:str|None = None, multiselect=False, show_index=False, **kwargs) -> str|list[str]|None:
     if show_index:
         choices = _enumerate_choices(choices)
 
@@ -452,7 +499,7 @@ class Spinner(LineClearingMixin):
         self.message = message
         self.finished_message = finished_message
         self.failed_message = failed_message
-        self.spinner_generator = itertools.cycle(["▰▱▱▱", "▰▰▱▱", "▰▰▰▱", "▰▰▰▰", "▱▰▰▰", "▱▱▰▰", "▱▱▱▰", "▱▱▱▱"])
+        self.spinner_generator = itertools.cycle(SPINNER_FRAMES)
         self.is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
         self.is_hex = isinstance(runtime_env, HexEnvironment)
         self.is_jupyter = isinstance(runtime_env, JupyterEnvironment)
@@ -565,12 +612,12 @@ class Spinner(LineClearingMixin):
         self.update(color="magenta", starting=True)
         # return control to the event loop briefly so stdout can be sure to flush:
         if self.delay:
-            time.sleep(0.25)
+            time.sleep(INITIAL_DISPLAY_DELAY)
         self.reset_cursor()
         if not self.delay:
             return self
         self.busy = True
-        threading.Thread(target=self.spinner_task).start()
+        threading.Thread(target=self.spinner_task, daemon=True).start()
         return self
 
     def __exit__(self, exception, value, _):
@@ -637,13 +684,23 @@ class SpanSpinner(Spinner):
 class TaskInfo:
     """Represents a single task with its state and metadata."""
     description: str
+    category: str = "other"
     completed: bool = False
     added_time: float = 0.0
+    completed_time: float = 0.0
+    hidden: bool = False
 
     def __post_init__(self):
         if self.added_time == 0.0:
             self.added_time = time.time()
 
+    def get_duration(self) -> float:
+        """Get the duration of this task in seconds."""
+        if not self.completed or self.completed_time == 0.0:
+            return 0.0
+
+        return self.completed_time - self.added_time
+
 
 class _TimerManager:
     """Manages all delayed operations for TaskProgress."""
@@ -670,6 +727,15 @@ class _TimerManager:
         self._operations[task_id] = ("delayed_removal", scheduled_time)
         self._start()
 
+    def schedule_task_hiding(self, task_id: str, delay: float | None = None):
+        """Schedule hiding of a completed task from display (but keep in data structure)."""
+        if delay is None:
+            delay = COMPLETION_DISPLAY_DURATION
+        scheduled_time = time.time() + delay
+        self._operations[task_id] = ("delayed_hiding", scheduled_time)
+        self._start()
+
+
     def _start(self):
         """Start the timer thread if not already running."""
         if self._thread is None or not self._thread.is_alive():
@@ -717,6 +783,14 @@ class _TimerManager:
             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
+        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()
 
     def stop(self):
         """Stop the timer manager."""
@@ -736,11 +810,6 @@ class TaskProgress:
     - Consistent task ordering with active tasks displayed above completed ones
     """
 
-    # Display symbols
-    SPINNER_FRAMES = ["▰▱▱▱", "▰▰▱▱", "▰▰▰▱", "▰▰▰▰", "▱▰▰▰", "▱▱▰▰", "▱▱▱▰", "▱▱▱▱"]
-    ARROW = "➜"
-    CHECK_MARK = "✓"
-
     def __init__(
         self,
         description: str = "",
@@ -750,6 +819,7 @@ class TaskProgress:
         trailing_newline: bool = False,
         transient: bool = False,
         hide_on_completion: bool = False,
+        show_duration_summary: bool = True,
     ):
         # Public configuration
         self.description = description
@@ -759,6 +829,7 @@ class TaskProgress:
         self.trailing_newline = trailing_newline
         self.transient = transient
         self.hide_on_completion = hide_on_completion
+        self.show_duration_summary = show_duration_summary
 
         # Detect CI environment to avoid cursor control issues
         from ..environments import CIEnvironment
@@ -775,6 +846,10 @@ class TaskProgress:
         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
 
@@ -839,7 +914,7 @@ class TaskProgress:
         elif self.main_completed:
             main_line = Text(f"{SUCCESS_ICON} ", style="green") + Text(self.description, style="green")
         else:
-            spinner_text = self.SPINNER_FRAMES[self.spinner_index]
+            spinner_text = SPINNER_FRAMES[self.spinner_index]
             main_line = Text(f"{spinner_text} ", style="magenta") + Text(self.description, style="magenta")
 
         # Build subtask lines
@@ -866,6 +941,9 @@ class TaskProgress:
         completed_tasks = []
 
         for task_id, task_info in self._tasks.items():
+            # Skip hidden tasks
+            if task_info.hidden:
+                continue
             if task_info.completed:
                 completed_tasks.append((task_id, task_info))
             else:
@@ -873,23 +951,23 @@ class TaskProgress:
 
         # Render incomplete tasks first
         for task_id, task_info in incomplete_tasks:
-            is_highlighted = (task_id in self._highlighted_tasks and 
+            is_highlighted = (task_id in self._highlighted_tasks and
                             current_time < self._highlighted_tasks[task_id])
 
             style = "yellow" if is_highlighted else "white"
-            line = Text(f"   {self.ARROW} ", style=style) + Text(task_info.description, style=style)
+            line = Text(f"   {ARROW} ", style=style) + Text(task_info.description, style=style)
             subtask_lines.append(line)
 
         # Render completed tasks
         for task_id, task_info in completed_tasks:
-            line = Text(f"   {self.CHECK_MARK} ", style="green") + Text(task_info.description, style="green")
+            line = Text(f"   {CHECK_MARK} ", style="green") + Text(task_info.description, style="green")
             subtask_lines.append(line)
 
         return subtask_lines
 
     def _advance_spinner(self):
         """Advance the spinner animation."""
-        self.spinner_index = (self.spinner_index + 1) % len(self.SPINNER_FRAMES)
+        self.spinner_index = (self.spinner_index + 1) % len(SPINNER_FRAMES)
 
     def _invalidate_cache(self):
         """Invalidate the render cache to force re-rendering."""
@@ -901,12 +979,13 @@ class TaskProgress:
         if self.live:
             self.live.update(self._render_display())
 
-    def add_sub_task(self, description: str, task_id: str | None = None) -> str:
+    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
@@ -915,7 +994,7 @@ class TaskProgress:
             task_id = self._generate_task_id()
 
         if task_id not in self._tasks:
-            self._tasks[task_id] = TaskInfo(description=description)
+            self._tasks[task_id] = TaskInfo(description=description, category=category)
             self._invalidate_cache()
             self._update_display()
 
@@ -940,19 +1019,24 @@ class TaskProgress:
             self._invalidate_cache()
             self._update_display()
 
-    def complete_sub_task(self, task_id: str) -> None:
+    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 removal after completion display duration
-            self._timer_manager.schedule_task_removal(task_id)
+            # 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."""
@@ -1006,13 +1090,136 @@ class TaskProgress:
         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 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.
+        """
+        if categories is None:
+            categories = DEFAULT_SUMMARY_CATEGORIES
+
+        # Get completed tasks by category and calculate durations
+        category_durations = {}
+        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
+
+            table = Table(show_header=False, box=None, padding=(0, 1))
+            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:
+                    table.add_row(
+                        f" {ARROW} {display_name}",
+                        format_duration(duration)
+                    )
+
+            # Add blank row for spacing
+            table.add_row("", "")
+
+            console = Console()
+            with console.capture() as capture:
+                console.print(table)
+            return capture.get()
+
+        except ImportError:
+            # Fallback to simple text format
+            lines = []
+            if total_duration > 0:
+                lines.append(f"{INITIALIZATION_COMPLETED_TEXT} {format_duration(total_duration)}")
+
+            for category_name, display_name in categories.items():
+                duration = category_durations[category_name]
+                if duration > MIN_CATEGORY_DURATION_SECONDS:
+                    lines.append(f" {ARROW} {display_name} {format_duration(duration)}")
+
+            if lines:
+                lines.append("")
+
+            return "\n".join(lines)
+
     def __enter__(self):
         if self.leading_newline:
             print()
 
         # Start the live display
         from rich.live import Live
-        self.live = Live(self._render_display(), console=self.console, refresh_per_second=10)
+        self.live = Live(self._render_display(), console=self.console, refresh_per_second=LIVE_REFRESH_RATE)
         self.live.start()
 
         # Start spinner animation
@@ -1051,19 +1258,19 @@ class TaskProgress:
         # Clear all tasks and update main task to show failure state
         self._clear_all_tasks()
         self.main_failed = True
-        
+
         # Update main task description to show failure message
         if self.failure_message:
             self.description = self.failure_message
         else:
             self.description = f"Failed: {exc_val}"
-        
+
         # Update the display to show the failure state before stopping
         if self.live:
             self.live.update(self._render_display())
             # Brief pause to show the failure state
-            time.sleep(0.1)
-        
+            time.sleep(BRIEF_PAUSE)
+
         if self.trailing_newline:
             print()
         self._cleanup()
@@ -1071,6 +1278,11 @@ class TaskProgress:
     def _handle_success(self):
         """Handle success case in context manager exit."""
         self.main_completed = True
+
+        # Generate summary before clearing tasks (so we have the timing data)
+        # Only generate if show_duration_summary flag is True
+        summary = self.generate_summary() if self.show_duration_summary else ""
+
         self._clear_all_tasks()
 
         # Update main task description to show success message
@@ -1083,6 +1295,12 @@ class TaskProgress:
             # Stop the live display
             self.live.stop()
 
+        # Print summary if available
+        if summary:
+            print()  # Blank line for separation
+            print(summary, end="")  # summary already has trailing newline
+            print()  # Add extra blank line after summary
+
         if self.trailing_newline:
             print()
         self._cleanup()
@@ -1101,33 +1319,55 @@ class TaskProgress:
             if not self.is_ci and sys.stdout.isatty():
                 print("\r\033[K", end="", flush=True)
 
+def _calculate_category_duration(category_name: str, tasks: Dict[str, TaskInfo]) -> float:
+    """Calculate duration for a category based on task type (parallel vs sequential)."""
+    if not tasks:
+        return 0.0
+
+    if category_name in PARALLEL_TASK_CATEGORIES:
+        # For parallel tasks, use time span (max completion - min start)
+        category_start_times = [task_info.added_time for task_info in tasks.values()]
+        category_completion_times = [
+            task_info.completed_time for task_info in tasks.values()
+            if task_info.completed_time > 0
+        ]
+        if category_start_times and category_completion_times:
+            return max(category_completion_times) - min(category_start_times)
+        else:
+            return 0.0
+    else:
+        # For sequential tasks, sum individual durations
+        return sum(task_info.get_duration() for task_info in tasks.values())
+
 
 def create_progress(description: str = "", success_message: str = "", failure_message: str = "",
-leading_newline: bool = False, trailing_newline: bool = False):
+leading_newline: bool = False, trailing_newline: bool = False, show_duration_summary: bool = True):
     """Factory function to create the appropriate progress component based on environment.
 
-    Automatically detects if we're in a Snowflake notebook or similar environment
+    Automatically detects if we're in a notebook environment (Snowflake, Jupyter, etc.)
     and returns the appropriate progress class.
     """
-    from ..environments import runtime_env, SnowbookEnvironment
+    from ..environments import runtime_env, SnowbookEnvironment, JupyterEnvironment
 
-    if isinstance(runtime_env, SnowbookEnvironment):
-        # Use NotebookTaskProgress for Snowflake notebooks
+    if isinstance(runtime_env, (SnowbookEnvironment, JupyterEnvironment)):
+        # Use NotebookTaskProgress for Snowflake and Jupyter notebooks
         return NotebookTaskProgress(
             description=description,
             success_message=success_message,
             failure_message=failure_message,
             leading_newline=leading_newline,
-            trailing_newline=trailing_newline
+            trailing_newline=trailing_newline,
+            show_duration_summary=show_duration_summary
         )
     else:
-        # Use TaskProgress for other environments
+        # Use TaskProgress for other environments (terminal, CI, etc.)
         return TaskProgress(
             description=description,
             success_message=success_message,
             failure_message=failure_message,
             leading_newline=leading_newline,
-            trailing_newline=trailing_newline
+            trailing_newline=trailing_newline,
+            show_duration_summary=show_duration_summary
         )
 
 
@@ -1166,14 +1406,24 @@ class NotebookTaskProgress:
         failure_message: str = "",
         leading_newline: bool = False,
         trailing_newline: bool = False,
+        show_duration_summary: bool = True,
     ):
         self.description = description
         self.success_message = success_message
         self.failure_message = failure_message
         self.leading_newline = leading_newline
         self.trailing_newline = trailing_newline
+        self.show_duration_summary = show_duration_summary
 
-        self.spinner_generator = itertools.cycle(["▰▱▱▱", "▰▰▱▱", "▰▰▰▱", "▰▰▰▰", "▱▰▰▰", "▱▱▰▰", "▱▱▱▰", "▱▱▱▱"])
+        # 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.spinner_generator = itertools.cycle(SPINNER_FRAMES)
 
         # Environment detection for notebook environments only
         self.is_snowflake_notebook = isinstance(runtime_env, SnowbookEnvironment)
@@ -1188,11 +1438,15 @@ class NotebookTaskProgress:
         self._update_lock = threading.Lock()
 
         # Add sub-task support for TaskProgress compatibility
-        self._tasks = {}  # Use same data structure as TaskProgress
-        self._next_task_id = 1
+        # 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."""
@@ -1219,7 +1473,7 @@ class NotebookTaskProgress:
     def get_message(self, starting=False):
         """Get the current message with spinner - notebook environments only."""
         # For notebook environments, use a reasonable default width
-        max_width = 80  # Default width for notebooks
+        max_width = DEFAULT_TERMINAL_WIDTH
         try:
             max_width = shutil.get_terminal_size().columns
         except (OSError, AttributeError):
@@ -1264,7 +1518,7 @@ class NotebookTaskProgress:
                 diff = width(self.last_message) - width(rich_string)
 
                 sys.stdout.write("\r")           # Move to beginning
-                sys.stdout.write(" " * 80)      # Clear with spaces (same as Spinner)
+                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
@@ -1298,12 +1552,13 @@ class NotebookTaskProgress:
             self._current_display = subtask_text
             # The spinner will now show the subtask instead of main task
 
-    def add_sub_task(self, description: str, task_id: str | None = None) -> str:
+    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
@@ -1312,7 +1567,7 @@ class NotebookTaskProgress:
             task_id = self._generate_task_id()
 
         if task_id not in self._tasks:
-            self._tasks[task_id] = TaskInfo(description=description)
+            self._tasks[task_id] = TaskInfo(description=description, category=category)
 
             # Show the subtask by updating the main task text
             self._update_subtask_display(description)
@@ -1326,9 +1581,12 @@ class NotebookTaskProgress:
             # Show the updated subtask by updating the main task text
             self._update_subtask_display(description)
 
-    def complete_sub_task(self, task_id: str) -> None:
+    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
@@ -1336,8 +1594,9 @@ class NotebookTaskProgress:
             self._current_display = ""
             # The spinner will now show the main task again
 
-            # Remove completed task immediately (no delay needed in notebooks)
-            del self._tasks[task_id]
+            # 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."""
@@ -1406,17 +1665,18 @@ class NotebookTaskProgress:
         self._current_display = ""
 
     def __enter__(self):
-        if self.leading_newline:
+        # Skip leading newline for Jupyter - it interferes with IPython display
+        if self.leading_newline and not self.is_jupyter:
             rich.print()
         self.update(starting=True)
         # return control to the event loop briefly so stdout can be sure to flush:
         if self.delay:
-            time.sleep(0.25)
+            time.sleep(INITIAL_DISPLAY_DELAY)
         self.reset_cursor()
         if not self.delay:
             return self
         self.busy = True
-        threading.Thread(target=self.spinner_task).start()
+        threading.Thread(target=self.spinner_task, daemon=True).start()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -1424,21 +1684,146 @@ class NotebookTaskProgress:
         if exc_type is not None:
             if self.failure_message is not None:
                 self.update(f"{self.failure_message} {exc_val}", file=sys.stderr)
-                # Use rich.print with explicit newline to ensure proper formatting
-                rich.print(file=sys.stderr)
+                # For non-Jupyter, add newline to ensure proper formatting
+                # For Jupyter, IPython display handles formatting
+                if not self.is_jupyter:
+                    rich.print(file=sys.stderr)
                 return True
             return False
         if self.delay: # will be None for non-interactive environments
             time.sleep(self.delay)
-        self.reset_cursor()
+
+        # Generate summary BEFORE clearing the spinner line (so we have timing data)
+        # Only generate if show_duration_summary flag is True
+        summary = self.generate_summary() if self.show_duration_summary else ""
+
+        # Clear the spinner line completely
+        self._clear_spinner_line()
+
         if self.success_message != "":
             final_message = f"{SUCCESS_ICON} {self.success_message}"
-            self.update(final_message)
-            # Use rich.print with explicit newline to ensure proper formatting
-            rich.print()
+            # For Jupyter, use update() to properly handle IPython display
+            # For Snowflake, use print() to get a new line
+            if self.is_jupyter:
+                self.update(final_message)
+            else:
+                # Print the success message on a clean line
+                print(final_message)
         elif self.success_message == "":
-            self.update("")
-            self.reset_cursor()
-        if self.trailing_newline:
+            # When there's no success message, clear the display for notebooks
+            # The summary will be printed below if available
+            if self.is_jupyter:
+                self.update("")
+            # For non-Jupyter notebooks, _clear_spinner_line() already handled it
+
+        # 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
+
+        # Skip trailing newline for Jupyter - it interferes with IPython display
+        if self.trailing_newline and not self.is_jupyter:
             rich.print()
         return True
+
+    def _clear_spinner_line(self):
+        """Clear the current spinner line completely."""
+        # Skip clearing for Jupyter notebooks - IPython display handles it
+        if self.is_jupyter:
+            return
+
+        # Write enough spaces to clear any content, then move to start of line
+        terminal_width = DEFAULT_TERMINAL_WIDTH
+        try:
+            terminal_width = shutil.get_terminal_size().columns
+        except (OSError, AttributeError):
+            pass
+
+        # Clear with spaces, carriage return, and newline to ensure we're on a fresh line
+        sys.stdout.write("\r" + " " * terminal_width + "\r\n")
+        sys.stdout.flush()
+
+    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 generate_summary(self, categories: dict[str, str] | None = None) -> str:
+        """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 = {}
+        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 ""
+
+        # Generate summary lines with proper alignment
+        summary_lines = []
+        label_width = 30  # Width for category labels
+        time_width = 10   # Width for time column (right-aligned)
+
+        # Add total time FIRST (at the top) - align with arrow lines
+        total_duration = self.get_total_duration()
+        if total_duration > 0:
+            formatted_total = format_duration(total_duration)
+            # Use the same format as arrow lines but with a different prefix
+            # This ensures perfect alignment with the time column
+            summary_lines.append(f" {INITIALIZATION_COMPLETED_TEXT:<{label_width-1}} {formatted_total:>{time_width}}")
+
+        # Add category breakdown
+        category_lines = []
+        for category_name, display_name in categories.items():
+            duration = category_durations[category_name]
+            if duration > MIN_CATEGORY_DURATION_SECONDS:  # Only show significant durations
+                formatted_duration = format_duration(duration)
+                # Use arrow for visual consistency with right-aligned time
+                category_lines.append(f" {ARROW} {display_name:<{label_width-4}} {formatted_duration:>{time_width}}")
+
+        # Only add category lines if there are any
+        if category_lines:
+            summary_lines.extend(category_lines)
+
+        # Add a visual separator line for Snowflake notebook environment
+        summary_lines.append("─" * SEPARATOR_WIDTH)
+
+        return "\n".join(summary_lines) + "\n"
+
+    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_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}