hud-python 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

hud/telemetry/instrumentation/mcp.py

@@ -1,259 +0,0 @@
-from __future__ import annotations
-
-import logging
-import time
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from collections.abc import Awaitable, Callable
-
-from wrapt import wrap_function_wrapper
-
-from hud.telemetry.context import (
-    create_notification_record,
-    create_request_record,
-    create_response_record,
-    get_current_task_run_id,
-)
-from hud.telemetry.mcp_models import DirectionType, MCPCallType, StatusType
-
-logger = logging.getLogger(__name__)
-
-
-class MCPInstrumentor:
-    """
-    Context-aware instrumentor for MCP calls.
-    Only instruments MCP methods when there's an active trace context.
-    """
-
-    def __init__(self) -> None:
-        self._installed = False
-
-    def install(self) -> None:
-        """Install instrumentation for MCP - but only activate when trace context exists."""
-        logger.debug("MCPInstrumentor: install() called (context-aware mode)")
-        if self._installed:
-            logger.debug("MCP instrumentation already installed")
-            return
-
-        try:
-            # Import and wrap the main session methods
-            import mcp.shared.session  # noqa: F401
-
-            wrap_function_wrapper(
-                "mcp.shared.session",
-                "BaseSession.send_request",
-                self._context_aware_send_request_wrapper,
-            )
-            wrap_function_wrapper(
-                "mcp.shared.session",
-                "BaseSession.send_notification",
-                self._context_aware_send_notification_wrapper,
-            )
-
-            logger.debug("Successfully wrapped BaseSession methods for context-aware telemetry")
-        except ImportError:
-            logger.debug("mcp.shared.session not available yet")
-        except Exception as e:
-            logger.warning("Failed to wrap BaseSession methods: %s", e)
-
-        self._installed = True
-        logger.debug("MCP instrumentation installed (context-aware mode)")
-
-    async def _context_aware_send_request_wrapper(
-        self,
-        wrapped: Callable[[Any, Any, Any, Any], Awaitable[Any]],
-        instance: Any,
-        args: Any,
-        kwargs: Any,
-    ) -> Any:
-        """Context-aware send request wrapper."""
-        hud_task_run_id = get_current_task_run_id()
-
-        if not hud_task_run_id:
-            # No active trace - pass through without instrumentation
-            return await wrapped(*args, **kwargs)
-
-        start_time = time.time()
-        request = args[0] if args else None
-
-        # Extract method name from the request
-        method_name = "unknown_method"
-        message_id = None
-        request_data = None
-
-        if request:
-            try:
-                # Get model dump for Pydantic models
-                if hasattr(request, "model_dump"):
-                    request_data = request.model_dump(mode="json", exclude_none=True)
-                    # Extract method from the dump
-                    if isinstance(request_data, dict) and "method" in request_data:
-                        method_name = request_data["method"]
-                # Fallback to direct method attribute
-                elif hasattr(request, "method"):
-                    method_name = str(request.method)
-                    request_data = {"method": method_name}
-            except Exception as e:
-                logger.debug("Failed to extract method: %s", e)
-
-        # Create the request record
-        create_request_record(
-            method=method_name,
-            message_id=message_id,
-            call_type=MCPCallType.SEND_REQUEST,
-            direction=DirectionType.SENT,
-            status=StatusType.STARTED,
-            start_time=start_time,
-            request_data=request_data,
-        )
-
-        try:
-            # Call the wrapped method and get the result
-            result = await wrapped(*args, **kwargs)
-
-            # Get the actual message_id from the session after sending
-            if hasattr(instance, "_request_id"):
-                # Current request ID minus 1 (since it was incremented)
-                message_id = getattr(instance, "_request_id", 1) - 1
-
-            # Update request record with completion
-            create_request_record(
-                method=method_name,
-                message_id=message_id,
-                call_type=MCPCallType.SEND_REQUEST,
-                direction=DirectionType.SENT,
-                status=StatusType.COMPLETED,
-                start_time=start_time,
-                end_time=time.time(),
-                duration=time.time() - start_time,
-                request_data=request_data,
-            )
-
-            # Capture the response
-            if result is not None:
-                response_data = None
-                try:
-                    if hasattr(result, "model_dump"):
-                        response_data = result.model_dump(mode="json", exclude_none=True)
-                    else:
-                        response_data = {"_type": type(result).__name__}
-                except Exception as e:
-                    logger.debug("Failed to serialize response data: %s", e)
-                    response_data = {"_type": type(result).__name__, "_error": str(e)}
-
-                create_response_record(
-                    method=method_name,
-                    related_request_id=message_id,
-                    is_error=False,
-                    message_id=message_id,
-                    direction=DirectionType.RECEIVED,
-                    call_type=MCPCallType.RECEIVE_RESPONSE,
-                    response_data=response_data,
-                    timestamp=time.time(),
-                )
-
-            return result
-
-        except Exception as e:
-            # Log the error
-            create_request_record(
-                method=method_name,
-                message_id=message_id,
-                call_type=MCPCallType.SEND_REQUEST,
-                direction=DirectionType.SENT,
-                status=StatusType.ERROR,
-                start_time=start_time,
-                end_time=time.time(),
-                duration=time.time() - start_time,
-                error=str(e),
-                error_type=type(e).__name__,
-                request_data=request_data,
-            )
-
-            # Also record error response
-            create_response_record(
-                method=method_name,
-                related_request_id=message_id,
-                is_error=True,
-                message_id=message_id,
-                direction=DirectionType.RECEIVED,
-                call_type=MCPCallType.RECEIVE_RESPONSE,
-                status=StatusType.ERROR,
-                error=str(e),
-                error_type=type(e).__name__,
-                timestamp=time.time(),
-            )
-
-            raise
-
-    async def _context_aware_send_notification_wrapper(
-        self,
-        wrapped: Callable[[Any, Any, Any, Any], Awaitable[Any]],
-        instance: Any,
-        args: Any,
-        kwargs: Any,
-    ) -> Any:
-        """Context-aware send notification wrapper."""
-        hud_task_run_id = get_current_task_run_id()
-
-        if not hud_task_run_id:
-            # No active trace - pass through without instrumentation
-            return await wrapped(*args, **kwargs)
-
-        start_time = time.time()
-        notification = args[0] if args else None
-
-        # Extract method name
-        method_name = "unknown_method"
-        notification_data = None
-
-        if notification:
-            try:
-                # Get model dump for Pydantic models
-                if hasattr(notification, "model_dump"):
-                    notification_data = notification.model_dump(mode="json", exclude_none=True)
-                    # Extract method from the dump
-                    if isinstance(notification_data, dict) and "method" in notification_data:
-                        method_name = notification_data["method"]
-                # Fallback to direct method attribute
-                elif hasattr(notification, "method"):
-                    method_name = str(notification.method)
-                    notification_data = {"method": method_name}
-            except Exception as e:
-                logger.debug("Failed to extract notification method: %s", e)
-
-        record_data_base = {
-            "method": method_name,
-            "message_id": None,  # Notifications don't have IDs
-            "call_type": MCPCallType.SEND_NOTIFICATION,
-            "direction": DirectionType.SENT,
-            "notification_data": notification_data,
-        }
-
-        try:
-            create_notification_record(
-                **record_data_base, status=StatusType.STARTED, start_time=start_time
-            )
-
-            result = await wrapped(*args, **kwargs)
-
-            create_notification_record(
-                **record_data_base,
-                status=StatusType.COMPLETED,
-                start_time=start_time,
-                end_time=time.time(),
-                duration=time.time() - start_time,
-            )
-            return result
-        except Exception as e:
-            create_notification_record(
-                **record_data_base,
-                status=StatusType.ERROR,
-                start_time=start_time,
-                end_time=time.time(),
-                duration=time.time() - start_time,
-                error=str(e),
-                error_type=type(e).__name__,
-            )
-            raise

