hud-python 0.3.0__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
@@ -13,14 +14,12 @@ from typing import (
     TypeVar,
 )
 
-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:
@@ -54,9 +53,104 @@ def _ensure_telemetry_initialized() -> None:
         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_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]:
@@ -75,12 +169,16 @@ def trace_open(
 
     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)
+    _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")
 
@@ -91,8 +189,39 @@ def trace_open(
     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
@@ -101,32 +230,65 @@ def trace_open(
 
         logger.debug("Finishing trace %s after %.2f seconds", task_run_id, duration)
 
+        # Update status for root traces
+        if is_root:
+            from hud.telemetry.exporter import (
+                TaskRunStatus,
+                submit_to_worker_loop,
+                update_task_run_status,
+            )
+
+            # 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
+                )
+                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)
+
+            # 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
+
+                    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)
+
+        # Export any remaining records before flushing
+        if is_root:
+            from hud.telemetry.context import export_incremental
+
+            export_incremental()
+
         # 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:
-            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)
 
         # 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)
+            _print_trace_complete_url(task_run_id)
 
 
 @contextmanager
 def trace(
     name: str | None = None,
+    agent_model: str | None = None,
     attributes: dict[str, Any] | None = None,
 ) -> Generator[str, None, None]:
     """
@@ -142,7 +304,7 @@ def trace(
     Returns:
         The generated task run ID (UUID string) used for this trace
     """
-    with trace_open(name=name, attributes=attributes) as task_run_id:
+    with trace_open(name=name, agent_model=agent_model, attributes=attributes) as task_run_id:
         yield task_run_id
 
     # Ensure telemetry is flushed synchronously
