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

hud/otel/context.py

@@ -1,570 +0,0 @@
-"""OpenTelemetry context utilities for HUD telemetry.
-
-This module provides internal utilities for managing OpenTelemetry contexts.
-User-facing APIs are in hud.telemetry.
-"""
-
-from __future__ import annotations
-
-import contextvars
-import logging
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any
-
-from opentelemetry import baggage, context
-from opentelemetry import trace as otel_trace
-from opentelemetry.trace import Status, StatusCode
-
-if TYPE_CHECKING:
-    from collections.abc import Generator
-    from types import TracebackType
-
-from hud.settings import settings
-from hud.shared import make_request, make_request_sync
-from hud.utils.async_utils import fire_and_forget
-
-logger = logging.getLogger(__name__)
-
-# Context variables for task tracking
-current_task_run_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
-    "current_task_run_id", default=None
-)
-is_root_trace_var: contextvars.ContextVar[bool] = contextvars.ContextVar(
-    "is_root_trace", default=False
-)
-
-# Step counters for different types
-current_base_mcp_steps: contextvars.ContextVar[int] = contextvars.ContextVar(
-    "current_base_mcp_steps", default=0
-)
-current_mcp_tool_steps: contextvars.ContextVar[int] = contextvars.ContextVar(
-    "current_mcp_tool_steps", default=0
-)
-current_agent_steps: contextvars.ContextVar[int] = contextvars.ContextVar(
-    "current_agent_steps", default=0
-)
-
-# Keys for OpenTelemetry baggage
-TASK_RUN_ID_KEY = "hud.task_run_id"
-IS_ROOT_TRACE_KEY = "hud.is_root_trace"
-BASE_MCP_STEPS_KEY = "hud.base_mcp_steps"
-MCP_TOOL_STEPS_KEY = "hud.mcp_tool_steps"
-AGENT_STEPS_KEY = "hud.agent_steps"
-
-
-def set_current_task_run_id(task_run_id: str | None) -> contextvars.Token:
-    """Set the current task run ID."""
-    return current_task_run_id.set(task_run_id)
-
-
-def get_current_task_run_id() -> str | None:
-    """Get current task_run_id from either contextvars or OTel baggage."""
-    # First try OTel baggage
-    task_run_id = baggage.get_baggage(TASK_RUN_ID_KEY)
-    if task_run_id and isinstance(task_run_id, str):
-        return task_run_id
-
-    # Fallback to contextvars
-    return current_task_run_id.get()
-
-
-def is_root_trace() -> bool:
-    """Check if current context is a root trace."""
-    # First try OTel baggage
-    is_root = baggage.get_baggage(IS_ROOT_TRACE_KEY)
-    if isinstance(is_root, str):
-        return is_root.lower() == "true"
-
-    # Fallback to contextvars
-    return is_root_trace_var.get()
-
-
-def get_base_mcp_steps() -> int:
-    """Get current base MCP step count from either contextvars or OTel baggage."""
-    # First try OTel baggage
-    step_count = baggage.get_baggage(BASE_MCP_STEPS_KEY)
-    if step_count and isinstance(step_count, str):
-        try:
-            return int(step_count)
-        except ValueError:
-            pass
-
-    # Fallback to contextvars
-    return current_base_mcp_steps.get()
-
-
-def get_mcp_tool_steps() -> int:
-    """Get current MCP tool step count from either contextvars or OTel baggage."""
-    # First try OTel baggage
-    step_count = baggage.get_baggage(MCP_TOOL_STEPS_KEY)
-    if step_count and isinstance(step_count, str):
-        try:
-            return int(step_count)
-        except ValueError:
-            pass
-
-    # Fallback to contextvars
-    return current_mcp_tool_steps.get()
-
-
-def get_agent_steps() -> int:
-    """Get current agent step count from either contextvars or OTel baggage."""
-    # First try OTel baggage
-    step_count = baggage.get_baggage(AGENT_STEPS_KEY)
-    if step_count and isinstance(step_count, str):
-        try:
-            return int(step_count)
-        except ValueError:
-            pass
-
-    # Fallback to contextvars
-    return current_agent_steps.get()
-
-
-def increment_base_mcp_steps() -> int:
-    """Increment the base MCP step count and update baggage.
-
-    Returns:
-        The new base MCP step count after incrementing
-    """
-    current = get_base_mcp_steps()
-    new_count = current + 1
-
-    # Update contextvar
-    current_base_mcp_steps.set(new_count)
-
-    # Update baggage for propagation
-    ctx = baggage.set_baggage(BASE_MCP_STEPS_KEY, str(new_count))
-    context.attach(ctx)
-
-    # Update current span if one exists
-    span = otel_trace.get_current_span()
-    if span and span.is_recording():
-        span.set_attribute("hud.base_mcp_steps", new_count)
-
-    return new_count
-
-
-def increment_mcp_tool_steps() -> int:
-    """Increment the MCP tool step count and update baggage.
-
-    Returns:
-        The new MCP tool step count after incrementing
-    """
-    current = get_mcp_tool_steps()
-    new_count = current + 1
-
-    # Update contextvar
-    current_mcp_tool_steps.set(new_count)
-
-    # Update baggage for propagation
-    ctx = baggage.set_baggage(MCP_TOOL_STEPS_KEY, str(new_count))
-    context.attach(ctx)
-
-    # Update current span if one exists
-    span = otel_trace.get_current_span()
-    if span and span.is_recording():
-        span.set_attribute("hud.mcp_tool_steps", new_count)
-
-    return new_count
-
-
-def increment_agent_steps() -> int:
-    """Increment the agent step count and update baggage.
-
-    Returns:
-        The new agent step count after incrementing
-    """
-    current = get_agent_steps()
-    new_count = current + 1
-
-    # Update contextvar
-    current_agent_steps.set(new_count)
-
-    # Update baggage for propagation
-    ctx = baggage.set_baggage(AGENT_STEPS_KEY, str(new_count))
-    context.attach(ctx)
-
-    # Update current span if one exists
-    span = otel_trace.get_current_span()
-    if span and span.is_recording():
-        span.set_attribute("hud.agent_steps", new_count)
-
-    return new_count
-
-
-@contextmanager
-def span_context(
-    name: str,
-    attributes: dict[str, Any] | None = None,
-    kind: otel_trace.SpanKind = otel_trace.SpanKind.INTERNAL,
-) -> Generator[otel_trace.Span, None, None]:
-    """Create a child span within the current trace context.
-
-    This is a simple wrapper around OpenTelemetry's span creation that
-    ensures the span inherits the current HUD context (task_run_id, etc).
-
-    Args:
-        name: Name for the span
-        attributes: Additional attributes to add to the span
-        kind: OpenTelemetry span kind
-
-    Example:
-        with span_context("process_data", {"items": 100}) as span:
-            # Process data...
-            span.set_attribute("processed", True)
-    """
-    tracer = otel_trace.get_tracer("hud-sdk")
-
-    # Current task_run_id will be added by HudEnrichmentProcessor
-    with tracer.start_as_current_span(
-        name,
-        attributes=attributes,
-        kind=kind,
-    ) as span:
-        yield span
-
-
-async def _update_task_status_async(
-    task_run_id: str,
-    status: str,
-    job_id: str | None = None,
-    error_message: str | None = None,
-    trace_name: str | None = None,
-    task_id: str | None = None,
-) -> None:
-    """Async task status update."""
-    if not settings.telemetry_enabled:
-        return
-
-    try:
-        data: dict[str, Any] = {"status": status}
-
-        # Resolve effective job_id from explicit param, OTel baggage, or current job context
-        effective_job_id: str | None = job_id
-        if not effective_job_id:
-            bj = baggage.get_baggage("hud.job_id")
-            if isinstance(bj, str) and bj:
-                effective_job_id = bj
-        if not effective_job_id:
-            try:
-                from hud.telemetry.job import get_current_job  # Local import to avoid cycles
-
-                current_job = get_current_job()
-                if current_job:
-                    effective_job_id = current_job.id
-            except Exception:
-                effective_job_id = None
-
-        if effective_job_id:
-            data["job_id"] = effective_job_id
-        if error_message:
-            data["error_message"] = error_message
-
-        # Build metadata with trace name and step counts
-        metadata = {}
-        if trace_name:
-            metadata["trace_name"] = trace_name
-
-        # Include all three step counts in metadata
-        metadata["base_mcp_steps"] = get_base_mcp_steps()
-        metadata["mcp_tool_steps"] = get_mcp_tool_steps()
-        metadata["agent_steps"] = get_agent_steps()
-
-        if metadata:
-            data["metadata"] = metadata
-
-        if task_id:
-            data["task_id"] = task_id
-
-        await make_request(
-            method="POST",
-            url=f"{settings.hud_telemetry_url}/trace/{task_run_id}/status",
-            json=data,
-            api_key=settings.api_key,
-        )
-        logger.debug("Updated task %s status to %s", task_run_id, status)
-    except Exception as e:
-        # Suppress warnings about interpreter shutdown
-        if "interpreter shutdown" not in str(e):
-            logger.warning("Failed to update task status: %s", e)
-
-
-def _fire_and_forget_status_update(
-    task_run_id: str,
-    status: str,
-    job_id: str | None = None,
-    error_message: str | None = None,
-    trace_name: str | None = None,
-    task_id: str | None = None,
-) -> None:
-    """Fire and forget status update - works in any context including Jupyter."""
-    fire_and_forget(
-        _update_task_status_async(task_run_id, status, job_id, error_message, trace_name, task_id),
-        f"update task {task_run_id} status to {status}",
-    )
-
-
-def _update_task_status_sync(
-    task_run_id: str,
-    status: str,
-    job_id: str | None = None,
-    error_message: str | None = None,
-    trace_name: str | None = None,
-    task_id: str | None = None,
-) -> None:
-    """Synchronous task status update."""
-    if not settings.telemetry_enabled:
-        return
-
-    try:
-        data: dict[str, Any] = {"status": status}
-
-        # Resolve effective job_id from explicit param, OTel baggage, or current job context
-        effective_job_id: str | None = job_id
-        if not effective_job_id:
-            bj = baggage.get_baggage("hud.job_id")
-            if isinstance(bj, str) and bj:
-                effective_job_id = bj
-        if not effective_job_id:
-            try:
-                from hud.telemetry.job import get_current_job  # Local import to avoid cycles
-
-                current_job = get_current_job()
-                if current_job:
-                    effective_job_id = current_job.id
-            except Exception:
-                effective_job_id = None
-
-        if effective_job_id:
-            data["job_id"] = effective_job_id
-        if error_message:
-            data["error_message"] = error_message
-
-        # Build metadata with trace name and step counts
-        metadata = {}
-        if trace_name:
-            metadata["trace_name"] = trace_name
-
-        # Include all three step counts in metadata
-        metadata["base_mcp_steps"] = get_base_mcp_steps()
-        metadata["mcp_tool_steps"] = get_mcp_tool_steps()
-        metadata["agent_steps"] = get_agent_steps()
-
-        if metadata:
-            data["metadata"] = metadata
-
-        if task_id:
-            data["task_id"] = task_id
-
-        make_request_sync(
-            method="POST",
-            url=f"{settings.hud_telemetry_url}/trace/{task_run_id}/status",
-            json=data,
-            api_key=settings.api_key,
-        )
-        logger.debug("Updated task %s status to %s", task_run_id, status)
-    except Exception as e:
-        # Suppress warnings about interpreter shutdown
-        if "interpreter shutdown" not in str(e):
-            logger.warning("Failed to update task status: %s", e)
-
-
-def _print_trace_url(task_run_id: str) -> None:
-    """Print the trace URL in a colorful box."""
-    # Only print HUD URL if HUD telemetry is enabled and has API key
-    if not (settings.telemetry_enabled and settings.api_key):
-        return
-
-    url = f"https://hud.so/trace/{task_run_id}"
-    header = "🚀 See your agent live at:"
-
-    # ANSI color codes
-    DIM = "\033[90m"  # Dim/Gray for border (visible on both light and dark terminals)
-    GOLD = "\033[33m"  # Gold/Yellow for URL
-    RESET = "\033[0m"
-    BOLD = "\033[1m"
-
-    # Calculate box width based on the longest line
-    box_width = max(len(url), len(header)) + 6
-
-    # Box drawing characters
-    top_border = "╔" + "═" * (box_width - 2) + "╗"
-    bottom_border = "╚" + "═" * (box_width - 2) + "╝"
-    divider = "╟" + "─" * (box_width - 2) + "╢"
-
-    # Center the content
-    header_padding = (box_width - len(header) - 2) // 2
-    url_padding = (box_width - len(url) - 2) // 2
-
-    # Print the box
-    print(f"\n{DIM}{top_border}{RESET}")  # noqa: T201
-    print(  # noqa: T201
-        f"{DIM}║{RESET}{' ' * header_padding}{header}{' ' * (box_width - len(header) - header_padding - 3)}{DIM}║{RESET}"  # noqa: E501
-    )
-    print(f"{DIM}{divider}{RESET}")  # noqa: T201
-    print(  # noqa: T201
-        f"{DIM}║{RESET}{' ' * url_padding}{BOLD}{GOLD}{url}{RESET}{' ' * (box_width - len(url) - url_padding - 2)}{DIM}║{RESET}"  # noqa: E501
-    )
-    print(f"{DIM}{bottom_border}{RESET}\n")  # noqa: T201
-
-
-def _print_trace_complete_url(task_run_id: str, error_occurred: bool = False) -> None:
-    """Print the trace completion URL with appropriate messaging."""
-    # Only print HUD URL if HUD telemetry is enabled and has API key
-    if not (settings.telemetry_enabled and settings.api_key):
-        return
-
-    url = f"https://hud.so/trace/{task_run_id}"
-
-    # ANSI color codes
-    GREEN = "\033[92m"
-    RED = "\033[91m"
-    GOLD = "\033[33m"
-    RESET = "\033[0m"
-    DIM = "\033[2m"
-    BOLD = "\033[1m"
-
-    if error_occurred:
-        print(  # noqa: T201
-            f"\n{RED}✗ Trace errored!{RESET} {DIM}More error details available at:{RESET} {BOLD}{GOLD}{url}{RESET}\n"  # noqa: E501
-        )
-    else:
-        print(f"\n{GREEN}✓ Trace complete!{RESET} {DIM}View at:{RESET} {BOLD}{GOLD}{url}{RESET}\n")  # noqa: T201
-
-
-class trace:
-    """Internal OpenTelemetry trace context manager.
-
-    This is the implementation class. Users should use hud.trace() instead.
-    """
-
-    def __init__(
-        self,
-        task_run_id: str,
-        is_root: bool = True,
-        span_name: str = "hud.task",
-        attributes: dict[str, Any] | None = None,
-        job_id: str | None = None,
-        task_id: str | None = None,
-    ) -> None:
-        self.task_run_id = task_run_id
-        self.job_id = job_id
-        self.task_id = task_id
-        self.is_root = is_root
-        self.span_name = span_name
-        self.attributes = attributes or {}
-        self._span: otel_trace.Span | None = None
-        self._span_manager: Any | None = None
-        self._otel_token: object | None = None
-        self._task_run_token = None
-        self._root_token = None
-
-    def __enter__(self) -> str:
-        """Enter the trace context and return the task_run_id."""
-        # Set context variables
-        self._task_run_token = set_current_task_run_id(self.task_run_id)
-        self._root_token = is_root_trace_var.set(self.is_root)
-
-        # Set OpenTelemetry baggage for propagation
-        ctx = baggage.set_baggage(TASK_RUN_ID_KEY, self.task_run_id)
-        ctx = baggage.set_baggage(IS_ROOT_TRACE_KEY, str(self.is_root), context=ctx)
-        if self.job_id:
-            ctx = baggage.set_baggage("hud.job_id", self.job_id, context=ctx)
-        if self.task_id:
-            ctx = baggage.set_baggage("hud.task_id", self.task_id, context=ctx)
-        self._otel_token = context.attach(ctx)
-
-        # Start a span as current
-        tracer = otel_trace.get_tracer("hud-sdk")
-        span_attrs = {
-            "hud.task_run_id": self.task_run_id,
-            "hud.is_root_trace": self.is_root,
-            **self.attributes,
-        }
-        if self.job_id:
-            span_attrs["hud.job_id"] = self.job_id
-        if self.task_id:
-            span_attrs["hud.task_id"] = self.task_id
-
-        # Use start_as_current_span context manager
-        self._span_manager = tracer.start_as_current_span(
-            self.span_name,
-            attributes=span_attrs,
-        )
-        self._span = self._span_manager.__enter__()
-
-        # Update task status to running if root (only for HUD backend)
-        if self.is_root and settings.telemetry_enabled and settings.api_key:
-            _fire_and_forget_status_update(
-                self.task_run_id,
-                "running",
-                job_id=self.job_id,
-                trace_name=self.span_name,
-                task_id=self.task_id,
-            )
-            # Print the nice trace URL box (only if not part of a job)
-            if not self.job_id:
-                _print_trace_url(self.task_run_id)
-
-        logger.debug("Started HUD trace context for task_run_id=%s", self.task_run_id)
-        return self.task_run_id
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> None:
-        """Exit the trace context."""
-        # Update task status if root (only for HUD backend)
-        if self.is_root and settings.telemetry_enabled and settings.api_key:
-            if exc_type is not None:
-                # Use synchronous update to ensure it completes before process exit
-                _update_task_status_sync(
-                    self.task_run_id,
-                    "error",
-                    job_id=self.job_id,
-                    error_message=str(exc_val),
-                    trace_name=self.span_name,
-                    task_id=self.task_id,
-                )
-                # Print error completion message (only if not part of a job)
-                if not self.job_id:
-                    _print_trace_complete_url(self.task_run_id, error_occurred=True)
-            else:
-                # Use synchronous update to ensure it completes before process exit
-                _update_task_status_sync(
-                    self.task_run_id,
-                    "completed",
-                    job_id=self.job_id,
-                    trace_name=self.span_name,
-                    task_id=self.task_id,
-                )
-                # Print success completion message (only if not part of a job)
-                if not self.job_id:
-                    _print_trace_complete_url(self.task_run_id, error_occurred=False)
-
-        # End the span
-        if self._span and self._span_manager is not None:
-            if exc_type is not None and exc_val is not None:
-                self._span.record_exception(exc_val)
-                self._span.set_status(Status(StatusCode.ERROR, str(exc_val)))
-            else:
-                self._span.set_status(Status(StatusCode.OK))
-            self._span_manager.__exit__(exc_type, exc_val, exc_tb)
-
-        # Detach OpenTelemetry context
-        if self._otel_token is not None:
-            try:
-                context.detach(self._otel_token)  # type: ignore[arg-type]
-            except Exception:
-                logger.warning("Failed to detach OpenTelemetry context")
-
-        # Reset context variables
-        if self._task_run_token is not None:
-            current_task_run_id.reset(self._task_run_token)  # type: ignore
-        if self._root_token is not None:
-            is_root_trace_var.reset(self._root_token)  # type: ignore
-
-        logger.debug("Ended HUD trace context for task_run_id=%s", self.task_run_id)