hud/telemetry/instrumentation/registry.py

@@ -1,59 +0,0 @@
-from __future__ import annotations
-
-import logging
-from typing import Any
-
-logger = logging.getLogger(__name__)
-
-
-class InstrumentorRegistry:
-    """Registry for telemetry instrumentors."""
-
-    def __init__(self) -> None:
-        self._instrumentors: dict[str, Any] = {}
-
-    def register(self, name: str, instrumentor: Any) -> None:
-        """Register an instrumentor.
-
-        Args:
-            name: Name of the instrumentor
-            instrumentor: The instrumentor instance
-        """
-        self._instrumentors[name] = instrumentor
-        logger.debug("Registered instrumentor: %s", name)
-
-    def _safe_install(self, name: str, instrumentor: Any) -> tuple[str, Exception | None]:
-        """Safely install an instrumentor and return result."""
-        try:
-            instrumentor.install()
-            return name, None
-        except Exception as e:
-            return name, e
-
-    def install_all(self) -> None:
-        """Install all registered instrumentors."""
-        # Use map to apply safe installation to all instrumentors
-        installation_results = [
-            self._safe_install(name, instrumentor)
-            for name, instrumentor in self._instrumentors.items()
-        ]
-
-        # Process results
-        for name, error in installation_results:
-            if error is None:
-                logger.debug("Installed instrumentor: %s", name)
-            else:
-                logger.warning("Failed to install instrumentor %s: %s", name, error)
-
-
-# Create a singleton registry
-registry = InstrumentorRegistry()
-
-# Try to register MCP instrumentor if available
-try:
-    from .mcp import MCPInstrumentor
-
-    registry.register("mcp", MCPInstrumentor())
-    logger.debug("MCP instrumentor registered")
-except Exception as e:
-    logger.debug("Could not register MCP instrumentor: %s", e)

