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

hud/telemetry/_trace.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+# ruff: noqa: T201
 import asyncio
 import logging
 import time
@@ -11,61 +12,173 @@ from typing import (
     Any,
     ParamSpec,
     TypeVar,
-    overload,
 )
 
-from hud.telemetry import exporter
 from hud.telemetry.context import (
     flush_buffer,
     get_current_task_run_id,
     is_root_trace,
     set_current_task_run_id,
 )
-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()
+
+
+def _detect_agent_model() -> str | None:
+    """
+    Try to auto-detect agent model from parent frames.
+    This is a best-effort approach and may not work in all cases.
+    """
+    import sys
+
+    try:
+        # Try different frame depths (2-3 typically covers most cases)
+        for depth in range(2, 3):
+            try:
+                frame = sys._getframe(depth)
+                # Check local variables for agent objects
+                for var_value in frame.f_locals.values():
+                    # Look for objects with model_name attribute
+                    if hasattr(var_value, "model_name") and hasattr(var_value, "run"):
+                        # Likely an agent object
+                        model_name = getattr(var_value, "model_name", None)
+                        if model_name:
+                            logger.debug(
+                                "Found agent with model_name in frame %d: %s", depth, model_name
+                            )
+                            return str(model_name)
+
+                # Also check self in case we're in a method
+                if "self" in frame.f_locals:
+                    self_obj = frame.f_locals["self"]
+                    if hasattr(self_obj, "model_name"):
+                        model_name = getattr(self_obj, "model_name", None)
+                        if model_name:
+                            logger.debug(
+                                "Found agent model_name in self at frame %d: %s", depth, model_name
+                            )
+                            return str(model_name)
+
+            except (ValueError, AttributeError):
+                # Frame doesn't exist at this depth or other issues
+                continue
+
+    except Exception as e:
+        logger.debug("Agent model detection failed: %s", e)
+
+    return None
+
+
+def _print_trace_url(task_run_id: str) -> None:
+    """Print the trace URL in a colorful box."""
+    url = f"https://app.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}")
+    print(
+        f"{DIM}║{RESET}{' ' * header_padding}{header}{' ' * (box_width - len(header) - header_padding - 3)}{DIM}║{RESET}"  # noqa: E501
+    )
+    print(f"{DIM}{divider}{RESET}")
+    print(
+        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")
+
+
+def _print_trace_complete_url(task_run_id: str) -> None:
+    """Print the trace completion URL in a simple colorful format."""
+    url = f"https://app.hud.so/trace/{task_run_id}"
+
+    # ANSI color codes
+    GREEN = "\033[92m"
+    GOLD = "\033[33m"
+    RESET = "\033[0m"
+    DIM = "\033[2m"
+    BOLD = "\033[1m"
+
+    print(f"\n{GREEN}✓ Trace complete!{RESET} {DIM}View at:{RESET} {BOLD}{GOLD}{url}{RESET}\n")
 
 
 @contextmanager
-def trace(
+def trace_open(
     name: str | None = None,
+    agent_model: 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())
+
+    _print_trace_url(task_run_id)
 
     local_attributes = attributes.copy() if attributes is not None else {}
     if name is not None:
         local_attributes["trace_name"] = name
 
+    # Auto-detect agent if not explicitly provided
+    if agent_model is None:
+        agent_model = _detect_agent_model()
+
     start_time = time.time()
     logger.debug("Starting trace %s (Name: %s)", task_run_id, name if name else "Unnamed")
 
@@ -76,96 +189,157 @@ def trace(
     is_root = previous_task_id is None
     is_root_trace.set(is_root)
 
+    # Update status to initializing for root traces
+    if is_root:
+        from hud.telemetry.exporter import (
+            TaskRunStatus,
+            submit_to_worker_loop,
+            update_task_run_status,
+        )
+        from hud.telemetry.job import get_current_job_id
+
+        # Include metadata in the initial status update
+        initial_metadata = local_attributes.copy()
+        initial_metadata["is_root_trace"] = is_root
+        if agent_model:
+            initial_metadata["agent_model"] = agent_model
+
+        # Get job_id if we're in a job context
+        job_id = get_current_job_id()
+
+        coro = update_task_run_status(
+            task_run_id, TaskRunStatus.INITIALIZING, metadata=initial_metadata, job_id=job_id
+        )
+        submit_to_worker_loop(coro)
+        logger.debug("Updated task run %s status to INITIALIZING with metadata", task_run_id)
+
+    error_occurred = False
+    error_message = None
+
     try:
         yield task_run_id
+    except Exception as e:
+        error_occurred = True
+        error_message = str(e)
+        raise
     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,
-        }
+        # Update status for root traces
+        if is_root:
+            from hud.telemetry.exporter import (
+                TaskRunStatus,
+                submit_to_worker_loop,
+                update_task_run_status,
+            )
 
-        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,
+            # Include final metadata with duration
+            final_metadata = local_attributes.copy()
+
+            if error_occurred:
+                coro = update_task_run_status(
+                    task_run_id, TaskRunStatus.ERROR, error_message, metadata=final_metadata
                 )
-                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)
+                logger.debug("Updated task run %s status to ERROR: %s", task_run_id, error_message)
+            else:
+                coro = update_task_run_status(
+                    task_run_id, TaskRunStatus.COMPLETED, metadata=final_metadata
+                )
+                logger.debug("Updated task run %s status to COMPLETED with metadata", task_run_id)
 
-        set_current_task_run_id(previous_task_id)
-        is_root_trace.set(was_root)
+            # Wait for the status update to complete
+            future = submit_to_worker_loop(coro)
+            if future:
+                try:
+                    # Wait up to 5 seconds for the status update
+                    import concurrent.futures
 
-        logger.debug(
-            "Ended trace %s (Name: %s) with %d MCP call(s)",
-            task_run_id,
-            name if name else "Unnamed",
-            len(mcp_calls),
-        )
+                    future.result(timeout=5.0)
+                    logger.debug("Status update completed successfully")
+                except concurrent.futures.TimeoutError:
+                    logger.warning("Timeout waiting for status update to complete")
+                except Exception as e:
+                    logger.error("Error waiting for status update: %s", e)
 
-        logger.info("View trace at https://app.hud.so/jobs/traces/%s", task_run_id)
+        # Export any remaining records before flushing
+        if is_root:
+            from hud.telemetry.context import export_incremental
 
+            export_incremental()
 
-P = ParamSpec("P")
-R = TypeVar("R")
+        # 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)
+
+        # Restore previous context
+        set_current_task_run_id(previous_task_id)
+        is_root_trace.set(was_root)
 
+        # Log at the end
+        if is_root:
+            _print_trace_complete_url(task_run_id)
 
-def register_trace(
-    name: str | None = None, attributes: dict[str, Any] | None = None
-) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+
+@contextmanager
+def trace(
+    name: str | None = None,
+    agent_model: str | None = None,
+    attributes: dict[str, Any] | None = None,
+) -> Generator[str, None, None]:
     """
-    Decorator to wrap a synchronous or asynchronous function call
-    within a hud._telemetry.trace context.
+    Synchronous context manager that traces and blocks until telemetry is sent.
+
+    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 the trace.
-        attributes: Optional dictionary of attributes for the trace.
+        name: Optional name for this trace
+        attributes: Optional attributes for the trace
+
+    Returns:
+        The generated task run ID (UUID string) used for this trace
     """
+    with trace_open(name=name, agent_model=agent_model, attributes=attributes) as task_run_id:
+        yield task_run_id
 
-    @overload
-    def decorator(
-        func: Callable[P, Coroutine[Any, Any, R]],
-    ) -> Callable[P, Coroutine[Any, Any, R]]: ...
+    # Ensure telemetry is flushed synchronously
+    from hud import flush
 
-    @overload
-    def decorator(func: Callable[P, R]) -> Callable[P, R]: ...
+    flush()
+
+
+def trace_decorator(
+    name: str | None = None,
+    agent_model: 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, agent_model=agent_model, 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, agent_model=agent_model, 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,8 +20,12 @@ 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)
+# Track the last exported index for each task_run_id
+_GLOBAL_EXPORT_INDICES: defaultdict[str, int] = defaultdict(int)
+# Track whether we've seen a non-init request for each task_run_id
+_GLOBAL_HAS_NON_INIT_REQUEST: defaultdict[str, bool] = defaultdict(bool)
 is_root_trace: contextvars.ContextVar[bool] = contextvars.ContextVar("is_root_trace", default=False)
 
 # Maximum buffer size before automatic flush
@@ -43,6 +46,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 +55,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)
@@ -67,6 +71,48 @@ def buffer_mcp_call(record: BaseMCPCall | dict[str, Any]) -> None:
         flush_buffer(export=True)
 
 
