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

hud/telemetry/_trace.py

@@ -11,7 +11,6 @@ from typing import (
     Any,
     ParamSpec,
     TypeVar,
-    overload,
 )
 
 from hud.telemetry import exporter
@@ -25,42 +24,58 @@ from hud.telemetry.exporter import submit_to_worker_loop
 from hud.telemetry.instrumentation.registry import registry
 
 if TYPE_CHECKING:
-    from collections.abc import (
-        Callable,
-        Coroutine,
-        Generator,
-    )
+    from collections.abc import Generator
 
-    from hud.telemetry.mcp_models import BaseMCPCall
 
 logger = logging.getLogger("hud.telemetry")
 T = TypeVar("T")
+P = ParamSpec("P")
+
+# Track whether telemetry has been initialized
+_telemetry_initialized = False
 
 
 def init_telemetry() -> None:
     """Initialize telemetry instrumentors and ensure worker is started if telemetry is active."""
+    global _telemetry_initialized
+    if _telemetry_initialized:
+        return
+
     registry.install_all()
     logger.info("Telemetry initialized.")
+    _telemetry_initialized = True
+
+
+def _ensure_telemetry_initialized() -> None:
+    """Ensure telemetry is initialized - called lazily by trace functions."""
+    from hud.settings import settings
+
+    if settings.telemetry_enabled and not _telemetry_initialized:
+        init_telemetry()
 
 
 @contextmanager
