hud-python 0.2.10__py3-none-any.whl → 0.3.0__py3-none-any.whl

hud/telemetry/instrumentation/mcp.py

@@ -2,494 +2,258 @@ from __future__ import annotations
 
 import logging
 import time
-from contextlib import asynccontextmanager
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Awaitable, Callable
+    from collections.abc import Awaitable, Callable
 
-# from dataclasses import dataclass # Likely unused now
-from mcp.shared.message import SessionMessage  # type: ignore
-
-# MCP type imports for type checking and attribute access safety
-from mcp.types import (  # type: ignore
-    JSONRPCError,
-    JSONRPCNotification,
-    JSONRPCRequest,
-    JSONRPCResponse,
-)
-from wrapt import register_post_import_hook, wrap_function_wrapper
+from wrapt import wrap_function_wrapper
 
 from hud.telemetry.context import (
-    buffer_mcp_call,
     create_notification_record,
     create_request_record,
     create_response_record,
     get_current_task_run_id,
 )
-from hud.telemetry.mcp_models import DirectionType, MCPCallType, MCPManualTestCall, StatusType
+from hud.telemetry.mcp_models import DirectionType, MCPCallType, StatusType
 
 logger = logging.getLogger(__name__)
 
 
 class MCPInstrumentor:
     """
-    Instrumentor for MCP calls.
+    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"""
-        logger.debug("MCPInstrumentor: install() called")
+        """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
 
-        # Updated hooks: removed _transport_wrapper and _base_session_init_wrapper related hooks
-        # if they become no-ops. For now, let's assume _transport_wrapper might still be used,
-        # but simplified. _base_session_init_wrapper hook is removed.
-        hooks = [
-            ("mcp.client.sse", "sse_client", self._transport_wrapper),
-            ("mcp.server.sse", "SseServerTransport.connect_sse", self._transport_wrapper),
-            ("mcp.client.stdio", "stdio_client", self._transport_wrapper),
-            ("mcp.server.stdio", "stdio_server", self._transport_wrapper),
-            ("mcp.shared.session", "BaseSession._receive_loop", self._receive_loop_wrapper),
-        ]
-
-        def create_hook(module: str, func_name: str, wrapper: Any) -> Any:
-            return lambda _: wrap_function_wrapper(module, func_name, wrapper)
-
-        for module, func_name, wrapper in hooks:
-            logger.debug(
-                "MCPInstrumentor: Registering post-import hook for %s.%s", module, func_name
-            )
-            register_post_import_hook(
-                create_hook(module, func_name, wrapper),
-                module,
-            )
-
-        # Removed hook for BaseSession.__init__ (it was for _base_session_init_wrapper)
-
-        logger.debug(
-            "MCPInstrumentor: Attempting immediate instrumentation of already imported modules."
-        )
-        for module, func_name, wrapper in hooks:
-            try:
-                mod = __import__(module, fromlist=[func_name.split(".")[0]])
-                target_obj = mod
-                parts = func_name.split(".")
-                for i, part in enumerate(parts):
-                    if hasattr(target_obj, part):
-                        if i == len(parts) - 1:
-                            target_obj = getattr(target_obj, part)
-                        else:
-                            target_obj = getattr(target_obj, part)
-                    else:
-                        target_obj = None
-                        break
-
-                if target_obj and callable(target_obj):
-                    logger.debug("MCPInstrumentor: Wrapping %s.%s", module, func_name)
-                    wrap_function_wrapper(module, func_name, wrapper)
-                else:
-                    logger.debug(
-                        "Function %s not found in %s during immediate instrumentation attempt",
-                        func_name,
-                        module,
-                    )
-            except ImportError:
-                logger.debug("Module %s not imported yet, post-import hook will cover it", module)
-            except Exception as e:
-                logger.warning("Failed to immediately wrap %s.%s: %s", module, func_name, e)
-
         try:
-            # Import only for testing availability, don't store reference
-            __import__("mcp.shared.session")
+            # Import and wrap the main session methods
+            import mcp.shared.session  # noqa: F401
+
             wrap_function_wrapper(
                 "mcp.shared.session",
                 "BaseSession.send_request",
-                self._send_request_telemetry_wrapper,
+                self._context_aware_send_request_wrapper,
             )
             wrap_function_wrapper(
                 "mcp.shared.session",
                 "BaseSession.send_notification",
-                self._send_notification_telemetry_wrapper,
+                self._context_aware_send_notification_wrapper,
             )
-            logger.debug("Successfully wrapped BaseSession methods for telemetry")
+
+            logger.debug("Successfully wrapped BaseSession methods for context-aware telemetry")
         except ImportError:
-            logger.debug("mcp.shared.session not imported yet, post-import hook will cover it")
+            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 (simplified)")
-
-    @asynccontextmanager
-    async def _transport_wrapper(
-        self, wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
-    ) -> AsyncGenerator[Any, None]:
-        """Wrap transport functions. Simplified: passes through original streams."""
-        original_func_name = f"{wrapped.__module__}.{wrapped.__name__}"
-        logger.debug(
-            "MCPInstrumentor: _transport_wrapper called for %s (passthrough)", original_func_name
-        )
+        logger.debug("MCP instrumentation installed (context-aware mode)")
 
-        # No OTel context or HUD Task ID manipulation here anymore for transport layer itself.
-        # Higher level wrappers (_send_request, _receive_loop) will handle HUD Task ID.
-
-        async with wrapped(*args, **kwargs) as (read_stream, write_stream):
-            logger.debug("_transport_wrapper: Yielding original streams for %s", original_func_name)
-            yield read_stream, write_stream  # Pass original streams directly
-            logger.debug(
-                "_transport_wrapper: Exited original stream context for %s", original_func_name
-            )
-
-    async def _receive_loop_wrapper(
+    async def _context_aware_send_request_wrapper(
         self,
         wrapped: Callable[[Any, Any, Any, Any], Awaitable[Any]],
         instance: Any,
         args: Any,
         kwargs: Any,
     ) -> Any:
-        """
-        Wrap _receive_loop to capture responses and handled incoming messages.
-        """
-        logger.debug("MCPInstrumentor: _receive_loop_wrapper called")
-
-        orig_handle_incoming = instance._handle_incoming
-
-        async def handle_incoming_with_telemetry(req_or_msg: Any) -> Any:
-            start_time = time.time()
-            hud_task_run_id = get_current_task_run_id()
-            method_name = "unknown_method"
-            message_id = None
-            is_response_or_error_flag = False
-            call_type_override = MCPCallType.HANDLE_INCOMING
-
-            actual_message_root = None
-            # SessionMessage
-            if hasattr(req_or_msg, "message") and hasattr(req_or_msg.message, "root"):
-                actual_message_root = req_or_msg.message.root
-            elif hasattr(req_or_msg, "root"):  # e.g. RequestResponder
-                actual_message_root = req_or_msg.root
-            elif isinstance(
-                req_or_msg, JSONRPCRequest | JSONRPCNotification | JSONRPCResponse | JSONRPCError
-            ):
-                actual_message_root = req_or_msg
-
-            if actual_message_root:
-                if isinstance(actual_message_root, JSONRPCRequest):
-                    method_name = actual_message_root.method
-                    message_id = actual_message_root.id
-                    # call_type_override can remain HANDLE_INCOMING or be more specific if desired
-                elif isinstance(actual_message_root, JSONRPCNotification):
-                    method_name = actual_message_root.method
-                    message_id = None  # Notifications don't have IDs
-                    # call_type_override can remain HANDLE_INCOMING
-                elif isinstance(actual_message_root, JSONRPCResponse | JSONRPCError):
-                    # This case implies _handle_incoming is processing a response/error directly
-                    # (e.g. an error encountered while trying to send/route a response previously)
-                    message_id = actual_message_root.id
-                    method_name = f"internal_response_handling_for_id_{message_id}"
-                    is_response_or_error_flag = True
-                    # Treat as receiving a response internally
-                    call_type_override = MCPCallType.RECEIVE_RESPONSE
-                else:
-                    # Fallback for other types, if any, that might appear here
-                    if hasattr(actual_message_root, "method"):
-                        method_name = actual_message_root.method
-                    if hasattr(actual_message_root, "id"):
-                        message_id = actual_message_root.id
-
-            record_data_base = {
-                "method": method_name,
-                "direction": DirectionType.RECEIVED,
-                "call_type": call_type_override,
-                "message_id": message_id,
-                "is_response_or_error": is_response_or_error_flag,  # For MCPRequestCall model
-            }
-
-            try:
-                if hud_task_run_id:
-                    create_request_record(
-                        **record_data_base,
-                        status=StatusType.STARTED,
-                        start_time=start_time,
-                        # request_data might be populated if needed for HANDLE_INCOMING
-                    )
-
-                result = await orig_handle_incoming(req_or_msg)
-
-                if hud_task_run_id:
-                    create_request_record(
-                        **record_data_base,
-                        status=StatusType.COMPLETED,
-                        start_time=start_time,
-                        end_time=time.time(),
-                        duration=time.time() - start_time,
-                        # result_data might be populated if needed
-                    )
-                return result
-            except Exception as e:
-                if hud_task_run_id:
-                    create_request_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
-
-        instance._handle_incoming = handle_incoming_with_telemetry
-
-        try:
-            logger.debug("MCPInstrumentor: Entering instrumented _receive_loop section")
-            # The original logic for distinguishing request/notification from response for
-            # telemetry:
-            # Ensure streams are context managed
-            async with instance._read_stream, instance._write_stream:
-                async for message in instance._read_stream:
-                    hud_task_run_id = get_current_task_run_id()  # Get ID for each message
-
-                    if isinstance(message, Exception):
-                        await instance._handle_incoming(message)  # Will be wrapped
-                        continue
-
-                    # Ensure we are dealing with SessionMessage
-                    if (
-                        not isinstance(message, SessionMessage)
-                        or not hasattr(message, "message")
-                        or not hasattr(message.message, "root")
-                    ):
-                        logger.warning(
-                            "Unexpected message type in _receive_loop: %s", type(message)
-                        )
-                        await instance._handle_incoming(
-                            RuntimeError(f"Unknown message type: {message}")
-                        )
-                        continue
-
-                    msg_root = message.message.root
-
-                    if isinstance(msg_root, JSONRPCRequest | JSONRPCNotification):
-                        # Let the (wrapped) _handle_incoming deal with these and record telemetry
-                        await instance._handle_incoming(message)
-                    elif isinstance(msg_root, JSONRPCResponse | JSONRPCError):
-                        # This is a direct response/error, record it here
-                        logger.debug(
-                            "MCPInstrumentor: Capturing direct response/error ID: %s", msg_root.id
-                        )
-                        if hud_task_run_id:
-                            is_error = isinstance(msg_root, JSONRPCError)
-                            error_message = None
-                            error_code = None
-                            if is_error and hasattr(msg_root, "error"):
-                                error_message = getattr(msg_root.error, "message", None)
-                                error_code = str(getattr(msg_root.error, "code", None))
-
-                            create_response_record(
-                                method=f"response_to_id_{msg_root.id}",  # Consistent method naming
-                                related_request_id=msg_root.id,
-                                response_id=msg_root.id,
-                                is_error=is_error,
-                                response_data=msg_root.model_dump(exclude_none=True),
-                                error=error_message,
-                                error_type=error_code,
-                                direction=DirectionType.RECEIVED,
-                                call_type=MCPCallType.RECEIVE_RESPONSE,
-                                timestamp=time.time(),  # Add timestamp here for accuracy
-                            )
-
-                        # Original logic to pass response to waiting stream
-                        stream = instance._response_streams.pop(msg_root.id, None)
-                        if stream:
-                            await stream.send(msg_root)
-                        else:
-                            # This case should ideally be handled by _handle_incoming if it's an
-                            # unsolicited response
-                            logger.warning(
-                                "Received response/error with unknown request ID %s and no "
-                                "response stream.",
-                                msg_root.id,
-                            )
-                            # Potentially pass to _handle_incoming if that's desired for unroutable
-                            # responses
-                            await instance._handle_incoming(message)  # Let wrapped handler decide
-                    else:
-                        logger.warning(
-                            "Unknown message root type in _receive_loop: %s", type(msg_root)
-                        )
-                        await instance._handle_incoming(message)  # Let wrapped handler decide
+        """Context-aware send request wrapper."""
+        hud_task_run_id = get_current_task_run_id()
 
-            logger.debug("MCPInstrumentor: Exiting instrumented _receive_loop section")
-        except Exception as e:
-            logger.error("Error in instrumented receive loop: %s, falling back to original", e)
-            instance._handle_incoming = orig_handle_incoming
-            # Re-raising to ensure original behavior if full loop fails, or call wrapped directly
-            # For simplicity now, we just log and let finally restore.
-            raise
-        finally:
-            instance._handle_incoming = orig_handle_incoming
-            logger.debug("MCPInstrumentor: Restored original _handle_incoming")
+        if not hud_task_run_id:
+            # No active trace - pass through without instrumentation
+            return await wrapped(*args, **kwargs)
 
-    async def _send_request_telemetry_wrapper(
-        self,
-        wrapped: Callable[[Any, Any, Any, Any], Awaitable[Any]],
-        instance: Any,
-        args: Any,
-        kwargs: Any,
-    ) -> Any:
         start_time = time.time()
-        hud_task_run_id = get_current_task_run_id()
+        request = args[0] if args else None
+
+        # Extract method name from the request
         method_name = "unknown_method"
-        request_obj_data = None
-        request_id = None
-
-        if args and len(args) > 0:
-            request_session_msg = args[0]  # Assuming SessionMessage
-            if hasattr(request_session_msg, "message") and hasattr(
-                request_session_msg.message, "root"
-            ):
-                actual_request = request_session_msg.message.root
-                if hasattr(actual_request, "method"):
-                    method_name = actual_request.method
-                if hasattr(actual_request, "id"):
-                    request_id = actual_request.id
-                if hasattr(actual_request, "model_dump"):
-                    try:
-                        request_obj_data = actual_request.model_dump(exclude_none=True)
-                    except Exception as e:
-                        logger.warning(
-                            "Could not dump request data for %s: %s",
-                            method_name,
-                            e,
-                        )
+        message_id = None
+        request_data = None
 
-        record_data_base = {
-            "method": method_name,
-            "direction": DirectionType.SENT,
-            "call_type": MCPCallType.SEND_REQUEST,  # More specific type
-            "request_id": request_id,
-            "message_id": request_id,
-            "request_data": request_obj_data,
-        }
+        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:
-            if hud_task_run_id:
-                create_request_record(
-                    **record_data_base, status=StatusType.STARTED, start_time=start_time
-                )
+            # 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,
+            )
 
-            result = await wrapped(*args, **kwargs)  # result here is usually None for send_request
-
-            if hud_task_run_id:
-                # For send_request, the 'result' is the response future, not the response itself.
-                # Completion means the request was successfully sent to the transport.
-                # The actual response is captured by _receive_loop_wrapper.
-                create_request_record(
-                    **record_data_base,
-                    status=StatusType.COMPLETED,  # Meaning: successfully dispatched
-                    start_time=start_time,
-                    end_time=time.time(),
-                    duration=time.time() - start_time,
+            # 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:
-            if hud_task_run_id:
-                create_request_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__,
-                )
+            # 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 _send_notification_telemetry_wrapper(
+    async def _context_aware_send_notification_wrapper(
         self,
         wrapped: Callable[[Any, Any, Any, Any], Awaitable[Any]],
         instance: Any,
         args: Any,
         kwargs: Any,
     ) -> Any:
-        start_time = time.time()
+        """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_obj_data = None
-
-        if args and len(args) > 0:
-            notification_session_msg = args[0]  # Assuming SessionMessage
-            if hasattr(notification_session_msg, "message") and hasattr(
-                notification_session_msg.message, "root"
-            ):
-                actual_notification = notification_session_msg.message.root
-                if hasattr(actual_notification, "method"):
-                    method_name = actual_notification.method
-                if hasattr(actual_notification, "model_dump"):
-                    try:
-                        notification_obj_data = actual_notification.model_dump(exclude_none=True)
-                    except Exception as e:
-                        logger.warning(
-                            "Could not dump notification data for %s: %s",
-                            method_name,
-                            e,
-                        )
+        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,
-            "call_type": MCPCallType.SEND_NOTIFICATION,  # More specific type
-            "notification_data": notification_obj_data,
+            "notification_data": notification_data,
         }
 
         try:
-            if hud_task_run_id:
-                create_notification_record(
-                    **record_data_base, status=StatusType.STARTED, start_time=start_time
-                )
+            create_notification_record(
+                **record_data_base, status=StatusType.STARTED, start_time=start_time
+            )
 
-            result = await wrapped(*args, **kwargs)  # Usually None
+            result = await wrapped(*args, **kwargs)
 
-            if hud_task_run_id:
-                create_notification_record(
-                    **record_data_base,
-                    status=StatusType.COMPLETED,
-                    start_time=start_time,
-                    end_time=time.time(),
-                    duration=time.time() - start_time,
-                )
+            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:
-            if hud_task_run_id:
-                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__,
-                )
+            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
-
-    def record_manual_test(self, **kwargs: Any) -> bool:
-        """Record a manual test telemetry entry"""
-        hud_task_run_id = get_current_task_run_id()
-        if not hud_task_run_id:
-            logger.warning("Manual test not recorded: No active task_run_id")
-            return False
-        try:
-            record = MCPManualTestCall.create(task_run_id=hud_task_run_id, **kwargs)
-            buffer_mcp_call(record)  # buffer_mcp_call handles Pydantic model directly
-            logger.debug("Manual test recorded: %s", record.model_dump(exclude_none=True))
-            return True
-        except Exception as e:
-            logger.warning("Manual test not recorded: %s", e)
-            return False