@@ -153,6 +315,7 @@ def trace(
 
 def trace_decorator(
     name: str | None = None,
+    agent_model: str | None = None,
     attributes: dict[str, Any] | None = None,
 ) -> Any:
     """
@@ -167,7 +330,7 @@ def trace_decorator(
             @wraps(func)
             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):
+                with trace_open(name=func_name, agent_model=agent_model, attributes=attributes):
                     return await func(*args, **kwargs)
 
             return async_wrapper
@@ -176,7 +339,7 @@ def trace_decorator(
             @wraps(func)
             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):
+                with trace_open(name=func_name, agent_model=agent_model, attributes=attributes):
                     return func(*args, **kwargs)
 
             return sync_wrapper

hud/telemetry/context.py

@@ -22,6 +22,10 @@ current_task_run_id: contextvars.ContextVar[str | None] = contextvars.ContextVar
 )
 # 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
@@ -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.
@@ -83,6 +129,10 @@ def flush_buffer(export: bool = False) -> list[BaseMCPCall]:
         return []
 
     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
 
 
@@ -95,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,
@@ -128,6 +203,10 @@ def create_response_record(
     )
 
     buffer_mcp_call(record)
+
+    # Trigger incremental export when we receive a response
+    export_incremental()
+
     return record
 
 

hud/telemetry/exporter.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures  # For run_coroutine_threadsafe return type
+import enum
 import json
 import logging
 import threading
@@ -13,6 +14,7 @@ if TYPE_CHECKING:
     from collections.abc import Coroutine
 
 import httpx
+from pydantic import BaseModel
 
 from hud.settings import settings
 
@@ -25,6 +27,41 @@ from hud.telemetry.mcp_models import (  # MCPResponseCall for isinstance check
 
 logger = logging.getLogger("hud.telemetry")
 
+
+# --- Task Run Status Models ---
+class TaskRunStatus(enum.StrEnum):
+    INITIALIZING = "initializing"
+    RUNNING = "running"
+    EVALUATING = "evaluating"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+
+class TaskRunStatusUpdateRequest(BaseModel):
+    """Request model for updating task run status."""
+
+    status: TaskRunStatus
+    error_message: str | None = None  # Optional error message if status is ERROR
+    metadata: dict[str, Any] | None = None  # Optional metadata for context
+    job_id: str | None = None  # Optional parent job ID
+
+
+# --- Job Status Models ---
+class JobStatus(enum.StrEnum):
+    RUNNING = "running"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+
+class JobStatusUpdateRequest(BaseModel):
+    """Request model for updating job status."""
+
+    status: JobStatus
+    error_message: str | None = None  # Optional error message if status is ERROR
+    metadata: dict[str, Any] | None = None  # Optional metadata for context
+    taskset_name: str | None = None  # Optional dataset/taskset name
+
+
 # --- Worker Thread and Event Loop Management ---
 _worker_thread: threading.Thread | None = None
 _worker_loop: asyncio.AbstractEventLoop | None = None
@@ -38,7 +75,8 @@ _export_lock_async = asyncio.Lock()  # Async lock for the async queue
 _export_task_async: asyncio.Task | None = None  # Async task for processing the queue
 
 # --- Constants ---
-EXPORT_INTERVAL = 5.0  # seconds
+EXPORT_INTERVAL = 5.0  # seconds - delay between non-incremental exports
+MIN_EXPORT_INTERVAL = 0.1  # seconds - minimum delay between any exports to avoid overwhelming
 # MAX_BATCH_SIZE removed as we send one trace payload at a time
 
 
@@ -265,12 +303,19 @@ async def _process_export_queue_async() -> None:
 
             if isinstance(payload_to_process, dict):  # Ensure it's a dict before processing as such
                 await _export_trace_payload_async(payload_to_process)
+
+                # Apply appropriate delay based on export type
+                is_incremental = payload_to_process.get("attributes", {}).get("incremental", False)
+                if is_incremental:
+                    # Small delay for incremental exports to avoid overwhelming the server
+                    await asyncio.sleep(MIN_EXPORT_INTERVAL)
+                else:
+                    # Longer delay for final exports
+                    await asyncio.sleep(EXPORT_INTERVAL)
             else:
                 # Should not happen if only dicts and sentinel are queued
                 logger.warning("Unexpected item in telemetry queue: %s", type(payload_to_process))
 
-            await asyncio.sleep(EXPORT_INTERVAL)
-
     except asyncio.CancelledError:
         logger.debug("Async telemetry export processing task cancelled.")
         _export_task_async = None
@@ -340,6 +385,119 @@ async def send_telemetry_to_server(task_run_id: str, data: dict[str, Any]) -> No
         logger.exception("Error exporting telemetry for task run %s: %s", task_run_id, e)
 
 
+async def update_task_run_status(
+    task_run_id: str,
+    status: TaskRunStatus,
+    error_message: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    job_id: str | None = None,
+) -> None:
+    """Update the status of a task run."""
+    if not settings.telemetry_enabled:
+        logger.debug("Status update skipped - telemetry not enabled")
+        return
+
+    status_url = f"{settings.base_url}/v2/task_runs/{task_run_id}/status"
+
+    try:
+        async with httpx.AsyncClient() as client:
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {settings.api_key}",
+            }
+
+            request_data = TaskRunStatusUpdateRequest(
+                status=status, error_message=error_message, metadata=metadata, job_id=job_id
+            )
+
+            logger.debug(
+                "Updating status for task run %s to %s",
+                task_run_id,
+                status,
+            )
+
+            response = await client.post(
+                status_url,
+                json=request_data.model_dump(exclude_none=True),
+                headers=headers,
+                timeout=10.0,
+            )
+
+            if response.status_code >= 200 and response.status_code < 300:
+                logger.debug(
+                    "Successfully updated status for task run %s to %s",
+                    task_run_id,
+                    status,
+                )
+            else:
+                logger.warning(
+                    "Failed to update status for task run %s: HTTP %s - %s",
+                    task_run_id,
+                    response.status_code,
+                    response.text,
+                )
+    except Exception as e:
+        logger.exception("Error updating status for task run %s: %s", task_run_id, e)
+
+
+async def update_job_status(
+    job_id: str,
+    status: JobStatus,
+    error_message: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    taskset_name: str | None = None,
+) -> None:
+    """Update the status of a job."""
+    if not settings.telemetry_enabled:
+        logger.debug("Job status update skipped - telemetry not enabled")
+        return
+
+    status_url = f"{settings.base_url}/v2/jobs/{job_id}/status"
+
+    try:
+        async with httpx.AsyncClient() as client:
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {settings.api_key}",
+            }
+
+            request_data = JobStatusUpdateRequest(
+                status=status,
+                error_message=error_message,
+                metadata=metadata,
+                taskset_name=taskset_name,
+            )
+
+            logger.debug(
+                "Updating status for job %s to %s",
+                job_id,
+                status,
+            )
+
+            response = await client.post(
+                status_url,
+                json=request_data.model_dump(exclude_none=True),
+                headers=headers,
+                timeout=10.0,
+            )
+
+            if response.status_code >= 200 and response.status_code < 300:
+                logger.debug(
+                    "Successfully updated status for job %s to %s",
+                    job_id,
+                    status,
+                )
+            else:
+                logger.warning(
+                    "Failed to update status for job %s: HTTP %s - %s",
+                    job_id,
+                    response.status_code,
+                    response.text,
+                )
+    except Exception as e:
+        logger.exception("Error updating status for job %s: %s", job_id, e)
+
+
 # --- Public Shutdown Function ---
 def flush(timeout: float = 10.0) -> None:
     """Flushes pending telemetry data and stops the worker thread."""
@@ -382,9 +540,10 @@ def flush(timeout: float = 10.0) -> None:
             time.sleep(0.1)
             # _export_task_async is set to None by _process_export_queue_async upon its exit.
         if _export_task_async is not None:
-            logger.warning(
-                "Telemetry processing task did not clear itself after sentinel. May still be "
-                "running or stuck."
+            # This is often a false positive due to race conditions during shutdown
+            logger.debug(
+                "Telemetry processing task did not clear itself after sentinel. "
+                "This is normal during shutdown."
             )
         else:
             logger.debug("Telemetry processing task appears to have completed after sentinel.")