webber/webber-api/src/shared/logging.py

"""
Centralized logging with temporal benchmarking.

Provides:
- @logged() decorator for automatic function timing
- trace_span() context manager for manual instrumentation
- Trace ID correlation across nested calls
- Configurable slow/warn thresholds
"""
import asyncio
import functools
import logging
import sys
import time
from collections.abc import Callable
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, ParamSpec, TypeVar, cast
from uuid import uuid4

# === Trace Context ===

@dataclass
class TraceSpan:
    """Represents a timed execution span."""
    name: str
    trace_id: str
    parent_id: str | None = None
    span_id: str = field(default_factory=lambda: uuid4().hex[:8])
    start_time: float = field(default_factory=time.perf_counter)
    end_time: float | None = None

    @property
    def duration_ms(self) -> float:
        """Get duration in milliseconds."""
        end = self.end_time or time.perf_counter()
        return (end - self.start_time) * 1000


_current_span: ContextVar[TraceSpan | None] = ContextVar('current_span', default=None)
_trace_id: ContextVar[str | None] = ContextVar('trace_id', default=None)


def get_current_trace_id() -> str | None:
    """Get current trace ID for log correlation."""
    return _trace_id.get()


def get_current_span() -> TraceSpan | None:
    """Get current trace span."""
    return _current_span.get()


# === Setup ===

def setup_logging(log_level: str = "INFO") -> None:
    """Configure application logging."""
    log_dir = Path("logs")
    log_dir.mkdir(exist_ok=True)

    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler(log_dir / "webber.log", encoding="utf-8")
        ]
    )

    # Quiet noisy libraries
    for name in ["httpx", "httpcore", "uvicorn.access", "uvicorn.error"]:
        logging.getLogger(name).setLevel(logging.WARNING)


def get_logger(name: str) -> logging.Logger:
    """Get a logger instance."""
    return logging.getLogger(name)


# === Decorator ===

# @logged wraps ~24 functions across this package. Untyped, its decorator
# erased every one of their signatures, so mypy saw `Any` coming back from
# annotated functions like `ToolResult.execute() -> ToolResult`. That surfaced
# as 33 no-any-return errors scattered across the tools and agents — each
# reading like a local annotation slip, all of them this one decorator.
P = ParamSpec("P")
R = TypeVar("R")


def logged(
    logger: logging.Logger | None = None,
    slow_threshold_ms: float = 100.0,
    warn_threshold_ms: float = 500.0,
    include_args: bool = False,
) -> Callable[[Callable[P, R]], Callable[P, R]]:
    """
    Decorator for automatic function logging with temporal benchmarking.

    Args:
        logger: Logger instance (defaults to module logger)
        slow_threshold_ms: Log INFO if execution exceeds this (default 100ms)
        warn_threshold_ms: Log WARNING if execution exceeds this (default 500ms)
        include_args: Include function arguments in log (careful with sensitive data)

    Usage:
        @logged()
        async def my_function(): ...

        @logged(slow_threshold_ms=50, warn_threshold_ms=200)
        def critical_path(): ...
    """
    def decorator(func: Callable[P, R]) -> Callable[P, R]:
        nonlocal logger
        if logger is None:
            logger = logging.getLogger(func.__module__)

        func_name = f"{func.__module__}.{func.__qualname__}"

        def _create_span() -> TraceSpan:
            parent = _current_span.get()
            trace_id = _trace_id.get() or uuid4().hex[:16]
            if _trace_id.get() is None:
                _trace_id.set(trace_id)
            return TraceSpan(
                name=func_name,
                trace_id=trace_id,
                parent_id=parent.span_id if parent else None,
            )

        def _log_completion(span: TraceSpan, error: Exception | None = None):
            span.end_time = time.perf_counter()
            duration = span.duration_ms
            tid = span.trace_id[:8]

            if error:
                logger.error(
                    f"[{tid}] {func_name} FAILED after {duration:.2f}ms: {error}",
                    exc_info=True
                )
            elif duration >= warn_threshold_ms:
                logger.warning(
                    f"[{tid}] {func_name} SLOW: {duration:.2f}ms (threshold: {warn_threshold_ms}ms)"
                )
            elif duration >= slow_threshold_ms:
                logger.info(f"[{tid}] {func_name} completed in {duration:.2f}ms")
            else:
                logger.debug(f"[{tid}] {func_name} completed in {duration:.2f}ms")

        @functools.wraps(func)
        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
            span = _create_span()
            token = _current_span.set(span)

            if include_args:
                logger.debug(f"[{span.trace_id[:8]}] -> {func_name}({args}, {kwargs})")
            else:
                logger.debug(f"[{span.trace_id[:8]}] -> {func_name}")

            try:
                result = await cast(Any, func(*args, **kwargs))
                _log_completion(span)
                return result
            except Exception as e:
                _log_completion(span, error=e)
                raise
            finally:
                _current_span.reset(token)

        @functools.wraps(func)
        def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            span = _create_span()
            token = _current_span.set(span)

            if include_args:
                logger.debug(f"[{span.trace_id[:8]}] -> {func_name}({args}, {kwargs})")
            else:
                logger.debug(f"[{span.trace_id[:8]}] -> {func_name}")

            try:
                result = func(*args, **kwargs)
                _log_completion(span)
                return result
            except Exception as e:
                _log_completion(span, error=e)
                raise
            finally:
                _current_span.reset(token)

        # The branch is chosen at decoration time; mypy cannot narrow R to a
        # coroutine on the strength of iscoroutinefunction, so the union is
        # asserted here once instead of at every call site.
        chosen = async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
        return cast(Callable[P, R], chosen)
    return decorator


# === Context Manager ===

class trace_span:
    """
    Context manager for manual span creation.

    Usage:
        with trace_span("database_query"):
            result = db.execute(query)

        async with trace_span("llm_call"):
            response = await agent.run(prompt)
    """

    def __init__(self, name: str, logger: logging.Logger | None = None):
        self.name = name
        self.logger = logger or logging.getLogger(__name__)
        self.span: TraceSpan | None = None
        self.token: Token[TraceSpan | None] | None = None

    def __enter__(self) -> TraceSpan:
        parent = _current_span.get()
        trace_id = _trace_id.get() or uuid4().hex[:16]
        if _trace_id.get() is None:
            _trace_id.set(trace_id)

        self.span = TraceSpan(
            name=self.name,
            trace_id=trace_id,
            parent_id=parent.span_id if parent else None,
        )
        self.token = _current_span.set(self.span)
        self.logger.debug(f"[{self.span.trace_id[:8]}] -> {self.name}")
        return self.span

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.span:
            self.span.end_time = time.perf_counter()
            duration = self.span.duration_ms
            tid = self.span.trace_id[:8]
            if exc_val:
                self.logger.error(f"[{tid}] {self.name} FAILED: {duration:.2f}ms")
            else:
                self.logger.debug(f"[{tid}] {self.name}: {duration:.2f}ms")
        if self.token:
            _current_span.reset(self.token)
        return False

    async def __aenter__(self) -> TraceSpan:
        return self.__enter__()

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return self.__exit__(exc_type, exc_val, exc_tb)