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

hud/telemetry/job.py

@@ -0,0 +1,141 @@
+"""Job context manager for grouping related traces."""
+
+from __future__ import annotations
+
+import logging
+import sys
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any
+
+from hud.telemetry.exporter import JobStatus, submit_to_worker_loop, update_job_status
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+    from typing import Self
+
+logger = logging.getLogger("hud.telemetry")
+
+# Context variables for current job
+current_job_id: ContextVar[str | None] = ContextVar("current_job_id", default=None)
+current_job_name: ContextVar[str | None] = ContextVar("current_job_name", default=None)
+
+
+class JobContext:
+    """Context manager for grouping traces under a job."""
+
+    def __init__(
+        self, name: str, taskset_name: str | None = None, metadata: dict[str, Any] | None = None
+    ) -> None:
+        self.id = str(uuid.uuid4())
+        self.name = name
+        self.metadata = metadata or {}
+        self.taskset_name: str | None = taskset_name
+
+    def __enter__(self) -> Self:
+        # Auto-detect dataset
+        if self.taskset_name is None:
+            self._detect_dataset()
+
+        # Set context variables
+        current_job_id.set(self.id)
+        current_job_name.set(self.name)
+
+        # Send initial status
+        job_metadata = {**self.metadata}
+        coro = update_job_status(
+            self.id, JobStatus.RUNNING, metadata=job_metadata, taskset_name=self.taskset_name
+        )
+        submit_to_worker_loop(coro)
+
+        logger.info("Started job %s (ID: %s)", self.name, self.id)
+        return self
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
+    ) -> None:
+        # Determine final status
+        if exc_type is not None:
+            # Job failed with exception
+            error_msg = f"{exc_type.__name__}: {exc_val}"
+            coro = update_job_status(
+                self.id, JobStatus.ERROR, error_message=error_msg, taskset_name=self.taskset_name
+            )
+        else:
+            # Job completed successfully
+            coro = update_job_status(self.id, JobStatus.COMPLETED, taskset_name=self.taskset_name)
+
+        submit_to_worker_loop(coro)
+
+        # Clear context
+        current_job_id.set(None)
+        current_job_name.set(None)
+
+        status = "failed" if exc_type else "completed"
+        logger.info("Job %s %s", self.name, status)
+
+    def _detect_dataset(self) -> None:
+        """Auto-detect HuggingFace dataset in parent scope."""
+        try:
+            # Check frames 2 and 3 (with statement and parent scope)
+            for frame_depth in [2, 3]:
+                try:
+                    frame = sys._getframe(frame_depth)
+
+                    # Search for Dataset objects
+                    for var_value in frame.f_locals.values():
+                        if hasattr(var_value, "info") and hasattr(var_value.info, "builder_name"):
+                            self.taskset_name = var_value.info.builder_name
+                            logger.debug(
+                                "Auto-detected dataset at frame %d: %s",
+                                frame_depth,
+                                self.taskset_name,
+                            )
+                            return
+                        elif hasattr(var_value, "builder_name"):
+                            # Older dataset format
+                            self.taskset_name = var_value.builder_name
+                            logger.debug(
+                                "Auto-detected dataset at frame %d: %s",
+                                frame_depth,
+                                self.taskset_name,
+                            )
+                            return
+                except ValueError:
+                    # Frame doesn't exist
+                    continue
+        except Exception as e:
+            logger.debug("Dataset auto-detection failed: %s", e)
+
+
+@contextmanager
+def job(
+    name: str, taskset_name: str | None = None, metadata: dict[str, Any] | None = None
+) -> Generator[JobContext, None, None]:
+    """
+    Create a job context for grouping related traces.
+
+    Args:
+        name: Name for the job
+        metadata: Optional metadata to include with the job
+
+    Example:
+        with hud.job("evaluation_run") as job:
+            for task in tasks:
+                with hud.trace(f"task_{task.id}"):
+                    # Trace automatically includes job_id
+                    result = await agent.run(task)
+    """
+    with JobContext(name, taskset_name, metadata) as ctx:
+        yield ctx
+
+
+def get_current_job_id() -> str | None:
+    """Get the current job ID if inside a job context."""
+    return current_job_id.get()
+
+
+def get_current_job_name() -> str | None:
+    """Get the current job name if inside a job context."""
+    return current_job_name.get()

hud/telemetry/mcp_models.py

@@ -28,16 +28,14 @@ class StatusType(str, Enum):
 
 
 class MCPCallType(str, Enum):
-    """Known MCP call types"""
+    """Enum for different types of MCP calls in telemetry."""
 
