mcghidra/src/ghydramcp/core/progress.py

"""Progress reporting utilities for long-running operations.

Provides async progress reporting using FastMCP's Context for
real-time progress notifications to MCP clients.
"""

from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from mcp.server.fastmcp import Context


class ProgressReporter:
    """Helper class for reporting progress during long operations.

    Provides throttled progress updates to avoid spamming the client
    with too many notifications.

    Example:
        async def long_operation(ctx: Context):
            progress = ProgressReporter(ctx, "Scanning", total=100)
            for i in range(100):
                await progress.update(message=f"Processing item {i}")
            await progress.complete("Scan finished")
    """

    def __init__(
        self,
        ctx: Optional["Context"],
        operation: str,
        total: int = 100
    ):
        """Initialize the progress reporter.

        Args:
            ctx: FastMCP context for progress reporting (may be None)
            operation: Name of the operation (used in log messages)
            total: Total number of steps (default: 100)
        """
        self.ctx = ctx
        self.operation = operation
        self.total = total
        self.current = 0
        self._last_reported = 0
        # Report every 5% at minimum
        self._report_threshold = max(1, total // 20)

    async def update(
        self,
        progress: Optional[int] = None,
        message: Optional[str] = None
    ) -> None:
        """Update progress, reporting to client if threshold reached.

        Args:
            progress: Current progress value (if None, increments by 1)
            message: Optional message to log with the progress update
        """
        if progress is not None:
            self.current = progress
        else:
            self.current += 1

        # Only report if we've crossed a threshold or reached the end
        should_report = (
            self.current - self._last_reported >= self._report_threshold
            or self.current >= self.total
        )

        if self.ctx and should_report:
            try:
                await self.ctx.report_progress(
                    progress=self.current,
                    total=self.total
                )
                if message:
                    await self.ctx.info(f"{self.operation}: {message}")
                self._last_reported = self.current
            except Exception:
                pass  # Silently ignore if context doesn't support progress

    async def info(self, message: str) -> None:
        """Send an info message to the client.

        Args:
            message: Message to send
        """
        if self.ctx:
            try:
                await self.ctx.info(f"{self.operation}: {message}")
            except Exception:
                pass

    async def complete(self, message: Optional[str] = None) -> None:
        """Mark operation as complete.

        Args:
            message: Optional completion message (supports format placeholders:
                     {count}, {total}, {operation})
        """
        self.current = self.total
        if self.ctx:
            try:
                await self.ctx.report_progress(
                    progress=self.total,
                    total=self.total
                )
                if message:
                    formatted = message.format(
                        count=self.current,
                        total=self.total,
                        operation=self.operation
                    )
                    await self.ctx.info(formatted)
            except Exception:
                pass


async def report_progress(
    ctx: Optional["Context"],
    progress: int,
    total: int,
    message: Optional[str] = None
) -> None:
    """Convenience function for one-off progress updates.

    Args:
        ctx: FastMCP context (may be None)
        progress: Current progress value
        total: Total progress value
        message: Optional message to log
    """
    if ctx:
        try:
            await ctx.report_progress(progress=progress, total=total)
            if message:
                await ctx.info(message)
        except Exception:
            pass


async def report_step(
    ctx: Optional["Context"],
    step: int,
    total_steps: int,
    description: str
) -> None:
    """Report a discrete step in a multi-step operation.

    Args:
        ctx: FastMCP context (may be None)
        step: Current step number (1-indexed)
        total_steps: Total number of steps
        description: Description of the current step
    """
    if ctx:
        try:
            await ctx.report_progress(progress=step, total=total_steps)
            await ctx.info(f"Step {step}/{total_steps}: {description}")
        except Exception:
            pass