diff --git a/posthog/test/tracing/__init__.py b/posthog/test/tracing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/posthog/test/tracing/test_ids.py b/posthog/test/tracing/test_ids.py new file mode 100644 index 00000000..b677cd9e --- /dev/null +++ b/posthog/test/tracing/test_ids.py @@ -0,0 +1,42 @@ +import re +from unittest import mock + +from posthog.tracing import _ids +from posthog.tracing._ids import ( + new_span_id, + new_trace_id, +) + +LOWER_HEX = re.compile(r"^[0-9a-f]+$") + + +class TestNewTraceId: + def test_is_32_lowercase_hex_characters(self): + trace_id = new_trace_id() + assert len(trace_id) == 32 + assert LOWER_HEX.match(trace_id) + + def test_is_never_all_zeros(self): + with mock.patch.object( + _ids.secrets, "token_hex", side_effect=["0" * 32, "0" * 32, "ab" * 16] + ): + assert new_trace_id() == "ab" * 16 + + def test_does_not_repeat(self): + assert len({new_trace_id() for _ in range(1000)}) == 1000 + + +class TestNewSpanId: + def test_is_16_lowercase_hex_characters(self): + span_id = new_span_id() + assert len(span_id) == 16 + assert LOWER_HEX.match(span_id) + + def test_is_never_all_zeros(self): + with mock.patch.object( + _ids.secrets, "token_hex", side_effect=["0" * 16, "cd" * 8] + ): + assert new_span_id() == "cd" * 8 + + def test_does_not_repeat(self): + assert len({new_span_id() for _ in range(1000)}) == 1000 diff --git a/posthog/test/tracing/test_traceparent.py b/posthog/test/tracing/test_traceparent.py new file mode 100644 index 00000000..3c6b06d4 --- /dev/null +++ b/posthog/test/tracing/test_traceparent.py @@ -0,0 +1,196 @@ +import pytest + +from posthog.tracing._traceparent import ( + RemoteSpanContext, + format_traceparent, + normalize_traceparent, + parse_traceparent, + sanitize_tracestate, + traceparent_header, +) + +TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736" +SPAN_ID = "00f067aa0ba902b7" + + +class TestParseTraceparent: + def test_parses_a_sampled_header(self): + assert parse_traceparent(f"00-{TRACE_ID}-{SPAN_ID}-01") == RemoteSpanContext( + TRACE_ID, SPAN_ID, "01" + ) + + def test_continues_a_sampled_out_trace_and_keeps_the_flag(self): + assert parse_traceparent(f"00-{TRACE_ID}-{SPAN_ID}-00") == RemoteSpanContext( + TRACE_ID, SPAN_ID, "00" + ) + + @pytest.mark.parametrize( + "inbound,expected", [("05", "01"), ("04", "00"), ("ff", "01")] + ) + def test_zeroes_flags_version_00_does_not_define(self, inbound, expected): + parsed = parse_traceparent(f"00-{TRACE_ID}-{SPAN_ID}-{inbound}") + assert parsed is not None + assert parsed.flags == expected + + def test_accepts_a_future_version_with_extra_fields(self): + parsed = parse_traceparent(f"01-{TRACE_ID}-{SPAN_ID}-01-extra") + assert parsed == RemoteSpanContext(TRACE_ID, SPAN_ID, "01") + + def test_ignores_surrounding_whitespace(self): + parsed = parse_traceparent(f" 00-{TRACE_ID}-{SPAN_ID}-01 ") + assert parsed == RemoteSpanContext(TRACE_ID, SPAN_ID, "01") + + def test_decodes_a_raw_asgi_header_value(self): + parsed = parse_traceparent(f"00-{TRACE_ID}-{SPAN_ID}-01".encode("ascii")) + assert parsed == RemoteSpanContext(TRACE_ID, SPAN_ID, "01") + + @pytest.mark.parametrize( + "value", + [ + f"00-{TRACE_ID.upper()}-{SPAN_ID}-01", + f"00-{TRACE_ID}-{SPAN_ID.upper()}-01", + f"00-{TRACE_ID}-{SPAN_ID}-0A", + ], + ) + def test_rejects_uppercase_rather_than_folding_it(self, value): + assert parse_traceparent(value) is None + + def test_rejects_a_version_00_header_with_trailing_fields(self): + assert parse_traceparent(f"00-{TRACE_ID}-{SPAN_ID}-01-extra") is None + + @pytest.mark.parametrize( + "value", + [ + "garbage", + "", + f"ff-{TRACE_ID}-{SPAN_ID}-01", + f"00-{'0' * 32}-{SPAN_ID}-01", + f"00-{TRACE_ID}-{'0' * 16}-01", + f"00-{TRACE_ID[:30]}-{SPAN_ID}-01", + f"00-{TRACE_ID}-{SPAN_ID}", + 42, + None, + [f"00-{TRACE_ID}-{SPAN_ID}-01"], + "00-caf\u00e9".encode("utf-8"), + ], + ) + def test_returns_none_for_malformed_input(self, value): + assert parse_traceparent(value) is None + + +class TestFormatTraceparent: + def test_sets_the_sampled_flag_on_a_trace_started_here(self): + assert format_traceparent(TRACE_ID, SPAN_ID) == f"00-{TRACE_ID}-{SPAN_ID}-01" + + def test_propagates_the_flags_byte_it_was_given(self): + assert ( + format_traceparent(TRACE_ID, SPAN_ID, "00") == f"00-{TRACE_ID}-{SPAN_ID}-00" + ) + + def test_round_trips_through_the_parser(self): + assert parse_traceparent(format_traceparent(TRACE_ID, SPAN_ID, "00")) == ( + RemoteSpanContext(TRACE_ID, SPAN_ID, "00") + ) + + +class TestNormalizeTraceparent: + def test_carries_version_and_flags_through_as_received(self): + assert ( + normalize_traceparent(f"01-{TRACE_ID}-{SPAN_ID}-05") + == f"01-{TRACE_ID}-{SPAN_ID}-05" + ) + + def test_decodes_a_raw_asgi_header_value(self): + header = f"00-{TRACE_ID}-{SPAN_ID}-01" + assert normalize_traceparent(header.encode("ascii")) == header + + def test_echoes_a_higher_version_whole_including_its_trailing_fields(self): + assert ( + normalize_traceparent(f" 01-{TRACE_ID}-{SPAN_ID}-01-what ") + == f"01-{TRACE_ID}-{SPAN_ID}-01-what" + ) + + @pytest.mark.parametrize( + "value", + [ + "garbage", + f"ff-{TRACE_ID}-{SPAN_ID}-01", + f"00-{'0' * 32}-{SPAN_ID}-01", + f"00-{TRACE_ID.upper()}-{SPAN_ID}-01", + f"00-{TRACE_ID}-{SPAN_ID}-01-extra", + f"01-{TRACE_ID}-{SPAN_ID}-01-x\rX-Injected: yes", + f"01-{TRACE_ID}-{SPAN_ID}-01-caf\u00e9", + f"01-{TRACE_ID}-{SPAN_ID}-01-" + "x" * 512, + ["a", "b"], + ], + ) + def test_rejects_malformed_input(self, value): + assert normalize_traceparent(value) is None + + +class TestTraceparentHeader: + def test_unwraps_a_one_element_list(self): + header = f"00-{TRACE_ID}-{SPAN_ID}-01" + assert traceparent_header([header]) == header + assert traceparent_header((header,)) == header + + def test_leaves_two_values_for_the_parser_to_reject(self): + values = [f"00-{TRACE_ID}-{SPAN_ID}-01", f"00-{TRACE_ID}-{SPAN_ID}-00"] + assert traceparent_header(values) == values + assert parse_traceparent(traceparent_header(values)) is None + + def test_passes_a_string_through(self): + assert traceparent_header("x") == "x" + + +class TestSanitizeTracestate: + def test_preserves_a_valid_vendor_list_unchanged(self): + assert sanitize_tracestate("vendor=abc,other=def") == "vendor=abc,other=def" + + def test_decodes_a_raw_asgi_header_value(self): + assert sanitize_tracestate(b"vendor=abc") == "vendor=abc" + + def test_trims_surrounding_whitespace(self): + assert sanitize_tracestate(" vendor=abc ") == "vendor=abc" + + def test_keeps_a_tab_separated_vendor_list(self): + assert sanitize_tracestate("vendor=abc,\tother=def") == "vendor=abc,\tother=def" + + def test_tolerates_empty_members(self): + assert sanitize_tracestate("vendor=abc,,other=def") == "vendor=abc,,other=def" + + @pytest.mark.parametrize( + "value", + [ + "", + " ", + "novalue", + "vendor=abc,novalue", + 42, + None, + ",".join(f"k{i}=v" for i in range(33)), + # One member, over 512 characters: nothing survives the trim. + "k=" + "v" * 600, + "vendor=abc\r\nother=def", + "vendor=\ud800", + "vendor=é", + ], + ) + def test_discards_invalid_values(self, value): + assert sanitize_tracestate(value) is None + + def test_trims_an_over_long_value_by_dropping_large_members_first(self): + # The spec's scenario: 600 characters, one member of 150. Dropping that + # member leaves 449 characters, kept unchanged. + large = "big=" + "x" * 146 + small = ["k{:02d}=".format(i) + "v" * 13 for i in range(25)] + value = ",".join(small[:10] + [large] + small[10:]) + assert len(large) == 150 and len(value) == 600 + trimmed = sanitize_tracestate(value) + assert trimmed == ",".join(small) + assert len(trimmed) == 449 + + def test_trims_from_the_right_once_no_large_member_is_left(self): + members = ["k{:02d}=".format(i) + "v" * 60 for i in range(10)] + # Ten 64-character members; seven fit in 512. + assert sanitize_tracestate(",".join(members)) == ",".join(members[:7]) diff --git a/posthog/tracing/__init__.py b/posthog/tracing/__init__.py new file mode 100644 index 00000000..7d17992e --- /dev/null +++ b/posthog/tracing/__init__.py @@ -0,0 +1,3 @@ +"""Distributed tracing: spans exported as OTLP JSON to PostHog.""" + +__all__: list[str] = [] diff --git a/posthog/tracing/_ids.py b/posthog/tracing/_ids.py new file mode 100644 index 00000000..7e885808 --- /dev/null +++ b/posthog/tracing/_ids.py @@ -0,0 +1,32 @@ +"""W3C Trace Context identifiers: 16-byte trace ids and 8-byte span ids, lowercase hex. + +The ingestion service zeroes an id of the wrong length rather than rejecting +it, which silently orphans the span. Ids are valid by construction: generated +here, or accepted from a header the parser matched in full. +""" + +import secrets + +TRACE_ID_BYTES = 16 +SPAN_ID_BYTES = 8 + +TRACE_ID_HEX = TRACE_ID_BYTES * 2 +SPAN_ID_HEX = SPAN_ID_BYTES * 2 + +INVALID_TRACE_ID = "0" * TRACE_ID_HEX +INVALID_SPAN_ID = "0" * SPAN_ID_HEX + + +def _random_hex_id(byte_length: int) -> str: + hex_id = secrets.token_hex(byte_length) + while hex_id == "0" * (byte_length * 2): + hex_id = secrets.token_hex(byte_length) + return hex_id + + +def new_trace_id() -> str: + return _random_hex_id(TRACE_ID_BYTES) + + +def new_span_id() -> str: + return _random_hex_id(SPAN_ID_BYTES) diff --git a/posthog/tracing/_traceparent.py b/posthog/tracing/_traceparent.py new file mode 100644 index 00000000..a30aa9dd --- /dev/null +++ b/posthog/tracing/_traceparent.py @@ -0,0 +1,154 @@ +"""W3C Trace Context interop: ``traceparent`` and ``tracestate`` handling.""" + +import re +from dataclasses import dataclass +from typing import List, Optional + +from ._ids import INVALID_SPAN_ID, INVALID_TRACE_ID + +TRACE_FLAGS_SAMPLED = "01" +TRACE_FLAGS_UNSAMPLED = "00" + +# A version above `00` may append fields after the first four. +_TRACEPARENT_RE = re.compile( + r"^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})(-.*)?$" +) + +# W3C bounds neither header; the caps keep a peer's header from being +# amplified onto every outbound call. +_TRACEPARENT_MAX_LENGTH = 512 +_TRACESTATE_MAX_MEMBERS = 32 +_TRACESTATE_MAX_LENGTH = 512 +# W3C: when trimming, drop members longer than this first. +_TRACESTATE_LARGE_MEMBER_LENGTH = 128 +_NON_PRINTABLE_ASCII_RE = re.compile(r"[^\x20-\x7e\t]") + + +@dataclass(frozen=True) +class RemoteSpanContext: + trace_id: str + span_id: str + # `01` sampled or `00` sampled out; other flag bits are zeroed. + flags: str + + +@dataclass(frozen=True) +class _TraceparentFields: + version: str + trace_id: str + span_id: str + flags: str + + +def _header_text(value: object) -> Optional[str]: + """A header value as text; a raw ASGI scope carries it as ASCII bytes.""" + if isinstance(value, bytes): + try: + return value.decode("ascii") + except UnicodeDecodeError: + return None + return value if isinstance(value, str) else None + + +def _match_traceparent(value: object) -> Optional[_TraceparentFields]: + value = _header_text(value) + if value is None: + return None + # Not lowercased: a conformant peer restarts the trace on uppercase hex. + stripped = value.strip() + if len(stripped) > _TRACEPARENT_MAX_LENGTH: + return None + match = _TRACEPARENT_RE.match(stripped) + if match is None: + return None + version, trace_id, span_id, flags, trailing = match.group(1, 2, 3, 4, 5) + if version == "ff" or (version == "00" and trailing): + return None + if trailing and _NON_PRINTABLE_ASCII_RE.search(trailing): + return None + if trace_id == INVALID_TRACE_ID or span_id == INVALID_SPAN_ID: + return None + return _TraceparentFields(version, trace_id, span_id, flags) + + +def parse_traceparent(value: object) -> Optional[RemoteSpanContext]: + """Parse an inbound ``traceparent``; ``None`` for anything malformed. + + A sampled-out (``00``) trace is still continued, and the flag is carried + onward so a downstream sampler sees the caller's decision. + """ + fields = _match_traceparent(value) + if fields is None: + return None + # W3C requires zeroing the flag bits version `00` does not define. + flags = ( + TRACE_FLAGS_SAMPLED if int(fields.flags, 16) & 0x01 else TRACE_FLAGS_UNSAMPLED + ) + return RemoteSpanContext(fields.trace_id, fields.span_id, flags) + + +def normalize_traceparent(value: object) -> Optional[str]: + """The inbound ``traceparent`` as received, or ``None`` when it is malformed. + + Echoed whole rather than rebuilt, so a higher version keeps the fields this + SDK does not read. + """ + text = _header_text(value) + if text is None or _match_traceparent(text) is None: + return None + return text.strip() + + +def traceparent_header(value: object) -> object: + """Unwrap the one-element list a multi-value header API returns. + + A longer list holds two different inbound values and is left to be rejected. + """ + if isinstance(value, (list, tuple)) and len(value) == 1: + return value[0] + return value + + +def format_traceparent( + trace_id: str, span_id: str, flags: str = TRACE_FLAGS_SAMPLED +) -> str: + return f"00-{trace_id}-{span_id}-{flags}" + + +def sanitize_tracestate(value: object) -> Optional[str]: + """Validate an inbound ``tracestate``; ``None`` when it is malformed. + + An invalid value is discarded without invalidating its traceparent. One + over 512 characters is trimmed by whole members. + """ + text = _header_text(value) + if text is None: + return None + trimmed = text.strip() + if not trimmed or _NON_PRINTABLE_ASCII_RE.search(trimmed): + return None + members = trimmed.split(",") + if len(members) > _TRACESTATE_MAX_MEMBERS: + return None + if any(member.strip() and "=" not in member for member in members): + return None + if len(trimmed) <= _TRACESTATE_MAX_LENGTH: + return trimmed + return _trim_to_length(members) + + +def _trim_to_length(members: List[str]) -> Optional[str]: + """Drop large members first, then from the right, keeping those nearest the caller.""" + kept = list(members) + + def joined_length() -> int: + return sum(len(member) for member in kept) + len(kept) - 1 + + for index in range(len(kept) - 1, -1, -1): + if joined_length() <= _TRACESTATE_MAX_LENGTH: + break + if len(kept[index]) > _TRACESTATE_LARGE_MEMBER_LENGTH: + del kept[index] + while kept and joined_length() > _TRACESTATE_MAX_LENGTH: + kept.pop() + return ",".join(kept) if kept else None diff --git a/pyproject.toml b/pyproject.toml index 54ad4b8a..a4525caa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,12 +113,14 @@ packages = [ "posthog.ai.claude_agent_sdk", "posthog.ai.otel", "posthog.mcp", + "posthog.tracing", "posthog.test", "posthog.test.ai", "posthog.test.ai.openai_agents", "posthog.test.ai.claude_agent_sdk", "posthog.test.ai.otel", "posthog.test.mcp", + "posthog.test.tracing", "posthog.integrations", ] diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 36b28e43..0650028b 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -1558,6 +1558,7 @@ module posthog.mcp.version module posthog.metrics_capture module posthog.poller module posthog.request +module posthog.tracing module posthog.types module posthog.utils module posthog.version