-    SEND_REQUEST = "mcp.shared.session.send_request"
-    SEND_NOTIFICATION = "mcp.shared.session.send_notification"
-    RECEIVE_RESPONSE = "mcp.shared.session.receive_response"
-    RECEIVE_REQUEST = "mcp.shared.session.receive_request"
-    STREAM_READ = "mcp.stream.read"
-    STREAM_WRITE = "mcp.stream.write"
-    HANDLE_INCOMING = "mcp.handle_incoming"
-    MANUAL_TEST = "manual.test"
+    # Requests and Notifications
+    SEND_REQUEST = "mcp.send_request"
+    SEND_NOTIFICATION = "mcp.send_notification"
+
+    # Responses
+    RECEIVE_RESPONSE = "mcp.receive_response"
 
 
 class BaseMCPCall(BaseModel):
@@ -87,6 +85,8 @@ class MCPRequestCall(BaseMCPCall):
     duration: float | None = None
     request_id: str | int | None = None
     request_data: dict[str, Any] | None = None
+    error: str | None = None
+    error_type: str | None = None
 
     @classmethod
     def from_jsonrpc_request(
@@ -191,6 +191,8 @@ class MCPNotificationCall(BaseMCPCall):
     end_time: float | None = None
     duration: float | None = None
     notification_data: dict[str, Any] | None = None
+    error: str | None = None
+    error_type: str | None = None
 
     @classmethod
     def from_jsonrpc_notification(
@@ -230,69 +232,6 @@ class MCPNotificationCall(BaseMCPCall):
         return None
 
 
-class MCPStreamEvent(BaseMCPCall):
-    """Record for an MCP stream event (read or write)"""
-
-    stream_event: bool = True
-    event_type: str = Field(..., description="Type of stream event: read or write")
-    item_type: str | None = None
-    is_response_or_error: bool = False
-    message_data: dict[str, Any] | None = None
-
-    @classmethod
-    def from_session_message(
-        cls, message: SessionMessage, task_run_id: str, event_type: str, **kwargs: Any
-    ) -> MCPStreamEvent:
-        """Create telemetry record for a stream event"""
-        method_name = "unknown_stream_operation"
-        is_response = False
-        item_type = "unknown"
-        message_data = None
-
-        if hasattr(message, "message") and hasattr(message.message, "root"):
-            msg_root = message.message.root
-            item_type = type(msg_root).__name__
-            message_data = msg_root.model_dump(exclude_none=True)
-
-            # Check type first before accessing attributes
-            if isinstance(msg_root, JSONRPCRequest | JSONRPCNotification) and hasattr(
-                msg_root, "method"
-            ):
-                method_name = msg_root.method
-            elif isinstance(msg_root, JSONRPCResponse | JSONRPCError) and hasattr(msg_root, "id"):
-                method_name = f"response_to_id_{msg_root.id}"
-                is_response = True
-
-        return cls(
-            task_run_id=task_run_id,
-            status=StatusType.COMPLETED,
-            method=method_name,
-            event_type=event_type,
-            item_type=item_type,
-            is_response_or_error=is_response,
-            message_data=message_data,
-            timestamp=datetime.now().timestamp(),
-            **kwargs,
-        )
-
-
-class MCPManualTestCall(BaseMCPCall):
-    """Record for a manual test record"""
-
-    call_type: str = MCPCallType.MANUAL_TEST
-    custom_data: dict[str, Any] = Field(default_factory=dict)
-
-    @classmethod
-    def create(cls, task_run_id: str, **custom_data: Any) -> MCPManualTestCall:
-        """Create a manual test record with custom data"""
-        return cls(
-            task_run_id=task_run_id,
-            status=StatusType.COMPLETED,
-            custom_data=custom_data,
-            timestamp=datetime.now().timestamp(),
-        )
-
-
 class MCPTelemetryRecord(BaseModel):
     """Container for a set of related MCP telemetry records"""
 
@@ -320,9 +259,9 @@ class MCPTelemetryRecord(BaseModel):
 
 
 class TrajectoryStep(BaseModel):
-    """Model representing a single step in a trajectory, for export."""
+    """Model for telemetry export format."""
 
-    type: str = Field(default="mcp-step")  # Default for MCP calls
+    type: str = Field(default="mcp-step")
     observation_url: str | None = None
     observation_text: str | None = None
     actions: list[dict[str, Any]] = Field(default_factory=list)

hud/telemetry/tests/test_context.py

@@ -64,21 +64,24 @@ class TestRootTraceContext:
 class TestMCPCallBuffer:
     """Test MCP call buffer management."""
 
-    def setUp(self):
+    def reset_context(self):
         """Clear buffer before each test."""
         # Flush any existing calls and reset context
+        set_current_task_run_id(None)
+        # Clear any existing buffers by setting a temporary task ID and flushing
+        set_current_task_run_id("temp-cleanup")
         flush_buffer()
         set_current_task_run_id(None)
 
     def test_flush_buffer_empty(self):
         """Test flushing empty buffer."""
-        self.setUp()
+        self.reset_context()
         result = flush_buffer()
         assert result == []
 
     def test_add_and_flush_mcp_call(self):
         """Test adding and flushing MCP calls."""
-        self.setUp()
+        self.reset_context()
 
         # Set active task run ID
         set_current_task_run_id("test-task")
@@ -101,7 +104,7 @@ class TestMCPCallBuffer:
 
     def test_add_multiple_mcp_calls(self):
         """Test adding multiple MCP calls."""
-        self.setUp()
+        self.reset_context()
 
         # Set active task run ID
         set_current_task_run_id("test-task")
@@ -122,7 +125,7 @@ class TestMCPCallBuffer:
 
     def test_buffer_isolation_per_task(self):
         """Test that MCP call buffers contain all calls regardless of task ID."""
-        self.setUp()
+        self.reset_context()
 
         # Set task run ID 1
         set_current_task_run_id("task-1")
@@ -150,7 +153,7 @@ class TestMCPCallBuffer:
 
     def test_buffer_mcp_call_without_task_id(self):
         """Test adding MCP call when no task run ID is set."""
-        self.setUp()
+        self.reset_context()
         set_current_task_run_id(None)
 
         mock_call = MagicMock(spec=BaseMCPCall)

hud/telemetry/tests/test_trace.py

@@ -8,21 +8,25 @@ import pytest
 
 from hud.telemetry._trace import (
     init_telemetry,
-    register_trace,
     trace,
+    trace_decorator,
+    trace_open,
 )
 from hud.telemetry.context import get_current_task_run_id as actual_get_current_task_run_id
 from hud.telemetry.context import is_root_trace as actual_is_root_trace
-from hud.telemetry.context import reset_context
 from hud.telemetry.context import set_current_task_run_id as actual_set_current_task_run_id
 
 
 @pytest.fixture(autouse=True)
 def reset_telemetry_context_fixture():
     """Ensures telemetry context is reset before and after each test in this file."""
-    reset_context()
+    # Reset context before test
+    actual_set_current_task_run_id(None)
+    actual_is_root_trace.set(False)
     yield
-    reset_context()
+    # Reset context after test
+    actual_set_current_task_run_id(None)
+    actual_is_root_trace.set(False)
 
 
 class TestInitTelemetry:
@@ -44,12 +48,12 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         initial_root_state = actual_is_root_trace.get()
 
-        with trace() as task_run_id:
+        with trace_open() as task_run_id:
             assert isinstance(task_run_id, str)
             uuid.UUID(task_run_id)
             assert actual_get_current_task_run_id() == task_run_id
@@ -58,7 +62,8 @@ class TestTrace:
         assert actual_get_current_task_run_id() is None
         assert actual_is_root_trace.get() == initial_root_state
         mock_flush.assert_called_once()
-        mock_submit_loop.assert_not_called()
+        # submit_to_worker_loop is now called for status updates
+        assert mock_submit_loop.call_count == 2  # INITIALIZING and COMPLETED
 
     def test_trace_with_name_and_attributes(self, mocker):
         """Test trace with name and attributes, checking they are passed on."""
@@ -67,17 +72,18 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         trace_name = "test_trace_with_data"
         attrs = {"key": "value", "number": 42}
 
-        with trace(name=trace_name, attributes=attrs) as task_run_id:
+        with trace_open(name=trace_name, attributes=attrs) as task_run_id:
             assert isinstance(task_run_id, str)
 
         mock_flush.assert_called_once()
-        mock_submit_loop.assert_called_once()
+        # submit_to_worker_loop is now called for status updates
+        assert mock_submit_loop.call_count == 2  # INITIALIZING and COMPLETED
 
     @pytest.mark.asyncio
     async def test_trace_with_mcp_calls_exports(self, mocker):
@@ -87,36 +93,32 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=mock_mcp_calls, autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         async def mock_export(*args, **kwargs):
             return None
 
-        mock_export_actual_coro = mocker.patch(
-            "hud.telemetry._trace.exporter.export_telemetry",
+        mocker.patch(
+            "hud.telemetry.exporter.export_telemetry",
             side_effect=mock_export,
         )
 
         test_attrs = {"custom_attr": "test_val"}
         test_name = "mcp_export_test"
 
-        with trace(name=test_name, attributes=test_attrs) as task_run_id:
+        with trace_open(name=test_name, attributes=test_attrs) as task_run_id:
             pass
 
         mock_flush.assert_called_once()
-        mock_submit_loop.assert_called_once()
-
-        mock_export_actual_coro.assert_called_once()
-        args, kwargs = mock_export_actual_coro.call_args
-        assert kwargs["task_run_id"] == task_run_id
-        assert kwargs["mcp_calls"] == mock_mcp_calls
-        assert kwargs["trace_attributes"]["trace_name"] == test_name
-        assert kwargs["trace_attributes"]["custom_attr"] == "test_val"
-        assert "start_time" in kwargs["trace_attributes"]
-        assert "end_time" in kwargs["trace_attributes"]
-        assert "duration" in kwargs["trace_attributes"]
-        assert kwargs["trace_attributes"]["is_root"] is True
+        # submit_to_worker_loop is now called for status updates and export
+        # The exact count may vary depending on whether export_incremental is called
+        assert mock_submit_loop.call_count >= 2  # At least INITIALIZING and COMPLETED
+
+        # With the new export flow, export_telemetry is submitted to worker loop
+        # so we can't directly assert on it being called synchronously
+        # Instead, verify that the trace completed successfully
+        assert task_run_id is not None
 
     def test_trace_nested(self, mocker):
         """Test nested traces, verifying context restoration and root trace logic."""
@@ -127,13 +129,13 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
         )
         mock_submit_loop_internal = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
         assert actual_get_current_task_run_id() is None
         assert actual_is_root_trace.get() is False
 
-        with trace(name="outer") as outer_id:
+        with trace_open(name="outer") as outer_id:
             assert actual_get_current_task_run_id() == outer_id
             assert actual_is_root_trace.get() is True
             with trace(name="inner") as inner_id:
@@ -146,7 +148,8 @@ class TestTrace:
         assert actual_get_current_task_run_id() is None
         assert actual_is_root_trace.get() is False
         assert mock_flush_internal.call_count == 2
-        mock_submit_loop_internal.assert_not_called()
+        # submit_to_worker_loop is now called for status updates
+        assert mock_submit_loop_internal.call_count == 2  # Only outer trace sends status updates
 
     def test_trace_exception_handling(self, mocker):
         """Test trace handles exceptions properly and restores context."""
@@ -159,10 +162,13 @@ class TestTrace:
             "hud.telemetry._trace.flush_buffer", return_value=[], autospec=True
         )
         mock_submit_loop = mocker.patch(
-            "hud.telemetry._trace.submit_to_worker_loop", return_value=MagicMock(), autospec=True
+            "hud.telemetry.exporter.submit_to_worker_loop", return_value=MagicMock(), autospec=True
         )
 
-        with pytest.raises(ValueError, match="Test exception"), trace(name="trace_with_exception"):
+        with (
+            pytest.raises(ValueError, match="Test exception"),
+            trace_open(name="trace_with_exception"),
+        ):
             assert actual_get_current_task_run_id() != initial_task_id_before_trace
             assert actual_is_root_trace.get() is False
             raise ValueError("Test exception")
@@ -173,98 +179,134 @@ class TestTrace:
         mock_submit_loop.assert_not_called()
 
 
-class TestRegisterTrace:
-    """Test the register_trace decorator."""
+class TestTraceSync:
+    """Test the trace_sync context manager."""
+
+    def test_trace_sync_basic(self, mocker):
+        """Test trace calls trace_open and flush."""
+        mock_flush = mocker.patch("hud.flush", autospec=True)
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open")
+        mock_trace_open.return_value.__enter__.return_value = "test-task-id"
+        mock_trace_open.return_value.__exit__.return_value = None
+
+        with trace(name="test_sync") as task_run_id:
+            assert task_run_id == "test-task-id"
+
+        mock_trace_open.assert_called_once_with(name="test_sync", agent_model=None, attributes=None)
+        mock_flush.assert_called_once()
+
+    def test_trace_sync_with_attributes(self, mocker):
+        """Test trace passes attributes correctly."""
+        mock_flush = mocker.patch("hud.flush", autospec=True)
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open")
+        mock_trace_open.return_value.__enter__.return_value = "test-task-id"
+        mock_trace_open.return_value.__exit__.return_value = None
+        attrs = {"key": "value"}
+
+        with trace(name="test_sync", attributes=attrs):
+            pass
+
+        mock_trace_open.assert_called_once_with(
+            name="test_sync", agent_model=None, attributes=attrs
+        )
+        mock_flush.assert_called_once()
+
 
-    def test_register_trace_sync_function(self, mocker):
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "mocked_task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+class TestTraceDecorator:
+    """Test the trace_decorator function decorator."""
 
-        @register_trace(name="test_func_sync")
+    def test_trace_decorator_sync_function(self, mocker):
+        """Test trace_decorator on synchronous functions."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "mocked_task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
+
+        @trace_decorator(name="test_func_sync")
         def sync_function(x, y):
             return x + y
 
         result = sync_function(1, 2)
         assert result == 3
-        mock_trace_context_manager.assert_called_once_with(name="test_func_sync", attributes=None)
+        mock_trace_open.assert_called_once_with(
+            name="test_func_sync", agent_model=None, attributes=None
+        )
 
-    def test_register_trace_async_function(self, mocker):
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "mocked_task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_decorator_async_function(self, mocker):
+        """Test trace_decorator on asynchronous functions."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "mocked_task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace(name="test_func_async")
+        @trace_decorator(name="test_func_async")
         async def async_function(x, y):
             return x + y
 
         async def run_test():
             result = await async_function(1, 2)
             assert result == 3
-            mock_trace_context_manager.assert_called_once_with(
-                name="test_func_async", attributes=None
+            mock_trace_open.assert_called_once_with(
+                name="test_func_async", agent_model=None, attributes=None
             )
 
         asyncio.run(run_test())
 
-    def test_register_trace_with_attributes(self, mocker):
-        """Test register_trace with attributes."""
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
+    def test_trace_decorator_with_attributes(self, mocker):
+        """Test trace_decorator with attributes."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        class _MockTraceContextManager:
-            def __enter__(self):
-                return "task_id"
+        attrs = {"operation": "multiply"}
 
-            def __exit__(self, exc_type, exc_value, traceback):
-                return None
-
-        mock_trace_context_manager.return_value = _MockTraceContextManager()
-
-        attrs = {"operation": "add"}
-
-        @register_trace(name="test_func", attributes=attrs)
+        @trace_decorator(name="test_func", attributes=attrs)
         def func_with_attrs(x):
             return x * 2
 
         result = func_with_attrs(5)
         assert result == 10
-        mock_trace_context_manager.assert_called_once_with(name="test_func", attributes=attrs)
+        mock_trace_open.assert_called_once_with(
+            name="test_func", agent_model=None, attributes=attrs
+        )
 
-    def test_register_trace_without_name(self, mocker):
-        """Test register_trace uses function name when name not provided."""
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_decorator_without_name(self, mocker):
+        """Test trace_decorator uses module.function name when name not provided."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace()
+        @trace_decorator()
         def my_function():
             return "result"
 
         result = my_function()
         assert result == "result"
-        mock_trace_context_manager.assert_called_once_with(name="my_function", attributes=None)
+        # Should use module.function name
+        expected_name = f"{my_function.__module__}.my_function"
+        mock_trace_open.assert_called_once_with(
+            name=expected_name, agent_model=None, attributes=None
+        )
 
-    def test_register_trace_preserves_function_metadata(self):
-        """Test register_trace preserves original function metadata."""
+    def test_trace_decorator_preserves_function_metadata(self):
+        """Test trace_decorator preserves original function metadata."""
 
-        @register_trace(name="test")
+        @trace_decorator(name="test")
         def original_function():
             """Original docstring."""
 
         assert original_function.__name__ == "original_function"
         assert original_function.__doc__ == "Original docstring."
 
-    def test_register_trace_exception_propagation(self, mocker):
-        """Test register_trace propagates exceptions."""
-        mock_trace_context_manager = mocker.patch("hud.telemetry._trace.trace", autospec=True)
-        mock_trace_context_manager.return_value.__enter__.return_value = "task_id"
-        mock_trace_context_manager.return_value.__exit__.return_value = None
+    def test_trace_decorator_exception_propagation(self, mocker):
+        """Test trace_decorator propagates exceptions."""
+        mock_trace_open = mocker.patch("hud.telemetry._trace.trace_open", autospec=True)
+        mock_trace_open.return_value.__enter__.return_value = "task_id"
+        mock_trace_open.return_value.__exit__.return_value = None
 
-        @register_trace()
+        @trace_decorator()
         def failing_function():
             raise RuntimeError("Test error")
 
         with pytest.raises(RuntimeError, match="Test error"):
             failing_function()
 
-        mock_trace_context_manager.assert_called_once()
+        mock_trace_open.assert_called_once()