mcp-hangar 0.2.0__py3-none-any.whl

mcp_hangar/progress.py

@@ -0,0 +1,542 @@
+"""Real-time operation feedback with progress events.
+
+This module provides streaming progress updates for long-running
+operations, giving users visibility into:
+
+- Cold starts and provider launches
+- Container initialization
+- Tool discovery
+- Network calls and execution
+
+Usage with callback::
+
+    from mcp_hangar import ProgressTracker, ProgressStage
+
+    def on_progress(stage: str, message: str, elapsed_ms: float):
+        print(f"⏳ [{stage}] {message} ({elapsed_ms:.0f}ms)")
+
+    tracker = ProgressTracker(callback=on_progress)
+    tracker.report(ProgressStage.LAUNCHING, "Starting...")
+    tracker.complete(result)
+
+See docs/guides/UX_IMPROVEMENTS.md for more examples.
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from enum import Enum
+import threading
+import time
+from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Optional
+
+from .logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class ProgressStage(str, Enum):
+    """Stages of operation progress."""
+
+    # Pre-execution stages
+    COLD_START = "cold_start"
+    LAUNCHING = "launching"
+    INITIALIZING = "initializing"
+    DISCOVERING_TOOLS = "discovering_tools"
+    CONNECTING = "connecting"
+
+    # Execution stages
+    READY = "ready"
+    EXECUTING = "executing"
+    PROCESSING = "processing"
+
+    # Completion stages
+    COMPLETE = "complete"
+    FAILED = "failed"
+    RETRYING = "retrying"
+
+
+class EventType(str, Enum):
+    """Types of events in the stream."""
+
+    PROGRESS = "progress"
+    RESULT = "result"
+    ERROR = "error"
+
+
+@dataclass
+class ProgressEvent:
+    """A single progress update event.
+
+    Attributes:
+        type: Event type (progress, result, error)
+        stage: Current operation stage
+        message: Human-readable description
+        elapsed_ms: Time since operation started
+        details: Additional context data
+        timestamp: Event timestamp
+    """
+
+    type: EventType = EventType.PROGRESS
+    stage: ProgressStage = ProgressStage.EXECUTING
+    message: str = ""
+    elapsed_ms: float = 0.0
+    details: Dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=time.time)
+
+    # For result/error events
+    data: Any = None
+    exception: Optional[Exception] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        result: Dict[str, Any] = {
+            "type": self.type.value,
+            "stage": self.stage.value,
+            "message": self.message,
+            "elapsed_ms": round(self.elapsed_ms, 2),
+            "timestamp": self.timestamp,
+        }
+        if self.details:
+            result["details"] = self.details
+        if self.data is not None:
+            result["data"] = self.data
+        if self.exception is not None:
+            result["error"] = str(self.exception)
+        return result
+
+
+# Type alias for progress callback
+ProgressCallback = Callable[[str, str, float], None]
+
+
+class ProgressTracker:
+    """Tracks and reports progress for a single operation.
+
+    Can be used with callbacks or as a streaming iterator.
+
+    Usage with callback:
+        tracker = ProgressTracker(callback=my_callback)
+        tracker.report(ProgressStage.COLD_START, "Launching container...")
+        tracker.report(ProgressStage.READY, "Provider ready")
+        tracker.complete(result)
+
+    Usage as iterator:
+        tracker = ProgressTracker()
+        # In one thread/task:
+        tracker.report(ProgressStage.LAUNCHING, "Starting...")
+        tracker.complete(result)
+        # In another:
+        for event in tracker:
+            logger.info("progress_event", event=event.to_dict())
+    """
+
+    def __init__(
+        self,
+        callback: Optional[ProgressCallback] = None,
+        provider: str = "",
+        operation: str = "",
+    ):
+        """Initialize progress tracker.
+
+        Args:
+            callback: Optional callback for progress updates
+            provider: Provider name for context
+            operation: Operation name for context
+        """
+        self._callback = callback
+        self._provider = provider
+        self._operation = operation
+        self._start_time = time.time()
+        self._events: List[ProgressEvent] = []
+        self._completed = False
+        self._result: Any = None
+        self._error: Optional[Exception] = None
+
+        # For iterator support
+        self._event_queue: asyncio.Queue = asyncio.Queue()
+        self._sync_queue: List[ProgressEvent] = []
+        self._lock = threading.Lock()
+        self._done = threading.Event()
+
+    @property
+    def elapsed_ms(self) -> float:
+        """Time elapsed since operation started."""
+        return (time.time() - self._start_time) * 1000
+
+    def report(
+        self,
+        stage: ProgressStage,
+        message: str,
+        details: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Report a progress update.
+
+        Args:
+            stage: Current progress stage
+            message: Human-readable description
+            details: Additional context
+        """
+        elapsed = self.elapsed_ms
+        event = ProgressEvent(
+            type=EventType.PROGRESS,
+            stage=stage,
+            message=message,
+            elapsed_ms=elapsed,
+            details=details or {},
+        )
+
+        with self._lock:
+            self._events.append(event)
+            self._sync_queue.append(event)
+
+        # Try to put in async queue (non-blocking)
+        try:
+            self._event_queue.put_nowait(event)
+        except asyncio.QueueFull:
+            pass
+
+        # Invoke callback if set
+        if self._callback:
+            try:
+                self._callback(stage.value, message, elapsed)
+            except Exception as e:
+                logger.debug("progress_callback_error", error=str(e))
+
+        logger.debug(
+            "progress_reported",
+            provider=self._provider,
+            operation=self._operation,
+            stage=stage.value,
+            message=message,
+            elapsed_ms=elapsed,
+        )
+
+    def complete(self, result: Any) -> None:
+        """Mark operation as complete with result.
+
+        Args:
+            result: The operation result
+        """
+        elapsed = self.elapsed_ms
+        event = ProgressEvent(
+            type=EventType.RESULT,
+            stage=ProgressStage.COMPLETE,
+            message=f"Operation completed (total: {elapsed:.0f}ms)",
+            elapsed_ms=elapsed,
+            data=result,
+        )
+
+        with self._lock:
+            self._events.append(event)
+            self._sync_queue.append(event)
+            self._completed = True
+            self._result = result
+
+        try:
+            self._event_queue.put_nowait(event)
+        except asyncio.QueueFull:
+            pass
+
+        self._done.set()
+
+        if self._callback:
+            try:
+                self._callback(ProgressStage.COMPLETE.value, event.message, elapsed)
+            except Exception:
+                pass
+
+        logger.debug(
+            "progress_complete",
+            provider=self._provider,
+            operation=self._operation,
+            elapsed_ms=elapsed,
+        )
+
+    def fail(self, error: Exception) -> None:
+        """Mark operation as failed with error.
+
+        Args:
+            error: The exception that caused failure
+        """
+        elapsed = self.elapsed_ms
+        event = ProgressEvent(
+            type=EventType.ERROR,
+            stage=ProgressStage.FAILED,
+            message=f"Operation failed: {str(error)[:100]}",
+            elapsed_ms=elapsed,
+            exception=error,
+        )
+
+        with self._lock:
+            self._events.append(event)
+            self._sync_queue.append(event)
+            self._completed = True
+            self._error = error
+
+        try:
+            self._event_queue.put_nowait(event)
+        except asyncio.QueueFull:
+            pass
+
+        self._done.set()
+
+        if self._callback:
+            try:
+                self._callback(ProgressStage.FAILED.value, event.message, elapsed)
+            except (TypeError, ValueError, RuntimeError) as e:
+                logger.debug("progress_callback_error", stage="failed", error=str(e))
+
+        logger.debug(
+            "progress_failed",
+            provider=self._provider,
+            operation=self._operation,
+            error=str(error)[:200],
+            elapsed_ms=elapsed,
+        )
+
+    def __iter__(self) -> Iterator[ProgressEvent]:
+        """Iterate over events synchronously."""
+        while True:
+            with self._lock:
+                if self._sync_queue:
+                    event = self._sync_queue.pop(0)
+                    yield event
+                    if event.type in (EventType.RESULT, EventType.ERROR):
+                        return
+                elif self._completed:
+                    return
+
+            # Wait a bit before checking again
+            if not self._done.wait(timeout=0.01):
+                continue
+
+    async def __aiter__(self) -> AsyncIterator[ProgressEvent]:
+        """Iterate over events asynchronously."""
+        while True:
+            try:
+                event = await asyncio.wait_for(
+                    self._event_queue.get(),
+                    timeout=0.1,
+                )
+                yield event
+                if event.type in (EventType.RESULT, EventType.ERROR):
+                    return
+            except asyncio.TimeoutError:
+                if self._completed:
+                    return
+
+    def get_all_events(self) -> List[ProgressEvent]:
+        """Get all recorded events."""
+        with self._lock:
+            return list(self._events)
+
+
+# =============================================================================
+# Progress-Aware Operation Wrapper
+# =============================================================================
+
+
+class ProgressOperation:
+    """Context manager for progress-tracked operations.
+
+    Usage:
+        async with ProgressOperation("math", "add", callback=my_cb) as progress:
+            progress.report(ProgressStage.LAUNCHING, "Starting provider...")
+            result = await do_work()
+            progress.complete(result)
+    """
+
+    def __init__(
+        self,
+        provider: str,
+        operation: str,
+        callback: Optional[ProgressCallback] = None,
+    ):
+        self.provider = provider
+        self.operation = operation
+        self.tracker = ProgressTracker(
+            callback=callback,
+            provider=provider,
+            operation=operation,
+        )
+
+    def __enter__(self) -> ProgressTracker:
+        return self.tracker
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_val is not None and not self.tracker._completed:
+            self.tracker.fail(exc_val)
+        return False
+
+    async def __aenter__(self) -> ProgressTracker:
+        return self.tracker
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if exc_val is not None and not self.tracker._completed:
+            self.tracker.fail(exc_val)
+        return False
+
+
+# =============================================================================
+# Standard Progress Messages
+# =============================================================================
+
+# Pre-defined messages for common operations
+PROGRESS_MESSAGES = {
+    ProgressStage.COLD_START: "Provider is cold, launching...",
+    ProgressStage.LAUNCHING: "Starting {mode} provider...",
+    ProgressStage.INITIALIZING: "Container started (PID: {pid}), initializing...",
+    ProgressStage.DISCOVERING_TOOLS: "Discovering available tools...",
+    ProgressStage.CONNECTING: "Connecting to provider...",
+    ProgressStage.READY: "Provider ready, executing request...",
+    ProgressStage.EXECUTING: "Calling tool '{tool}'...",
+    ProgressStage.PROCESSING: "Processing response...",
+    ProgressStage.COMPLETE: "Operation completed (total: {elapsed_ms:.0f}ms)",
+    ProgressStage.FAILED: "Operation failed: {error}",
+    ProgressStage.RETRYING: "Retrying (attempt {attempt}/{max_attempts})...",
+}
+
+
+def get_stage_message(
+    stage: ProgressStage,
+    **kwargs,
+) -> str:
+    """Get formatted message for a progress stage.
+
+    Args:
+        stage: The progress stage
+        **kwargs: Format arguments for message template
+
+    Returns:
+        Formatted progress message
+    """
+    template = PROGRESS_MESSAGES.get(stage, stage.value)
+    try:
+        return template.format(**kwargs)
+    except KeyError:
+        return template
+
+
+# =============================================================================
+# Event Bus Integration
+# =============================================================================
+
+
+class ProgressEventHandler:
+    """Event handler that forwards domain events to progress tracker.
+
+    This bridges the gap between domain events and progress reporting,
+    allowing existing code to emit progress updates without modification.
+    """
+
+    # Map domain event types to progress stages
+    EVENT_STAGE_MAP = {
+        "ProviderStarted": ProgressStage.READY,
+        "ProviderStopped": ProgressStage.COMPLETE,
+        "ProviderStateChanged": ProgressStage.INITIALIZING,
+        "ToolInvocationRequested": ProgressStage.EXECUTING,
+        "ToolInvocationCompleted": ProgressStage.COMPLETE,
+        "ToolInvocationFailed": ProgressStage.FAILED,
+    }
+
+    def __init__(self):
+        self._trackers: Dict[str, ProgressTracker] = {}
+        self._lock = threading.Lock()
+
+    def register_tracker(self, correlation_id: str, tracker: ProgressTracker) -> None:
+        """Register a tracker for a correlation ID."""
+        with self._lock:
+            self._trackers[correlation_id] = tracker
+
+    def unregister_tracker(self, correlation_id: str) -> None:
+        """Unregister a tracker."""
+        with self._lock:
+            self._trackers.pop(correlation_id, None)
+
+    def handle(self, event: Any) -> None:
+        """Handle a domain event and forward to appropriate tracker."""
+        event_type = type(event).__name__
+        stage = self.EVENT_STAGE_MAP.get(event_type)
+
+        if not stage:
+            return
+
+        # Try to find correlation ID
+        correlation_id = getattr(event, "correlation_id", None)
+        if not correlation_id:
+            return
+
+        with self._lock:
+            tracker = self._trackers.get(correlation_id)
+
+        if tracker:
+            message = self._format_event_message(event, event_type)
+            tracker.report(stage, message, self._extract_details(event))
+
+    def _format_event_message(self, event: Any, event_type: str) -> str:
+        """Format event into progress message."""
+        if event_type == "ProviderStarted":
+            return f"Provider started ({getattr(event, 'tools_count', 0)} tools)"
+        elif event_type == "ProviderStateChanged":
+            new_state = getattr(event, "new_state", "unknown")
+            return f"Provider state: {new_state}"
+        elif event_type == "ToolInvocationRequested":
+            tool = getattr(event, "tool_name", "unknown")
+            return f"Invoking tool: {tool}"
+        elif event_type == "ToolInvocationCompleted":
+            duration = getattr(event, "duration_ms", 0)
+            return f"Tool completed ({duration:.0f}ms)"
+        elif event_type == "ToolInvocationFailed":
+            error = getattr(event, "error_message", "unknown error")
+            return f"Tool failed: {error[:50]}"
+        return event_type
+
+    def _extract_details(self, event: Any) -> Dict[str, Any]:
+        """Extract relevant details from event."""
+        details = {}
+        for attr in ["provider_id", "tool_name", "duration_ms", "error_type"]:
+            if hasattr(event, attr):
+                details[attr] = getattr(event, attr)
+        return details
+
+
+# Global progress event handler instance
+_progress_handler: Optional[ProgressEventHandler] = None
+
+
+def get_progress_handler() -> ProgressEventHandler:
+    """Get or create the global progress event handler."""
+    global _progress_handler
+    if _progress_handler is None:
+        _progress_handler = ProgressEventHandler()
+    return _progress_handler
+
+
+def create_progress_tracker(
+    provider: str = "",
+    operation: str = "",
+    callback: Optional[ProgressCallback] = None,
+    correlation_id: Optional[str] = None,
+) -> ProgressTracker:
+    """Create a progress tracker with optional event bus integration.
+
+    Args:
+        provider: Provider name
+        operation: Operation name
+        callback: Optional progress callback
+        correlation_id: Optional correlation ID for event bus integration
+
+    Returns:
+        New ProgressTracker instance
+    """
+    tracker = ProgressTracker(
+        callback=callback,
+        provider=provider,
+        operation=operation,
+    )
+
+    if correlation_id:
+        handler = get_progress_handler()
+        handler.register_tracker(correlation_id, tracker)
+
+    return tracker