hud-python 0.4.45__py3-none-any.whl → 0.5.1__py3-none-any.whl

hud/otel/exporters.py

@@ -1,369 +0,0 @@
-"""Custom OpenTelemetry exporter that sends spans to the existing HUD telemetry
-HTTP endpoint (/trace/<id>/telemetry-upload).
-
-The exporter groups spans by ``hud.task_run_id`` baggage / attribute so we keep
-exactly the same semantics the old async worker in ``hud.telemetry.exporter``
-implemented.
-
-This exporter is *synchronous* (derives from :class:`SpanExporter`).  We rely on
-``hud.shared.make_request_sync`` which already contains retry & auth logic.
-"""
-
-from __future__ import annotations
-
-import contextlib
-import json
-import logging
-import time
-from collections import defaultdict
-from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Any
-
-from mcp.types import ClientRequest, ServerResult
-from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
-from pydantic import BaseModel, ConfigDict, Field
-
-from hud.shared import make_request_sync
-from hud.types import TraceStep as HudSpanAttributes
-
-if TYPE_CHECKING:
-    from opentelemetry.sdk.trace import ReadableSpan
-
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# Models
-# ---------------------------------------------------------------------------
-
-
-class HudSpan(BaseModel):
-    """A telemetry span ready for export."""
-
-    name: str
-    trace_id: str = Field(pattern=r"^[0-9a-fA-F]{32}$")
-    span_id: str = Field(pattern=r"^[0-9a-fA-F]{16}$")
-    parent_span_id: str | None = Field(None, pattern=r"^[0-9a-fA-F]{16}$")
-
-    start_time: str  # ISO format
-    end_time: str  # ISO format
-
-    status_code: str  # "UNSET", "OK", "ERROR"
-    status_message: str | None = None
-
-    attributes: HudSpanAttributes
-    exceptions: list[dict[str, Any]] | None = None
-
-    model_config = ConfigDict(extra="forbid")
-
-
-def extract_span_attributes(
-    attrs: dict[str, Any], method_name: str | None = None, span_name: str | None = None
-) -> HudSpanAttributes:
-    """Extract and parse span attributes into typed model.
-
-    This handles:
-    - Detecting span type (MCP vs Agent)
-    - Renaming verbose OpenTelemetry semantic conventions
-    - Parsing JSON strings to MCP types
-    """
-    # Start with core attributes - map to TraceStep field names
-    result_attrs = {
-        "task_run_id": attrs.get(
-            "hud.task_run_id"
-        ),  # TraceStep expects task_run_id, not hud.task_run_id
-        "job_id": attrs.get("hud.job_id"),  # TraceStep expects job_id, not hud.job_id
-        "type": attrs.get("span.kind", "CLIENT"),  # TraceStep expects type, not span.kind
-    }
-
-    # Determine span type based on presence of agent or MCP attributes
-    # Note: The input attrs might already have "category" set
-    existing_category = attrs.get("category")
-
-    if existing_category:
-        # Use the explicit category if provided
-        result_attrs["category"] = existing_category
-    elif span_name and span_name.startswith("agent."):
-        # Legacy support for spans named "agent.*"
-        result_attrs["category"] = "agent"
-    else:
-        result_attrs["category"] = "mcp"  # Default to MCP
-
-    # No special processing needed for different categories
-    # The backend will handle them based on the category field
-
-    # Add method_name and request_id for MCP spans
-    if result_attrs["category"] == "mcp":
-        if method_name:
-            result_attrs["method_name"] = method_name
-        # Check for request_id with and without semconv_ai prefix
-        request_id = attrs.get("semconv_ai.mcp.request_id") or attrs.get("mcp.request.id")
-        if request_id:
-            result_attrs["request_id"] = request_id
-
-    # Parse input/output - check both with and without semconv_ai prefix
-    input_str = attrs.get("semconv_ai.traceloop.entity.input") or attrs.get(
-        "traceloop.entity.input"
-    )
-    output_str = attrs.get("semconv_ai.traceloop.entity.output") or attrs.get(
-        "traceloop.entity.output"
-    )
-
-    logger.debug(
-        "Category: %s, has input: %s, has output: %s",
-        result_attrs.get("category"),
-        bool(input_str),
-        bool(output_str),
-    )
-
-    # Check for direct request/result attributes first
-    if "request" in attrs and not result_attrs.get("request"):
-        req = attrs["request"]
-        if isinstance(req, str):
-            with contextlib.suppress(json.JSONDecodeError):
-                req = json.loads(req)
-        result_attrs["request"] = req
-
-    if "result" in attrs and not result_attrs.get("result"):
-        res = attrs["result"]
-        if isinstance(res, str):
-            with contextlib.suppress(json.JSONDecodeError):
-                res = json.loads(res)
-        result_attrs["result"] = res
-
-    # Process input/output from MCP instrumentation
-    if input_str and not result_attrs.get("request"):
-        try:
-            input_data = json.loads(input_str) if isinstance(input_str, str) else input_str
-
-            # For MCP category, try to parse as ClientRequest to extract the root
-            if result_attrs["category"] == "mcp" and isinstance(input_data, dict):
-                try:
-                    if "method" in input_data and "params" in input_data:
-                        client_request = ClientRequest.model_validate(input_data)
-                        result_attrs["request"] = client_request.root
-                    else:
-                        result_attrs["request"] = input_data
-                except Exception:
-                    result_attrs["request"] = input_data
-            else:
-                # For all other categories, just store the data
-                result_attrs["request"] = input_data
-        except Exception as e:
-            logger.debug("Failed to parse request JSON: %s", e)
-
-    if output_str and not result_attrs.get("result"):
-        try:
-            output_data = json.loads(output_str) if isinstance(output_str, str) else output_str
-
-            # For MCP category, try to parse as ServerResult to extract the root
-            if result_attrs["category"] == "mcp" and isinstance(output_data, dict):
-                # Check for error
-                if "error" in output_data:
-                    result_attrs["mcp_error"] = True
-                try:
-                    server_result = ServerResult.model_validate(output_data)
-                    result_attrs["result"] = server_result.root
-                    # Check for isError in the result
-                    if getattr(server_result.root, "isError", False):
-                        result_attrs["mcp_error"] = True
-                except Exception:
-                    result_attrs["result"] = output_data
-            else:
-                # For all other categories, just store the data
-                result_attrs["result"] = output_data
-        except Exception as e:
-            logger.debug("Failed to parse result JSON: %s", e)
-
-    # Don't include the verbose attributes or ones we've already processed
-    exclude_keys = {
-        "hud.task_run_id",
-        "hud.job_id",
-        "span.kind",
-        "semconv_ai.mcp.method_name",
-        "mcp.method.name",  # Also exclude non-prefixed version
-        "semconv_ai.mcp.request_id",
-        "mcp.request.id",  # Also exclude non-prefixed version
-        "semconv_ai.traceloop.entity.input",
-        "semconv_ai.traceloop.entity.output",
-        "traceloop.entity.input",  # Also exclude non-prefixed versions
-        "traceloop.entity.output",
-        "mcp_request",  # Exclude to prevent overwriting parsed values
-        "mcp_result",  # Exclude to prevent overwriting parsed values
-        "request",  # Exclude to prevent overwriting parsed values
-        "result",  # Exclude to prevent overwriting parsed values
-        "category",  # Already handled above
-    }
-
-    # Add any extra attributes
-    for key, value in attrs.items():
-        if key not in exclude_keys:
-            result_attrs[key] = value  # noqa: PERF403
-
-    logger.debug(
-        """Final result_attrs before creating HudSpanAttributes:
-        request=%s,
-        result=%s""",
-        result_attrs.get("request"),
-        result_attrs.get("result"),
-    )
-    return HudSpanAttributes(**result_attrs)
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _ts_ns_to_iso(ts_ns: int) -> str:
-    """Convert a ``Span`` timestamp (nanoseconds) to ISO-8601 string."""
-    # OpenTelemetry times are epoch nanoseconds
-    dt = datetime.fromtimestamp(ts_ns / 1_000_000_000, tz=UTC)
-    return dt.isoformat().replace("+00:00", "Z")
-
-
-def _span_to_dict(span: ReadableSpan) -> dict[str, Any]:
-    """Convert an OpenTelemetry span to a dict using typed models."""
-
-    attrs = dict(span.attributes or {})
-
-    # Extract method name from span name if not in attributes
-    # Check both with and without semconv_ai prefix
-    raw_method = attrs.get("semconv_ai.mcp.method_name") or attrs.get("mcp.method.name")
-    method_name: str | None = None
-    if isinstance(raw_method, str):
-        method_name = raw_method
-    if method_name is None and isinstance(span.name, str) and span.name.endswith(".mcp"):
-        method_name = span.name[:-4]  # Remove .mcp suffix
-
-    # Create typed attributes
-    typed_attrs = extract_span_attributes(attrs, method_name, str(span.name))
-
-    # Record span kind as extra attribute (TraceStep allows extras)
-    try:
-        typed_attrs.span_kind = span.kind.name  # type: ignore[attr-defined]
-    except Exception:
-        logger.warning("Failed to set span kind attribute")
-
-    # Build typed span
-    # Guard context/parent/timestamps
-    context = getattr(span, "context", None)
-    trace_id_hex = (
-        format(context.trace_id, "032x") if context and hasattr(context, "trace_id") else "0" * 32
-    )
-    span_id_hex = (
-        format(context.span_id, "016x") if context and hasattr(context, "span_id") else "0" * 16
-    )
-    parent = getattr(span, "parent", None)
-    parent_id_hex = (
-        format(parent.span_id, "016x") if parent and hasattr(parent, "span_id") else None
-    )
-    start_ns = span.start_time or 0
-    end_ns = span.end_time or start_ns
-
-    typed_span = HudSpan(
-        name=span.name,
-        trace_id=trace_id_hex,
-        span_id=span_id_hex,
-        parent_span_id=parent_id_hex,
-        start_time=_ts_ns_to_iso(int(start_ns)),
-        end_time=_ts_ns_to_iso(int(end_ns)),
-        status_code=span.status.status_code.name if span.status else "UNSET",
-        status_message=span.status.description if span.status else None,
-        attributes=typed_attrs,
-        exceptions=None,
-    )
-
-    # Add error information if present
-    if span.events:
-        exceptions = []
-        exceptions = [
-            {
-                "timestamp": _ts_ns_to_iso(event.timestamp),
-                "attributes": dict(event.attributes or {}),
-            }
-            for event in span.events
-        ]
-        if exceptions:
-            typed_span.exceptions = exceptions
-
-    # Convert to dict for export
-    return typed_span.model_dump(mode="json", by_alias=True, exclude_none=True)
-
-
-# ---------------------------------------------------------------------------
-# Exporter
-# ---------------------------------------------------------------------------
-
-
-class HudSpanExporter(SpanExporter):
-    """Exporter that forwards spans to HUD backend using existing endpoint."""
-
-    def __init__(self, *, telemetry_url: str, api_key: str) -> None:
-        super().__init__()
-        self._telemetry_url = telemetry_url.rstrip("/")
-        self._api_key = api_key
-
-    # ------------------------------------------------------------------
-    # Core API
-    # ------------------------------------------------------------------
-    def export(self, spans: list[ReadableSpan]) -> SpanExportResult:  # type: ignore[override]
-        if not spans:
-            return SpanExportResult.SUCCESS
-
-        # Group spans by hud.task_run_id attribute
-        grouped: dict[str, list[ReadableSpan]] = defaultdict(list)
-        for span in spans:
-            run_id = span.attributes.get("hud.task_run_id") if span.attributes else None
-            if not run_id:
-                # Skip spans that are outside HUD traces
-                continue
-            grouped[str(run_id)].append(span)
-
-        # Send each group synchronously (retry inside make_request_sync)
-        for run_id, span_batch in grouped.items():
-            try:
-                url = f"{self._telemetry_url}/trace/{run_id}/telemetry-upload"
-                telemetry_spans = [_span_to_dict(s) for s in span_batch]
-                # Include current step count in metadata
-                metadata = {}
-                # Get the HIGHEST step count from the batch (most recent)
-                step_count = 0
-                for span in span_batch:
-                    if span.attributes and "hud.step_count" in span.attributes:
-                        current_step = span.attributes["hud.step_count"]
-                        if isinstance(current_step, int) and current_step > step_count:
-                            step_count = current_step
-
-                payload = {
-                    "metadata": metadata,
-                    "telemetry": telemetry_spans,
-                }
-
-                # Only include step_count if we found any steps
-                if step_count > 0:
-                    payload["step_count"] = step_count
-
-                logger.debug("HUD exporter sending %d spans to %s", len(span_batch), url)
-                make_request_sync(
-                    method="POST",
-                    url=url,
-                    json=payload,
-                    api_key=self._api_key,
-                )
-            except Exception as exc:
-                logger.exception("HUD exporter failed to send spans for task %s: %s", run_id, exc)
-                # If *any* group fails we return FAILURE so the OTEL SDK can retry
-                return SpanExportResult.FAILURE
-
-        return SpanExportResult.SUCCESS
-
-    def shutdown(self) -> None:  # type: ignore[override]
-        # Nothing to cleanup, httpx handled inside make_request_sync
-        pass
-
-    def force_flush(self, timeout_millis: int | None = None) -> bool:  # type: ignore[override]
-        if timeout_millis:
-            time.sleep(timeout_millis / 1000)
-        # Synchronous export, nothing buffered here
-        return True

