From c36506108425bc910cbe9d510735e987c5c2f2ae Mon Sep 17 00:00:00 2001
From: Robert Witt <robert_witt@icloud.com>
Date: Thu, 1 Oct 2026 10:45:27 +0200
Subject: [PATCH] feat(openai-agents): map group_id to session.id

---
 .../openai_agents/_processor.py               | 14 ++++--
 .../tests/test_processor_span_attributes.py   | 45 +++++++++++++++++++
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/python/instrumentation/openinference-instrumentation-openai-agents/src/openinference/instrumentation/openai_agents/_processor.py b/python/instrumentation/openinference-instrumentation-openai-agents/src/openinference/instrumentation/openai_agents/_processor.py
index 6ebb8372d3..46a659a32f 100644
--- a/python/instrumentation/openinference-instrumentation-openai-agents/src/openinference/instrumentation/openai_agents/_processor.py
+++ b/python/instrumentation/openinference-instrumentation-openai-agents/src/openinference/instrumentation/openai_agents/_processor.py
@@ -129,11 +129,18 @@ def on_trace_start(self, trace: Trace) -> None:
         Args:
             trace: The trace that started.
         """
+        attributes: dict[str, AttributeValue] = {
+            OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.AGENT.value,
+        }
+        # The SDK's Trace.group_id carries the conversation/thread grouping
+        # (e.g. set via trace(group_id=...)). Map it to session.id so multiple
+        # traces in one conversation correlate.
+        group_id = getattr(trace, "group_id", None)
+        if group_id is not None:
+            attributes[SESSION_ID] = group_id
         otel_span = self._tracer.start_span(
             name=trace.name,
-            attributes={
-                OPENINFERENCE_SPAN_KIND: OpenInferenceSpanKindValues.AGENT.value,
-            },
+            attributes=attributes,
         )
         self._root_spans[trace.trace_id] = otel_span
 
@@ -1188,6 +1195,7 @@ def _get_span_status(obj: Span[Any]) -> Status:
 GRAPH_NODE_ID = SpanAttributes.GRAPH_NODE_ID
 GRAPH_NODE_PARENT_ID = SpanAttributes.GRAPH_NODE_PARENT_ID
 AGENT_NAME = SpanAttributes.AGENT_NAME
+SESSION_ID = SpanAttributes.SESSION_ID
 
 IMAGE_URL = ImageAttributes.IMAGE_URL
 
diff --git a/python/instrumentation/openinference-instrumentation-openai-agents/tests/test_processor_span_attributes.py b/python/instrumentation/openinference-instrumentation-openai-agents/tests/test_processor_span_attributes.py
index d839dcacf9..6178489dbc 100644
--- a/python/instrumentation/openinference-instrumentation-openai-agents/tests/test_processor_span_attributes.py
+++ b/python/instrumentation/openinference-instrumentation-openai-agents/tests/test_processor_span_attributes.py
@@ -1143,3 +1143,48 @@ def model_dump(self, **kwargs: Any) -> dict[str, Any]:
     assert attrs["llm.input_messages.1.message.tool_call_id"] == "call-x"
     assert "_Undumpable" in str(attrs["llm.input_messages.1.message.content"])
     assert "_Undumpable" in str(attrs["input.value"])
+
+
+# --- Trace.group_id maps to session.id on the root span ----------------------------
+
+
+class _FakeTraceWithGroupId(_FakeTrace):
+    def __init__(self, group_id: Any, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.group_id = group_id
+
+
+def _root_span(exporter: InMemorySpanExporter) -> ReadableSpan:
+    """Return the root AGENT span (the one produced by on_trace_start)."""
+    spans = exporter.get_finished_spans()
+    roots = [s for s in spans if s.parent is None]
+    assert len(roots) == 1, f"Expected 1 root span, got {[s.name for s in roots]}"
+    return roots[0]
+
+
+def test_group_id_mapped_to_session_id() -> None:
+    """group_id present → root span carries session.id."""
+    processor, exporter = _make_processor()
+    _run(processor, _FakeTraceWithGroupId("conv-abc"), [])
+    assert _root_span(exporter).attributes["session.id"] == "conv-abc"
+
+
+def test_none_group_id_omits_session_id() -> None:
+    """group_id=None → session.id must not be set on the root span."""
+    processor, exporter = _make_processor()
+    _run(processor, _FakeTraceWithGroupId(None), [])
+    assert "session.id" not in (_root_span(exporter).attributes or {})
+
+
+def test_empty_group_id_mapped_to_session_id() -> None:
+    """group_id='' is a present value → mapped verbatim, not dropped."""
+    processor, exporter = _make_processor()
+    _run(processor, _FakeTraceWithGroupId(""), [])
+    assert _root_span(exporter).attributes["session.id"] == ""
+
+
+def test_missing_group_id_attr_no_crash() -> None:
+    """Trace without a group_id attribute (older SDK) → no crash, session.id omitted."""
+    processor, exporter = _make_processor()
+    _run(processor, _FakeTrace(), [])  # _FakeTrace has no group_id attr
+    assert "session.id" not in (_root_span(exporter).attributes or {})