diff --git a/.sampo/changesets/span-processor-request-context.md b/.sampo/changesets/span-processor-request-context.md new file mode 100644 index 000000000..2d3b5f637 --- /dev/null +++ b/.sampo/changesets/span-processor-request-context.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +`PostHogSpanProcessor` now adds the active PostHog context to the AI spans it exports: the distinct ID from `identify_context`, the session ID from `set_context_session` (sent as `$session_id`), and context tags as properties. Spans started inside `posthog.new_context()` reach PostHog under the right person, so one process can trace many users. Tags set by the request middleware (for example `email`, `$ip`, `$user_agent` and `$current_url`) now appear on AI events too. Attributes already on a span are kept, and other processors on the provider never see the context. Apps already using PostHog contexts (for example the Django, Flask or ASGI request middleware) will see their AI events attributed to the context's person after upgrading. diff --git a/posthog/ai/otel/_request_context.py b/posthog/ai/otel/_request_context.py new file mode 100644 index 000000000..b34bf4c61 --- /dev/null +++ b/posthog/ai/otel/_request_context.py @@ -0,0 +1,146 @@ +"""Captures the active PostHog context for a span and applies it to an exported copy.""" + +import datetime +import decimal +import json +import logging +import math +import uuid +from dataclasses import dataclass +from typing import Any, Dict, Optional, Sequence + +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.util.types import AttributeValue + +from ...contexts import get_context_distinct_id, get_context_session_id, get_tags + +DISTINCT_ID_ATTRIBUTE = "posthog.distinct_id" +SESSION_ID_ATTRIBUTE = "$session_id" +# Ingestion prefers posthog.distinct_id, so adding the context's would override the span's explicit user. +IDENTITY_ATTRIBUTES = ( + DISTINCT_ID_ATTRIBUTE, + "user.id", + "ai.telemetry.metadata.posthog_distinct_id", + "ai.settings.context.posthog_distinct_id", +) +_MAX_INT64 = 2**63 - 1 +_MIN_INT64 = -(2**63) +_PRIMITIVE_TYPES = (str, bool, int, float) + +log = logging.getLogger("posthog") + + +@dataclass(frozen=True) +class RequestContext: + distinct_id: Optional[str] + session_id: Optional[str] + tags: Dict[str, Any] + + +def capture_request_context() -> Optional[RequestContext]: + """Snapshot the active context's distinct ID, session ID and tags, or None if all are empty.""" + snapshot = RequestContext( + distinct_id=get_context_distinct_id() or None, + session_id=get_context_session_id() or None, + tags=dict(get_tags()), + ) + if not (snapshot.distinct_id or snapshot.session_id or snapshot.tags): + return None + return snapshot + + +def request_context_attributes(context: RequestContext) -> Dict[str, AttributeValue]: + """Return the snapshot as span attributes. Tags win over the context IDs.""" + values: Dict[str, Any] = { + key: value + for key, value in ( + (DISTINCT_ID_ATTRIBUTE, context.distinct_id), + (SESSION_ID_ATTRIBUTE, context.session_id), + ) + if value + } + values.update(context.tags) + + attributes: Dict[str, AttributeValue] = {} + for key, value in values.items(): + try: + attribute = to_attribute_value(value) + except Exception: + log.debug("Failed to read PostHog context value %r", key, exc_info=True) + continue + if attribute is not None: + attributes[key] = attribute + return attributes + + +def with_request_context(span: ReadableSpan, context: RequestContext) -> ReadableSpan: + """Return a copy of ``span`` with the context's values added, keeping attributes it already has.""" + own = span.attributes or {} + has_identity = any(key in own for key in IDENTITY_ATTRIBUTES) + extra = { + key: value + for key, value in request_context_attributes(context).items() + if key not in own and not (has_identity and key == DISTINCT_ID_ATTRIBUTE) + } + if not extra: + return span + return ReadableSpan( + name=span.name, + context=span.get_span_context(), + parent=span.parent, + resource=span.resource, + attributes={**own, **extra}, + events=span.events, + links=span.links, + kind=span.kind, + status=span.status, + start_time=span.start_time, + end_time=span.end_time, + instrumentation_scope=span.instrumentation_scope, + ) + + +def to_attribute_value(value: Any) -> Optional[AttributeValue]: + """Return ``value`` as an OpenTelemetry attribute value, JSON-encoding anything OTel can't hold.""" + if value is None or (isinstance(value, str) and not value): + return None + if isinstance(value, float) and not math.isfinite(value): + return None + if _is_out_of_int64_range(value): + return str(value) + if isinstance(value, _PRIMITIVE_TYPES): + return value + if isinstance(value, (list, tuple)) and _is_homogeneous(value): + return list(value) + try: + encoded = json.dumps(value, allow_nan=False, default=_encode_string_like) + except (TypeError, ValueError): + return None + if encoded.startswith('"'): + return json.loads(encoded) or None + return encoded + + +def _encode_string_like(value: Any) -> str: + if isinstance(value, (datetime.datetime, datetime.date, datetime.time)): + return value.isoformat() + if isinstance(value, (uuid.UUID, decimal.Decimal)): + return str(value) + raise TypeError(f"{type(value).__name__} is not JSON serializable") + + +def _is_out_of_int64_range(value: Any) -> bool: + return ( + isinstance(value, int) + and not isinstance(value, bool) + and not _MIN_INT64 <= value <= _MAX_INT64 + ) + + +def _is_homogeneous(values: Sequence[Any]) -> bool: + return ( + all(isinstance(v, _PRIMITIVE_TYPES) for v in values) + and not any(_is_out_of_int64_range(v) for v in values) + and all(math.isfinite(v) for v in values if isinstance(v, float)) + and len({type(v) for v in values}) <= 1 + ) diff --git a/posthog/ai/otel/exporter.py b/posthog/ai/otel/exporter.py index 02c1c22e3..ec6e387ff 100644 --- a/posthog/ai/otel/exporter.py +++ b/posthog/ai/otel/exporter.py @@ -20,6 +20,8 @@ class PostHogTraceExporter(SpanExporter): Wraps an OTLPSpanExporter configured for PostHog's OTLP endpoint. Spans that are not AI-related are silently dropped, returning SUCCESS immediately. + Exporters only see finished spans, so they can't add per-request PostHog + context; use ``PostHogSpanProcessor`` for that. Usage:: diff --git a/posthog/ai/otel/processor.py b/posthog/ai/otel/processor.py index 4f520591b..31073c54b 100644 --- a/posthog/ai/otel/processor.py +++ b/posthog/ai/otel/processor.py @@ -5,7 +5,11 @@ for setups using TracerProvider.add_span_processor(). """ -from typing import Optional +import logging +import os +import threading +import weakref +from typing import Dict, Optional, Tuple from opentelemetry.context import Context from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor @@ -13,9 +17,43 @@ from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter from ..gateway import warn_if_posthog_ai_gateway_otel_attributes +from ._request_context import ( + RequestContext, + capture_request_context, + with_request_context, +) from .spans import DEFAULT_HOST, is_ai_span +log = logging.getLogger("posthog") + + +def _span_key(span: ReadableSpan) -> Tuple[int, int]: + # The SDK passes on_end a ReadableSpan snapshot, not the Span object on_start saw. + context = span.get_span_context() + return (context.trace_id, context.span_id) + + +_processors: "weakref.WeakSet[PostHogSpanProcessor]" = weakref.WeakSet() + + +def _discard_abandoned( + processor_ref: "weakref.ref[PostHogSpanProcessor]", key: Tuple[int, int] +) -> None: + processor = processor_ref() + if processor is not None: + processor._discard_captured(key) + + +def _reset_processors_after_fork() -> None: + for processor in list(_processors): + processor._reset_after_fork() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_processors_after_fork) + + class PostHogSpanProcessor(SpanProcessor): """Span processor that filters AI spans and exports them to PostHog. @@ -32,6 +70,19 @@ class PostHogSpanProcessor(SpanProcessor): provider.add_span_processor( PostHogSpanProcessor(api_key="phc_...") ) + + # Per request: AI spans started inside a context export with its user, session and tags. + with posthog.new_context(): + posthog.identify_context(user_id) + posthog.tag("plan", "pro") + handle_request() + + Attributes already on a span win over explicit tags, which win over the + context's distinct ID and session ID (exported as ``$session_id``). Threads that don't copy contextvars + carry no user. + + The context is captured when a span starts, but tag values are read when the + span ends. Mutating a tagged object while a span is open shows up in the export. """ def __init__( @@ -52,15 +103,50 @@ def __init__( headers={"Authorization": f"Bearer {self._api_key}"}, ) self._processor = BatchSpanProcessor(exporter) + self._captured: Dict[ + Tuple[int, int], Tuple[RequestContext, weakref.finalize] + ] = {} + # Reentrant: cyclic GC can run a span finalizer on a thread that already holds the lock. + self._captured_lock = threading.RLock() + _processors.add(self) + + def _reset_after_fork(self) -> None: + # A lock held by another thread at fork time would never release in the child. + self._captured_lock = threading.RLock() + self._captured = {} + + def _discard_captured(self, key: Tuple[int, int]) -> None: + with self._captured_lock: + self._captured.pop(key, None) + + def _take_captured(self, key: Tuple[int, int]) -> Optional[RequestContext]: + with self._captured_lock: + entry = self._captured.pop(key, None) + if entry is None: + return None + captured, finalizer = entry + finalizer.detach() + return captured def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None: """ - Handle span start notifications. + Remember the active PostHog context for the span. - This processor does not need to do work at span start; filtering happens - in ``on_end``. + The span itself is left untouched, so other processors on the provider + never see the context. The values are added to the copy exported to + PostHog in ``on_end``. """ - pass + try: + captured = capture_request_context() + if captured: + key = _span_key(span) + finalizer = weakref.finalize( + span, _discard_abandoned, weakref.ref(self), key + ) + with self._captured_lock: + self._captured[key] = (captured, finalizer) + except Exception: + log.debug("Failed to capture PostHog context for span", exc_info=True) def on_end(self, span: ReadableSpan) -> None: """ @@ -69,14 +155,24 @@ def on_end(self, span: ReadableSpan) -> None: Args: span: The ended OpenTelemetry span. """ + captured = self._take_captured(_span_key(span)) if not is_ai_span(span): return warn_if_posthog_ai_gateway_otel_attributes(span.attributes) + if captured: + try: + span = with_request_context(span, captured) + except Exception: + log.debug("Failed to add PostHog context to span", exc_info=True) self._processor.on_end(span) def shutdown(self) -> None: """Shut down the underlying batch span processor.""" self._processor.shutdown() + with self._captured_lock: + for _, finalizer in self._captured.values(): + finalizer.detach() + self._captured.clear() def force_flush(self, timeout_millis: Optional[int] = None) -> bool: """ diff --git a/posthog/test/ai/otel/test_processor.py b/posthog/test/ai/otel/test_processor.py index c2e2c2694..06deb61a2 100644 --- a/posthog/test/ai/otel/test_processor.py +++ b/posthog/test/ai/otel/test_processor.py @@ -1,7 +1,17 @@ +import gc +import json +import threading import unittest +import uuid +from datetime import datetime, timezone +from decimal import Decimal +from concurrent.futures import ThreadPoolExecutor from unittest.mock import MagicMock, patch -from posthog.ai.otel.processor import PostHogSpanProcessor +from opentelemetry.sdk.trace import TracerProvider + +from posthog import contexts +from posthog.ai.otel.processor import PostHogSpanProcessor, _reset_processors_after_fork from posthog.test.ai.otel.conftest import make_span @@ -108,3 +118,294 @@ def test_force_flush_without_timeout(self, mock_batch_cls, mock_otlp_cls): inner.force_flush.return_value = result self.assertIs(processor.force_flush(), result) inner.force_flush.assert_called_once_with() + + +def new_request_context(): + return contexts.new_context(fresh=True, capture_exceptions=False) + + +class TestPostHogSpanProcessorRequestContext(unittest.TestCase): + def setUp(self): + exporter_patch = patch("posthog.ai.otel.processor.OTLPSpanExporter") + batch_patch = patch("posthog.ai.otel.processor.BatchSpanProcessor") + exporter_patch.start() + self.inner = batch_patch.start().return_value + self.addCleanup(exporter_patch.stop) + self.addCleanup(batch_patch.stop) + provider = TracerProvider() + self.processor = PostHogSpanProcessor(api_key="phc_test") + provider.add_span_processor(self.processor) + self.other_processor = MagicMock() + provider.add_span_processor(self.other_processor) + self.tracer = provider.get_tracer("test") + + def _emit_span(self, name="gen_ai.chat", attributes=None): + with self.tracer.start_as_current_span(name, attributes=attributes): + pass + + def _exported_attributes(self): + return [ + dict(call.args[0].attributes) for call in self.inner.on_end.call_args_list + ] + + def test_exports_distinct_id_session_and_tags(self): + with new_request_context(): + contexts.identify_context("alice") + contexts.set_context_session("s_alice") + contexts.tag("plan", "pro") + self._emit_span() + + self.assertEqual( + self._exported_attributes(), + [ + { + "posthog.distinct_id": "alice", + "$session_id": "s_alice", + "plan": "pro", + } + ], + ) + + def test_other_processors_do_not_see_the_context(self): + with new_request_context(): + contexts.identify_context("alice") + contexts.tag("plan", "pro") + self._emit_span(attributes={"gen_ai.system": "openai"}) + + started = self.other_processor.on_start.call_args.args[0] + ended = self.other_processor.on_end.call_args.args[0] + self.assertEqual(dict(started.attributes), {"gen_ai.system": "openai"}) + self.assertEqual(dict(ended.attributes), {"gen_ai.system": "openai"}) + + def test_drops_non_ai_spans_started_in_a_context(self): + with new_request_context(): + contexts.identify_context("alice") + self._emit_span(name="http.request") + + self.inner.on_end.assert_not_called() + self.assertEqual(self.processor._captured, {}) + + def test_shutdown_clears_captured_contexts(self): + with new_request_context(): + contexts.identify_context("alice") + with self.tracer.start_as_current_span("gen_ai.chat"): + self.assertEqual(len(self.processor._captured), 1) + self.processor.shutdown() + self.assertEqual(self.processor._captured, {}) + + def test_keeps_attributes_already_on_the_span(self): + with new_request_context(): + contexts.set_context_session("s_alice") + self._emit_span(attributes={"$session_id": "conv_42"}) + + self.assertEqual(self._exported_attributes()[0]["$session_id"], "conv_42") + + def test_skips_context_distinct_id_when_span_has_an_identity(self): + identity_keys = ( + "posthog.distinct_id", + "user.id", + "ai.telemetry.metadata.posthog_distinct_id", + "ai.settings.context.posthog_distinct_id", + ) + for key in identity_keys: + with self.subTest(key=key): + self.inner.on_end.reset_mock() + with new_request_context(): + contexts.identify_context("alice") + self._emit_span(attributes={key: "bob"}) + + attributes = self._exported_attributes()[0] + self.assertEqual(attributes[key], "bob") + if key != "posthog.distinct_id": + self.assertNotIn("posthog.distinct_id", attributes) + + def test_skips_distinct_id_tag_when_span_has_an_identity(self): + with new_request_context(): + contexts.tag("posthog.distinct_id", "alice") + self._emit_span(attributes={"user.id": "bob"}) + + attributes = self._exported_attributes()[0] + self.assertEqual(attributes, {"user.id": "bob"}) + + def test_encodes_structured_tags_as_json(self): + with new_request_context(): + contexts.tag("account", {"tier": "pro", "seats": 3}) + contexts.tag("arms", ["short-prompt", "new-router"]) + contexts.tag("mixed", ["a", 1]) + contexts.tag("nested_not_finite", {"x": float("inf")}) + self._emit_span() + + attributes = self._exported_attributes()[0] + self.assertEqual(json.loads(attributes["account"]), {"tier": "pro", "seats": 3}) + self.assertEqual(list(attributes["arms"]), ["short-prompt", "new-router"]) + self.assertEqual(json.loads(attributes["mixed"]), ["a", 1]) + self.assertNotIn("nested_not_finite", attributes) + + def test_skips_non_finite_numbers(self): + with new_request_context(): + contexts.tag("nan", float("nan")) + contexts.tag("inf", float("-inf")) + contexts.tag("list_with_inf", [1.0, float("inf")]) + contexts.tag("plan", "pro") + self._emit_span() + + self.assertEqual(self._exported_attributes(), [{"plan": "pro"}]) + + def test_string_like_values_arrive_as_plain_strings(self): + moment = datetime(2026, 10, 14, 8, 30, tzinfo=timezone.utc) + identifier = uuid.UUID("12345678-1234-5678-1234-567812345678") + with new_request_context(): + contexts.tag("when", moment) + contexts.tag("id", identifier) + contexts.tag("price", Decimal("1.50")) + self._emit_span() + + self.assertEqual( + self._exported_attributes(), + [{"when": moment.isoformat(), "id": str(identifier), "price": "1.50"}], + ) + + def test_skips_values_json_cannot_represent(self): + with new_request_context(): + contexts.tag("tags", {"a", "b"}) + contexts.tag("raw", b"bytes") + contexts.tag("plan", "pro") + self._emit_span() + + self.assertEqual(self._exported_attributes(), [{"plan": "pro"}]) + + def test_tags_win_over_context_ids_and_survive_an_empty_context(self): + with new_request_context(): + contexts.tag("$session_id", "conv_42") + self._emit_span() + with new_request_context(): + contexts.set_context_session("s_alice") + contexts.tag("$session_id", "conv_42") + self._emit_span() + + self.assertEqual( + [a["$session_id"] for a in self._exported_attributes()], + ["conv_42", "conv_42"], + ) + + def test_explicit_ai_session_id_tag_passes_through(self): + with new_request_context(): + contexts.set_context_session("s_alice") + contexts.tag("$ai_session_id", "conv_42") + self._emit_span() + + attributes = self._exported_attributes()[0] + self.assertEqual(attributes["$ai_session_id"], "conv_42") + self.assertEqual(attributes["$session_id"], "s_alice") + + def test_context_session_does_not_set_ai_session_id(self): + with new_request_context(): + contexts.set_context_session("s_alice") + self._emit_span() + + self.assertNotIn("$ai_session_id", self._exported_attributes()[0]) + + def test_integers_outside_signed_64_bit_range_are_sent_as_strings(self): + with new_request_context(): + contexts.tag("big", 2**63) + contexts.tag("small", -(2**63) - 1) + contexts.tag("edge", 2**63 - 1) + contexts.tag("bigs", [2**64, 1]) + self._emit_span() + + attributes = self._exported_attributes()[0] + self.assertEqual(attributes["big"], str(2**63)) + self.assertEqual(attributes["small"], str(-(2**63) - 1)) + self.assertEqual(attributes["edge"], 2**63 - 1) + self.assertEqual(json.loads(attributes["bigs"]), [2**64, 1]) + + def test_abandoned_span_entry_is_removed_when_garbage_collected(self): + with new_request_context(): + contexts.identify_context("alice") + span = self.tracer.start_span("gen_ai.chat") + self.assertEqual(len(self.processor._captured), 1) + + self.other_processor.reset_mock() + del span + gc.collect() + + self.assertEqual(self.processor._captured, {}) + + def test_collecting_an_abandoned_span_does_not_deadlock_on_the_held_lock(self): + def collect_while_holding_lock(): + with new_request_context(): + contexts.identify_context("alice") + holder = [] + span = self.tracer.start_span("gen_ai.chat") + holder.append(span) + holder.append(holder) + del span, holder + self.other_processor.reset_mock() + with self.processor._captured_lock: + gc.collect() + + worker = threading.Thread(target=collect_while_holding_lock, daemon=True) + worker.start() + worker.join(timeout=5) + + self.assertFalse(worker.is_alive(), "gc deadlocked on the captured lock") + self.assertEqual(self.processor._captured, {}) + + def test_finalizer_is_detached_once_the_span_ends(self): + with new_request_context(): + contexts.identify_context("alice") + span = self.tracer.start_span("gen_ai.chat") + _, finalizer = next(iter(self.processor._captured.values())) + self.assertTrue(finalizer.alive) + + span.end() + + self.assertFalse(finalizer.alive) + self.assertEqual(self.processor._captured, {}) + + def test_after_fork_resets_lock_and_store(self): + with new_request_context(): + contexts.identify_context("alice") + span = self.tracer.start_span("gen_ai.chat") + old_lock = self.processor._captured_lock + + _reset_processors_after_fork() + + self.assertIsNot(self.processor._captured_lock, old_lock) + self.assertEqual(self.processor._captured, {}) + span.end() + + def test_unstampable_tag_does_not_drop_the_other_tags(self): + class Unstampable(list): + def __iter__(self): + raise RuntimeError("cannot iterate") + + with new_request_context(): + contexts.tag("bad", Unstampable([1])) + contexts.tag("plan", "pro") + self._emit_span() + + self.assertEqual(self._exported_attributes(), [{"plan": "pro"}]) + + def test_adds_nothing_without_a_context(self): + self._emit_span() + + self.assertEqual(self._exported_attributes(), [{}]) + + def test_overlapping_requests_export_their_own_values(self): + both_started = threading.Barrier(2) + + def request(user): + with new_request_context(): + contexts.identify_context(user) + span = self.tracer.start_span("gen_ai.chat") + both_started.wait(timeout=5) + span.end() + + with ThreadPoolExecutor(max_workers=2) as pool: + list(pool.map(request, ["alice", "bob"])) + + self.assertEqual( + sorted(a["posthog.distinct_id"] for a in self._exported_attributes()), + ["alice", "bob"], + ) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index f8ce5837d..b66a265f1 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -568,6 +568,7 @@ attribute posthog.ai.openai.openai_providers.AzureOpenAI.embeddings: WrappedEmbe attribute posthog.ai.openai.openai_providers.AzureOpenAI.responses: WrappedResponses attribute posthog.ai.openai.wrapper_utils.log = logging.getLogger('posthog') attribute posthog.ai.openai_agents.processor.log = logging.getLogger('posthog') +attribute posthog.ai.otel.processor.log = logging.getLogger('posthog') attribute posthog.ai.otel.spans.AI_SPAN_PREFIXES = ('gen_ai.', 'llm.', 'ai.', 'traceloop.') attribute posthog.ai.otel.spans.DEFAULT_HOST = 'https://us.i.posthog.com' attribute posthog.ai.prompts.APP_ENDPOINT = 'https://us.posthog.com'