feat(ai): add per-request context to span processor

.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.

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
+    )

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::
 

posthog/ai/otel/processor.py

@@ -5,17 +5,55 @@
 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
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 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:
         """