From 6c2e02b017aba72dbbfdcdce79595392eb55c994 Mon Sep 17 00:00:00 2001 From: Graham Neubig Date: Tue, 9 Jun 2026 00:36:18 +0000 Subject: [PATCH] Add conversation attribution to Laminar traces --- .../openhands/sdk/conversation/base.py | 29 ++++++++++++++--- .../conversation/impl/local_conversation.py | 2 +- .../conversation/impl/remote_conversation.py | 4 ++- .../openhands/sdk/observability/laminar.py | 12 +++++-- .../conversation/test_base_span_management.py | 31 +++++++++++++++++++ .../conversation/test_conversation_factory.py | 20 ++++++++++++ tests/sdk/observability/test_laminar.py | 24 ++++++++++++++ 7 files changed, 114 insertions(+), 8 deletions(-) diff --git a/openhands-sdk/openhands/sdk/conversation/base.py b/openhands-sdk/openhands/sdk/conversation/base.py index c6f56da5ee..2915fc9f5f 100644 --- a/openhands-sdk/openhands/sdk/conversation/base.py +++ b/openhands-sdk/openhands/sdk/conversation/base.py @@ -28,6 +28,14 @@ from openhands.sdk.workspace.base import BaseWorkspace +def _conversation_tag_attributes( + tags: Mapping[str, str] | None, +) -> dict[str, str] | None: + if not tags: + return None + return {f"conversation.tags.{key}": value for key, value in tags.items()} + + if TYPE_CHECKING: from openhands.sdk.agent.base import AgentBase from openhands.sdk.conversation.state import ConversationExecutionStatus @@ -128,22 +136,35 @@ def __init__(self) -> None: self._observability_root_span: RootSpan | None = None def _start_observability_span( - self, session_id: str, user_id: str | None = None + self, + session_id: str, + user_id: str | None = None, + tags: Mapping[str, str] | None = None, ) -> None: """Start a per-conversation observability root span. Args: session_id: The session ID to associate with the trace user_id: Optional user ID to associate with the trace + tags: Optional conversation tags to add as root span attributes """ if not should_enable_observability(): return if self._observability_root_span is not None: # Idempotent: never start two roots for one conversation. return - self._observability_root_span = start_root_span( - "conversation", session_id=session_id, user_id=user_id - ) + attributes = _conversation_tag_attributes(tags) + if attributes: + self._observability_root_span = start_root_span( + "conversation", + session_id=session_id, + user_id=user_id, + attributes=attributes, + ) + else: + self._observability_root_span = start_root_span( + "conversation", session_id=session_id, user_id=user_id + ) def _end_observability_span(self) -> None: """End the observability span if it hasn't been ended already.""" diff --git a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py index 83901ce35d..cf0c4a8285 100644 --- a/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py +++ b/openhands-sdk/openhands/sdk/conversation/impl/local_conversation.py @@ -341,7 +341,7 @@ def _default_callback(e): self.update_secrets(secret_values) atexit.register(self.close) - self._start_observability_span(str(desired_id), user_id=user_id) + self._start_observability_span(str(desired_id), user_id=user_id, tags=tags) self.delete_on_close = delete_on_close @property diff --git a/openhands-sdk/openhands/sdk/conversation/impl/remote_conversation.py b/openhands-sdk/openhands/sdk/conversation/impl/remote_conversation.py index 146388f31e..c25276ad81 100644 --- a/openhands-sdk/openhands/sdk/conversation/impl/remote_conversation.py +++ b/openhands-sdk/openhands/sdk/conversation/impl/remote_conversation.py @@ -768,6 +768,8 @@ def __init__( # Include tags if provided "tags": tags or {}, } + if user_id: + payload["user_id"] = user_id if stuck_detection_thresholds is not None: # Convert to StuckDetectionThresholds if dict, then serialize if isinstance(stuck_detection_thresholds, Mapping): @@ -920,7 +922,7 @@ def run_complete_callback(event: Event) -> None: secret_values: dict[str, SecretValue] = {k: v for k, v in secrets.items()} self.update_secrets(secret_values) - self._start_observability_span(str(self._id), user_id=user_id) + self._start_observability_span(str(self._id), user_id=user_id, tags=tags) # All hooks (including SessionStart/SessionEnd) are executed server-side. # hook_config is sent in the creation payload. self.delete_on_close = delete_on_close diff --git a/openhands-sdk/openhands/sdk/observability/laminar.py b/openhands-sdk/openhands/sdk/observability/laminar.py index fd9701eccb..0816249c1f 100644 --- a/openhands-sdk/openhands/sdk/observability/laminar.py +++ b/openhands-sdk/openhands/sdk/observability/laminar.py @@ -2,7 +2,7 @@ import functools import inspect import sys -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal from openhands.sdk.logger import get_logger @@ -253,12 +253,17 @@ def __init__( name: str, session_id: str | None = None, user_id: str | None = None, + attributes: Mapping[str, str] | None = None, ) -> None: from lmnr import Laminar # ``start_span`` returns a span without attaching it as the current # OTel context; we'll restore it on every entry point via ``use_span``. self.span = Laminar.start_span(name) + if attributes: + with contextlib.suppress(Exception): + for key, value in attributes.items(): + self.span.set_attribute(key, value) if session_id or user_id: # ``set_trace_session_id`` / ``set_trace_user_id`` require an # active span; briefly enter the span context to apply them. @@ -285,6 +290,7 @@ def start_root_span( name: str, session_id: str | None = None, user_id: str | None = None, + attributes: Mapping[str, str] | None = None, ) -> RootSpan | None: """Create a long-lived root span for an owning object. @@ -293,7 +299,9 @@ def start_root_span( if not should_enable_observability(): return None try: - return RootSpan(name, session_id=session_id, user_id=user_id) + return RootSpan( + name, session_id=session_id, user_id=user_id, attributes=attributes + ) except Exception: logger.debug("Failed to create observability root span", exc_info=True) return None diff --git a/tests/sdk/conversation/test_base_span_management.py b/tests/sdk/conversation/test_base_span_management.py index 933056dbd0..933993636e 100644 --- a/tests/sdk/conversation/test_base_span_management.py +++ b/tests/sdk/conversation/test_base_span_management.py @@ -116,6 +116,37 @@ def test_base_conversation_span_management(): assert conversation._span_ended is True +def test_base_conversation_span_includes_tag_attributes(): + """Conversation tags are attached to the root span for trace filtering.""" + conversation = MockConversation() + + with ( + patch( + "openhands.sdk.conversation.base.should_enable_observability" + ) as mock_should_enable, + patch("openhands.sdk.conversation.base.start_root_span") as mock_start_span, + ): + mock_should_enable.return_value = True + fake_root = MagicMock(name="root-span") + mock_start_span.return_value = fake_root + + conversation._start_observability_span( + "test-session-id", + user_id="user-42", + tags={"automationid": "auto-1", "automationrunid": "run-1"}, + ) + + mock_start_span.assert_called_once_with( + "conversation", + session_id="test-session-id", + user_id="user-42", + attributes={ + "conversation.tags.automationid": "auto-1", + "conversation.tags.automationrunid": "run-1", + }, + ) + + def test_base_conversation_span_management_disabled(): """Test that BaseConversation doesn't perform span operations when observability is disabled.""" # noqa: E501 diff --git a/tests/sdk/conversation/test_conversation_factory.py b/tests/sdk/conversation/test_conversation_factory.py index 244c0d112e..9473394059 100644 --- a/tests/sdk/conversation/test_conversation_factory.py +++ b/tests/sdk/conversation/test_conversation_factory.py @@ -102,6 +102,26 @@ def test_conversation_factory_forwards_remote_parameters( assert conversation.max_iteration_per_run == 200 +@patch("openhands.sdk.conversation.impl.remote_conversation.WebSocketCallbackClient") +def test_conversation_factory_forwards_remote_user_id_to_create_payload( + mock_ws_client, agent, remote_workspace +): + """RemoteConversation must send user_id to the agent-server create request.""" + conversation = Conversation( + agent=agent, + workspace=remote_workspace, + tags={"automationid": "auto-1"}, + user_id="user-42", + ) + + assert isinstance(conversation, RemoteConversation) + create_call = remote_workspace._client.request.call_args_list[0] + assert create_call.args[:2] == ("POST", "/api/conversations") + payload = create_call.kwargs["json"] + assert payload["user_id"] == "user-42" + assert payload["tags"] == {"automationid": "auto-1"} + + def test_conversation_factory_string_workspace_creates_local(agent): """Test that string workspace creates LocalConversation.""" conversation = Conversation(agent=agent, workspace="") diff --git a/tests/sdk/observability/test_laminar.py b/tests/sdk/observability/test_laminar.py index 19edcdfcaf..e27dd22829 100644 --- a/tests/sdk/observability/test_laminar.py +++ b/tests/sdk/observability/test_laminar.py @@ -438,6 +438,30 @@ def fake_use_span(span, *args, **kwargs): os.environ.pop("LMNR_PROJECT_API_KEY", None) +def test_root_span_sets_attributes(): + """RootSpan must attach provided attributes to the underlying span.""" + os.environ["LMNR_PROJECT_API_KEY"] = "test-key" + try: + from lmnr import Laminar + + from openhands.sdk.observability import laminar as lam + + mock_span = MagicMock(name="span") + + with patch.object(Laminar, "start_span", return_value=mock_span): + lam._observability_enabled = True + root = lam.RootSpan( + "conversation", + attributes={"conversation.tags.automationid": "auto-1"}, + ) + assert root.span is mock_span + mock_span.set_attribute.assert_called_once_with( + "conversation.tags.automationid", "auto-1" + ) + finally: + os.environ.pop("LMNR_PROJECT_API_KEY", None) + + def test_two_concurrent_conversations_do_not_collide(): """Each conversation must own its own root span (no global stack).