+def export_incremental() -> list[BaseMCPCall]:
+    """
+    Export only new MCP calls since last export without clearing the buffer.
+
+    Returns:
+        The list of newly exported MCP calls
+    """
+    task_run_id = get_current_task_run_id()
+    if not task_run_id or not is_root_trace.get():
+        return []
+
+    buffer = _GLOBAL_MCP_CALL_BUFFERS.get(task_run_id, [])
+    last_exported_idx = _GLOBAL_EXPORT_INDICES.get(task_run_id, 0)
+
+    # Get only the new records since last export
+    new_records = buffer[last_exported_idx:]
+
+    if new_records:
+        # Update the export index
+        _GLOBAL_EXPORT_INDICES[task_run_id] = len(buffer)
+
+        # Trigger export
+        from hud.telemetry import exporter
+        from hud.telemetry.exporter import submit_to_worker_loop
+
+        # Get current trace attributes if available
+        attributes = {"incremental": True}
+
+        coro = exporter.export_telemetry(
+            task_run_id=task_run_id,
+            trace_attributes=attributes,
+            mcp_calls=new_records.copy(),  # Copy to avoid modification during export
+        )
+        submit_to_worker_loop(coro)
+
+        logger.debug(
+            "Incremental export: %d new MCP calls for trace %s", len(new_records), task_run_id
+        )
+
+    return new_records
+
+
 def flush_buffer(export: bool = False) -> list[BaseMCPCall]:
     """
     Clear the MCP calls buffer and return its contents.
@@ -82,11 +128,12 @@ 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, [])
+    # Clean up export index when buffer is flushed
+    _GLOBAL_EXPORT_INDICES.pop(task_run_id, None)
+    # Clean up non-init request tracking
+    _GLOBAL_HAS_NON_INIT_REQUEST.pop(task_run_id, None)
+    return buffer_for_task
 
 
 def create_request_record(
@@ -98,6 +145,31 @@ def create_request_record(
         logger.warning("No active task_run_id, request record will not be created")
         raise ValueError("No active task_run_id")
 
+    # Check if this is the first non-init request and update status
+    if is_root_trace.get() and not _GLOBAL_HAS_NON_INIT_REQUEST[task_run_id]:
+        # Common initialization method patterns
+        init_methods = {"initialize", "session/new", "init", "setup", "connect"}
+        method_lower = method.lower()
+
+        # Check if this is NOT an initialization method
+        if not any(init_pattern in method_lower for init_pattern in init_methods):
+            _GLOBAL_HAS_NON_INIT_REQUEST[task_run_id] = True
+
+            # Update status to running
+            from hud.telemetry.exporter import (
+                TaskRunStatus,
+                submit_to_worker_loop,
+                update_task_run_status,
+            )
+
+            coro = update_task_run_status(task_run_id, TaskRunStatus.RUNNING)
+            submit_to_worker_loop(coro)
+            logger.debug(
+                "Updated task run %s status to RUNNING on first non-init request: %s",
+                task_run_id,
+                method,
+            )
+
     record = MCPRequestCall(
         task_run_id=task_run_id,
         method=method,
@@ -118,16 +190,23 @@ 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,
     )
 
     buffer_mcp_call(record)
+
+    # Trigger incremental export when we receive a response
+    export_incremental()
+
     return record
 
 
@@ -149,21 +228,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)