hud-python 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

hud/otel/instrumentation.py

@@ -1,97 +1,97 @@
-"""MCP instrumentation support for HUD.
-
-This module provides functions to enable MCP OpenTelemetry instrumentation
-for automatic tracing of MCP protocol communication.
-"""
-
-from __future__ import annotations
-
-import logging
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Callable
-
-    from opentelemetry.trace import TracerProvider
-
-logger = logging.getLogger(__name__)
-
-LIFECYCLE_TOOLS = {"setup", "evaluate"}
-
-
-def install_mcp_instrumentation(provider: TracerProvider) -> None:
-    """Enable community MCP OpenTelemetry instrumentation if present.
-
-    Args:
-        provider: The TracerProvider to use for instrumentation
-    """
-    import logging
-
-    logger = logging.getLogger(__name__)
-
-    try:
-        from opentelemetry.instrumentation.mcp.instrumentation import (
-            McpInstrumentor,
-        )
-
-        # First, patch the instrumentation to handle 3-value transports correctly
-        _patch_mcp_instrumentation()
-
-        McpInstrumentor().instrument(tracer_provider=provider)
-        logger.debug("MCP instrumentation installed with fastmcp compatibility patch")
-    except ImportError:
-        logger.debug("opentelemetry-instrumentation-mcp not available, skipping")
-    except Exception as exc:
-        logger.warning("Failed to install MCP instrumentation: %s", exc)
-
-
-def _patch_mcp_instrumentation() -> None:
-    """Patch MCP instrumentation to handle 3-value transport yields correctly."""
-    from contextlib import asynccontextmanager
-
-    try:
-        from opentelemetry.instrumentation.mcp.instrumentation import McpInstrumentor
-
-        def patched_transport_wrapper(self: Any, tracer: Any) -> Callable[..., Any]:
-            @asynccontextmanager
-            async def traced_method(
-                wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
-            ) -> AsyncGenerator[Any, None]:
-                async with wrapped(*args, **kwargs) as result:
-                    # Check if we got a tuple with 3 values
-                    if isinstance(result, tuple) and len(result) == 3:
-                        read_stream, write_stream, third_value = result
-                        # Import here to avoid circular imports
-                        from opentelemetry.instrumentation.mcp.instrumentation import (
-                            InstrumentedStreamReader,
-                            InstrumentedStreamWriter,
-                        )
-
-                        yield (
-                            InstrumentedStreamReader(read_stream, tracer),
-                            InstrumentedStreamWriter(write_stream, tracer),
-                            third_value,
-                        )
-                    else:
-                        # Fall back to 2-value case
-                        read_stream, write_stream = result
-                        from opentelemetry.instrumentation.mcp.instrumentation import (
-                            InstrumentedStreamReader,
-                            InstrumentedStreamWriter,
-                        )
-
-                        yield (
-                            InstrumentedStreamReader(read_stream, tracer),
-                            InstrumentedStreamWriter(write_stream, tracer),
-                        )
-
-            return traced_method
-
-        # Apply the patch
-        McpInstrumentor._transport_wrapper = patched_transport_wrapper
-
-    except Exception as e:
-        import logging
-
-        logger = logging.getLogger(__name__)
-        logger.warning("Failed to patch MCP instrumentation: %s", e)
+"""MCP instrumentation support for HUD.
+
+This module provides functions to enable MCP OpenTelemetry instrumentation
+for automatic tracing of MCP protocol communication.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Callable
+
+    from opentelemetry.trace import TracerProvider
+
+logger = logging.getLogger(__name__)
+
+LIFECYCLE_TOOLS = {"setup", "evaluate"}
+
+
+def install_mcp_instrumentation(provider: TracerProvider) -> None:
+    """Enable community MCP OpenTelemetry instrumentation if present.
+
+    Args:
+        provider: The TracerProvider to use for instrumentation
+    """
+    import logging
+
+    logger = logging.getLogger(__name__)
+
+    try:
+        from opentelemetry.instrumentation.mcp.instrumentation import (
+            McpInstrumentor,
+        )
+
+        # First, patch the instrumentation to handle 3-value transports correctly
+        _patch_mcp_instrumentation()
+
+        McpInstrumentor().instrument(tracer_provider=provider)
+        logger.debug("MCP instrumentation installed with fastmcp compatibility patch")
+    except ImportError:
+        logger.debug("opentelemetry-instrumentation-mcp not available, skipping")
+    except Exception as exc:
+        logger.warning("Failed to install MCP instrumentation: %s", exc)
+
+
+def _patch_mcp_instrumentation() -> None:
+    """Patch MCP instrumentation to handle 3-value transport yields correctly."""
+    from contextlib import asynccontextmanager
+
+    try:
+        from opentelemetry.instrumentation.mcp.instrumentation import McpInstrumentor
+
+        def patched_transport_wrapper(self: Any, tracer: Any) -> Callable[..., Any]:
+            @asynccontextmanager
+            async def traced_method(
+                wrapped: Callable[..., Any], instance: Any, args: Any, kwargs: Any
+            ) -> AsyncGenerator[Any, None]:
+                async with wrapped(*args, **kwargs) as result:
+                    # Check if we got a tuple with 3 values
+                    if isinstance(result, tuple) and len(result) == 3:
+                        read_stream, write_stream, third_value = result
+                        # Import here to avoid circular imports
+                        from opentelemetry.instrumentation.mcp.instrumentation import (
+                            InstrumentedStreamReader,
+                            InstrumentedStreamWriter,
+                        )
+
+                        yield (
+                            InstrumentedStreamReader(read_stream, tracer),
+                            InstrumentedStreamWriter(write_stream, tracer),
+                            third_value,
+                        )
+                    else:
+                        # Fall back to 2-value case
+                        read_stream, write_stream = result
+                        from opentelemetry.instrumentation.mcp.instrumentation import (
+                            InstrumentedStreamReader,
+                            InstrumentedStreamWriter,
+                        )
+
+                        yield (
+                            InstrumentedStreamReader(read_stream, tracer),
+                            InstrumentedStreamWriter(write_stream, tracer),
+                        )
+
+            return traced_method
+
+        # Apply the patch
+        McpInstrumentor._transport_wrapper = patched_transport_wrapper
+
+    except Exception as e:
+        import logging
+
+        logger = logging.getLogger(__name__)
+        logger.warning("Failed to patch MCP instrumentation: %s", e)

hud/otel/processors.py

@@ -1,118 +1,118 @@
-from __future__ import annotations
-
-import logging
-from typing import Any
-
-from opentelemetry import baggage
-from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
-
-from .context import (
-    get_agent_steps,
-    get_base_mcp_steps,
-    get_mcp_tool_steps,
-    increment_agent_steps,
-    increment_base_mcp_steps,
-    increment_mcp_tool_steps,
-)
-
-logger = logging.getLogger(__name__)
-
-
-class HudEnrichmentProcessor(SpanProcessor):
-    """Span processor that enriches every span with HUD-specific context.
-
-    • Adds ``hud.task_run_id`` attribute if available.
-    • Adds ``hud.job_id`` attribute if available in baggage.
-    • Adds ``hud.step_count`` attribute if available in baggage.
-    """
-
-    def __init__(self) -> None:
-        # No state, everything comes from context vars
-        super().__init__()
-
-    # --- callback hooks -------------------------------------------------
-    def on_start(self, span: Span, parent_context: Any) -> None:  # type: ignore[override]
-        try:
-            # Get task_run_id from baggage in parent context
-            run_id = baggage.get_baggage("hud.task_run_id", context=parent_context)
-            if run_id and span.is_recording():
-                span.set_attribute("hud.task_run_id", str(run_id))
-
-            # Get job_id from baggage if available
-            job_id = baggage.get_baggage("hud.job_id", context=parent_context)
-            if job_id and span.is_recording():
-                span.set_attribute("hud.job_id", str(job_id))
-
-            # Check what type of step this is and increment appropriate counters
-            if span.is_recording():
-                step_type = self._get_step_type(span)
-
-                if step_type == "agent":
-                    # Increment agent steps
-                    new_agent_count = increment_agent_steps()
-                    span.set_attribute("hud.agent_steps", new_agent_count)
-                    logger.debug("Incremented agent steps to %d", new_agent_count)
-
-                elif step_type == "base_mcp":
-                    # Increment base MCP steps
-                    new_base_count = increment_base_mcp_steps()
-                    span.set_attribute("hud.base_mcp_steps", new_base_count)
-                    logger.debug("Incremented base MCP steps to %d", new_base_count)
-
-                elif step_type == "mcp_tool":
-                    # Increment both base MCP and MCP tool steps
-                    new_base_count = increment_base_mcp_steps()
-                    new_tool_count = increment_mcp_tool_steps()
-                    span.set_attribute("hud.base_mcp_steps", new_base_count)
-                    span.set_attribute("hud.mcp_tool_steps", new_tool_count)
-                    logger.debug(
-                        "Incremented MCP steps to base=%d, tool=%d", new_base_count, new_tool_count
-                    )
-
-                # Always set all current step counts on the span
-                span.set_attribute("hud.base_mcp_steps", get_base_mcp_steps())
-                span.set_attribute("hud.mcp_tool_steps", get_mcp_tool_steps())
-                span.set_attribute("hud.agent_steps", get_agent_steps())
-
-        except Exception as exc:  # defensive; never fail the tracer
-            logger.debug("HudEnrichmentProcessor.on_start error: %s", exc, exc_info=False)
-
-    def _get_step_type(self, span: Span) -> str | None:
-        """Determine what type of step this span represents.
-
-        Returns:
-            'base_mcp' for any MCP span
-            'mcp_tool' for MCP tool calls (tools/call.mcp)
-            'agent' for agent spans
-            None if not a step
-        """
-        # Check span attributes
-        attrs = span.attributes or {}
-        span_name = span.name
-
-        # Check for agent steps (instrumented with span_type="agent")
-        if attrs.get("category") == "agent":
-            return "agent"
-
-        # Check span name pattern for MCP calls
-        if span_name:
-            # tools/call.mcp is an mcp_tool step
-            if span_name == "tools/call.mcp":
-                return "mcp_tool"
-
-            # Any other .mcp suffixed span is a base MCP step
-            elif span_name.endswith(".mcp"):
-                return "base_mcp"
-
-        return None
-
-    def on_end(self, span: ReadableSpan) -> None:
-        # Nothing to do enrichment is on_start only
-        pass
-
-    # Required to fully implement abstract base, but we don't batch spans
-    def shutdown(self) -> None:  # type: ignore[override]
-        pass
-
-    def force_flush(self, timeout_millis: int | None = None) -> bool:  # type: ignore[override]
-        return True
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from opentelemetry import baggage
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+
+from .context import (
+    get_agent_steps,
+    get_base_mcp_steps,
+    get_mcp_tool_steps,
+    increment_agent_steps,
+    increment_base_mcp_steps,
+    increment_mcp_tool_steps,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class HudEnrichmentProcessor(SpanProcessor):
+    """Span processor that enriches every span with HUD-specific context.
+
+    • Adds ``hud.task_run_id`` attribute if available.
+    • Adds ``hud.job_id`` attribute if available in baggage.
+    • Adds ``hud.step_count`` attribute if available in baggage.
+    """
+
+    def __init__(self) -> None:
+        # No state, everything comes from context vars
+        super().__init__()
+
+    # --- callback hooks -------------------------------------------------
+    def on_start(self, span: Span, parent_context: Any) -> None:  # type: ignore[override]
+        try:
+            # Get task_run_id from baggage in parent context
+            run_id = baggage.get_baggage("hud.task_run_id", context=parent_context)
+            if run_id and span.is_recording():
+                span.set_attribute("hud.task_run_id", str(run_id))
+
+            # Get job_id from baggage if available
+            job_id = baggage.get_baggage("hud.job_id", context=parent_context)
+            if job_id and span.is_recording():
+                span.set_attribute("hud.job_id", str(job_id))
+
+            # Check what type of step this is and increment appropriate counters
+            if span.is_recording():
+                step_type = self._get_step_type(span)
+
+                if step_type == "agent":
+                    # Increment agent steps
+                    new_agent_count = increment_agent_steps()
+                    span.set_attribute("hud.agent_steps", new_agent_count)
+                    logger.debug("Incremented agent steps to %d", new_agent_count)
+
+                elif step_type == "base_mcp":
+                    # Increment base MCP steps
+                    new_base_count = increment_base_mcp_steps()
+                    span.set_attribute("hud.base_mcp_steps", new_base_count)
+                    logger.debug("Incremented base MCP steps to %d", new_base_count)
+
+                elif step_type == "mcp_tool":
+                    # Increment both base MCP and MCP tool steps
+                    new_base_count = increment_base_mcp_steps()
+                    new_tool_count = increment_mcp_tool_steps()
+                    span.set_attribute("hud.base_mcp_steps", new_base_count)
+                    span.set_attribute("hud.mcp_tool_steps", new_tool_count)
+                    logger.debug(
+                        "Incremented MCP steps to base=%d, tool=%d", new_base_count, new_tool_count
+                    )
+
+                # Always set all current step counts on the span
+                span.set_attribute("hud.base_mcp_steps", get_base_mcp_steps())
+                span.set_attribute("hud.mcp_tool_steps", get_mcp_tool_steps())
+                span.set_attribute("hud.agent_steps", get_agent_steps())
+
+        except Exception as exc:  # defensive; never fail the tracer
+            logger.debug("HudEnrichmentProcessor.on_start error: %s", exc, exc_info=False)
+
+    def _get_step_type(self, span: Span) -> str | None:
+        """Determine what type of step this span represents.
+
+        Returns:
+            'base_mcp' for any MCP span
+            'mcp_tool' for MCP tool calls (tools/call.mcp)
+            'agent' for agent spans
+            None if not a step
+        """
+        # Check span attributes
+        attrs = span.attributes or {}
+        span_name = span.name
+
+        # Check for agent steps (instrumented with span_type="agent")
+        if attrs.get("category") == "agent":
+            return "agent"
+
+        # Check span name pattern for MCP calls
+        if span_name:
+            # tools/call.mcp is an mcp_tool step
+            if span_name == "tools/call.mcp":
+                return "mcp_tool"
+
+            # Any other .mcp suffixed span is a base MCP step
+            elif span_name.endswith(".mcp"):
+                return "base_mcp"
+
+        return None
+
+    def on_end(self, span: ReadableSpan) -> None:
+        # Nothing to do enrichment is on_start only
+        pass
+
+    # Required to fully implement abstract base, but we don't batch spans
+    def shutdown(self) -> None:  # type: ignore[override]
+        pass
+
+    def force_flush(self, timeout_millis: int | None = None) -> bool:  # type: ignore[override]
+        return True

hud/otel/tests/__init__.py

@@ -1 +1 @@
-"""Tests for OpenTelemetry integration."""
+"""Tests for OpenTelemetry integration."""