-def trace(
+def trace_open(
     name: str | None = None,
+    run_id: str | None = None,
     attributes: dict[str, Any] | None = None,
 ) -> Generator[str, None, None]:
     """
     Context manager for tracing a block of code.
-    The task_run_id is always generated internally as a UUID.
-    Telemetry export is handled by a background worker thread.
 
     Args:
-        attributes: Optional dictionary of attributes to associate with this trace
         name: Optional name for this trace, will be added to attributes.
+        attributes: Optional dictionary of attributes to associate with this trace
 
     Returns:
         The generated task run ID (UUID string) used for this trace
     """
-    task_run_id = str(uuid.uuid4())
+    # Lazy initialization - only initialize telemetry when trace() is actually called
+    _ensure_telemetry_initialized()
+
+    task_run_id = run_id or str(uuid.uuid4())
+
+    logger.info("See your agent live at https://app.hud.so/trace/%s", task_run_id)
 
     local_attributes = attributes.copy() if attributes is not None else {}
     if name is not None:
@@ -81,91 +96,87 @@ def trace(
     finally:
         end_time = time.time()
         duration = end_time - start_time
+        local_attributes["duration_seconds"] = duration
+        local_attributes["is_root_trace"] = is_root
 
-        mcp_calls: list[BaseMCPCall] = flush_buffer()
+        logger.debug("Finishing trace %s after %.2f seconds", task_run_id, duration)
 
-        trace_attributes_final = {
-            **local_attributes,
-            "start_time": start_time,
-            "end_time": end_time,
-            "duration": duration,
-            "is_root": is_root,
-        }
+        # Always flush the buffer for the current task
+        mcp_calls = flush_buffer(export=True)
+        logger.debug("Flushed %d MCP calls for trace %s", len(mcp_calls), task_run_id)
 
+        # Submit the telemetry payload to the worker queue
         if is_root and mcp_calls:
-            try:
-                coro_to_submit = exporter.export_telemetry(
-                    task_run_id=task_run_id,
-                    trace_attributes=trace_attributes_final,
-                    mcp_calls=mcp_calls,
-                )
-                future = submit_to_worker_loop(coro_to_submit)
-                if future:
-                    logger.debug(
-                        "Telemetry for trace %s submitted to background worker.", task_run_id
-                    )
-                else:
-                    logger.warning(
-                        "Failed to submit telemetry for trace %s to"
-                        "background worker (loop not available).",
-                        task_run_id,
-                    )
-            except Exception as e:
-                logger.warning("Failed to submit telemetry for trace %s: %s", task_run_id, e)
-
+            coro = exporter.export_telemetry(
+                task_run_id=task_run_id,
+                trace_attributes=local_attributes,
+                mcp_calls=mcp_calls,
+            )
+            submit_to_worker_loop(coro)
+
+        # Restore previous context
         set_current_task_run_id(previous_task_id)
         is_root_trace.set(was_root)
 
-        logger.debug(
-            "Ended trace %s (Name: %s) with %d MCP call(s)",
-            task_run_id,
-            name if name else "Unnamed",
-            len(mcp_calls),
-        )
+        # Log at the end
+        if is_root:
+            view_url = f"https://app.hud.so/trace/{task_run_id}"
+            logger.info("View trace at %s", view_url)
 
-        logger.info("View trace at https://app.hud.so/jobs/traces/%s", task_run_id)
 
+@contextmanager
+def trace(
+    name: str | None = None,
+    attributes: dict[str, Any] | None = None,
+) -> Generator[str, None, None]:
+    """
+    Synchronous context manager that traces and blocks until telemetry is sent.
 
-P = ParamSpec("P")
-R = TypeVar("R")
+    This is the "worry-free" option when you want to ensure telemetry is
+    sent immediately before continuing, rather than relying on background workers.
 
+    Args:
+        name: Optional name for this trace
+        attributes: Optional attributes for the trace
 
-def register_trace(
-    name: str | None = None, attributes: dict[str, Any] | None = None
-) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    Returns:
+        The generated task run ID (UUID string) used for this trace
     """
-    Decorator to wrap a synchronous or asynchronous function call
-    within a hud._telemetry.trace context.
+    with trace_open(name=name, attributes=attributes) as task_run_id:
+        yield task_run_id
 
-    Args:
-        name: Optional name for the trace.
-        attributes: Optional dictionary of attributes for the trace.
-    """
+    # Ensure telemetry is flushed synchronously
+    from hud import flush
+
+    flush()
 
-    @overload
-    def decorator(
-        func: Callable[P, Coroutine[Any, Any, R]],
-    ) -> Callable[P, Coroutine[Any, Any, R]]: ...
 
-    @overload
-    def decorator(func: Callable[P, R]) -> Callable[P, R]: ...
+def trace_decorator(
+    name: str | None = None,
+    attributes: dict[str, Any] | None = None,
+) -> Any:
+    """
+    Decorator for tracing functions.
+
+    Can be used on both sync and async functions.
+    """
 
-    def decorator(func: Callable[P, Any]) -> Callable[P, Any]:
+    def decorator(func: Any) -> Any:
         if asyncio.iscoroutinefunction(func):
 
             @wraps(func)
-            async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
-                effective_name = name if name else func.__name__
-                with trace(name=effective_name, attributes=attributes):
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                func_name = name or f"{func.__module__}.{func.__name__}"
+                with trace_open(name=func_name, attributes=attributes):
                     return await func(*args, **kwargs)
 
             return async_wrapper
         else:
 
             @wraps(func)
-            def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
-                effective_name = name if name else func.__name__
-                with trace(name=effective_name, attributes=attributes):
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                func_name = name or f"{func.__module__}.{func.__name__}"
+                with trace_open(name=func_name, attributes=attributes):
                     return func(*args, **kwargs)
 
             return sync_wrapper

hud/telemetry/context.py

@@ -8,7 +8,6 @@ from typing import Any, TypeVar
 
 from hud.telemetry.mcp_models import (
     BaseMCPCall,
-    MCPManualTestCall,
     MCPNotificationCall,
     MCPRequestCall,
     MCPResponseCall,
@@ -21,7 +20,7 @@ logger = logging.getLogger("hud.telemetry")
 current_task_run_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
     "current_task_run_id", default=None
 )
-# NEW: Global dictionary for buffering, keyed by task_run_id
+# Global dictionary for buffering, keyed by task_run_id
 _GLOBAL_MCP_CALL_BUFFERS: defaultdict[str, list[BaseMCPCall]] = defaultdict(list)
 is_root_trace: contextvars.ContextVar[bool] = contextvars.ContextVar("is_root_trace", default=False)
 
@@ -43,6 +42,7 @@ def set_current_task_run_id(task_run_id: str | None) -> None:
 
 
 def buffer_mcp_call(record: BaseMCPCall | dict[str, Any]) -> None:
+    """Buffer an MCP call record for the current trace."""
     task_run_id = get_current_task_run_id()
 
     if not task_run_id:
@@ -51,7 +51,7 @@ def buffer_mcp_call(record: BaseMCPCall | dict[str, Any]) -> None:
         )
         return
 
-    # Ensure 'record' is a Pydantic model instance from here
+    # Ensure 'record' is a Pydantic model instance
     if isinstance(record, dict):
         try:
             record_model = BaseMCPCall.from_dict(record)
@@ -82,11 +82,8 @@ def flush_buffer(export: bool = False) -> list[BaseMCPCall]:
         logger.warning("FLUSH_BUFFER: No current task_run_id. Cannot flush.")
         return []
 
-    buffer_for_task = _GLOBAL_MCP_CALL_BUFFERS.pop(
-        task_run_id, []
-    )  # Get and remove the list for this task
-
-    return buffer_for_task  # Return the flushed items
+    buffer_for_task = _GLOBAL_MCP_CALL_BUFFERS.pop(task_run_id, [])
+    return buffer_for_task
 
 
 def create_request_record(
@@ -118,10 +115,13 @@ def create_response_record(
         logger.warning("No active task_run_id, response record will not be created")
         raise ValueError("No active task_run_id")
 
+    # Default to COMPLETED status if not provided
+    if "status" not in kwargs:
+        kwargs["status"] = StatusType.COMPLETED
+
     record = MCPResponseCall(
         task_run_id=task_run_id,
         method=method,
-        status=StatusType.COMPLETED,
         related_request_id=related_request_id,
         is_error=is_error,
         **kwargs,
@@ -149,21 +149,3 @@ def create_notification_record(
     )
     buffer_mcp_call(record)
     return record
-
-
-def create_manual_test_record(**custom_data: Any) -> MCPManualTestCall | None:
-    """Create and buffer a manual test record"""
-    task_run_id = get_current_task_run_id()
-    if not task_run_id:
-        logger.warning("No active task_run_id, manual test record will not be created")
-        return None
-
-    record = MCPManualTestCall.create(task_run_id=task_run_id, **custom_data)
-    buffer_mcp_call(record)
-    return record
-
-
-def reset_context() -> None:
-    """Reset all telemetry context variables. Useful for test isolation."""
-    set_current_task_run_id(None)
-    is_root_trace.set(False)

hud/telemetry/exporter.py

@@ -6,7 +6,7 @@ import json
 import logging
 import threading
 import time
-from datetime import datetime, timezone  # For ISO timestamp conversion
+from datetime import UTC, datetime  # For ISO timestamp conversion
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
@@ -157,7 +157,7 @@ async def export_telemetry(
         actual_start_time_float = getattr(mcp_call_model, "start_time", None)
         if actual_start_time_float:
             start_ts_iso = (
-                datetime.fromtimestamp(actual_start_time_float, timezone.utc)
+                datetime.fromtimestamp(actual_start_time_float, UTC)
                 .isoformat()
                 .replace("+00:00", "Z")
             )
@@ -170,7 +170,7 @@ async def export_telemetry(
 
         if effective_end_timestamp_float:
             end_ts_iso = (
-                datetime.fromtimestamp(effective_end_timestamp_float, timezone.utc)
+                datetime.fromtimestamp(effective_end_timestamp_float, UTC)
                 .isoformat()
                 .replace("+00:00", "Z")
             )
@@ -375,8 +375,9 @@ def flush(timeout: float = 10.0) -> None:
         # This check is racy, but it's the best we can do without more complex inter-thread
         # sync for task completion. Give some time for the task to process the sentinel and
         # clear itself.
-        # Max wait for task to clear
-        attempt_timeout = time.time() + (timeout / 2 if timeout else 2.0)
+        # Max wait for task to clear - should be longer than EXPORT_INTERVAL to ensure
+        # the task has time to wake from sleep and process the sentinel
+        attempt_timeout = time.time() + (timeout / 2 if timeout else 2.0) + EXPORT_INTERVAL + 1.0
         while _export_task_async is not None and time.time() < attempt_timeout:
             time.sleep(0.1)
             # _export_task_async is set to None by _process_export_queue_async upon its exit.