Merge pull request #49 from context-labs/integrate-session-id-in-tracing

integrate session id in tracing

Author: francescov1

engine/main.py

@@ -59,10 +59,9 @@ async def stream_engine_async(
     telemetry_handle = setup_telemetry(enable=telemetry, run_id=run_id)
     try:
         # Root + subagent share ``agent_id="halo"`` so Catalyst's Agents
-        # view collapses every HALO run into one identity row; the
-        # span-level ``agent.name`` ("halo-root" vs "halo-sub") still
-        # distinguishes the two for span-tree rendering.
-        with halo_agent_span(name="halo-root", agent_id="halo", system="openai"):
+        # view collapses every HALO run into one identity row; the span
+        # name still distinguishes root vs subagent work in the trace tree.
+        with halo_agent_span(span_name="halo-root.run", agent_id="halo", system="openai"):
             configure_default_sdk_client(engine_config.model_provider)
             sandbox = Sandbox.get()
 

engine/telemetry/tracing.py

@@ -22,36 +22,48 @@
 
 from __future__ import annotations
 
+import os
 from collections.abc import Iterator
 from contextlib import contextmanager
 
 from inference_catalyst_tracing import AgentSpanHandle, agent_span
 from opentelemetry import trace
 
 _TRACER_NAME = "halo-engine"
+_SESSION_ID_ENV = "CATALYST_TRACING_CONVERSATION_ID"
+
+
+def _session_id_from_env() -> str | None:
+    value = os.environ.get(_SESSION_ID_ENV, "").strip()
+    if not value:
+        return None
+    return value
 
 
 @contextmanager
 def halo_agent_span(
     *,
-    name: str,
+    span_name: str,
     system: str = "openai",
     agent_id: str | None = None,
 ) -> Iterator[AgentSpanHandle]:
     """Open an OpenInference AGENT span around a chunk of HALO agent work.
 
-    ``name`` becomes the span's ``agent.name`` and the prefix of its
-    span name (``f"{name}.run"``). ``agent_id`` becomes the span's
-    ``agent.id`` — Catalyst groups the Agents view on that field with
-    ``agent.name`` as a fallback, so passing a stable ``agent_id`` (e.g.
-    ``"halo"`` for both root + subagent) collapses every HALO run under
-    a single Agents-tab row regardless of which name the span carries.
-    ``system`` becomes ``gen_ai.system``.
+    ``span_name`` is the OTel operation name. ``agent_id`` is the
+    stable agent identity Catalyst groups by, so passing ``"halo"``
+    for both root and subagent spans collapses every HALO run under a
+    single Agents-tab row. ``system`` becomes ``gen_ai.system``.
 
     Yields the ``AgentSpanHandle`` from catalyst-tracing. When telemetry
     is off or the local backend is active, the handle wraps a
     no-op ``NonRecordingSpan`` and all attribute setters silently no-op.
     """
     tracer = trace.get_tracer(_TRACER_NAME)
-    with agent_span(tracer, name=name, system=system, agent_id=agent_id) as span:
+    with agent_span(
+        tracer,
+        span_name=span_name,
+        system=system,
+        agent_id=agent_id,
+        session_id=_session_id_from_env(),
+    ) as span:
         yield span

engine/tools/subagent_tool_factory.py

@@ -272,12 +272,10 @@ async def _run_streamed(*, agent, input, context):
 
             # ``agent_id="halo"`` matches the root span (see
             # ``engine/main.py``) so Catalyst groups root + every
-            # subagent invocation under one Agents-tab identity. The
-            # span ``agent.name`` is fixed at ``"halo-sub"`` for
-            # display; ``engine_config.subagent.name`` ("sub") still
-            # drives the HALO run page UI via
-            # ``child_execution.agent_name``.
-            with halo_agent_span(name="halo-sub", agent_id="halo", system="openai"):
+            # subagent invocation under one Agents-tab identity.
+            # ``engine_config.subagent.name`` ("sub") still drives the
+            # HALO run page UI via ``child_execution.agent_name``.
+            with halo_agent_span(span_name="halo-sub.run", agent_id="halo", system="openai"):
                 try:
                     await runner.run(
                         sdk_agent=child_agent,

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     # by ``CATALYST_OTLP_TOKEN``. When the token is unset and the local
     # JSONL backend is in use, OTel's proxy tracer makes the wrapper a
     # no-op without any extra branching in engine code.
-    "inference-catalyst-tracing[openai-agents]==0.0.13",
+    "inference-catalyst-tracing[openai-agents]==0.0.19",
 ]
 
 [project.scripts]

pyrightconfig.json

@@ -12,7 +12,7 @@
   "pythonVersion": "3.10",
   "venvPath": ".",
   "venv": ".venv",
-  "typeCheckingMode": "basic",
+  "typeCheckingMode": "standard",
   "reportMissingTypeStubs": false,
   "executionEnvironments": [
     {

tests/unit/test_telemetry.py

@@ -4,6 +4,8 @@
 
 import logging
 import os
+from collections.abc import Iterator
+from contextlib import contextmanager
 
 import pytest
 
@@ -42,6 +44,46 @@ def test_setup_returns_none_when_disabled(monkeypatch) -> None:
     assert cleared == [[]]
 
 
+def test_halo_agent_span_passes_conversation_id_as_session_id(monkeypatch) -> None:
+    from engine.telemetry import tracing
+
+    captured: dict[str, object] = {}
+
+    class _FakeSpan:
+        pass
+
+    @contextmanager
+    def _fake_agent_span(
+        tracer: object,
+        *,
+        span_name: str,
+        system: str,
+        agent_id: str | None,
+        session_id: str | None,
+    ) -> Iterator[_FakeSpan]:
+        captured["tracer"] = tracer
+        captured["span_name"] = span_name
+        captured["system"] = system
+        captured["agent_id"] = agent_id
+        captured["session_id"] = session_id
+        yield _FakeSpan()
+
+    monkeypatch.setenv("CATALYST_TRACING_CONVERSATION_ID", "  conv-123  ")
+    monkeypatch.setattr(tracing.trace, "get_tracer", lambda name: f"tracer:{name}")
+    monkeypatch.setattr(tracing, "agent_span", _fake_agent_span)
+
+    with tracing.halo_agent_span(span_name="halo-root.run", agent_id="halo"):
+        pass
+
+    assert captured == {
+        "tracer": "tracer:halo-engine",
+        "span_name": "halo-root.run",
+        "system": "openai",
+        "agent_id": "halo",
+        "session_id": "conv-123",
+    }
+
+
 def test_setup_attaches_local_processor(monkeypatch, tmp_path) -> None:
     """setup_telemetry attaches the InferenceOtlpFileProcessor to the
     openai-agents SDK at the path indicated by HALO_TELEMETRY_PATH."""