feat(openai-agents): per-run identity via RunConfig trace_metadata

posthog/ai/openai_agents/__init__.py

@@ -34,11 +34,18 @@ def instrument(
 
     Args:
         client: Optional PostHog client instance. If not provided, uses the default client.
-        distinct_id: Optional distinct ID to associate with all traces.
-            Can also be a callable that takes a trace and returns a distinct ID.
+        distinct_id: Optional default distinct ID for all traces. Only suitable
+            as a static value when one process serves one user (a CLI or
+            worker); servers should pass identity per run via
+            ``RunConfig(trace_metadata={"posthog_distinct_id": ...})``, which
+            takes precedence. Can also be a callable that takes a trace and
+            returns a distinct ID.
         privacy_mode: If True, redacts input/output content from events.
         groups: Optional PostHog groups to associate with events.
         properties: Optional additional properties to include with all events.
+            Per-run properties can be passed via
+            ``RunConfig(trace_metadata={"posthog_properties": {...}})`` and
+            override these defaults.
 
     Returns:
         PostHogTracingProcessor: The registered processor instance.
@@ -47,20 +54,25 @@ def instrument(
         ```python
         from posthog.ai.openai_agents import instrument
 
-        # Simple setup
+        # One-user process (CLI/worker): a static distinct ID is fine
         instrument(distinct_id="user@example.com")
 
-        # With custom properties
-        instrument(
-            distinct_id="user@example.com",
-            privacy_mode=True,
-            properties={"environment": "production"}
-        )
+        # Server: pass identity and session per run instead
+        instrument()
 
-        # Now run agents as normal - traces automatically sent to PostHog
-        from agents import Agent, Runner
+        from agents import Agent, Runner, RunConfig
         agent = Agent(name="Assistant", instructions="You are helpful.")
-        result = Runner.run_sync(agent, "Hello!")
+        result = Runner.run_sync(
+            agent,
+            "Hello!",
+            run_config=RunConfig(
+                group_id=conversation_id,  # becomes $ai_session_id
+                trace_metadata={
+                    "posthog_distinct_id": user_id,
+                    "posthog_properties": {"plan": "scale"},
+                },
+            ),
+        )
         ```
     """
     from agents.tracing import add_trace_processor

posthog/ai/openai_agents/processor.py

@@ -39,6 +39,41 @@ def _parse_iso_timestamp(iso_str: Optional[str]) -> Optional[float]:
         return None
 
 
+# PostHog control keys read from RunConfig(trace_metadata=...). Same names as
+# the per-call kwargs on the OpenAI/Anthropic wrappers, carried by the Agents
+# SDK's per-run surface instead.
+_METADATA_DISTINCT_ID_KEY = "posthog_distinct_id"
+_METADATA_PROPERTIES_KEY = "posthog_properties"
+
+
+def _extract_posthog_metadata(
+    metadata: Optional[Dict[str, Any]],
+) -> tuple[Optional[str], Dict[str, Any], Optional[Dict[str, Any]]]:
+    """Split PostHog control keys out of RunConfig trace_metadata.
+
+    Returns (distinct_id, properties, remaining_metadata). The control keys are
+    removed from the metadata that lands in $ai_trace_metadata.
+    """
+    if not metadata or not isinstance(metadata, dict):
+        return None, {}, metadata
+
+    distinct_id = metadata.get(_METADATA_DISTINCT_ID_KEY)
+    properties = metadata.get(_METADATA_PROPERTIES_KEY)
+    if not isinstance(properties, dict):
+        properties = {}
+
+    remaining = {
+        k: v
+        for k, v in metadata.items()
+        if k not in (_METADATA_DISTINCT_ID_KEY, _METADATA_PROPERTIES_KEY)
+    }
+    return (
+        str(distinct_id) if distinct_id else None,
+        properties,
+        remaining or None,
+    )
+
+
 class PostHogTracingProcessor(TracingProcessor):
     """
     A tracing processor that sends OpenAI Agents SDK traces to PostHog.
@@ -80,6 +115,8 @@ def __init__(
             client: Optional PostHog client instance. If not provided, uses the default client.
             distinct_id: Either a string distinct ID or a callable that takes a Trace
                 and returns a distinct ID. If not provided, uses the trace_id.
+                A per-run ``RunConfig(trace_metadata={"posthog_distinct_id": ...})``
+                takes precedence over this default.
             privacy_mode: If True, redacts input/output content from events.
             groups: Optional PostHog groups to associate with all events.
             properties: Optional additional properties to include with all events.
@@ -173,9 +210,18 @@ def _capture_event(
             ):
                 return
 
+            # Per-run posthog_properties (from RunConfig trace_metadata) merge
+            # last so they win over the instrument()-level defaults.
+            trace_id = properties.get("$ai_trace_id")
+            run_properties = (
+                self._trace_metadata.get(trace_id, {}).get("run_properties")
+                if trace_id
+                else None
+            )
             final_properties = {
                 **properties,
                 **self._properties,
+                **(run_properties or {}),
             }
 
             _capture_ai_event(
@@ -195,16 +241,22 @@ def on_trace_start(self, trace: Trace) -> None:
             trace_id = trace.trace_id
             trace_name = trace.name
             group_id = getattr(trace, "group_id", None)
-            metadata = getattr(trace, "metadata", None)
+            raw_metadata = getattr(trace, "metadata", None)
 
-            distinct_id = self._get_distinct_id(trace)
+            # Per-run values from RunConfig(trace_metadata=...) win over the
+            # instrument()-level defaults, which are process-global.
+            run_distinct_id, run_properties, metadata = _extract_posthog_metadata(
+                raw_metadata
+            )
+            distinct_id = run_distinct_id or self._get_distinct_id(trace)
 
             # Store trace metadata for later (used by spans and on_trace_end)
             self._trace_metadata[trace_id] = {
                 "name": trace_name,
                 "group_id": group_id,
                 "metadata": metadata,
                 "distinct_id": distinct_id,
+                "run_properties": run_properties,
                 "start_time": time.time(),
             }
         except Exception as e:
@@ -215,11 +267,26 @@ def on_trace_end(self, trace: Trace) -> None:
         try:
             trace_id = trace.trace_id
 
-            # Pop stored metadata (also cleans up)
-            trace_info = self._trace_metadata.pop(trace_id, {})
+            # Read stored metadata; popped after the capture below so
+            # _capture_event can still look up the trace's run_properties.
+            trace_info = self._trace_metadata.get(trace_id)
+            if trace_info is None:
+                # Evicted or never started — re-extract from the trace itself
+                # and seed the store so _capture_event finds run_properties.
+                (
+                    fallback_distinct_id,
+                    fallback_run_properties,
+                    fallback_metadata,
+                ) = _extract_posthog_metadata(getattr(trace, "metadata", None))
+                trace_info = {
+                    "metadata": fallback_metadata,
+                    "distinct_id": fallback_distinct_id,
+                    "run_properties": fallback_run_properties,
+                }
+                self._trace_metadata[trace_id] = trace_info
             trace_name = trace_info.get("name") or trace.name
             group_id = trace_info.get("group_id") or getattr(trace, "group_id", None)
-            metadata = trace_info.get("metadata") or getattr(trace, "metadata", None)
+            metadata = trace_info.get("metadata")
             distinct_id = trace_info.get("distinct_id") or self._get_distinct_id(trace)
 
             # Calculate trace-level latency
@@ -255,6 +322,7 @@ def on_trace_end(self, trace: Trace) -> None:
                 distinct_id=distinct_id or trace_id,
                 properties=properties,
             )
+            self._trace_metadata.pop(trace_id, None)
         except Exception as e:
             log.debug(f"Error in on_trace_end: {e}")
 

posthog/test/ai/openai_agents/test_processor.py

@@ -943,6 +943,109 @@ def test_eviction_of_stale_entries(self, mock_client):
         assert len(processor._trace_metadata) <= 10
 
 
+class TestPerRunMetadata:
+    """Tests for per-run identity/properties via RunConfig(trace_metadata=...)."""
+
+    def test_trace_metadata_distinct_id_overrides_instrument_default(
+        self, processor, mock_client, mock_trace
+    ):
+        """posthog_distinct_id in trace_metadata wins over instrument()-level."""
+        mock_trace.metadata = {"posthog_distinct_id": "run-user"}
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["distinct_id"] == "run-user"
+
+    def test_trace_metadata_distinct_id_used_for_spans(
+        self, processor, mock_client, mock_trace, mock_span
+    ):
+        """Spans inherit the per-run distinct_id resolved at trace start."""
+        mock_trace.metadata = {"posthog_distinct_id": "run-user"}
+        processor.on_trace_start(mock_trace)
+        mock_client.capture.reset_mock()
+
+        mock_span.span_data = GenerationSpanData(model="gpt-4o")
+        processor.on_span_start(mock_span)
+        processor.on_span_end(mock_span)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["distinct_id"] == "run-user"
+
+    def test_trace_metadata_properties_merge_and_override_defaults(
+        self, mock_client, mock_trace
+    ):
+        """posthog_properties merge into events, winning over instrument()-level."""
+        processor = PostHogTracingProcessor(
+            client=mock_client,
+            distinct_id="test-user",
+            properties={"env": "instrument", "keep": "yes"},
+        )
+        mock_trace.metadata = {"posthog_properties": {"env": "run", "extra": 1}}
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        properties = mock_client.capture.call_args[1]["properties"]
+        assert properties["env"] == "run"
+        assert properties["keep"] == "yes"
+        assert properties["extra"] == 1
+
+    def test_control_keys_stripped_from_trace_metadata_property(
+        self, processor, mock_client, mock_trace
+    ):
+        """posthog_* control keys never land in $ai_trace_metadata."""
+        mock_trace.metadata = {
+            "posthog_distinct_id": "run-user",
+            "posthog_properties": {"plan": "scale"},
+            "batch": "nightly",
+        }
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        properties = mock_client.capture.call_args[1]["properties"]
+        assert properties["$ai_trace_metadata"] == {"batch": "nightly"}
+
+    def test_only_control_keys_omits_trace_metadata_property(
+        self, processor, mock_client, mock_trace
+    ):
+        """Metadata that is all control keys emits no $ai_trace_metadata."""
+        mock_trace.metadata = {"posthog_distinct_id": "run-user"}
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        properties = mock_client.capture.call_args[1]["properties"]
+        assert "$ai_trace_metadata" not in properties
+
+    def test_without_control_keys_falls_back_to_instrument_default(
+        self, processor, mock_client, mock_trace
+    ):
+        """No posthog_* keys → instrument()-level distinct_id still applies."""
+        mock_trace.metadata = {"batch": "nightly"}
+
+        processor.on_trace_start(mock_trace)
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["distinct_id"] == "test-user"
+        assert call_kwargs["properties"]["$ai_trace_metadata"] == {"batch": "nightly"}
+
+    def test_trace_end_without_start_still_reads_control_keys(
+        self, processor, mock_client, mock_trace
+    ):
+        """A trace evicted before end still resolves per-run identity."""
+        mock_trace.metadata = {"posthog_distinct_id": "run-user"}
+
+        processor.on_trace_end(mock_trace)
+
+        call_kwargs = mock_client.capture.call_args[1]
+        assert call_kwargs["distinct_id"] == "run-user"
+        assert mock_trace.trace_id not in processor._trace_metadata
+
+
 class TestEnsureSerializableCycleGuard:
     def test_self_referencing_dict_returns_circular_marker(self):
         node = {"a": 1}