hud/telemetry/mcp_models.py

@@ -1,270 +0,0 @@
-from __future__ import annotations
-
-from datetime import datetime
-from enum import Enum
-from typing import TYPE_CHECKING, Any, ClassVar
-
-# Import MCP types
-from mcp.types import JSONRPCError, JSONRPCNotification, JSONRPCRequest, JSONRPCResponse
-from pydantic import BaseModel, Field, field_validator
-
-if TYPE_CHECKING:
-    from mcp.shared.message import SessionMessage
-
-
-class DirectionType(str, Enum):
-    """Direction of an MCP message"""
-
-    SENT = "sent"
-    RECEIVED = "received"
-
-
-class StatusType(str, Enum):
-    """Status of an MCP operation"""
-
-    STARTED = "started"
-    COMPLETED = "completed"
-    ERROR = "error"
-
-
-class MCPCallType(str, Enum):
-    """Enum for different types of MCP calls in telemetry."""
-
-    # Requests and Notifications
-    SEND_REQUEST = "mcp.send_request"
-    SEND_NOTIFICATION = "mcp.send_notification"
-
-    # Responses
-    RECEIVE_RESPONSE = "mcp.receive_response"
-
-
-class BaseMCPCall(BaseModel):
-    """Base model for all MCP telemetry records"""
-
-    task_run_id: str
-    call_type: str
-    timestamp: float = Field(default_factory=lambda: datetime.now().timestamp())
-    method: str = "unknown_method"
-    status: StatusType
-    direction: DirectionType | None = None
-    # Additional data that might be useful for any call
-    message_id: str | int | None = None
-
-    # Mapping of call types to model classes - to be populated by subclasses
-    _call_type_mapping: ClassVar[dict[str, type["BaseMCPCall"]]] = {}
-
-    @field_validator("call_type")
-    @classmethod
-    def validate_call_type(cls, v: str) -> str:
-        """Allow any string but preferably from MCPCallType"""
-        return v
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> BaseMCPCall:
-        """Create a record from a dictionary, using the appropriate subclass"""
-        call_type = data.get("call_type", "")
-        record_cls = cls._call_type_mapping.get(call_type, BaseMCPCall)
-        return record_cls.model_validate(data)
-
-    def __init_subclass__(cls, **kwargs: Any) -> None:
-        """Register subclasses in the mapping by their default call_type"""
-        super().__init_subclass__(**kwargs)
-        if hasattr(cls, "__annotations__") and "call_type" in cls.__annotations__:
-            default_call_type = getattr(cls, "call_type", None)
-            if isinstance(default_call_type, str):
-                BaseMCPCall._call_type_mapping[default_call_type] = cls
-
-
-class MCPRequestCall(BaseMCPCall):
-    """Record for an MCP request"""
-
-    direction: DirectionType = DirectionType.SENT
-    call_type: str = MCPCallType.SEND_REQUEST
-    start_time: float
-    end_time: float | None = None
-    duration: float | None = None
-    request_id: str | int | None = None
-    request_data: dict[str, Any] | None = None
-    error: str | None = None
-    error_type: str | None = None
-
-    @classmethod
-    def from_jsonrpc_request(
-        cls,
-        request: JSONRPCRequest,
-        task_run_id: str,
-        status: StatusType = StatusType.STARTED,
-        **kwargs: Any,
-    ) -> MCPRequestCall:
-        """Create telemetry record from a JSONRPCRequest"""
-        return cls(
-            task_run_id=task_run_id,
-            status=status,
-            request_id=request.id,
-            message_id=request.id,
-            method=request.method,
-            request_data=request.model_dump(exclude_none=True),
-            start_time=datetime.now().timestamp(),
-            **kwargs,
-        )
-
-    @classmethod
-    def from_session_message(
-        cls,
-        message: SessionMessage,
-        task_run_id: str,
-        status: StatusType = StatusType.STARTED,
-        **kwargs: Any,
-    ) -> MCPRequestCall | None:
-        """Create telemetry record from a SessionMessage containing a JSONRPCRequest"""
-        if (
-            hasattr(message, "message")
-            and hasattr(message.message, "root")
-            and isinstance(message.message.root, JSONRPCRequest)
-        ):
-            return cls.from_jsonrpc_request(
-                message.message.root, task_run_id=task_run_id, status=status, **kwargs
-            )
-        return None
-
-
-class MCPResponseCall(BaseMCPCall):
-    """Record for an MCP response"""
-
-    direction: DirectionType = DirectionType.RECEIVED
-    call_type: str = MCPCallType.RECEIVE_RESPONSE
-    is_response_or_error: bool = True
-    is_error: bool = False
-    response_id: str | int | None = None
-    related_request_id: str | int | None = None
-    response_data: dict[str, Any] | None = None
-    error: str | None = None
-    error_type: str | None = None
-
-    @classmethod
-    def from_jsonrpc_response(
-        cls, response: JSONRPCResponse | JSONRPCError, task_run_id: str, **kwargs: Any
-    ) -> MCPResponseCall:
-        """Create telemetry record from a JSONRPCResponse or JSONRPCError"""
-        is_error = isinstance(response, JSONRPCError)
-
-        result = cls(
-            task_run_id=task_run_id,
-            status=StatusType.COMPLETED,
-            response_id=response.id,
-            message_id=response.id,
-            related_request_id=response.id,  # In MCP, response ID matches request ID
-            is_error=is_error,
-            method=f"response_to_id_{response.id}",
-            response_data=response.model_dump(exclude_none=True),
-            **kwargs,
-        )
-
-        if is_error and hasattr(response, "error"):
-            result.error = response.error.message
-            result.error_type = str(response.error.code)
-
-        return result
-
-    @classmethod
-    def from_session_message(
-        cls, message: SessionMessage, task_run_id: str, **kwargs: Any
-    ) -> MCPResponseCall | None:
-        """Create telemetry record from a SessionMessage containing a response or error"""
-        if (
-            hasattr(message, "message")
-            and hasattr(message.message, "root")
-            and isinstance(message.message.root, JSONRPCResponse | JSONRPCError)
-        ):
-            return cls.from_jsonrpc_response(
-                message.message.root, task_run_id=task_run_id, **kwargs
-            )
-        return None
-
-
-class MCPNotificationCall(BaseMCPCall):
-    """Record for an MCP notification"""
-
-    direction: DirectionType = DirectionType.SENT
-    call_type: str = MCPCallType.SEND_NOTIFICATION
-    start_time: float
-    end_time: float | None = None
-    duration: float | None = None
-    notification_data: dict[str, Any] | None = None
-    error: str | None = None
-    error_type: str | None = None
-
-    @classmethod
-    def from_jsonrpc_notification(
-        cls,
-        notification: JSONRPCNotification,
-        task_run_id: str,
-        status: StatusType = StatusType.STARTED,
-        **kwargs: Any,
-    ) -> MCPNotificationCall:
-        """Create telemetry record from a JSONRPCNotification"""
-        return cls(
-            task_run_id=task_run_id,
-            status=status,
-            method=notification.method,
-            notification_data=notification.model_dump(exclude_none=True),
-            start_time=datetime.now().timestamp(),
-            **kwargs,
-        )
-
-    @classmethod
-    def from_session_message(
-        cls,
-        message: SessionMessage,
-        task_run_id: str,
-        status: StatusType = StatusType.STARTED,
-        **kwargs: Any,
-    ) -> MCPNotificationCall | None:
-        """Create telemetry record from a SessionMessage containing a JSONRPCNotification"""
-        if (
-            hasattr(message, "message")
-            and hasattr(message.message, "root")
-            and isinstance(message.message.root, JSONRPCNotification)
-        ):
-            return cls.from_jsonrpc_notification(
-                message.message.root, task_run_id=task_run_id, status=status, **kwargs
-            )
-        return None
-
-
-class MCPTelemetryRecord(BaseModel):
-    """Container for a set of related MCP telemetry records"""
-
-    task_run_id: str
-    records: list[BaseMCPCall]
-    timestamp: float = Field(default_factory=lambda: datetime.now().timestamp())
-
-    @property
-    def count_by_type(self) -> dict[str, int]:
-        """Count records by call_type"""
-        result: dict[str, int] = {}
-        for record in self.records:
-            result[record.call_type] = result.get(record.call_type, 0) + 1
-        return result
-
-    @property
-    def count_by_direction(self) -> dict[str, int]:
-        """Count records by direction"""
-        result: dict[str, int] = {}
-        for record in self.records:
-            if record.direction:
-                direction = record.direction.value
-                result[direction] = result.get(direction, 0) + 1
-        return result
-
-
-class TrajectoryStep(BaseModel):
-    """Model for telemetry export format."""
-
-    type: str = Field(default="mcp-step")
-    observation_url: str | None = None
-    observation_text: str | None = None
-    actions: list[dict[str, Any]] = Field(default_factory=list)
-    start_timestamp: str | None = None  # ISO 8601 format
-    end_timestamp: str | None = None  # ISO 8601 format
-    metadata: dict[str, Any] = Field(default_factory=dict)

hud/telemetry/tests/__init__.py

@@ -1 +0,0 @@
-# Tests for hud.telemetry module