hud/otel/instrumentation.py

@@ -1,135 +0,0 @@
-"""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__)
-
-
-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:
-        # First, patch the _instruments to use our fork
-        import opentelemetry.instrumentation.mcp.instrumentation as mcp_inst
-
-        mcp_inst._instruments = ("hud-mcp-python-sdk >= 3.13.1",)
-
-        from opentelemetry.instrumentation.mcp.instrumentation import (
-            McpInstrumentor,
-        )
-
-        # Then, 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
-
-        # First, patch the get_error_type function to handle invalid HTTP status codes
-        _patch_get_error_type()
-
-        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)
-
-
-def _patch_get_error_type() -> None:
-    """Patch get_error_type to handle invalid HTTP status codes gracefully."""
-    import re
-    from http import HTTPStatus
-
-    try:
-        import opentelemetry.instrumentation.mcp.instrumentation as mcp_inst
-
-        def patched_get_error_type(error_message: str) -> str | None:
-            """Extract HTTP status from error message, handling invalid codes."""
-            if not isinstance(error_message, str):
-                return None
-            match = re.search(r"\b(4\d{2}|5\d{2})\b", error_message)
-            if match:
-                num = int(match.group())
-                try:
-                    # Only return if it's a valid HTTPStatus
-                    if 400 <= num <= 599:
-                        return HTTPStatus(num).name
-                except ValueError:
-                    # Not a valid HTTP status code
-                    logger.debug("Ignoring invalid HTTP status code: %s", num)
-            return None
-
-        # Apply the patch
-        mcp_inst.get_error_type = patched_get_error_type
-        logger.debug("Patched get_error_type to handle invalid HTTP status codes")
-
-    except Exception as e:
-        logger.warning("Failed to patch get_error_type: %s", e)

hud/otel/processors.py

@@ -1,121 +0,0 @@
-from __future__ import annotations
-
-import logging
-import time
-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]
-        if timeout_millis:
-            time.sleep(timeout_millis / 1000)
-        return True

hud/otel/tests/__init__.py

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