diff --git a/posthog/ai/openai_agents/__init__.py b/posthog/ai/openai_agents/__init__.py index 2e1611cc..ca3bb19a 100644 --- a/posthog/ai/openai_agents/__init__.py +++ b/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 diff --git a/posthog/ai/openai_agents/processor.py b/posthog/ai/openai_agents/processor.py index 054f023c..c0da73ff 100644 --- a/posthog/ai/openai_agents/processor.py +++ b/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,9 +241,14 @@ 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] = { @@ -205,6 +256,7 @@ def on_trace_start(self, trace: Trace) -> None: "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}") diff --git a/posthog/test/ai/openai_agents/test_processor.py b/posthog/test/ai/openai_agents/test_processor.py index d1c38186..6449ed87 100644 --- a/posthog/test/ai/openai_agents/test_processor.py +++ b/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}