From 3fa7ba9ebe53e3b92aba886184815bad9882bcfd Mon Sep 17 00:00:00 2001 From: "posthog-eu[bot]" <226701856+posthog-eu[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 11:25:17 +0000 Subject: [PATCH] Harden AI engine PostHog telemetry against silent, permanent halts The AI engine translates pydantic-ai OpenTelemetry spans into PostHog `$ai_generation`/`$ai_trace` events via a custom SpanProcessor. Two gaps in that path let LLM telemetry stop silently and stay at zero: - `PostHogSpanProcessor.on_end` was unguarded. OpenTelemetry calls it synchronously from `Span.end()` with no exception handling of its own, and pydantic-ai ends the model span inside a `with` block, so any error while translating a span propagates into the model request itself. A telemetry defect could therefore break the agent run and/or halt emission for every subsequent span. It now translates and captures defensively and logs-and-drops on failure. - The PostHog client was created without an `on_error` callback. Its `capture` is decorated `@no_throw` and delivery happens on a background consumer thread, so a failing or stalled delivery pipeline was completely invisible. We now log delivery failures and log the enabled/disabled state at startup, so a halt is noticeable in minutes instead of going unnoticed. Adds unit tests covering the happy path, per-trace dedup, non-chat spans, and that `on_end` never raises on delivery failure or a malformed span. Generated-By: PostHog Code Task-Id: 505d85dd-9cd2-4258-b76a-e027f4240277 --- engine/src/stirling/services/tracking.py | 54 ++++++++-- engine/tests/test_tracking.py | 132 +++++++++++++++++++++++ 2 files changed, 175 insertions(+), 11 deletions(-) create mode 100644 engine/tests/test_tracking.py diff --git a/engine/src/stirling/services/tracking.py b/engine/src/stirling/services/tracking.py index 0064611bbf..ed895cb24e 100644 --- a/engine/src/stirling/services/tracking.py +++ b/engine/src/stirling/services/tracking.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import logging from collections import OrderedDict from collections.abc import Mapping from contextvars import ContextVar @@ -29,6 +30,8 @@ from posthog.client import Client as PostHogClient from stirling.config import AppSettings from stirling.models import UserId +logger = logging.getLogger(__name__) + # Per-request user ID, set by middleware from the X-User-Id header. # When not set, PostHog generates a random ID and marks the event as personless. current_user_id: ContextVar[UserId | None] = ContextVar("current_user_id", default=None) @@ -127,17 +130,27 @@ class PostHogSpanProcessor(SpanProcessor): pass def on_end(self, span: ReadableSpan) -> None: - attrs = dict(span.attributes or {}) - if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value: - return + # OpenTelemetry calls this synchronously from ``Span.end()`` with no + # exception handling of its own (see SynchronousMultiSpanProcessor), + # and pydantic-ai ends the model span inside a ``with`` block. So any + # exception raised here propagates straight into the model request and + # breaks the agent run. A telemetry defect must never take down the + # assistant or silently halt emission for every subsequent span, so we + # translate and capture defensively and swallow (but log) any failure. + try: + attrs = dict(span.attributes or {}) + if attrs.get(GEN_AI_OPERATION_NAME) != GenAiOperationNameValues.CHAT.value: + return - properties = self._build_generation_properties(span, attrs) - self._maybe_emit_trace_event(span, attrs, properties) - self._client.capture( - distinct_id=current_user_id.get(), - event="$ai_generation", - properties=properties, - ) + properties = self._build_generation_properties(span, attrs) + self._maybe_emit_trace_event(span, attrs, properties) + self._client.capture( + distinct_id=current_user_id.get(), + event="$ai_generation", + properties=properties, + ) + except Exception: + logger.exception("Failed to emit $ai_generation for span; dropping this event") def _build_generation_properties(self, span: ReadableSpan, attrs: Mapping[str, Any]) -> dict[str, object]: """Build the $ai_generation event properties from span data.""" @@ -231,6 +244,19 @@ class PostHogSpanProcessor(SpanProcessor): return True +def _log_delivery_error(error: Exception, _batch: Any) -> None: + """PostHog client ``on_error`` callback. + + The PostHog SDK delivers events from a background consumer thread and its + ``capture`` method is decorated ``@no_throw``, so without this hook a + failing (or dead) delivery pipeline is completely invisible: events simply + stop arriving with nothing in our logs. Surfacing the error lets a halt be + noticed in minutes instead of sitting silently at zero. Kept lightweight so + a transient upload error can't spam at capture volume. + """ + logger.warning("PostHog event delivery failed: %s", error) + + def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None: """Configure OpenTelemetry with a PostHog span processor for LLM analytics. @@ -238,11 +264,17 @@ def setup_posthog_tracking(settings: AppSettings) -> TracerProvider | None: or None when tracking is disabled. """ if not settings.posthog_enabled or not settings.posthog_api_key: + logger.info("PostHog LLM tracking is disabled (posthog_enabled=%s)", settings.posthog_enabled) return None - client = PostHogClient(project_api_key=settings.posthog_api_key, host=settings.posthog_host) + client = PostHogClient( + project_api_key=settings.posthog_api_key, + host=settings.posthog_host, + on_error=_log_delivery_error, + ) processor = PostHogSpanProcessor(client) provider = TracerProvider() provider.add_span_processor(processor) + logger.info("PostHog LLM tracking enabled (host=%s)", settings.posthog_host) return provider diff --git a/engine/tests/test_tracking.py b/engine/tests/test_tracking.py new file mode 100644 index 0000000000..55c8f51357 --- /dev/null +++ b/engine/tests/test_tracking.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any + +from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import ( + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_SYSTEM, + GenAiOperationNameValues, +) + +from stirling.config import AppSettings +from stirling.services.tracking import PostHogSpanProcessor, current_user_id, setup_posthog_tracking + + +@dataclass +class _FakeContext: + trace_id: int + span_id: int + + +@dataclass +class _FakeParent: + span_id: int + + +class _FakeSpan: + """Minimal stand-in for the OpenTelemetry ReadableSpan fields the processor reads.""" + + def __init__(self, attributes: dict[str, Any], trace_id: int = 0x1234, span_id: int = 0x9) -> None: + self.attributes = attributes + self.context = _FakeContext(trace_id=trace_id, span_id=span_id) + self.parent = _FakeParent(span_id=0x1) + self.start_time = 1_000_000_000 + self.end_time = 2_000_000_000 + + +class _RecordingClient: + def __init__(self) -> None: + self.events: list[tuple[str, dict[str, Any]]] = [] + + def capture(self, *, distinct_id: str | None, event: str, properties: dict[str, Any]) -> None: + self.events.append((event, properties)) + + +class _RaisingClient: + def capture(self, **_kwargs: Any) -> None: + raise RuntimeError("boom") + + +def _processor(client: object) -> PostHogSpanProcessor: + # The processor only calls ``capture`` on its client; the test doubles model + # that surface without being real PostHog clients (which start threads). + return PostHogSpanProcessor(client) # type: ignore[arg-type] + + +def _on_end(processor: PostHogSpanProcessor, span: object) -> None: + processor.on_end(span) # type: ignore[arg-type] + + +def _chat_attrs() -> dict[str, Any]: + return { + GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CHAT.value, + GEN_AI_SYSTEM: "anthropic", + GEN_AI_INPUT_MESSAGES: json.dumps([{"role": "user", "parts": [{"type": "text", "content": "hello"}]}]), + } + + +def test_chat_span_emits_generation_and_trace() -> None: + client = _RecordingClient() + + _on_end(_processor(client), _FakeSpan(_chat_attrs())) + + emitted = [event for event, _ in client.events] + assert "$ai_generation" in emitted + assert "$ai_trace" in emitted + + +def test_trace_event_deduplicated_per_trace() -> None: + client = _RecordingClient() + processor = _processor(client) + + _on_end(processor, _FakeSpan(_chat_attrs(), trace_id=0x42, span_id=0x1)) + _on_end(processor, _FakeSpan(_chat_attrs(), trace_id=0x42, span_id=0x2)) + + trace_events = [event for event, _ in client.events if event == "$ai_trace"] + generation_events = [event for event, _ in client.events if event == "$ai_generation"] + assert len(trace_events) == 1 + assert len(generation_events) == 2 + + +def test_non_chat_span_is_ignored() -> None: + client = _RecordingClient() + + _on_end(_processor(client), _FakeSpan({GEN_AI_OPERATION_NAME: "embeddings"})) + + assert client.events == [] + + +def test_on_end_never_raises_when_delivery_fails() -> None: + # A telemetry failure must not propagate into Span.end() and break the + # model request, nor wedge emission for every subsequent span. + _on_end(_processor(_RaisingClient()), _FakeSpan(_chat_attrs())) # must not raise + + +def test_on_end_never_raises_on_malformed_span() -> None: + processor = _processor(_RecordingClient()) + + broken = _FakeSpan(_chat_attrs()) + broken.context = None # type: ignore[assignment] + + _on_end(processor, broken) # must not raise + + +def test_distinct_id_pulled_from_context_var() -> None: + client = _RecordingClient() + processor = _processor(client) + + token = current_user_id.set("user-123") # type: ignore[arg-type] + try: + _on_end(processor, _FakeSpan(_chat_attrs())) + finally: + current_user_id.reset(token) + + assert client.events, "expected at least one captured event" + + +def test_setup_returns_none_when_disabled(app_settings: AppSettings) -> None: + # conftest builds settings with posthog disabled; setup must no-op. + assert setup_posthog_tracking(app_settings) is None