From 698b3a42f322ec2ce3c4436d687968b5ace33db4 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 6 Oct 2026 12:02:58 -0700 Subject: [PATCH 01/14] chore: add major changeset for posthog 8.0 capture v1 release --- .sampo/changesets/capture-v1-major.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .sampo/changesets/capture-v1-major.md diff --git a/.sampo/changesets/capture-v1-major.md b/.sampo/changesets/capture-v1-major.md new file mode 100644 index 000000000..9271aee97 --- /dev/null +++ b/.sampo/changesets/capture-v1-major.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: major +--- + +Capture v1 is the only capture path. Events and AI events send to the capture v1 endpoints, and the legacy v0 capture path is removed. See the migration guide for breaking changes. From 5b43ec7775172f24ace706db90a4dae02f353919 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Sat, 10 Oct 2026 15:37:37 -0700 Subject: [PATCH 02/14] feat!: send capture through v1 by default; capture_ai to /i/v1/ai/events (#1017) feat!: send capture through v1 by default; capture_ai posts to /i/v1/ai/events Capture v1 is now the default capture mode. The AI lane always posts capture v1 to /i/v1/ai/events, uncompressed, regardless of the analytics capture_mode. Tests move to a shared capture v1 send seam (posthog/test/capture_helpers.py); v0-only tests are pinned to CaptureMode.V0 until v0 removal. --- AGENTS.md | 2 +- posthog/__init__.py | 2 +- posthog/capture_mode.py | 17 +- posthog/capture_v1.py | 19 +- posthog/client.py | 44 +-- posthog/consumer.py | 19 +- posthog/test/capture_helpers.py | 58 ++++ posthog/test/test_ai_capture_lane.py | 143 +++++----- posthog/test/test_async_client.py | 123 ++++----- posthog/test/test_async_feature_flags.py | 13 +- posthog/test/test_before_send.py | 66 ++--- posthog/test/test_capture_mode.py | 24 +- posthog/test/test_capture_v1.py | 9 +- posthog/test/test_client.py | 297 +++++++++++---------- posthog/test/test_code_variables.py | 10 +- posthog/test/test_consumer.py | 53 ++-- posthog/test/test_exception_capture.py | 12 +- posthog/test/test_gevent_compat.py | 9 +- posthog/test/test_module.py | 7 +- posthog/test/test_release_id.py | 7 +- posthog/test/tracing/test_client_traces.py | 5 +- references/public_api_snapshot.txt | 3 +- sdk_compliance_adapter/adapter.py | 9 +- sdk_compliance_adapter/test_adapter.py | 2 +- 24 files changed, 534 insertions(+), 419 deletions(-) create mode 100644 posthog/test/capture_helpers.py diff --git a/AGENTS.md b/AGENTS.md index 745e91ad7..e05e0a255 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -23,7 +23,7 @@ Follow [Public API changes](./CONTRIBUTING.md#public-api-changes). As an agent, Before changing capture configuration, serialization, routing, or retries, read the relevant implementation and tests. -Preserve v0 defaults/compatibility; strictly typed v1 options and `$set`/`$set_once` relocation; v1-only compression (zlib-wrapped deflate, optional zstd); partial-only per-event retries with stable identity; accumulated drop reporting even on 2xx; terminal v1 `429`; `Retry-After` as a minimum bounded by the shared 30s ceiling; and inline blocking retries with `sync_mode=True`. +Preserve capture v1 as the default (v0 is opt-in for analytics only; `capture_ai` always posts v1 to `/i/v1/ai/events`); strictly typed v1 options and `$set`/`$set_once` relocation; v1-only compression (zlib-wrapped deflate, optional zstd); partial-only per-event retries with stable identity; accumulated drop reporting even on 2xx; terminal v1 `429`; `Retry-After` as a minimum bounded by the shared 30s ceiling; and inline blocking retries with `sync_mode=True`. ## Mirror and build safety diff --git a/posthog/__init__.py b/posthog/__init__.py index 5d6ae1193..ccbc8db58 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -425,7 +425,7 @@ def get_tags() -> Dict[str, Any]: enable_local_evaluation = True # type: bool flag_definition_cache_provider = None # type: Optional[FlagDefinitionCacheProvider] # Capture wire protocol for the global client. None defers to POSTHOG_CAPTURE_MODE -# then CaptureMode.V0. See posthog.capture_mode.CaptureMode. +# then CaptureMode.V1. See posthog.capture_mode.CaptureMode. capture_mode = None # type: Optional[CaptureMode] # Routes AI SDK wrapper events through the dedicated AI capture lane, skips # truncation, and passes media unredacted. `privacy_mode` always wins. diff --git a/posthog/capture_mode.py b/posthog/capture_mode.py index 7ea3696f2..5f72e3702 100644 --- a/posthog/capture_mode.py +++ b/posthog/capture_mode.py @@ -13,10 +13,9 @@ class CaptureMode(str, Enum): """Selects the capture wire protocol used for event ingestion. - ``V0`` is the legacy ``POST /batch/`` endpoint and the default, so upgrading - is transparent to existing callers. ``V1`` opts into - ``POST /i/v1/analytics/events`` (Bearer auth, per-event results, partial - retry). Inheriting from ``str`` keeps the members directly comparable to and + ``V1`` is ``POST /i/v1/analytics/events`` (Bearer auth, per-event results, + partial retry) and the default. ``V0`` opts back into the legacy + ``POST /batch/`` endpoint. Inheriting from ``str`` keeps the members directly comparable to and serializable as their ``"v0"`` / ``"v1"`` values. """ @@ -61,15 +60,15 @@ def _resolve_capture_mode( """Resolve the effective capture mode. Precedence: explicit ``capture_mode`` argument > ``POSTHOG_CAPTURE_MODE`` env - var > ``CaptureMode.V0``. An unrecognized env value logs a warning and falls - back to ``V0`` so a typo never silently flips the wire protocol. + var > ``CaptureMode.V1``. An unrecognized env value logs a warning and falls + back to ``V1`` so a typo never silently flips the wire protocol. """ if capture_mode is not None: return _coerce_explicit(capture_mode) raw = os.environ.get(CAPTURE_MODE_ENV_VAR) if raw is None or raw.strip() == "": - return CaptureMode.V0 + return CaptureMode.V1 resolved = _ALIASES.get(raw.strip().lower()) if resolved is None: @@ -77,8 +76,8 @@ def _resolve_capture_mode( "Unrecognized %s=%r; falling back to %s. Expected one of %s.", CAPTURE_MODE_ENV_VAR, raw, - CaptureMode.V0.value, + CaptureMode.V1.value, sorted(_ALIASES), ) - return CaptureMode.V0 + return CaptureMode.V1 return resolved diff --git a/posthog/capture_v1.py b/posthog/capture_v1.py index 478175484..438aaf011 100644 --- a/posthog/capture_v1.py +++ b/posthog/capture_v1.py @@ -1,9 +1,10 @@ """Serialization and transport for the Capture V1 wire protocol. -This module owns everything specific to ``POST /i/v1/analytics/events``: the -*transform* layer (legacy-shaped queued message -> v1 wire event + batch -envelope) and the *transport* layer (a single HTTP attempt, response parsing, -and the partial-retry send loop). +This module owns everything specific to the capture v1 endpoints +(``POST /i/v1/analytics/events`` and ``POST /i/v1/ai/events``, which share one +wire contract): the *transform* layer (legacy-shaped queued message -> v1 wire +event + batch envelope) and the *transport* layer (a single HTTP attempt, +response parsing, and the partial-retry send loop). The v1 contract (see ``rust/capture/src/v1/analytics/types.rs``) differs from the legacy ``/batch/`` shape in a few load-bearing ways that this module @@ -66,6 +67,7 @@ __all__ = ["CaptureV1Error"] _CAPTURE_V1_PATH = "/i/v1/analytics/events" +_CAPTURE_AI_V1_PATH = "/i/v1/ai/events" # Required request/response headers for the v1 endpoint. Defined here as the # single source of truth; the transport layer builds requests from them. @@ -350,8 +352,9 @@ def _post_v1( timeout: int = 15, sdk_info: str = USER_AGENT, session: Optional["requests.Session"] = None, + path: str = _CAPTURE_V1_PATH, ) -> "requests.Response": - """Perform a single ``POST /i/v1/analytics/events`` attempt. + """Perform a single capture v1 ``POST`` to ``path``. Bearer-authed (no ``api_key`` in the body) with the required v1 headers. ``attempt`` (1-based) and the stable ``request_id`` are echoed via @@ -361,7 +364,7 @@ def _post_v1( the caller. """ trimmed_host = remove_trailing_slash(normalize_host(host)) - url = trimmed_host + _CAPTURE_V1_PATH + url = trimmed_host + path data = json.dumps(batch_body, cls=DatetimeSerializer) headers = { "Content-Type": "application/json", @@ -473,8 +476,9 @@ def _send_v1_batch( historical_migration: bool = False, sdk_info: str = USER_AGENT, session: Optional["requests.Session"] = None, + path: str = _CAPTURE_V1_PATH, ) -> None: - """Deliver ``batch`` to the v1 endpoint with partial retry. + """Deliver ``batch`` to the v1 endpoint at ``path`` with partial retry. The v1 sibling of ``Consumer._send``: it loops up to ``max_retries + 1`` attempts, but unlike v0 it shrinks the batch to only the events the server @@ -525,6 +529,7 @@ def _send_v1_batch( timeout=timeout, sdk_info=sdk_info, session=session, + path=path, ) except Exception as e: # Transport-level failure (connection/timeout): retry like v0 does. diff --git a/posthog/client.py b/posthog/client.py index 2bb16c492..d2d705d52 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -32,7 +32,11 @@ _resolve_capture_compression, ) from posthog.capture_mode import CaptureMode, _resolve_capture_mode -from posthog.capture_v1 import _send_v1_batch +from posthog.capture_v1 import ( + _CAPTURE_AI_V1_PATH, + _CAPTURE_V1_PATH, + _send_v1_batch, +) from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE, Consumer, _DrainSignal from posthog.contexts import ( _get_current_context, @@ -85,7 +89,6 @@ from posthog.poller import Poller from posthog.release_id import _resolve_release_id from posthog.request import ( - AI_EVENTS_ENDPOINT, EVENTS_ENDPOINT, USER_AGENT as _USER_AGENT, APIError, @@ -841,10 +844,10 @@ def __init__( exception_autocapture_refill_interval_seconds: Seconds between token refills for autocaptured exception rate limiting. capture_mode: Capture wire protocol to use. Defaults to - ``CaptureMode.V0`` (legacy ``/batch/``). Set ``CaptureMode.V1`` - (or pass the string ``"v1"``) to opt into - ``/i/v1/analytics/events``. When omitted, the - ``POSTHOG_CAPTURE_MODE`` env var is consulted, then ``V0``. + ``CaptureMode.V1`` (``/i/v1/analytics/events``). Set + ``CaptureMode.V0`` (or pass the string ``"v0"``) to opt back + into the legacy ``/batch/`` endpoint. When omitted, the + ``POSTHOG_CAPTURE_MODE`` env var is consulted, then ``V1``. capture_compression: Request-body compression for capture-v1 uploads (ignored in V0, which uses ``gzip``). ``CaptureCompression.GZIP`` or ``DEFLATE`` (or the strings ``"gzip"``/``"deflate"``). When @@ -1085,24 +1088,27 @@ def __init__( self._analytics_lane = _Lane( name="analytics", **lane_defaults, - endpoint=EVENTS_ENDPOINT, + endpoint=( + _CAPTURE_V1_PATH + if self.capture_mode == CaptureMode.V1 + else EVENTS_ENDPOINT + ), max_msg_size=MAX_MSG_SIZE, capture_mode=self.capture_mode, capture_compression=self.capture_compression, eager_start=not sync_mode, ) - # The AI lane is pinned to the v0 submitter: the AI endpoint has no v1 - # form, and this keeps multi-MB AI events away from capture v1's - # smaller caps. The `capture_compression` pin is inert on v0 — its wire - # compression is the `gzip` flag, inherited from client config. Lazy - # start, so the many clients that never emit AI events pay for no - # extra threads. + # The AI lane always posts capture v1 to the AI endpoint, whatever the + # analytics `capture_mode`, so multi-MB AI events stay off the + # analytics endpoint's smaller caps. It sends uncompressed. Lazy start, + # so the many clients that never emit AI events pay for no extra + # threads. self._ai_lane = _Lane( name="ai", **lane_defaults, - endpoint=AI_EVENTS_ENDPOINT, + endpoint=_CAPTURE_AI_V1_PATH, max_msg_size=AI_MAX_MSG_SIZE, - capture_mode=CaptureMode.V0, + capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.NONE, eager_start=False, ) @@ -2438,19 +2444,19 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): self.log.debug("enqueued with blocking %s.", msg["event"]) def send_sync() -> None: - # Sync mode bypasses the lane's queue but keeps its wire config: - # the AI lane is pinned to v0, so its events post to the AI - # endpoint regardless of `capture_mode`. + # Sync mode bypasses the lane's queue but keeps its wire config, + # so AI events post to the AI endpoint whatever `capture_mode`. if lane.capture_mode == CaptureMode.V1: _send_v1_batch( self.api_key, self.host, [msg], - compression=self.capture_compression, + compression=lane.capture_compression, timeout=self.timeout, max_retries=self.max_retries, historical_migration=self.historical_migration, sdk_info=self._sdk_info, + path=lane.endpoint, ) return diff --git a/posthog/consumer.py b/posthog/consumer.py index c604b9756..8a69844cf 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -7,7 +7,7 @@ from posthog._logging import _configure_posthog_logging from posthog.capture_compression import CaptureCompression from posthog.capture_mode import CaptureMode -from posthog.capture_v1 import _backoff, _send_v1_batch +from posthog.capture_v1 import _CAPTURE_V1_PATH, _backoff, _send_v1_batch from posthog.request import ( EVENTS_ENDPOINT, USER_AGENT as _USER_AGENT, @@ -113,9 +113,9 @@ def __init__( retries=10, timeout=15, historical_migration=False, - endpoint=EVENTS_ENDPOINT, + endpoint=None, max_msg_size=MAX_MSG_SIZE, - capture_mode=CaptureMode.V0, + capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.NONE, ): """Create a consumer thread.""" @@ -129,6 +129,12 @@ def __init__( self.on_error = on_error self.queue = queue self.gzip = gzip + # Without an explicit endpoint, post to the analytics path of the + # selected protocol. + if endpoint is None: + endpoint = ( + _CAPTURE_V1_PATH if capture_mode == CaptureMode.V1 else EVENTS_ENDPOINT + ) self.endpoint = endpoint self.max_msg_size = max_msg_size self.capture_mode = capture_mode @@ -288,10 +294,10 @@ def next(self): return items def request(self, batch): - """Upload the batch via the wire protocol selected by `capture_mode`. + """Upload the batch to this consumer's `endpoint` via the wire protocol + selected by `capture_mode`. - V1 uses the partial-retry submitter (which posts to its own path); V0 - posts the batch to this consumer's `endpoint`. + V1 uses the partial-retry submitter; V0 posts the whole batch. """ if self.capture_mode == CaptureMode.V1: _send_v1_batch( @@ -303,6 +309,7 @@ def request(self, batch): max_retries=self.retries, historical_migration=self.historical_migration, sdk_info=self._sdk_info, + path=self.endpoint, ) return self._send(batch, self.endpoint) diff --git a/posthog/test/capture_helpers.py b/posthog/test/capture_helpers.py new file mode 100644 index 000000000..542d3e10f --- /dev/null +++ b/posthog/test/capture_helpers.py @@ -0,0 +1,58 @@ +"""Intercept capture uploads at the batch submitter for client-level tests. + +Patching the submitter (not the HTTP layer) lets tests assert on the event +dicts the SDK built, before the wire encoding in ``capture_v1``. Wire shape is +covered by ``test_capture_v1``. +""" + +import json +from unittest import mock + +from requests import Response + +_SUBMITTER = "_send_v1_batch" + + +def offline_v1_post(url: str, data=None, **kwargs) -> Response: + """Stand-in for ``requests.Session.post`` that accepts every v1 event. + + For subprocess tests with no server. Prints the uncompressed request body, + because the SDK never logs payloads, and answers ``ok`` for each event. + """ + print(f"capture request body: {data}", flush=True) # noqa: T201 + events = json.loads(data)["batch"] + response = Response() + response.status_code = 200 + response._content = json.dumps( + {"results": {event["uuid"]: {"result": "ok"} for event in events}} + ).encode() + return response + + +def patch_capture_send(site: str = "client", **kwargs) -> "mock._patch": + """Patch the submitter where ``posthog.`` imported it. + + ``site="client"`` sees ``sync_mode`` uploads; ``site="consumer"`` sees + background consumer uploads. + """ + return mock.patch(f"posthog.{site}.{_SUBMITTER}", **kwargs) + + +def patch_async_capture_send(**kwargs) -> "mock._patch": + """Patch the submitter the ``AsyncPosthog`` consumer awaits.""" + return mock.patch("posthog._async_consumer.async_send_v1_batch", **kwargs) + + +def sent_batch(send_mock: mock.Mock, call_index: int = -1) -> list[dict]: + """Return the event batch from one recorded upload (default: the last).""" + call = send_mock.call_args_list[call_index] + return call.args[2] if len(call.args) > 2 else call.kwargs["batch"] + + +def sent_events(send_mock: mock.Mock) -> list[dict]: + """Return every event uploaded through ``send_mock``, in send order.""" + return [ + event + for index in range(len(send_mock.call_args_list)) + for event in sent_batch(send_mock, index) + ] diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 3ab2f9f87..91b8c7256 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -6,18 +6,21 @@ import posthog from posthog.ai.utils import _capture_ai_event, finalize_ai_content, with_privacy_mode +from posthog.capture_compression import CaptureCompression from posthog.capture_mode import CaptureMode from posthog.client import Client from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE -from posthog.request import AI_EVENTS_ENDPOINT, EVENTS_ENDPOINT +from posthog.capture_v1 import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH +from posthog.request import EVENTS_ENDPOINT from posthog.version import VERSION +from posthog.test.capture_helpers import patch_capture_send, sent_batch from posthog.test.test_utils import TEST_API_KEY -def _events_by_path(mock_post): +def _events_by_path(mock_send): by_path: dict[str, list] = {} - for call in mock_post.call_args_list: - by_path.setdefault(call.kwargs["path"], []).extend(call.kwargs["batch"]) + for index, call in enumerate(mock_send.call_args_list): + by_path.setdefault(call.kwargs["path"], []).extend(sent_batch(mock_send, index)) return by_path @@ -29,24 +32,24 @@ def _client(self, **kwargs): def test_capture_ai_and_capture_ride_separate_lanes(self): client = self._client() - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client.capture("button_clicked", distinct_id="d") client.capture_ai("$ai_generation", distinct_id="d") client.flush() by_path = _events_by_path(mock_post) - self.assertEqual(set(by_path), {EVENTS_ENDPOINT, AI_EVENTS_ENDPOINT}) + self.assertEqual(set(by_path), {_CAPTURE_V1_PATH, _CAPTURE_AI_V1_PATH}) self.assertEqual( - [e["event"] for e in by_path[EVENTS_ENDPOINT]], ["button_clicked"] + [e["event"] for e in by_path[_CAPTURE_V1_PATH]], ["button_clicked"] ) self.assertEqual( - [e["event"] for e in by_path[AI_EVENTS_ENDPOINT]], ["$ai_generation"] + [e["event"] for e in by_path[_CAPTURE_AI_V1_PATH]], ["$ai_generation"] ) - for call in mock_post.call_args_list: - events = {e["event"] for e in call.kwargs["batch"]} + for index, call in enumerate(mock_post.call_args_list): + events = {e["event"] for e in sent_batch(mock_post, index)} expected = ( {"$ai_generation"} - if call.kwargs["path"] == AI_EVENTS_ENDPOINT + if call.kwargs["path"] == _CAPTURE_AI_V1_PATH else {"button_clicked"} ) self.assertEqual(events, expected) @@ -55,13 +58,13 @@ def test_capture_does_not_reroute_ai_named_events(self): # The two-lane rule: `capture()` never special-cases AI events, no # matter their name. Only `capture_ai()` reaches the AI lane. client = self._client() - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client.capture("$ai_generation", distinct_id="d") client.flush() self.assertEqual( [call.kwargs["path"] for call in mock_post.call_args_list], - [EVENTS_ENDPOINT], + [_CAPTURE_V1_PATH], ) def test_capture_ai_returns_event_uuid_like_capture(self): @@ -71,21 +74,21 @@ def test_capture_ai_returns_event_uuid_like_capture(self): def test_sync_mode_capture_ai_posts_single_event_batch_to_ai_endpoint(self): client = Client(TEST_API_KEY, sync_mode=True) - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client.capture_ai("$ai_generation", distinct_id="d") mock_post.assert_called_once() - self.assertEqual(mock_post.call_args.kwargs["path"], AI_EVENTS_ENDPOINT) - batch = mock_post.call_args.kwargs["batch"] + self.assertEqual(mock_post.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) + batch = sent_batch(mock_post) self.assertEqual([e["event"] for e in batch], ["$ai_generation"]) def test_multimodal_client_routes_wrapper_captures_to_ai_lane(self): client = self._client(enable_full_ai_capture=True) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: _capture_ai_event(client, "$ai_generation", distinct_id="d") client.flush() by_path = _events_by_path(mock_post) - self.assertEqual(set(by_path), {AI_EVENTS_ENDPOINT}) + self.assertEqual(set(by_path), {_CAPTURE_AI_V1_PATH}) def test_disabled_client_never_starts_ai_lane(self): client = Client(TEST_API_KEY, disabled=True) @@ -122,7 +125,7 @@ def test_analytics_consumers_keep_todays_parameters(self): self.assertEqual(len(consumers), 2) for consumer in consumers: self.assertIs(consumer.queue, client.queue) - self.assertEqual(consumer.endpoint, EVENTS_ENDPOINT) + self.assertEqual(consumer.endpoint, _CAPTURE_V1_PATH) self.assertEqual(consumer.max_msg_size, MAX_MSG_SIZE) self.assertEqual(consumer.flush_at, 7) self.assertEqual(consumer.flush_interval, 0.5) @@ -136,15 +139,15 @@ def test_analytics_consumers_keep_todays_parameters(self): def test_analytics_traffic_posts_to_single_endpoint(self): client = Client(TEST_API_KEY, flush_interval=0.05) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client.capture("event_a", distinct_id="d") client.capture("event_b", distinct_id="d") client.flush() by_path = _events_by_path(mock_post) - self.assertEqual(set(by_path), {EVENTS_ENDPOINT}) + self.assertEqual(set(by_path), {_CAPTURE_V1_PATH}) self.assertEqual( - sorted(e["event"] for e in by_path[EVENTS_ENDPOINT]), + sorted(e["event"] for e in by_path[_CAPTURE_V1_PATH]), ["event_a", "event_b"], ) client.join() @@ -152,9 +155,9 @@ def test_analytics_traffic_posts_to_single_endpoint(self): def test_sync_mode_analytics_path_unchanged(self): client = Client(TEST_API_KEY, sync_mode=True) self.assertIsNone(client.consumers) - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client.capture("button_clicked", distinct_id="d") - self.assertEqual(mock_post.call_args.kwargs["path"], EVENTS_ENDPOINT) + self.assertEqual(mock_post.call_args.kwargs["path"], _CAPTURE_V1_PATH) class TestLaneSizeCaps(unittest.TestCase): @@ -192,48 +195,68 @@ def test_analytics_lane_rejects_events_over_900kib(self): self.assertTrue(client.queue.empty()) -class TestAiLaneV0Pinned(unittest.TestCase): - """The AI endpoint has no v1 form: the AI lane ignores `capture_mode=v1`.""" +class TestAiLaneAlwaysV1(unittest.TestCase): + """The AI lane posts capture v1 to the AI endpoint whatever `capture_mode`.""" - def test_ai_lane_consumers_pin_v0_and_ai_endpoint(self): - client = Client(TEST_API_KEY, send=False, capture_mode="v1", thread=2) + def test_ai_lane_consumers_use_v1_and_ai_endpoint(self): + client = Client(TEST_API_KEY, send=False, capture_mode="v0", thread=2) client._ai_lane.start() self.assertEqual(len(client._ai_lane.consumers), 2) for consumer in client._ai_lane.consumers: self.assertIs(consumer.queue, client._ai_lane.queue) - self.assertEqual(consumer.endpoint, AI_EVENTS_ENDPOINT) + self.assertEqual(consumer.endpoint, _CAPTURE_AI_V1_PATH) self.assertEqual(consumer.max_msg_size, AI_MAX_MSG_SIZE) - self.assertEqual(consumer.capture_mode, CaptureMode.V0) + self.assertEqual(consumer.capture_mode, CaptureMode.V1) + self.assertEqual(consumer.capture_compression, CaptureCompression.NONE) - def test_async_ai_events_use_v0_even_with_capture_mode_v1(self): - client = Client(TEST_API_KEY, capture_mode="v1", flush_interval=0.05) + def test_async_ai_events_use_v1_with_capture_mode_v0(self): + client = Client( + TEST_API_KEY, + capture_mode="v0", + capture_compression="gzip", + flush_interval=0.05, + ) with ( mock.patch("posthog.consumer.batch_post") as mock_post, - mock.patch("posthog.consumer._send_v1_batch") as mock_v1, + patch_capture_send("consumer") as mock_v1, ): client.capture_ai("$ai_generation", distinct_id="d") client.capture("button_clicked", distinct_id="d") client.flush() - mock_v1.assert_called() self.assertEqual( [call.kwargs["path"] for call in mock_post.call_args_list], - [AI_EVENTS_ENDPOINT], + [EVENTS_ENDPOINT], + ) + mock_v1.assert_called_once() + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) + self.assertEqual( + mock_v1.call_args.kwargs["compression"], CaptureCompression.NONE ) + self.assertEqual([e["event"] for e in sent_batch(mock_v1)], ["$ai_generation"]) client.join() - def test_sync_ai_events_use_v0_even_with_capture_mode_v1(self): - client = Client(TEST_API_KEY, sync_mode=True, capture_mode="v1") + def test_sync_ai_events_use_v1_with_capture_mode_v0(self): + client = Client( + TEST_API_KEY, + sync_mode=True, + capture_mode="v0", + capture_compression="gzip", + ) with ( mock.patch("posthog.client.batch_post") as mock_post, - mock.patch("posthog.client._send_v1_batch") as mock_v1, + patch_capture_send("client") as mock_v1, ): client.capture_ai("$ai_generation", distinct_id="d") client.capture("button_clicked", distinct_id="d") mock_post.assert_called_once() - self.assertEqual(mock_post.call_args.kwargs["path"], AI_EVENTS_ENDPOINT) + self.assertEqual(mock_post.call_args.kwargs["path"], EVENTS_ENDPOINT) mock_v1.assert_called_once() + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) + self.assertEqual( + mock_v1.call_args.kwargs["compression"], CaptureCompression.NONE + ) class TestAiLaneLazyStart(unittest.TestCase): @@ -241,12 +264,12 @@ def test_no_ai_consumers_until_first_capture_ai(self): client = Client(TEST_API_KEY, flush_interval=0.05) self.assertEqual(client._ai_lane.consumers, []) - with mock.patch("posthog.consumer.batch_post"): + with patch_capture_send("consumer"): client.capture("button_clicked", distinct_id="d") client.flush() self.assertEqual(client._ai_lane.consumers, []) - with mock.patch("posthog.consumer.batch_post"): + with patch_capture_send("consumer"): client.capture_ai("$ai_generation", distinct_id="d") self.assertEqual(len(client._ai_lane.consumers), 1) self.assertTrue(client._ai_lane.consumers[0].is_alive()) @@ -262,7 +285,7 @@ def fire(): client.capture_ai("$ai_generation", distinct_id="d") threads = [threading.Thread(target=fire) for _ in range(8)] - with mock.patch("posthog.consumer.batch_post"): + with patch_capture_send("consumer"): for thread in threads: thread.start() for thread in threads: @@ -285,7 +308,7 @@ def test_fork_rebuild_restarts_analytics_and_resets_ai(self): client = Client( TEST_API_KEY, flush_interval=0.05, enable_local_evaluation=False ) - with mock.patch("posthog.consumer.batch_post"): + with patch_capture_send("consumer"): client.capture_ai("$ai_generation", distinct_id="d") client.flush() self.assertEqual(len(client._ai_lane.consumers), 1) @@ -305,14 +328,14 @@ def test_fork_rebuild_restarts_analytics_and_resets_ai(self): self.assertTrue(client._analytics_lane.consumers[0].is_alive()) self.assertEqual(client._ai_lane.consumers, []) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client.capture("button_clicked", distinct_id="d") client.capture_ai("$ai_generation", distinct_id="d") client.flush() self.assertEqual(len(client._ai_lane.consumers), 1) self.assertEqual( - set(_events_by_path(mock_post)), {EVENTS_ENDPOINT, AI_EVENTS_ENDPOINT} + set(_events_by_path(mock_post)), {_CAPTURE_V1_PATH, _CAPTURE_AI_V1_PATH} ) client.join() @@ -332,7 +355,7 @@ class TestCaptureAiEventHelper(unittest.TestCase): def test_opted_in_routes_through_ai_lane(self): client = Client(TEST_API_KEY, flush_interval=0.05, enable_full_ai_capture=True) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: _capture_ai_event( client, "$ai_generation", @@ -343,19 +366,19 @@ def test_opted_in_routes_through_ai_lane(self): self.assertEqual( [call.kwargs["path"] for call in mock_post.call_args_list], - [AI_EVENTS_ENDPOINT], + [_CAPTURE_AI_V1_PATH], ) client.join() def test_default_keeps_capture_path(self): client = Client(TEST_API_KEY, flush_interval=0.05) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: _capture_ai_event(client, "$ai_generation", distinct_id="d") client.flush() self.assertEqual( [call.kwargs["path"] for call in mock_post.call_args_list], - [EVENTS_ENDPOINT], + [_CAPTURE_V1_PATH], ) self.assertEqual(client._ai_lane.consumers, []) client.join() @@ -437,7 +460,7 @@ class TestLanesRefuseWorkAfterShutdown(unittest.TestCase): def test_late_ai_capture_after_shutdown_starts_nothing_and_sends_nothing(self): client = Client(TEST_API_KEY, enable_full_ai_capture=True, flush_interval=0.05) client.shutdown() - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: _capture_ai_event(client, "$ai_generation", distinct_id="d") client._ai_lane.queue.join() self.assertEqual(client._ai_lane.consumers, []) @@ -499,10 +522,10 @@ def _client(self, **kwargs): def test_new_flag_routes_wrapper_captures_to_ai_lane(self): client = self._client(enable_full_ai_capture=True) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: _capture_ai_event(client, "$ai_generation", distinct_id="d") client.flush() - self.assertEqual(set(_events_by_path(mock_post)), {AI_EVENTS_ENDPOINT}) + self.assertEqual(set(_events_by_path(mock_post)), {_CAPTURE_AI_V1_PATH}) def test_deprecated_kwargs_map_to_new_flag(self): for kwargs in ({"_use_ai_lane": True}, {"_enable_multimodal_capture": True}): @@ -570,24 +593,24 @@ def _client(self, **kwargs): def test_returned_uuid_matches_the_wire_event_uuid(self): client = self._client() - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: returned_uuid = client.capture_ai("$ai_generation", distinct_id="d") client.flush() - batch = mock_post.call_args.kwargs["batch"] + batch = sent_batch(mock_post) self.assertEqual(batch[0]["uuid"], returned_uuid) def test_supplied_uuid_is_preserved_end_to_end(self): client = self._client() supplied_uuid = str(uuid.uuid4()) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: returned_uuid = client.capture_ai( "$ai_generation", distinct_id="d", uuid=supplied_uuid ) client.flush() self.assertEqual(returned_uuid, supplied_uuid) - batch = mock_post.call_args.kwargs["batch"] + batch = sent_batch(mock_post) self.assertEqual(batch[0]["uuid"], supplied_uuid) def test_returned_uuid_reflects_before_send_replacement(self): @@ -598,12 +621,12 @@ def replace_uuid(event): return event client = self._client(before_send=replace_uuid) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: returned_uuid = client.capture_ai("$ai_generation", distinct_id="d") client.flush() self.assertEqual(returned_uuid, replacement_uuid) - batch = mock_post.call_args.kwargs["batch"] + batch = sent_batch(mock_post) self.assertEqual(batch[0]["uuid"], replacement_uuid) def test_returned_uuid_is_regenerated_when_before_send_removes_it(self): @@ -612,12 +635,12 @@ def drop_uuid(event): return event client = self._client(before_send=drop_uuid) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: returned_uuid = client.capture_ai("$ai_generation", distinct_id="d") client.flush() self.assertIsNotNone(returned_uuid) - batch = mock_post.call_args.kwargs["batch"] + batch = sent_batch(mock_post) self.assertEqual(batch[0]["uuid"], returned_uuid) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index fcb6e2298..9c100bc76 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -21,6 +21,7 @@ set_code_variables_mask_url_credentials_context, ) from posthog.request import APIError +from posthog.test.capture_helpers import patch_async_capture_send @pytest.mark.asyncio @@ -45,10 +46,10 @@ async def test_async_posthog_is_the_customer_facing_async_client(): async def test_capture_is_a_synchronous_queue_write_and_flushes(): batches = [] - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): async with AsyncPosthog("test-key", flush_at=100, flush_interval=30) as client: event_uuid = client.capture( "async event", @@ -76,16 +77,15 @@ def test_capture_from_worker_thread_wakes_loop_bound_queue(): import threading from unittest import mock from posthog import AsyncPosthog +from posthog.test.capture_helpers import patch_async_capture_send async def main(): delivered = [] - async def batch_post(*args, **kwargs): - delivered.extend(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + delivered.extend(batch) - with mock.patch( - "posthog._async_consumer.async_batch_post", side_effect=batch_post - ): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", flush_interval=30) client._ensure_workers_started() while not client._queue._getters: @@ -162,10 +162,10 @@ async def before_send(event): event["properties"]["from_before_send"] = True return event - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): async with AsyncPosthog( "test-key", before_send=before_send, flush_interval=30 ) as client: @@ -185,10 +185,10 @@ def before_send(event): event["properties"]["request_id"] = request_id.get() return event - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", before_send=before_send, flush_interval=30) token = request_id.set("request-A") client.capture("event-A", distinct_id="user-1") @@ -217,7 +217,7 @@ def before_send(event): callback_thread = threading.get_ident() return event - with mock.patch("posthog._async_consumer.async_batch_post", new=mock.AsyncMock()): + with patch_async_capture_send(new=mock.AsyncMock()): client = AsyncPosthog("test-key", before_send=before_send) result = await client.capture_immediate("event", distinct_id="user-1") await client.shutdown() @@ -231,9 +231,7 @@ async def test_capture_drops_event_when_before_send_raises(): async def before_send(_event): raise RuntimeError("callback failed") - with mock.patch( - "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() - ) as batch_post: + with patch_async_capture_send(new=mock.AsyncMock()) as send_batch: async with AsyncPosthog( "test-key", before_send=before_send, flush_interval=0.01 ) as client: @@ -241,18 +239,18 @@ async def before_send(_event): await client.flush(timeout_seconds=1) assert accepted_uuid is not None - batch_post.assert_not_awaited() + send_batch.assert_not_awaited() @pytest.mark.asyncio async def test_capture_immediate_waits_for_delivery(): delivered = asyncio.Event() - async def batch_post(*args, **kwargs): - assert kwargs["batch"][0]["event"] == "immediate event" + async def send_batch(api_key, host, batch, **kwargs): + assert batch[0]["event"] == "immediate event" delivered.set() - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key") event_uuid = await client.capture_immediate( "immediate event", distinct_id="user-1" @@ -271,10 +269,10 @@ async def before_send(event): event["properties"]["processed"] = True return event - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", before_send=before_send) result = await client.capture_immediate("event", distinct_id="user-1") await client.shutdown() @@ -289,15 +287,13 @@ def before_send(event): event["properties"]["user_input"] = "x" * MAX_MSG_SIZE return event - with mock.patch( - "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() - ) as batch_post: + with patch_async_capture_send(new=mock.AsyncMock()) as send_batch: client = AsyncPosthog("test-key", before_send=before_send) result = await client.capture_immediate("event", distinct_id="user-1") await client.shutdown() assert result is None - batch_post.assert_not_awaited() + send_batch.assert_not_awaited() @pytest.mark.asyncio @@ -329,7 +325,7 @@ async def test_missing_async_extra_does_not_accept_undeliverable_events(): "posthog.async_client._require_httpx", side_effect=RuntimeError("install posthog[async]"), ): - client = AsyncPosthog("test-key") + client = AsyncPosthog("test-key", capture_mode=CaptureMode.V0) assert client.capture("event", distinct_id="user-1") is None assert ( client.set(distinct_id="user-1", properties={"email": "a@example.com"}) @@ -382,10 +378,10 @@ async def test_identify_methods_enqueue_events( ): batches = [] - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): async with AsyncPosthog("test-key", flush_interval=30) as client: method = getattr(client, method_name) assert method(**method_kwargs) is not None @@ -409,12 +405,12 @@ async def test_capture_after_shutdown_is_dropped_without_restarting_workers(): async def test_batch_size_overflow_event_is_sent_in_the_next_batch(): batches = [] - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) with ( mock.patch("posthog._async_consumer.BATCH_SIZE_LIMIT", 800), - mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post), + patch_async_capture_send(side_effect=send_batch), ): client = AsyncPosthog("test-key", flush_at=10, flush_interval=30) client.capture("first", distinct_id="user-1", properties={"value": "a" * 400}) @@ -440,10 +436,10 @@ async def before_send(event): await allow_slow_callback.wait() return event - async def batch_post(*args, **kwargs): - delivered.extend(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + delivered.extend(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog( "test-key", thread=2, @@ -469,12 +465,12 @@ async def test_shutdown_waits_for_an_in_flight_batch_instead_of_cancelling_it(): allow_upload = asyncio.Event() delivered = [] - async def batch_post(*args, **kwargs): + async def send_batch(api_key, host, batch, **kwargs): upload_started.set() await allow_upload.wait() - delivered.extend(kwargs["batch"]) + delivered.extend(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", flush_at=1) client.capture("event", distinct_id="user-1") await upload_started.wait() @@ -498,9 +494,7 @@ async def before_send(event): callback_finished.set() return event - with mock.patch( - "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() - ) as batch_post: + with patch_async_capture_send(new=mock.AsyncMock()) as send_batch: client = AsyncPosthog("test-key", before_send=before_send, flush_at=1) client.capture("event", distinct_id="user-1") await asyncio.wait_for(callback_finished.wait(), timeout=1) @@ -511,7 +505,7 @@ async def wait_until_closed(): await asyncio.wait_for(wait_until_closed(), timeout=1) - batch_post.assert_awaited_once() + send_batch.assert_awaited_once() assert client.capture("after shutdown", distinct_id="user-1") is None @@ -526,10 +520,10 @@ async def before_send(event): await allow_callback.wait() return event - async def batch_post(*args, **kwargs): - delivered.extend(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + delivered.extend(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", before_send=before_send, flush_at=1) client.capture("event", distinct_id="user-1") await callback_started.wait() @@ -563,7 +557,7 @@ async def test_shutdown_waits_for_immediate_operation_not_its_long_lived_caller( allow_upload = asyncio.Event() shutdown_task = None - async def batch_post(*args, **kwargs): + async def send_batch(api_key, host, batch, **kwargs): upload_started.set() await allow_upload.wait() @@ -573,7 +567,7 @@ async def capture_then_await_shutdown(client): assert shutdown_task is not None await shutdown_task - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key") caller = asyncio.create_task(capture_then_await_shutdown(client)) await upload_started.wait() @@ -587,11 +581,11 @@ async def test_shutdown_waits_for_in_flight_immediate_capture(): upload_started = asyncio.Event() allow_upload = asyncio.Event() - async def batch_post(*args, **kwargs): + async def send_batch(api_key, host, batch, **kwargs): upload_started.set() await allow_upload.wait() - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key") capture = asyncio.create_task( client.capture_immediate("event", distinct_id="user-1") @@ -618,7 +612,7 @@ async def test_reuses_and_closes_instance_owned_http_client(): "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() ) as batch_post, ): - client = AsyncPosthog("test-key") + client = AsyncPosthog("test-key", capture_mode=CaptureMode.V0) await client.capture_immediate("first", distinct_id="user-1") await client.capture_immediate("second", distinct_id="user-1") await client.shutdown() @@ -636,13 +630,11 @@ def test_capture_before_loop_starts_is_flushed_when_loop_runs(): client = AsyncPosthog("test-key", flush_interval=30) assert client.capture("event", distinct_id="user-1") is not None - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) async def flush_and_close(): - with mock.patch( - "posthog._async_consumer.async_batch_post", side_effect=batch_post - ): + with patch_async_capture_send(side_effect=send_batch): await client.shutdown() asyncio.run(flush_and_close()) @@ -699,7 +691,7 @@ async def test_capture_exception_uses_context_code_variable_settings(): @pytest.mark.asyncio async def test_queued_payload_is_not_written_to_debug_logs(caplog): caplog.set_level(logging.DEBUG, logger="posthog") - with mock.patch("posthog._async_consumer.async_batch_post", new=mock.AsyncMock()): + with patch_async_capture_send(new=mock.AsyncMock()): async with AsyncPosthog("test-key", flush_interval=30) as client: client.capture( "event", @@ -720,8 +712,7 @@ def on_error(error, batch): nonlocal callback_thread callback_thread = threading.get_ident() - with mock.patch( - "posthog._async_consumer.async_batch_post", + with patch_async_capture_send( side_effect=APIError(400, "failed"), ): client = AsyncPosthog("test-key", on_error=on_error, max_retries=0) @@ -738,9 +729,7 @@ async def test_failed_capture_does_not_log_server_response_detail(caplog, immedi caplog.set_level(logging.DEBUG, logger="posthog") server_error = APIError(400, "password=server-secret") - with mock.patch( - "posthog._async_consumer.async_batch_post", side_effect=server_error - ): + with patch_async_capture_send(side_effect=server_error): client = AsyncPosthog("test-key", flush_at=1, max_retries=0) if immediate: await client.capture_immediate("event", distinct_id="user-1") @@ -759,11 +748,11 @@ async def test_flush_timeout_reports_unfinished_items(caplog): upload_started = asyncio.Event() allow_upload = asyncio.Event() - async def batch_post(*args, **kwargs): + async def send_batch(api_key, host, batch, **kwargs): upload_started.set() await allow_upload.wait() - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", flush_at=1) client.capture("event", distinct_id="user-1") await upload_started.wait() diff --git a/posthog/test/test_async_feature_flags.py b/posthog/test/test_async_feature_flags.py index 75c8b207e..da35bf17d 100644 --- a/posthog/test/test_async_feature_flags.py +++ b/posthog/test/test_async_feature_flags.py @@ -11,6 +11,7 @@ from posthog.client import _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES from posthog.contexts import new_context, set_context_device_id from posthog.request import APIError +from posthog.test.capture_helpers import patch_async_capture_send def flags_response(*, minimal=False): @@ -138,15 +139,15 @@ async def test_empty_flag_keys_skips_remote_request(): async def test_snapshot_access_captures_one_flag_called_event_and_attaches_flags(): batches = [] - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) with ( mock.patch( "posthog.async_client._async_flags", new=mock.AsyncMock(return_value=flags_response()), ), - mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post), + patch_async_capture_send(side_effect=send_batch), ): client = AsyncPosthog("project-key", flush_interval=30) snapshot = await client.evaluate_flags("user-1") @@ -171,15 +172,15 @@ async def batch_post(*args, **kwargs): async def test_minimal_flag_called_event_uses_strict_property_allowlist(): batches = [] - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) with ( mock.patch( "posthog.async_client._async_flags", new=mock.AsyncMock(return_value=flags_response(minimal=True)), ), - mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post), + patch_async_capture_send(side_effect=send_batch), ): client = AsyncPosthog("project-key", flush_interval=30) snapshot = await client.evaluate_flags("user-1", flag_keys=["beta"]) diff --git a/posthog/test/test_before_send.py b/posthog/test/test_before_send.py index ee048290e..3d7441d97 100644 --- a/posthog/test/test_before_send.py +++ b/posthog/test/test_before_send.py @@ -6,14 +6,19 @@ from posthog.client import Client from posthog.test.test_utils import FAKE_TEST_API_KEY +from posthog.test.capture_helpers import ( + patch_capture_send, + sent_batch, + sent_events, +) class TestClient(unittest.TestCase): @classmethod def setUpClass(cls): # This ensures no real HTTP POST requests are made - cls.client_post_patcher = mock.patch("posthog.client.batch_post") - cls.consumer_post_patcher = mock.patch("posthog.consumer.batch_post") + cls.client_post_patcher = patch_capture_send("client") + cls.consumer_post_patcher = patch_capture_send("consumer") cls.client_post_patcher.start() cls.consumer_post_patcher.start() @@ -42,7 +47,7 @@ def my_before_send(event): event["properties"]["processed_by_before_send"] = True return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -57,7 +62,7 @@ def my_before_send(event): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertEqual( @@ -75,7 +80,7 @@ def replace_uuid(event): event["uuid"] = replacement_uuid return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -87,7 +92,7 @@ def replace_uuid(event): self.assertEqual(msg_uuid, replacement_uuid) mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertEqual(enqueued_msg["uuid"], replacement_uuid) @@ -98,7 +103,7 @@ def remove_uuid(event): del event["uuid"] return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -110,7 +115,7 @@ def remove_uuid(event): self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertEqual(enqueued_msg["uuid"], msg_uuid) @@ -122,7 +127,7 @@ def drop_test_events(event): return None return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -140,7 +145,7 @@ def drop_test_events(event): # Check the enqueued message mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertEqual(enqueued_msg["event"], "keep_me") @@ -152,7 +157,7 @@ def buggy_before_send(event): event["uuid"] = "invalid" raise ValueError("Oops!") - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -187,7 +192,7 @@ def buggy_before_send(event): ) try: with ( - mock.patch("posthog.consumer.batch_post") as mock_post, + patch_capture_send("consumer") as mock_post, mock.patch.object( client._analytics_lane, "enqueue", @@ -210,7 +215,7 @@ def add_unsupported_value(event): event["properties"]["marker"] = marker return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, before_send=add_unsupported_value, @@ -218,12 +223,12 @@ def add_unsupported_value(event): ) self.assertIsNotNone(client.capture("recleaned", distinct_id="user1")) - sent_event = mock_post.call_args.kwargs["batch"][0] + sent_event = sent_batch(mock_post)[0] self.assertIsNone(sent_event["properties"]["marker"]) def test_before_send_callback_non_dict_output_drops_event(self): with ( - mock.patch("posthog.client.batch_post") as mock_post, + patch_capture_send("client") as mock_post, mock.patch("posthog.client.Client.log.exception") as mock_log, ): client = Client( @@ -251,7 +256,7 @@ def add_invalid_mapping_key(event): flush_interval=0.01, ) with ( - mock.patch("posthog.consumer.batch_post") as mock_post, + patch_capture_send("consumer") as mock_post, self.assertLogs("posthog", level="ERROR") as logs, ): client.capture("malformed", distinct_id="user1") @@ -259,8 +264,7 @@ def add_invalid_mapping_key(event): client.shutdown() mock_post.assert_called_once() - sent_batch = mock_post.call_args.kwargs["batch"] - self.assertEqual([event["event"] for event in sent_batch], ["valid"]) + self.assertEqual([event["event"] for event in sent_batch(mock_post)], ["valid"]) self.assertEqual(client.queue.unfinished_tasks, 0) self.assertTrue(all(not consumer.is_alive() for consumer in client.consumers)) self.assertNotIn("private-key", "\n".join(logs.output)) @@ -275,7 +279,7 @@ def add_marker(event): event["properties"]["marked"] = True return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -293,14 +297,12 @@ def add_marker(event): # Check all events were marked self.assertEqual(mock_post.call_count, 2) - for call in mock_post.call_args_list: - batch_data = call[1]["batch"] - enqueued_msg = batch_data[0] - self.assertTrue(enqueued_msg["properties"]["marked"]) + for event in sent_events(mock_post): + self.assertTrue(event["properties"]["marked"]) def test_before_send_callback_disabled_when_none(self): """Test that client works normally when before_send is None.""" - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -312,7 +314,7 @@ def test_before_send_callback_disabled_when_none(self): # Check the event was sent normally mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertEqual(enqueued_msg["event"], "normal_event") @@ -336,7 +338,7 @@ def scrub_pii(event): return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -357,7 +359,7 @@ def scrub_pii(event): # Check the enqueued message was scrubbed mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertEqual(enqueued_msg["properties"]["email"], "***@example.com") @@ -379,7 +381,7 @@ def my_before_send(event): event["properties"]["module_level_before_send"] = True return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: posthog.api_key = FAKE_TEST_API_KEY posthog.before_send = my_before_send posthog.sync_mode = True @@ -390,7 +392,7 @@ def my_before_send(event): self.assertIs(posthog.default_client.before_send, my_before_send) mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) enqueued_msg = batch_data[0] self.assertTrue(enqueued_msg["properties"]["module_level_before_send"]) @@ -399,7 +401,7 @@ def my_before_send(event): event["properties"]["updated_after_init"] = True return event - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: posthog.api_key = FAKE_TEST_API_KEY posthog.sync_mode = True @@ -413,8 +415,8 @@ def my_before_send(event): self.assertIs(posthog.default_client.before_send, my_before_send) self.assertEqual(mock_post.call_count, 2) - first_batch = mock_post.call_args_list[0][1]["batch"] - second_batch = mock_post.call_args_list[1][1]["batch"] + first_batch = sent_batch(mock_post, 0) + second_batch = sent_batch(mock_post, 1) self.assertNotIn("updated_after_init", first_batch[0]["properties"]) self.assertTrue(second_batch[0]["properties"]["updated_after_init"]) diff --git a/posthog/test/test_capture_mode.py b/posthog/test/test_capture_mode.py index 473e1a886..e63705d45 100644 --- a/posthog/test/test_capture_mode.py +++ b/posthog/test/test_capture_mode.py @@ -16,10 +16,10 @@ class TestResolveCaptureMode(unittest.TestCase): - def test_defaults_to_v0_with_no_kwarg_and_no_env(self) -> None: + def test_defaults_to_v1_with_no_kwarg_and_no_env(self) -> None: with mock.patch.dict(os.environ, {}, clear=False): os.environ.pop(CAPTURE_MODE_ENV_VAR, None) - self.assertIs(_resolve_capture_mode(None), CaptureMode.V0) + self.assertIs(_resolve_capture_mode(None), CaptureMode.V1) @parameterized.expand( [ @@ -62,14 +62,14 @@ def test_env_var_resolution(self, _name, env_value, expected) -> None: self.assertIs(_resolve_capture_mode(None), expected) @parameterized.expand([("empty", ""), ("whitespace", " ")]) - def test_blank_env_var_defaults_to_v0(self, _name, env_value) -> None: + def test_blank_env_var_defaults_to_v1(self, _name, env_value) -> None: with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: env_value}): - self.assertIs(_resolve_capture_mode(None), CaptureMode.V0) + self.assertIs(_resolve_capture_mode(None), CaptureMode.V1) - def test_unrecognized_env_var_warns_and_defaults_to_v0(self) -> None: + def test_unrecognized_env_var_warns_and_defaults_to_v1(self) -> None: with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: "bogus"}): with capture_message_only_logs() as stream: - self.assertIs(_resolve_capture_mode(None), CaptureMode.V0) + self.assertIs(_resolve_capture_mode(None), CaptureMode.V1) self.assertIn("bogus", stream.getvalue()) @parameterized.expand([("bad_str", "bogus"), ("wrong_type", 1)]) @@ -79,11 +79,11 @@ def test_invalid_explicit_kwarg_raises(self, _name, value) -> None: class TestCaptureModePlumbing(unittest.TestCase): - def test_client_resolves_and_stores_default_v0(self) -> None: + def test_client_resolves_and_stores_default_v1(self) -> None: with mock.patch.dict(os.environ, {}, clear=False): os.environ.pop(CAPTURE_MODE_ENV_VAR, None) client = Client(TEST_API_KEY, sync_mode=True) - self.assertIs(client.capture_mode, CaptureMode.V0) + self.assertIs(client.capture_mode, CaptureMode.V1) @parameterized.expand( [ @@ -99,11 +99,11 @@ def test_client_kwarg_sets_mode(self, _name, kwarg, expected) -> None: def test_client_propagates_mode_to_consumers(self) -> None: # Async (non-sync) client builds Consumer threads; assert each carries # the resolved mode. - client = Client(TEST_API_KEY, capture_mode=CaptureMode.V1, send=False, thread=2) + client = Client(TEST_API_KEY, capture_mode=CaptureMode.V0, send=False, thread=2) self.assertEqual(len(client.consumers), 2) for consumer in client.consumers: - self.assertIs(consumer.capture_mode, CaptureMode.V1) + self.assertIs(consumer.capture_mode, CaptureMode.V0) - def test_consumer_defaults_to_v0(self) -> None: + def test_consumer_defaults_to_v1(self) -> None: consumer = Consumer(None, TEST_API_KEY) - self.assertIs(consumer.capture_mode, CaptureMode.V0) + self.assertIs(consumer.capture_mode, CaptureMode.V1) diff --git a/posthog/test/test_capture_v1.py b/posthog/test/test_capture_v1.py index f4809d501..e11e74c7c 100644 --- a/posthog/test/test_capture_v1.py +++ b/posthog/test/test_capture_v1.py @@ -10,6 +10,7 @@ from posthog.capture_compression import CaptureCompression from posthog.capture_v1 import ( + _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, _HEADER_ATTEMPT, _HEADER_REQUEST_ID, @@ -84,9 +85,11 @@ def __call__( timeout=15, sdk_info=USER_AGENT, session=None, + path=_CAPTURE_V1_PATH, ): self.calls.append( { + "path": path, "attempt": attempt, "request_id": request_id, "compression": compression, @@ -410,6 +413,8 @@ def _post(self, response, **kwargs): def test_url_uses_v1_path_and_trims_host(self) -> None: call = self._post(_results_response({})) self.assertEqual(call["url"], "https://app.posthog.com" + _CAPTURE_V1_PATH) + call = self._post(_results_response({}), path=_CAPTURE_AI_V1_PATH) + self.assertEqual(call["url"], "https://app.posthog.com" + _CAPTURE_AI_V1_PATH) def test_required_headers_present(self) -> None: headers = self._post(_results_response({}))["headers"] @@ -585,13 +590,15 @@ def test_request_id_and_created_at_stable_attempt_increments(self) -> None: self.assertEqual(stub.calls[0]["created_at"], stub.calls[1]["created_at"]) self.assertEqual([c["attempt"] for c in stub.calls], [1, 2]) - def test_compression_forwarded_to_post_v1(self) -> None: + def test_compression_and_path_forwarded_to_post_v1(self) -> None: stub = self._run( [_msg("u-1")], [_results_response({"u-1": "ok"})], compression=CaptureCompression.DEFLATE, + path=_CAPTURE_AI_V1_PATH, ) self.assertEqual(stub.calls[0]["compression"], CaptureCompression.DEFLATE) + self.assertEqual(stub.calls[0]["path"], _CAPTURE_AI_V1_PATH) def test_drop_on_2xx_surfaces_via_error(self) -> None: # A server-chosen drop is terminal: even on an all-ok-otherwise 2xx with diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 531bbf1c9..faef807c4 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -18,6 +18,7 @@ import pytest from posthog.capture_compression import CaptureCompression +from posthog.capture_v1 import _CAPTURE_V1_PATH from posthog.client import Client from posthog.contexts import get_context_session_id, new_context, set_context_session from posthog.request import APIError, GetResponse @@ -26,6 +27,7 @@ from posthog.types import FeatureFlag, FeatureFlagResult, LegacyFlagMetadata from posthog.version import VERSION from posthog.contexts import tag +from posthog.test.capture_helpers import patch_capture_send, sent_batch # Legacy single-flag behavior remains covered here; warning emission itself is @@ -53,8 +55,8 @@ class TestClient(unittest.TestCase): @classmethod def setUpClass(cls): # This ensures no real HTTP POST requests are made - cls.client_post_patcher = mock.patch("posthog.client.batch_post") - cls.consumer_post_patcher = mock.patch("posthog.consumer.batch_post") + cls.client_post_patcher = patch_capture_send("client") + cls.consumer_post_patcher = patch_capture_send("consumer") cls.client_post_patcher.start() cls.consumer_post_patcher.start() @@ -292,7 +294,7 @@ def test_empty_flush(self): self.client.flush() def test_empty_flush_does_not_drain_a_later_event(self): - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client = Client(FAKE_TEST_API_KEY, flush_at=100, flush_interval=0.5) client.flush() @@ -321,7 +323,7 @@ def test_flush_timeout_returns_when_queue_does_not_drain(self): def test_flush_does_not_wait_for_flush_interval(self): # flush() must attempt delivery now rather than letting the consumer sit # on a below-flush_at batch until flush_interval elapses. - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client = Client(FAKE_TEST_API_KEY, flush_interval=30) client.capture("event", distinct_id="distinct_id") @@ -335,7 +337,7 @@ def test_flush_does_not_wait_for_flush_interval(self): def test_flush_delivers_when_flush_interval_exceeds_the_flush_timeout(self): # Waiting out flush_interval meant a flush_interval longer than the # flush timeout delivered nothing at all. - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client = Client(FAKE_TEST_API_KEY, flush_interval=30) client.capture("event", distinct_id="distinct_id") @@ -346,7 +348,7 @@ def test_flush_delivers_when_flush_interval_exceeds_the_flush_timeout(self): def test_flush_keeps_batches_whole(self): # Draining early must not turn a full queue into one request per event. - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: client = Client(FAKE_TEST_API_KEY, flush_at=10, flush_interval=30) for _ in range(30): client.capture("event", distinct_id="distinct_id") @@ -355,7 +357,8 @@ def test_flush_keeps_batches_whole(self): self.assertTrue(client.queue.empty()) batch_sizes = [ - len(call.kwargs["batch"]) for call in mock_post.call_args_list + len(sent_batch(mock_post, index)) + for index in range(mock_post.call_count) ] self.assertEqual(sum(batch_sizes), 30) self.assertLessEqual(len(batch_sizes), 5) @@ -378,7 +381,7 @@ def test_flush_logs_and_returns_on_unexpected_error(self): client.queue.task_done() def test_basic_capture(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.capture("python test event", distinct_id="distinct_id") self.assertIsNotNone(msg_uuid) @@ -386,7 +389,7 @@ def test_basic_capture(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -403,7 +406,7 @@ def test_basic_capture(self): assert msg["properties"]["$os_version"] == mock.ANY def test_capture_omits_is_server_when_disabled(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -413,12 +416,12 @@ def test_capture_omits_is_server_when_disabled(self): client.capture("python test event", distinct_id="distinct_id") self.assertFalse(self.failed) - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["properties"]["$lib"], "posthog-python") self.assertNotIn("$is_server", msg["properties"]) def test_is_server_not_overridden_by_super_properties(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -428,11 +431,11 @@ def test_is_server_not_overridden_by_super_properties(self): client.capture("python test event", distinct_id="distinct_id") self.assertFalse(self.failed) - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["properties"]["$is_server"], True) def test_basic_capture_with_uuid(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) uuid = str(uuid4()) msg_uuid = client.capture( @@ -443,7 +446,7 @@ def test_basic_capture_with_uuid(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -454,7 +457,7 @@ def test_basic_capture_with_uuid(self): self.assertEqual(msg["properties"]["$lib_version"], VERSION) def test_basic_capture_with_uuid_object(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) uuid = UUID("00000000-0000-4000-8000-000000000002") msg_uuid = client.capture( @@ -464,7 +467,7 @@ def test_basic_capture_with_uuid_object(self): self.assertFalse(self.failed) mock_post.assert_called_once() - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["uuid"], str(uuid)) @parameterized.expand( @@ -478,7 +481,7 @@ def test_basic_capture_with_uuid_object(self): def test_capture_with_invalid_uuid_logs_and_falls_back_to_generated_uuid( self, _name, invalid_uuid ): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with self.assertLogs("posthog", level="ERROR") as logs: msg_uuid = client.capture( @@ -488,7 +491,7 @@ def test_capture_with_invalid_uuid_logs_and_falls_back_to_generated_uuid( self.assertIsNotNone(msg_uuid) UUID(msg_uuid) mock_post.assert_called_once() - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["uuid"], msg_uuid) self.assertNotEqual(msg["uuid"], str(invalid_uuid)) self.assertTrue( @@ -509,7 +512,7 @@ def test_capture_with_invalid_uuid_logs_and_falls_back_to_generated_uuid( ] ) def test_capture_with_invalid_uuid_falls_back_in_debug(self, _name, invalid_uuid): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, debug=True, sync_mode=True) with self.assertLogs("posthog", level="ERROR"): msg_uuid = client.capture( @@ -521,7 +524,7 @@ def test_capture_with_invalid_uuid_falls_back_in_debug(self, _name, invalid_uuid mock_post.assert_called_once() def test_basic_capture_with_project_api_key(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( project_api_key=FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -534,7 +537,7 @@ def test_basic_capture_with_project_api_key(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -544,7 +547,7 @@ def test_basic_capture_with_project_api_key(self): self.assertEqual(msg["properties"]["$lib_version"], VERSION) def test_basic_super_properties(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, super_properties={"source": "repo-name"}, @@ -556,7 +559,7 @@ def test_basic_super_properties(self): # Check the enqueued message mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -611,13 +614,13 @@ def test_capture_exception_uses_current_otel_span_context( ) with ( - mock.patch("posthog.client.batch_post") as mock_post, + patch_capture_send("client") as mock_post, use_span(NonRecordingSpan(span_context)), ): client = Client(FAKE_TEST_API_KEY, sync_mode=True) client.capture_exception(Exception("test exception"), properties=properties) - event = mock_post.call_args.kwargs["batch"][0] + event = sent_batch(mock_post)[0] if expected_trace_id is None: self.assertNotIn("$trace_id", event["properties"]) self.assertNotIn("$span_id", event["properties"]) @@ -685,7 +688,7 @@ def test_capture_uses_current_otel_span_context_when_enabled( ) with ( - mock.patch("posthog.client.batch_post") as mock_post, + patch_capture_send("client") as mock_post, use_span(NonRecordingSpan(span_context)), ): client = Client( @@ -694,7 +697,7 @@ def test_capture_uses_current_otel_span_context_when_enabled( capture = getattr(client, entrypoint) capture("$ai_event", distinct_id="distinct_id", properties=properties) - event = mock_post.call_args.kwargs["batch"][0] + event = sent_batch(mock_post)[0] if expected_trace_id is None: self.assertNotIn("$trace_id", event["properties"]) self.assertNotIn("$span_id", event["properties"]) @@ -712,14 +715,14 @@ def test_capture_does_not_attach_otel_span_context_by_default(self, entrypoint): ) with ( - mock.patch("posthog.client.batch_post") as mock_post, + patch_capture_send("client") as mock_post, use_span(NonRecordingSpan(span_context)), ): client = Client(FAKE_TEST_API_KEY, sync_mode=True) capture = getattr(client, entrypoint) capture("$ai_event", distinct_id="distinct_id") - event = mock_post.call_args.kwargs["batch"][0] + event = sent_batch(mock_post)[0] self.assertNotIn("$trace_id", event["properties"]) self.assertNotIn("$span_id", event["properties"]) @@ -856,7 +859,7 @@ def test_capture_exception_logs_when_enabled(self): def test_basic_capture_with_feature_flags(self, patch_flags): patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}} - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -871,7 +874,7 @@ def test_basic_capture_with_feature_flags(self, patch_flags): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -979,7 +982,7 @@ def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_flags): }, } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -996,7 +999,7 @@ def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_flags): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1017,7 +1020,7 @@ def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_flags): self.assertEqual(patch_flags.call_count, 0) # test that flags are not evaluated without local evaluation - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1031,7 +1034,7 @@ def test_basic_capture_with_locally_evaluated_feature_flags(self, patch_flags): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] assert "$feature/beta-feature" not in msg["properties"] @@ -1149,7 +1152,7 @@ def test_dont_override_capture_with_local_flags(self, patch_flags): }, } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1169,7 +1172,7 @@ def test_dont_override_capture_with_local_flags(self, patch_flags): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1199,7 +1202,7 @@ def test_basic_capture_with_feature_flags_returns_active_only(self, patch_flags) } } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1214,7 +1217,7 @@ def test_basic_capture_with_feature_flags_returns_active_only(self, patch_flags) # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1283,7 +1286,7 @@ def test_basic_capture_with_feature_flags_and_disable_geoip_returns_correctly( } } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, host="https://app.posthog.com", @@ -1304,7 +1307,7 @@ def test_basic_capture_with_feature_flags_and_disable_geoip_returns_correctly( # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1343,7 +1346,7 @@ def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them( ): patch_flags.return_value = {"featureFlags": {"beta-feature": "random-variant"}} - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1358,7 +1361,7 @@ def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them( # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1421,7 +1424,7 @@ def test_capture_with_send_feature_flags_false_and_local_evaluation_doesnt_send_ }, } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1440,7 +1443,7 @@ def test_capture_with_send_feature_flags_false_and_local_evaluation_doesnt_send_ # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1503,7 +1506,7 @@ def test_capture_with_send_feature_flags_true_and_local_evaluation_uses_local_fl }, } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1522,7 +1525,7 @@ def test_capture_with_send_feature_flags_true_and_local_evaluation_uses_local_fl # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -1549,7 +1552,7 @@ def test_capture_with_send_feature_flags_options_only_evaluate_locally_true( self, patch_flags ): """Test that SendFeatureFlagsOptions with only_evaluate_locally=True uses local evaluation""" - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1591,7 +1594,7 @@ def test_capture_with_send_feature_flags_options_only_evaluate_locally_true( # Check the message includes the local flag mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$feature/local-flag"], True) @@ -1604,7 +1607,7 @@ def test_capture_with_send_feature_flags_options_only_evaluate_locally_false( """Test that SendFeatureFlagsOptions with only_evaluate_locally=False forces remote evaluation""" patch_flags.return_value = {"featureFlags": {"remote-flag": "remote-value"}} - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1638,7 +1641,7 @@ def test_capture_with_send_feature_flags_options_only_evaluate_locally_false( # Check the message includes the remote flag mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$feature/remote-flag"], "remote-value") @@ -1650,7 +1653,7 @@ def test_capture_with_send_feature_flags_options_default_behavior( """Test that SendFeatureFlagsOptions without only_evaluate_locally defaults to remote evaluation""" patch_flags.return_value = {"featureFlags": {"default-flag": "default-value"}} - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1676,7 +1679,7 @@ def test_capture_with_send_feature_flags_options_default_behavior( # Check the message includes the flag mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual( @@ -1688,7 +1691,7 @@ def test_capture_exception_with_send_feature_flags_options(self, patch_flags): """Test that capture_exception also supports SendFeatureFlagsOptions""" patch_flags.return_value = {"featureFlags": {"exception-flag": True}} - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -1718,7 +1721,7 @@ def test_capture_exception_with_send_feature_flags_options(self, patch_flags): # Check the message includes the flag mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "$exception") @@ -1727,7 +1730,7 @@ def test_capture_exception_with_send_feature_flags_options(self, patch_flags): def test_stringifies_distinct_id(self): # A large number that loses precision in node: # node -e "console.log(157963456373623802 + 1)" > 157963456373623800 - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.capture( "python test event", distinct_id=157963456373623802 @@ -1737,13 +1740,13 @@ def test_stringifies_distinct_id(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["distinct_id"], "157963456373623802") def test_advanced_capture(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.capture( "python test event", @@ -1757,7 +1760,7 @@ def test_advanced_capture(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00") @@ -1770,7 +1773,7 @@ def test_advanced_capture(self): self.assertTrue("$groups" not in msg["properties"]) def test_capture_converts_aware_timestamp_to_utc_without_changing_instant(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) client.capture( "python test event", @@ -1780,11 +1783,11 @@ def test_capture_converts_aware_timestamp_to_utc_without_changing_instant(self): ), ) - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00") def test_capture_converts_parseable_timestamp_string_to_utc(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) client.capture( "python test event", @@ -1792,14 +1795,14 @@ def test_capture_converts_parseable_timestamp_string_to_utc(self): timestamp="2014-09-03T05:30:00+05:30", ) - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00") @parameterized.expand(["2026-06-27", "not-an-iso-timestamp"]) def test_capture_replaces_invalid_timestamp_with_current_utc_time(self, timestamp): now = datetime(2026, 6, 27, 12, 30, tzinfo=timezone.utc) with ( - mock.patch("posthog.client.batch_post") as mock_post, + patch_capture_send("client") as mock_post, mock.patch("posthog.client.datetime", wraps=datetime) as mock_datetime, mock.patch.object(Client.log, "warning") as mock_warning, ): @@ -1812,7 +1815,7 @@ def test_capture_replaces_invalid_timestamp_with_current_utc_time(self, timestam ) self.assertIsNotNone(result) - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["timestamp"], "2026-06-27T12:30:00+00:00") mock_warning.assert_called_once_with( "Invalid timestamp %r. Falling back to the current UTC time.", timestamp @@ -1822,7 +1825,7 @@ def test_capture_does_not_normalize_datetime_properties(self): property_value = datetime( 2014, 9, 3, 5, 30, tzinfo=timezone(timedelta(hours=5, minutes=30)) ) - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) client.capture( "python test event", @@ -1830,11 +1833,11 @@ def test_capture_does_not_normalize_datetime_properties(self): properties={"caller_datetime": property_value}, ) - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertIs(msg["properties"]["caller_datetime"], property_value) def test_groups_capture(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.capture( "test_event", @@ -1846,7 +1849,7 @@ def test_groups_capture(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual( @@ -1855,7 +1858,7 @@ def test_groups_capture(self): ) def test_basic_set(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.set( distinct_id="distinct_id", properties={"trait": "value"} @@ -1865,7 +1868,7 @@ def test_basic_set(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["$set"]["trait"], "value") @@ -1874,7 +1877,7 @@ def test_basic_set(self): self.assertEqual(msg["distinct_id"], "distinct_id") def test_advanced_set(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.set( distinct_id="distinct_id", @@ -1887,7 +1890,7 @@ def test_advanced_set(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00") @@ -1899,7 +1902,7 @@ def test_advanced_set(self): self.assertEqual(msg["distinct_id"], "distinct_id") def test_basic_set_once(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.set_once( distinct_id="distinct_id", properties={"trait": "value"} @@ -1909,7 +1912,7 @@ def test_basic_set_once(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["$set_once"]["trait"], "value") @@ -1918,7 +1921,7 @@ def test_basic_set_once(self): self.assertEqual(msg["distinct_id"], "distinct_id") def test_advanced_set_once(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.set_once( distinct_id="distinct_id", @@ -1931,7 +1934,7 @@ def test_advanced_set_once(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00") @@ -1943,7 +1946,7 @@ def test_advanced_set_once(self): self.assertEqual(msg["distinct_id"], "distinct_id") def test_basic_group_identify(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.group_identify("organization", "id:5") @@ -1951,7 +1954,7 @@ def test_basic_group_identify(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "$groupidentify") @@ -1971,7 +1974,7 @@ def test_basic_group_identify(self): self.assertIsNotNone(msg.get("uuid")) def test_basic_group_identify_with_distinct_id(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.group_identify( "organization", "id:5", distinct_id="distinct_id" @@ -1980,7 +1983,7 @@ def test_basic_group_identify_with_distinct_id(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "$groupidentify") @@ -2001,7 +2004,7 @@ def test_basic_group_identify_with_distinct_id(self): self.assertIsNotNone(msg.get("uuid")) def test_advanced_group_identify(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.group_identify( "organization", @@ -2015,7 +2018,7 @@ def test_advanced_group_identify(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "$groupidentify") @@ -2034,7 +2037,7 @@ def test_advanced_group_identify(self): self.assertEqual(msg["timestamp"], "2014-09-03T00:00:00+00:00") def test_advanced_group_identify_with_distinct_id(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.group_identify( "organization", @@ -2049,7 +2052,7 @@ def test_advanced_group_identify_with_distinct_id(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "$groupidentify") @@ -2076,7 +2079,7 @@ def test_advanced_group_identify_with_distinct_id(self): ] ) def test_group_identify_without_group_type_is_dropped(self, _name, group_type): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with self.assertLogs("posthog", level="WARNING") as logs: msg_uuid = client.group_identify(group_type, "id:5") @@ -2092,7 +2095,7 @@ def test_group_identify_without_group_type_is_dropped(self, _name, group_type): ] ) def test_group_identify_without_group_key_is_dropped(self, _name, group_key): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with self.assertLogs("posthog", level="WARNING") as logs: msg_uuid = client.group_identify("organization", group_key) @@ -2102,18 +2105,18 @@ def test_group_identify_without_group_key_is_dropped(self, _name, group_key): self.assertIn("group_key", logs.output[0]) def test_group_identify_accepts_falsy_non_string_group_key(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.group_identify("organization", 0) self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] # The group key is validated, not normalized - it goes out as passed. self.assertEqual(msg["properties"]["$group_key"], 0) def test_basic_alias(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.alias("previousId", "distinct_id") self.assertIsNotNone(msg_uuid) @@ -2121,7 +2124,7 @@ def test_basic_alias(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["distinct_id"], "previousId") self.assertEqual(msg["properties"]["alias"], "distinct_id") @@ -2133,7 +2136,7 @@ def test_basic_alias(self): ] ) def test_alias_without_previous_id_is_dropped(self, _name, previous_id): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with self.assertLogs("posthog", level="WARNING") as logs: msg_uuid = client.alias(previous_id, "distinct_id") @@ -2143,18 +2146,18 @@ def test_alias_without_previous_id_is_dropped(self, _name, previous_id): self.assertIn("previous_id", logs.output[0]) def test_alias_accepts_non_string_previous_id(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) msg_uuid = client.alias(0, "distinct_id") self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() - msg = mock_post.call_args[1]["batch"][0] + msg = sent_batch(mock_post)[0] self.assertEqual(msg["distinct_id"], "0") self.assertEqual(msg["properties"]["distinct_id"], "0") def test_alias_without_distinct_id_is_dropped(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with self.assertLogs("posthog", level="WARNING") as logs: msg_uuid = client.alias("previousId", None) @@ -2190,7 +2193,7 @@ def test_alias_without_distinct_id_is_dropped(self): def test_capture_with_session_id_variations( self, test_name, session_id, additional_properties, expected_properties ): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) properties = {"$session_id": session_id, **additional_properties} @@ -2203,7 +2206,7 @@ def test_capture_with_session_id_variations( # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -2217,7 +2220,7 @@ def test_capture_with_session_id_variations( self.assertEqual(msg["properties"][key], value) def test_session_id_preserved_with_groups(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) session_id = "group-session-101" @@ -2232,7 +2235,7 @@ def test_session_id_preserved_with_groups(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$session_id"], session_id) @@ -2242,7 +2245,7 @@ def test_session_id_preserved_with_groups(self): ) def test_session_id_with_anonymous_event(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) session_id = "anonymous-session-202" @@ -2259,7 +2262,7 @@ def test_session_id_with_anonymous_event(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$session_id"], session_id) @@ -2339,7 +2342,7 @@ def test_session_id_with_different_event_types( additional_properties, expected_additional_properties, ): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) properties = {"$session_id": session_id, **additional_properties} @@ -2351,7 +2354,7 @@ def test_session_id_with_different_event_types( # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], event_name) @@ -2406,7 +2409,7 @@ def test_session_id_with_super_properties_variations( expected_session_id, expected_super_props, ): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, super_properties=super_properties, sync_mode=True ) @@ -2421,7 +2424,7 @@ def test_session_id_with_super_properties_variations( # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$session_id"], expected_session_id) @@ -2712,9 +2715,7 @@ def test_sync_send_failure_does_not_invoke_async_on_error(self): on_error=on_error, ) - with mock.patch( - "posthog.client.batch_post", side_effect=Exception("upload failed") - ): + with patch_capture_send("client", side_effect=Exception("upload failed")): result = client.capture("event", distinct_id="distinct_id") self.assertIsNone(result) @@ -3401,7 +3402,7 @@ def blocking_post(*args, **kwargs): target=lambda: (client.shutdown(), shutdown_done.set()) ) - with mock.patch("posthog.client.batch_post", side_effect=blocking_post) as post: + with patch_capture_send("client", side_effect=blocking_post) as post: capture_thread.start() self.assertTrue(send_started.wait(2)) shutdown_thread.start() @@ -3424,7 +3425,7 @@ def blocking_post(*args, **kwargs): post.assert_called_once() def test_synchronous(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, sync_mode=True) msg_uuid = client.capture("test event", distinct_id="distinct_id") @@ -3516,14 +3517,12 @@ def test_user_defined_flush_at(self): FAKE_TEST_API_KEY, on_error=self.fail, flush_at=10, flush_interval=3 ) - def mock_post_fn(*args, **kwargs): - self.assertEqual(len(kwargs["batch"]), 10) + def mock_post_fn(api_key, host, batch, **kwargs): + self.assertEqual(len(batch), 10) # the post function should be called 2 times, with a batch size of 10 # each time. - with mock.patch( - "posthog.consumer.batch_post", side_effect=mock_post_fn - ) as mock_post: + with patch_capture_send("consumer", side_effect=mock_post_fn) as mock_post: for _ in range(20): client.capture( "event", distinct_id="distinct_id", properties={"trait": "value"} @@ -3576,7 +3575,7 @@ def test_disabled_with_feature_flags(self, patch_flags): self.assertTrue(client.queue.empty()) def test_enabled_to_disabled(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -3590,7 +3589,7 @@ def test_enabled_to_disabled(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["event"], "python test event") @@ -3600,7 +3599,7 @@ def test_enabled_to_disabled(self): self.assertFalse(self.failed) def test_disable_geoip_default_on_events(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -3612,12 +3611,12 @@ def test_disable_geoip_default_on_events(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) capture_msg = batch_data[0] self.assertEqual(capture_msg["properties"]["$geoip_disable"], True) def test_disable_geoip_override_on_events(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -3643,17 +3642,17 @@ def test_disable_geoip_override_on_events(self): self.assertEqual(mock_post.call_count, 2) # Check set event - set_batch = mock_post.call_args_list[0][1]["batch"] + set_batch = sent_batch(mock_post, 0) capture_msg = set_batch[0] self.assertEqual(capture_msg["properties"]["$geoip_disable"], True) # Check page event - page_batch = mock_post.call_args_list[1][1]["batch"] + page_batch = sent_batch(mock_post, 1) identify_msg = page_batch[0] self.assertEqual("$geoip_disable" not in identify_msg["properties"], True) def test_disable_geoip_method_overrides_init_on_events(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -3667,7 +3666,7 @@ def test_disable_geoip_method_overrides_init_on_events(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertTrue("$geoip_disable" not in msg["properties"]) @@ -4119,7 +4118,7 @@ def test_get_decide_returns_normalized_decide_response(self, patch_flags): } def test_set_context_session_with_capture(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with new_context(): set_context_session("context-session-123") @@ -4134,7 +4133,7 @@ def test_set_context_session_with_capture(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual( @@ -4143,7 +4142,7 @@ def test_set_context_session_with_capture(self): @parameterized.expand([("new_context",), ("scoped",)]) def test_client_context_helpers_apply_to_capture(self, context_helper): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) def capture_in_context(): @@ -4171,7 +4170,7 @@ def scoped_capture(): self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["distinct_id"], "context-user") @@ -4182,7 +4181,7 @@ def scoped_capture(): self.assertEqual(client.get_tags(), {}) def test_client_scoped_context_helpers_apply_to_capture_async(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) @client.scoped(fresh=True) @@ -4197,7 +4196,7 @@ async def scoped_capture(): self.assertIsNotNone(msg_uuid) mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["distinct_id"], "async-scoped-user") @@ -4213,7 +4212,7 @@ async def scoped_capture(): self.assertEqual(client.get_tags(), {}) def test_set_context_session_with_page_explicit_properties(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with new_context(): set_context_session("page-explicit-session-789") @@ -4230,7 +4229,7 @@ def test_set_context_session_with_page_explicit_properties(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual( @@ -4241,7 +4240,7 @@ def test_set_context_session_override_in_capture(self): """Test that explicit session ID overrides context session ID in capture""" from posthog.contexts import new_context, set_context_session - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with new_context(): set_context_session("context-session-override") @@ -4259,7 +4258,7 @@ def test_set_context_session_override_in_capture(self): # Get the enqueued message from the mock mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual( @@ -4462,7 +4461,7 @@ def test_capture_with_send_feature_flags_flag_keys_filter(self, patch_flags): } } - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client( FAKE_TEST_API_KEY, on_error=self.set_fail, @@ -4490,7 +4489,7 @@ def test_capture_with_send_feature_flags_flag_keys_filter(self, patch_flags): # Check the message includes only the filtered flags mock_post.assert_called_once() - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$feature/flag1"], "value1") @@ -4498,7 +4497,7 @@ def test_capture_with_send_feature_flags_flag_keys_filter(self, patch_flags): # flag2 should not be included since it wasn't requested self.assertNotIn("$feature/flag2", msg["properties"]) - @mock.patch("posthog.client.batch_post") + @patch_capture_send("client") def test_get_feature_flag_result_with_empty_string_payload(self, patch_batch_post): """Test that get_feature_flag_result returns a FeatureFlagResult when payload is empty string""" client = Client( @@ -4549,7 +4548,7 @@ def test_get_feature_flag_result_with_empty_string_payload(self, patch_batch_pos self.assertEqual(result.get_value(), "empty-variant") self.assertIsNone(result.payload) - @mock.patch("posthog.client.batch_post") + @patch_capture_send("client") def test_get_all_flags_and_payloads_with_empty_string(self, patch_batch_post): """Test that get_all_flags_and_payloads includes flags with empty string payloads""" client = Client( @@ -4605,14 +4604,14 @@ def test_get_all_flags_and_payloads_with_empty_string(self, patch_batch_post): ) def test_context_tags_added(self): - with mock.patch("posthog.client.batch_post") as mock_post: + with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) with new_context(): tag("random_tag", 12345) client.capture("python test event", distinct_id="distinct_id") - batch_data = mock_post.call_args[1]["batch"] + batch_data = sent_batch(mock_post) msg = batch_data[0] self.assertEqual(msg["properties"]["$context_tags"], ["random_tag"]) @@ -4720,8 +4719,9 @@ def _client(self, **kwargs): @parameterized.expand( [ - ("v0", None, False), + ("default", None, True), ("v1", "v1", True), + ("v0", "v0", False), ] ) def test_capture_mode_selects_sync_submitter(self, _name, capture_mode, expects_v1): @@ -4734,9 +4734,10 @@ def test_capture_mode_selects_sync_submitter(self, _name, capture_mode, expects_ if expects_v1: mock_post.assert_not_called() mock_v1.assert_called_once() - sent_batch = mock_v1.call_args.args[2] - self.assertEqual(len(sent_batch), 1) - self.assertEqual(sent_batch[0]["event"], "evt") + batch = mock_v1.call_args.args[2] + self.assertEqual(len(batch), 1) + self.assertEqual(batch[0]["event"], "evt") + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) else: mock_v1.assert_not_called() mock_post.assert_called_once() @@ -4780,6 +4781,6 @@ def test_v1_sync_ai_named_event_through_capture_uses_v1(self): client.capture("$ai_generation", distinct_id="d") mock_post.assert_not_called() mock_v1.assert_called_once() - sent_batch = mock_v1.call_args.args[2] - self.assertEqual(len(sent_batch), 1) - self.assertEqual(sent_batch[0]["event"], "$ai_generation") + batch = mock_v1.call_args.args[2] + self.assertEqual(len(batch), 1) + self.assertEqual(batch[0]["event"], "$ai_generation") diff --git a/posthog/test/test_code_variables.py b/posthog/test/test_code_variables.py index 6263b4b8d..21d527469 100644 --- a/posthog/test/test_code_variables.py +++ b/posthog/test/test_code_variables.py @@ -112,17 +112,11 @@ def extract( import os import posthog from posthog import Posthog -from requests import Response import posthog.request +from posthog.test.capture_helpers import offline_v1_post -def offline_post(url, **kwargs): - response = Response() - response.status_code = 200 - return response - - -posthog.request._session.post = offline_post +posthog.request._session.post = offline_v1_post def make_client(**options): diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index c2c54cf52..1199d2da1 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -16,8 +16,10 @@ from posthog.capture_compression import CaptureCompression from posthog.capture_mode import CaptureMode +from posthog.capture_v1 import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.consumer import MAX_MSG_SIZE, Consumer, _DrainSignal from posthog.request import AI_EVENTS_ENDPOINT, EVENTS_ENDPOINT, APIError +from posthog.test.capture_helpers import patch_capture_send, sent_batch from posthog.test.logging_helpers import capture_message_only_logs from posthog.test.test_utils import TEST_API_KEY @@ -212,11 +214,11 @@ def test_upload(self) -> None: consumer = Consumer(q, TEST_API_KEY, flush_at=1) event = _track_event() q.put(event) - with mock.patch("posthog.consumer.batch_post") as post: + with patch_capture_send("consumer") as post: success = consumer.upload() self.assertTrue(success) post.assert_called_once() - self.assertEqual(post.call_args.kwargs["batch"], [event]) + self.assertEqual(sent_batch(post), [event]) self.assertEqual(q.unfinished_tasks, 0) self.assertTrue(q.empty()) @@ -307,12 +309,12 @@ def test_multiple_uploads_per_interval(self) -> None: delivered = threading.Event() batches = [] - def record_batch(*args, **kwargs): - batches.append(kwargs["batch"]) + def record_batch(api_key, host, batch, **kwargs): + batches.append(batch) if len(batches) == 2: delivered.set() - with mock.patch("posthog.consumer.batch_post", side_effect=record_batch): + with patch_capture_send("consumer", side_effect=record_batch): consumer.start() try: events = [ @@ -328,7 +330,7 @@ def record_batch(*args, **kwargs): self.assertFalse(consumer.is_alive()) def test_request(self) -> None: - consumer = Consumer(None, TEST_API_KEY) + consumer = Consumer(None, TEST_API_KEY, capture_mode=CaptureMode.V0) batch = [_track_event()] with mock.patch("posthog.consumer.batch_post") as post: consumer.request(batch) @@ -357,7 +359,9 @@ def mock_post(*args: Any, **kwargs: Any) -> None: if call_count[0] <= exception_count: raise exception - consumer = Consumer(None, TEST_API_KEY, retries=retries) + consumer = Consumer( + None, TEST_API_KEY, retries=retries, capture_mode=CaptureMode.V0 + ) batch = [_track_event()] with ( mock.patch("posthog.consumer.batch_post", side_effect=mock_post) as post, @@ -400,7 +404,7 @@ def test_request_fails_when_exceptions_exceed_retries(self) -> None: ) def test_negative_retries_still_attempts_delivery_once(self) -> None: - consumer = Consumer(None, TEST_API_KEY, retries=-1) + consumer = Consumer(None, TEST_API_KEY, retries=-1, capture_mode=CaptureMode.V0) with mock.patch("posthog.consumer.batch_post") as mock_post: consumer.request([_track_event()]) @@ -602,7 +606,7 @@ def test_max_batch_size(self) -> None: # Let's capture 8MB of data to trigger two batches n_msgs = int(8_000_000 / msg_size) - with mock.patch.object(consumer, "_send") as mock_send: + with mock.patch.object(consumer, "request") as mock_send: consumer.start() try: for _ in range(0, n_msgs + 2): @@ -629,7 +633,7 @@ def mock_post(*args: Any, **kwargs: Any) -> None: if call_count[0] <= 1: raise error - consumer = Consumer(None, TEST_API_KEY, retries=3) + consumer = Consumer(None, TEST_API_KEY, retries=3, capture_mode=CaptureMode.V0) with ( mock.patch("posthog.consumer.batch_post", side_effect=mock_post), mock.patch("posthog.consumer.time.sleep") as mock_sleep, @@ -646,7 +650,7 @@ def mock_post(*args: Any, **kwargs: Any) -> None: if call_count[0] <= 3: raise error - consumer = Consumer(None, TEST_API_KEY, retries=3) + consumer = Consumer(None, TEST_API_KEY, retries=3, capture_mode=CaptureMode.V0) with ( mock.patch("posthog.consumer.batch_post", side_effect=mock_post), mock.patch("posthog.consumer.time.sleep") as mock_sleep, @@ -687,7 +691,7 @@ def test_request_bounds_retry_after_without_reducing_attempts( session = mock.Mock() session.post.side_effect = [retry_response, retry_response, success_response] - consumer = Consumer(None, TEST_API_KEY, retries=2) + consumer = Consumer(None, TEST_API_KEY, retries=2, capture_mode=CaptureMode.V0) with ( mock.patch("posthog.request._get_session", return_value=session), mock.patch("posthog.consumer.time.sleep") as mock_sleep, @@ -707,7 +711,7 @@ def mock_post(*args: Any, **kwargs: Any) -> None: if call_count[0] <= 1: raise APIError(408, "Request Timeout") - consumer = Consumer(None, TEST_API_KEY, retries=3) + consumer = Consumer(None, TEST_API_KEY, retries=3, capture_mode=CaptureMode.V0) with ( mock.patch("posthog.consumer.batch_post", side_effect=mock_post), mock.patch("posthog.consumer.time.sleep"), @@ -752,10 +756,11 @@ def _ai_event(event_name: str = "$ai_generation") -> dict[str, str]: class TestConsumerCaptureModeRouting(unittest.TestCase): - """`capture_mode` selects the submitter; V0 posts to the consumer's `endpoint`.""" + """`capture_mode` selects the submitter; both post to the consumer's `endpoint`.""" @parameterized.expand( [ + ("default", None, True), ("v0", CaptureMode.V0, False), ("v1", CaptureMode.V1, True), ] @@ -763,7 +768,8 @@ class TestConsumerCaptureModeRouting(unittest.TestCase): def test_capture_mode_selects_analytics_submitter( self, _name, mode, expects_v1 ) -> None: - consumer = Consumer(None, TEST_API_KEY, capture_mode=mode) + kwargs = {"capture_mode": mode} if mode else {} + consumer = Consumer(None, TEST_API_KEY, **kwargs) batch = [_track_event()] with ( mock.patch("posthog.consumer.batch_post") as mock_post, @@ -774,6 +780,7 @@ def test_capture_mode_selects_analytics_submitter( mock_post.assert_not_called() mock_v1.assert_called_once() self.assertEqual(mock_v1.call_args.args[2], batch) + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) else: mock_v1.assert_not_called() mock_post.assert_called_once() @@ -800,8 +807,22 @@ def test_v1_forwards_consumer_config_to_submitter(self) -> None: self.assertEqual(kwargs["max_retries"], 4) self.assertEqual(kwargs["historical_migration"], True) + def test_v1_posts_to_configured_endpoint(self) -> None: + consumer = Consumer(None, TEST_API_KEY, endpoint=_CAPTURE_AI_V1_PATH) + batch = [_ai_event()] + with patch_capture_send("consumer") as mock_v1: + consumer.request(batch) + mock_v1.assert_called_once() + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) + self.assertEqual(sent_batch(mock_v1), batch) + def test_v0_posts_to_configured_endpoint(self) -> None: - consumer = Consumer(None, TEST_API_KEY, endpoint=AI_EVENTS_ENDPOINT) + consumer = Consumer( + None, + TEST_API_KEY, + endpoint=AI_EVENTS_ENDPOINT, + capture_mode=CaptureMode.V0, + ) batch = [_ai_event()] with mock.patch("posthog.consumer.batch_post") as mock_post: consumer.request(batch) diff --git a/posthog/test/test_exception_capture.py b/posthog/test/test_exception_capture.py index ef83c433f..b12ff88fa 100644 --- a/posthog/test/test_exception_capture.py +++ b/posthog/test/test_exception_capture.py @@ -278,15 +278,10 @@ def test_excepthook(tmpdir): dedent( """ from posthog import Posthog - from requests import Response import posthog.request + from posthog.test.capture_helpers import offline_v1_post - def offline_post(url, **kwargs): - response = Response() - response.status_code = 200 - return response - - posthog.request._session.post = offline_post + posthog.request._session.post = offline_v1_post posthog = Posthog('phc_x', host='https://eu.i.posthog.com', enable_exception_autocapture=True, debug=True, on_error=lambda e, batch: print('error handling batch: ', e, batch)) # frame_value = "LOL" @@ -305,7 +300,8 @@ def offline_post(url, **kwargs): assert b"ZeroDivisionError" in output assert b"LOL" in output - assert b"DEBUG:posthog:[PostHog] data uploaded successfully" in output + assert b"[PostHog] capture v1 response" in output + assert b" ok=1 " in output assert ( b'"$exception_list": [{"mechanism": {"type": "generic", "handled": true}, "module": null, "type": "ZeroDivisionError", "value": "division by zero", "stacktrace": {"frames": [{"platform": "python", "filename": "app.py", "abs_path"' in output diff --git a/posthog/test/test_gevent_compat.py b/posthog/test/test_gevent_compat.py index 34e0655bf..2948e9b2c 100644 --- a/posthog/test/test_gevent_compat.py +++ b/posthog/test/test_gevent_compat.py @@ -10,6 +10,7 @@ from posthog.client import Client, _new_lane_queue from posthog.test.test_utils import FAKE_TEST_API_KEY +from posthog.test.capture_helpers import patch_capture_send class TestLaneQueueFallback(unittest.TestCase): @@ -80,7 +81,7 @@ def test_incompatible_queue_does_not_disable_queue_independent_capabilities(self with ( mock.patch("posthog.client.Queue", return_value=incompatible_queue), mock.patch.dict(sys.modules, {"gevent.monkey": None}), - mock.patch("posthog.client.batch_post") as mock_post, + patch_capture_send("client") as mock_post, mock.patch( "posthog.client.flags", return_value={"featureFlags": {"beta-feature": True}}, @@ -113,9 +114,9 @@ def test_capture_and_flush_after_monkey_patching(self): "is not being exercised" ) - from unittest import mock + from posthog.test.capture_helpers import patch_capture_send - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_post: from posthog.client import Client client = Client("phc_test", flush_at=1, flush_interval=60) @@ -137,7 +138,7 @@ def test_capture_and_flush_after_monkey_patching(self): client.flush(timeout_seconds=10) client.join() - assert mock_post.called, "batch_post was never called" + assert mock_post.called, "capture send was never called" assert client.queue.empty(), "flush did not drain the queue" """ ) diff --git a/posthog/test/test_module.py b/posthog/test/test_module.py index dc875e674..6ebd8fe4e 100644 --- a/posthog/test/test_module.py +++ b/posthog/test/test_module.py @@ -7,6 +7,7 @@ import posthog from posthog import Posthog +from posthog.test.capture_helpers import patch_capture_send, sent_batch class TestModule(unittest.TestCase): @@ -17,7 +18,7 @@ def _assert_enqueue_result(self, result): self.assertTrue(result) def setUp(self): - patcher = mock.patch("posthog.consumer.batch_post") + patcher = patch_capture_send("consumer") self.transport = patcher.start() self.addCleanup(patcher.stop) self.on_error = mock.Mock() @@ -32,7 +33,7 @@ def test_track(self): self._assert_enqueue_result(res) self.posthog.flush() self.transport.assert_called_once() - event = self.transport.call_args.kwargs["batch"][0] + event = sent_batch(self.transport)[0] self.assertEqual(event["event"], "python module event") self.assertEqual(event["distinct_id"], "distinct_id") self.assertEqual(event["uuid"], res) @@ -42,7 +43,7 @@ def test_alias(self): self._assert_enqueue_result(res) self.posthog.flush() self.transport.assert_called_once() - event = self.transport.call_args.kwargs["batch"][0] + event = sent_batch(self.transport)[0] self.assertEqual(event["event"], "$create_alias") self.assertEqual(event["distinct_id"], "previousId") self.assertEqual(event["properties"]["alias"], "distinct_id") diff --git a/posthog/test/test_release_id.py b/posthog/test/test_release_id.py index 12b8e7e24..109f79f44 100644 --- a/posthog/test/test_release_id.py +++ b/posthog/test/test_release_id.py @@ -10,6 +10,7 @@ from posthog.client import _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, Client from posthog.release_id import RELEASE_ID_ENV_VAR, _resolve_release_id from posthog.test.test_utils import FAKE_TEST_API_KEY +from posthog.test.capture_helpers import patch_async_capture_send # (name, call, expected event): one row per public event-producing method, shared # by the sync and async clients. Each call builds its own arguments, because a @@ -155,10 +156,10 @@ async def _async_events(env_value, send_events): """Build an async client under `env_value`, run `send_events`, return the batch.""" batches = [] - async def batch_post(*args, **kwargs): - batches.append(kwargs["batch"]) + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) - with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + with patch_async_capture_send(side_effect=send_batch): with _release_id_env(env_value): client = AsyncPosthog("test-key", flush_interval=30) async with client: diff --git a/posthog/test/tracing/test_client_traces.py b/posthog/test/tracing/test_client_traces.py index ab9d0ee0e..f3b2b48c9 100644 --- a/posthog/test/tracing/test_client_traces.py +++ b/posthog/test/tracing/test_client_traces.py @@ -15,6 +15,7 @@ from posthog.tracing._transport import OK from posthog.tracing._span import NOOP_SPAN, RecordingSpan, Span from posthog.version import VERSION +from posthog.test.capture_helpers import patch_capture_send, sent_batch FAKE_API_KEY = "phc_test_key" @@ -671,10 +672,10 @@ def test_the_span_exit_drain_leaves_sync_mode_events_alone(self): client.start_span("x").end() with mock.patch("posthog.client._atexit_deadline", None): client._atexit_spans() - with mock.patch("posthog.client.batch_post") as batch_post: + with patch_capture_send("client") as batch_post: client.capture("after-exit", distinct_id="d") batch_post.assert_called_once() - assert batch_post.call_args[1]["batch"][0]["event"] == "after-exit" + assert sent_batch(batch_post)[0]["event"] == "after-exit" client.shutdown() @pytest.mark.parametrize("traces", [{}, None]) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 51ee00d9d..319451d6a 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -232,7 +232,6 @@ alias posthog.ai.utils.tag -> posthog.tag alias posthog.ai.utils.warn_if_posthog_ai_gateway -> posthog.ai.gateway.warn_if_posthog_ai_gateway alias posthog.args.FeatureFlagEvaluations -> posthog.feature_flag_evaluations.FeatureFlagEvaluations alias posthog.args.SendFeatureFlagsOptions -> posthog.types.SendFeatureFlagsOptions -alias posthog.client.AI_EVENTS_ENDPOINT -> posthog.request.AI_EVENTS_ENDPOINT alias posthog.client.AI_MAX_MSG_SIZE -> posthog.consumer.AI_MAX_MSG_SIZE alias posthog.client.APIError -> posthog.request.APIError alias posthog.client.CaptureCompression -> posthog.capture_compression.CaptureCompression @@ -1186,7 +1185,7 @@ class posthog.capture_compression.CaptureCompression class posthog.capture_mode.CaptureMode class posthog.capture_v1.CaptureV1Error(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) class posthog.client.Client(project_api_key: str, host=None, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, gzip=False, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) -class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, gzip=False, retries=10, timeout=15, historical_migration=False, endpoint=EVENTS_ENDPOINT, max_msg_size=MAX_MSG_SIZE, capture_mode=CaptureMode.V0, capture_compression=CaptureCompression.NONE) +class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, gzip=False, retries=10, timeout=15, historical_migration=False, endpoint=None, max_msg_size=MAX_MSG_SIZE, capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) class posthog.exception_utils.AnnotatedValue(value, metadata) diff --git a/sdk_compliance_adapter/adapter.py b/sdk_compliance_adapter/adapter.py index 0d803ae52..2a68dbf56 100644 --- a/sdk_compliance_adapter/adapter.py +++ b/sdk_compliance_adapter/adapter.py @@ -15,6 +15,7 @@ from posthog import Client from posthog.capture_compression import CaptureCompression +from posthog.capture_v1 import _CAPTURE_V1_PATH from posthog.capture_v1 import _post_v1 as original_post_v1 from posthog.request import EVENTS_ENDPOINT, USER_AGENT from posthog.request import batch_post as original_batch_post @@ -239,6 +240,7 @@ def patched_post_v1( timeout: int = 15, sdk_info: str = USER_AGENT, session: Any = None, + path: str = _CAPTURE_V1_PATH, ): """Patched version of _post_v1 that records requests for /state assertions. @@ -257,6 +259,7 @@ def patched_post_v1( timeout=timeout, sdk_info=sdk_info, session=session, + path=path, ) except Exception as e: status_code = getattr(e, "status", 0) @@ -305,10 +308,10 @@ def patched_post_v1( @app.route("/health", methods=["GET"]) def health(): """Health check endpoint""" + # No AI capture capability: `capture_ai` posts capture v1 to + # /i/v1/ai/events, which this harness version has no suite for. capabilities = ( - ["capture_v1", "capture_ai_v0", "encoding_gzip"] - if is_v1() - else ["capture_v0", "capture_ai_v0", "encoding_gzip"] + ["capture_v1", "encoding_gzip"] if is_v1() else ["capture_v0", "encoding_gzip"] ) capabilities.append("feature_flags_local_evaluation_v1") return jsonify( diff --git a/sdk_compliance_adapter/test_adapter.py b/sdk_compliance_adapter/test_adapter.py index 8acc78386..ac1ddc294 100644 --- a/sdk_compliance_adapter/test_adapter.py +++ b/sdk_compliance_adapter/test_adapter.py @@ -87,7 +87,7 @@ def test_health_opts_into_local_evaluation_without_losing_capture( capabilities = adapter.app.test_client().get("/health").json["capabilities"] assert "feature_flags_local_evaluation_v1" in capabilities assert capability in capabilities - assert "capture_ai_v0" in capabilities + assert "capture_ai_v0" not in capabilities def test_init_enables_explicit_definitions_loading_without_polling(adapter): From 8da46d847f6d42489589a249137a8402f55b6a17 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Sat, 10 Oct 2026 15:38:24 -0700 Subject: [PATCH 03/14] feat!: remove capture v0 and legacy transport options (#1018) Capture v1 is the only capture path. Remove CaptureMode, the capture_mode kwarg and POSTHOG_CAPTURE_MODE, the gzip kwarg, request.batch_post and the v0 endpoint constants, and the backoff dependency. Client and AsyncPosthog constructors are keyword-only after host. --- .github/workflows/sdk-compliance.yml | 14 +- AGENTS.md | 2 +- posthog/__init__.py | 5 - posthog/_async_consumer.py | 67 +---- posthog/_async_request.py | 112 +------- posthog/async_client.py | 23 +- posthog/capture_compression.py | 16 +- posthog/capture_mode.py | 83 ------ posthog/client.py | 82 ++---- posthog/consumer.py | 95 +------ posthog/request.py | 38 +-- posthog/test/mcp/test_posthog_mcp.py | 15 +- ...cy_event_family.json => event_family.json} | 33 ++- posthog/test/snapshots/exception_event.json | 16 +- posthog/test/test_ai_capture_lane.py | 73 ++--- posthog/test/test_async_client.py | 33 +-- posthog/test/test_async_consumer.py | 129 --------- posthog/test/test_async_request.py | 144 ---------- posthog/test/test_capture_compression.py | 49 +--- posthog/test/test_capture_mode.py | 109 -------- posthog/test/test_client.py | 115 +++----- posthog/test/test_consumer.py | 256 ++---------------- posthog/test/test_flag_definition_cache.py | 5 +- posthog/test/test_request.py | 118 +------- posthog/test/test_server_payload_snapshots.py | 30 +- pyproject.toml | 1 - references/public_api_snapshot.txt | 36 +-- sdk_compliance_adapter/Dockerfile.v1 | 25 -- sdk_compliance_adapter/README.md | 2 +- sdk_compliance_adapter/adapter.py | 124 ++------- sdk_compliance_adapter/docker-compose.yml | 12 +- sdk_compliance_adapter/test_adapter.py | 11 +- uv.lock | 11 - 33 files changed, 231 insertions(+), 1653 deletions(-) delete mode 100644 posthog/capture_mode.py rename posthog/test/snapshots/{legacy_event_family.json => event_family.json} (77%) delete mode 100644 posthog/test/test_capture_mode.py delete mode 100644 sdk_compliance_adapter/Dockerfile.v1 diff --git a/.github/workflows/sdk-compliance.yml b/.github/workflows/sdk-compliance.yml index 3891e105c..55cfc8ebe 100644 --- a/.github/workflows/sdk-compliance.yml +++ b/.github/workflows/sdk-compliance.yml @@ -28,21 +28,11 @@ jobs: run: python -m pytest sdk_compliance_adapter/test_adapter.py --timeout=30 compliance: - name: PostHog SDK compliance tests (capture v0) + name: PostHog SDK compliance tests uses: PostHog/posthog-sdk-test-harness/.github/workflows/test-sdk-action.yml@4593de8b423f61fa222115da592e5c18dc82ad3c # 1.11.0 with: adapter-dockerfile: "sdk_compliance_adapter/Dockerfile" adapter-context: "." test-harness-version: "1.1.1" continue-on-error: false - report-name: "sdk-compliance-report-v0" - - compliance-v1: - name: PostHog SDK compliance tests (capture v1) - uses: PostHog/posthog-sdk-test-harness/.github/workflows/test-sdk-action.yml@4593de8b423f61fa222115da592e5c18dc82ad3c # 1.11.0 - with: - adapter-dockerfile: "sdk_compliance_adapter/Dockerfile.v1" - adapter-context: "." - test-harness-version: "1.1.1" - continue-on-error: false - report-name: "sdk-compliance-report-v1" + report-name: "sdk-compliance-report" diff --git a/AGENTS.md b/AGENTS.md index e05e0a255..6b7a8d74f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -23,7 +23,7 @@ Follow [Public API changes](./CONTRIBUTING.md#public-api-changes). As an agent, Before changing capture configuration, serialization, routing, or retries, read the relevant implementation and tests. -Preserve capture v1 as the default (v0 is opt-in for analytics only; `capture_ai` always posts v1 to `/i/v1/ai/events`); strictly typed v1 options and `$set`/`$set_once` relocation; v1-only compression (zlib-wrapped deflate, optional zstd); partial-only per-event retries with stable identity; accumulated drop reporting even on 2xx; terminal v1 `429`; `Retry-After` as a minimum bounded by the shared 30s ceiling; and inline blocking retries with `sync_mode=True`. +Capture v1 is the only capture protocol (`capture` posts to `/i/v1/analytics/events`, `capture_ai` to `/i/v1/ai/events`); strictly typed v1 options and `$set`/`$set_once` relocation; compression (gzip, zlib-wrapped deflate, optional zstd, default none); partial-only per-event retries with stable identity; accumulated drop reporting even on 2xx; terminal v1 `429`; `Retry-After` as a minimum bounded by the shared 30s ceiling; and inline blocking retries with `sync_mode=True`. ## Mirror and build safety diff --git a/posthog/__init__.py b/posthog/__init__.py index ccbc8db58..9e28ec5bf 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -10,7 +10,6 @@ OptionalSetArgs, ) from posthog.capture_compression import CaptureCompression as CaptureCompression -from posthog.capture_mode import CaptureMode as CaptureMode from posthog.client import Client from posthog.tracing.span import Span as Span from posthog.async_client import AsyncClient as AsyncClient @@ -424,9 +423,6 @@ def get_tags() -> Dict[str, Any]: # We recommend setting this to False if you are only using the personalApiKey for evaluating remote config payloads via `get_remote_config_payload` and not using local evaluation. enable_local_evaluation = True # type: bool flag_definition_cache_provider = None # type: Optional[FlagDefinitionCacheProvider] -# Capture wire protocol for the global client. None defers to POSTHOG_CAPTURE_MODE -# then CaptureMode.V1. See posthog.capture_mode.CaptureMode. -capture_mode = None # type: Optional[CaptureMode] # Routes AI SDK wrapper events through the dedicated AI capture lane, skips # truncation, and passes media unredacted. `privacy_mode` always wins. enable_full_ai_capture = False # type: bool @@ -1423,7 +1419,6 @@ def setup() -> Client: exception_autocapture_bucket_size=exception_autocapture_bucket_size, exception_autocapture_refill_rate=exception_autocapture_refill_rate, exception_autocapture_refill_interval_seconds=exception_autocapture_refill_interval_seconds, - capture_mode=capture_mode, ) # Always set in case user changes it. Preserve Client's auto-disabled state diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 04eeb082b..9352139ee 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -9,11 +9,10 @@ from dataclasses import dataclass from typing import Any, Optional -from ._async_request import async_batch_post, async_send_v1_batch +from ._async_request import async_send_v1_batch from .capture_compression import CaptureCompression -from .capture_mode import CaptureMode from .consumer import BATCH_SIZE_LIMIT, MAX_MSG_SIZE -from .request import APIError, DatetimeSerializer, EVENTS_ENDPOINT +from .request import DatetimeSerializer _STOP = object() _PROCESSING_EVENT = contextvars.ContextVar( @@ -68,13 +67,10 @@ def __init__( process_event: Callable[[dict[str, Any]], Awaitable[Optional[dict[str, Any]]]], flush_at: int, flush_interval: float, - gzip: bool, retries: int, timeout: int, historical_migration: bool, - capture_mode: CaptureMode, capture_compression: CaptureCompression, - http_client: Optional[Any], ) -> None: self.queue = queue self.api_key = api_key @@ -83,13 +79,10 @@ def __init__( self.process_event = process_event self.flush_at = flush_at self.flush_interval = flush_interval - self.gzip = gzip self.retries = max(0, retries) self.timeout = timeout self.historical_migration = historical_migration - self.capture_mode = capture_mode self.capture_compression = capture_compression - self.http_client = http_client self._carryover: Optional[tuple[dict[str, Any], int]] = None self._flush_event = asyncio.Event() @@ -238,50 +231,12 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: return items, stop async def request(self, batch: list[dict[str, Any]]) -> None: - if self.capture_mode == CaptureMode.V1: - await async_send_v1_batch( - self.api_key, - self.host, - batch, - compression=self.capture_compression, - timeout=self.timeout, - max_retries=self.retries, - historical_migration=self.historical_migration, - ) - return - - last_error: Optional[Exception] = None - for attempt in range(self.retries + 1): - try: - await async_batch_post( - self.api_key, - self.host, - batch=batch, - path=EVENTS_ENDPOINT, - gzip=self.gzip, - timeout=self.timeout, - historical_migration=self.historical_migration, - client=self.http_client, - ) - return - except Exception as error: - last_error = error - if not self._is_retryable(error) or attempt >= self.retries: - raise - retry_after = getattr(error, "retry_after", None) - delay = max( - min(2**attempt, 30), - min(retry_after, 30) if retry_after and retry_after > 0 else 0, - ) - await asyncio.sleep(delay) - - if last_error is not None: # pragma: no cover - loop always raises first - raise last_error - - @staticmethod - def _is_retryable(error: Exception) -> bool: - if not isinstance(error, APIError): - return True - if not isinstance(error.status, int): - return False - return not (400 <= error.status < 500 and error.status not in (408, 429)) + await async_send_v1_batch( + self.api_key, + self.host, + batch, + compression=self.capture_compression, + timeout=self.timeout, + max_retries=self.retries, + historical_migration=self.historical_migration, + ) diff --git a/posthog/_async_request.py b/posthog/_async_request.py index ae750de9a..15b218f74 100644 --- a/posthog/_async_request.py +++ b/posthog/_async_request.py @@ -2,13 +2,9 @@ import asyncio import json -import logging -import zlib from datetime import datetime, timezone -from gzip import GzipFile -from io import BytesIO from typing import Any, Optional -from urllib.parse import quote, urljoin, urlsplit +from urllib.parse import quote from .capture_compression import CaptureCompression from .capture_v1 import _parse_retry_after, _send_v1_batch @@ -41,54 +37,6 @@ def _build_client(host: Optional[str] = None): return httpx_module.AsyncClient(base_url=base_url, follow_redirects=False) -def _serialize_v0_body( - api_key: str, gzip_enabled: bool, body: dict[str, Any] -) -> tuple[str | bytes, dict[str, str]]: - payload = { - **body, - "sent_at": datetime.now(tz=timezone.utc).isoformat(), - "api_key": api_key, - } - serialized = json.dumps(payload, cls=DatetimeSerializer) - data: str | bytes = serialized - headers = {"Content-Type": "application/json", "User-Agent": USER_AGENT} - - if gzip_enabled: - try: - buf = BytesIO() - with GzipFile(fileobj=buf, mode="w") as gz: - gz.write(serialized.encode("utf-8")) - data = buf.getvalue() - headers["Content-Encoding"] = "gzip" - except (OSError, zlib.error) as exc: - logging.getLogger("posthog").warning( - "failed to gzip async request body, sending uncompressed: %s", exc - ) - - return data, headers - - -def _origin(url: str) -> tuple[str, str, Optional[int]]: - parsed = urlsplit(url) - port = parsed.port - if port is None: - port = 443 if parsed.scheme.lower() == "https" else 80 - return parsed.scheme.lower(), (parsed.hostname or "").lower(), port - - -def _same_origin_redirect_url( - base_url: str, current_url: str, location: str -) -> Optional[str]: - target = urlsplit(urljoin(current_url, location)) - if _origin(target.geturl()) != _origin(base_url): - return None - return ( - urlsplit(base_url) - ._replace(path=target.path or "/", query=target.query, fragment="") - .geturl() - ) - - def _serialize_flags_body( project_api_key: str, body: dict[str, Any] ) -> tuple[str, dict[str, str]]: @@ -196,64 +144,6 @@ async def async_remote_config( await http_client.aclose() -async def async_batch_post( - api_key: str, - host: Optional[str], - *, - batch: list[dict[str, Any]], - path: str, - gzip: bool = False, - timeout: int = 15, - historical_migration: bool = False, - client: Optional[Any] = None, -) -> None: - """Post one legacy capture batch without blocking the event loop.""" - if not path.startswith("/") or "://" in path: - raise ValueError("async capture paths must be relative") - - data, headers = await asyncio.to_thread( - _serialize_v0_body, - api_key, - gzip, - { - "batch": batch, - "historical_migration": historical_migration, - }, - ) - - owns_client = client is None - http_client = client or _build_client(host) - try: - logging.getLogger("posthog").debug("making async capture request") - base_url = remove_trailing_slash(normalize_host(host)) - # Absolute URLs avoid reapplying an HTTPX base_url path on redirects. - request_url = f"{base_url}{path}" - for redirect_count in range(6): - response = await http_client.post( - request_url, content=data, headers=headers, timeout=timeout - ) - if response.status_code not in (307, 308): - _process_response(response) - return - - location = response.headers.get("Location") or response.headers.get( - "location" - ) - redirect_url = ( - _same_origin_redirect_url(base_url, request_url, location) - if location - else None - ) - if redirect_url is None: - raise APIError(400, "Cross-origin or invalid redirect blocked") - if redirect_count >= 5: - raise APIError(400, "Too many capture redirects") - request_url = redirect_url - finally: - if owns_client: - await http_client.aclose() - - async def async_send_v1_batch( api_key: str, host: Optional[str], diff --git a/posthog/async_client.py b/posthog/async_client.py index 3f429e81a..e46f0e494 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -26,7 +26,6 @@ ) from ._async_request import ( _build_client, - _require_httpx, async_flags as _async_flags, async_remote_config as _async_remote_config, ) @@ -35,7 +34,6 @@ CaptureCompression, _resolve_capture_compression, ) -from .capture_mode import CaptureMode, _resolve_capture_mode from .client import ( MAX_DICT_SIZE as _MAX_DICT_SIZE, _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, @@ -98,13 +96,13 @@ def __init__( self, project_api_key: str, host: Optional[str] = None, + *, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, - gzip: bool = False, max_retries: int = 3, timeout: int = 15, thread: int = 1, @@ -122,7 +120,6 @@ def __init__( code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, - capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, secret_key: Optional[str] = None, @@ -141,7 +138,6 @@ def __init__( self.debug = debug self.send = send self.on_error = on_error - self.gzip = gzip self.max_retries = max(0, max_retries) self.timeout = timeout self.disabled = disabled or not self.api_key @@ -150,10 +146,7 @@ def __init__( self.historical_migration = historical_migration self.super_properties = super_properties self._release_id = _resolve_release_id() - self.capture_mode = _resolve_capture_mode(capture_mode) - self.capture_compression = _resolve_capture_compression( - capture_compression, gzip_fallback=gzip - ) + self.capture_compression = _resolve_capture_compression(capture_compression) self.capture_trace_context = capture_trace_context if personal_api_key is not None and secret_key is None: warnings.warn( @@ -310,9 +303,6 @@ def _get_http_client(self): return self._http_client def _new_consumer(self) -> _AsyncConsumer: - http_client = ( - self._get_http_client() if self.capture_mode == CaptureMode.V0 else None - ) return _AsyncConsumer( self._queue, self.api_key, @@ -321,19 +311,12 @@ def _new_consumer(self) -> _AsyncConsumer: process_event=self._process_event, flush_at=self._flush_at, flush_interval=self._flush_interval, - gzip=self.gzip, retries=self.max_retries, timeout=self.timeout, historical_migration=self.historical_migration, - capture_mode=self.capture_mode, capture_compression=self.capture_compression, - http_client=http_client, ) - def _validate_transport_available(self) -> None: - if self.capture_mode == CaptureMode.V0: - _require_httpx() - def _ensure_workers_started(self) -> None: if self.disabled or not self.send or self._closed or self._worker_tasks: return @@ -523,7 +506,6 @@ def capture( if not self.send: return sent_uuid - self._validate_transport_available() if not self._enqueue_prepared_event(prepared): return None self.log.debug("queued async event %s", event) @@ -745,7 +727,6 @@ def _enqueue_built_event( return None if not self.send: return sent_uuid - self._validate_transport_available() if not self._enqueue_prepared_event(prepared): return None return sent_uuid diff --git a/posthog/capture_compression.py b/posthog/capture_compression.py index 659c12fab..52a1d4c2d 100644 --- a/posthog/capture_compression.py +++ b/posthog/capture_compression.py @@ -19,10 +19,9 @@ class CaptureCompression(str, Enum): - """Selects the request-body compression for capture-v1 uploads. + """Selects the request-body compression for capture uploads. - Only honored when ``capture_mode`` is ``V1``; the legacy ``/batch/`` path - keeps using its own ``gzip`` flag. ``NONE`` sends the body uncompressed. + ``NONE`` sends the body uncompressed. ``GZIP`` and ``DEFLATE`` (zlib, RFC 1950) are both stdlib / zero-dependency; ``ZSTD`` is faster and compresses better but needs the optional zstandard package (``pip install posthog[zstd]``) until stdlib support lands in @@ -76,15 +75,12 @@ def _coerce_explicit( def _resolve_capture_compression( capture_compression: Optional[Union[CaptureCompression, str]] = None, - *, - gzip_fallback: bool = False, ) -> CaptureCompression: - """Resolve the effective v1 compression. + """Resolve the effective capture compression. Precedence: explicit ``capture_compression`` argument > - ``POSTHOG_CAPTURE_COMPRESSION`` env var > the legacy ``gzip`` flag - (``GZIP`` when set) > ``NONE``. An unrecognized env value logs a warning and - falls back to the ``gzip`` flag, so a typo never silently changes encoding. + ``POSTHOG_CAPTURE_COMPRESSION`` env var > ``NONE``. An unrecognized env + value logs a warning and falls back to ``NONE``. ``ZSTD`` requires the optional zstandard package: explicitly requesting it without the package raises ``ValueError`` (programming error, fail loud), @@ -100,7 +96,7 @@ def _resolve_capture_compression( ) return resolved - fallback = CaptureCompression.GZIP if gzip_fallback else CaptureCompression.NONE + fallback = CaptureCompression.NONE raw = os.environ.get(CAPTURE_COMPRESSION_ENV_VAR) if raw is None or raw.strip() == "": diff --git a/posthog/capture_mode.py b/posthog/capture_mode.py deleted file mode 100644 index 5f72e3702..000000000 --- a/posthog/capture_mode.py +++ /dev/null @@ -1,83 +0,0 @@ -import logging -import os -from enum import Enum -from typing import Optional, Union - -__all__ = ["CAPTURE_MODE_ENV_VAR", "CaptureMode"] - -log = logging.getLogger("posthog") - -CAPTURE_MODE_ENV_VAR = "POSTHOG_CAPTURE_MODE" - - -class CaptureMode(str, Enum): - """Selects the capture wire protocol used for event ingestion. - - ``V1`` is ``POST /i/v1/analytics/events`` (Bearer auth, per-event results, - partial retry) and the default. ``V0`` opts back into the legacy - ``POST /batch/`` endpoint. Inheriting from ``str`` keeps the members directly comparable to and - serializable as their ``"v0"`` / ``"v1"`` values. - """ - - V0 = "v0" - V1 = "v1" - - -# Accepted spellings for both the explicit kwarg and the env var. Aliases mirror -# the posthog-go naming (``legacy`` / ``analytics_v1``) so the two SDKs are -# configured with the same vocabulary. -_ALIASES: dict[str, CaptureMode] = { - "v0": CaptureMode.V0, - "legacy": CaptureMode.V0, - "v1": CaptureMode.V1, - "analytics_v1": CaptureMode.V1, -} - - -def _coerce_explicit(value: Union[CaptureMode, str]) -> CaptureMode: - """Normalize an explicitly-supplied capture mode to a ``CaptureMode``. - - Accepts a ``CaptureMode`` or one of the string aliases. An explicit but - unrecognized value is a programming error, so it raises ``ValueError`` rather - than silently defaulting (unlike the env var, which is operator-supplied and - defaults defensively). - """ - if isinstance(value, CaptureMode): - return value - if isinstance(value, str): - resolved = _ALIASES.get(value.strip().lower()) - if resolved is not None: - return resolved - raise ValueError( - f"invalid capture_mode {value!r}; expected a CaptureMode or one of " - f"{sorted(_ALIASES)}" - ) - - -def _resolve_capture_mode( - capture_mode: Optional[Union[CaptureMode, str]] = None, -) -> CaptureMode: - """Resolve the effective capture mode. - - Precedence: explicit ``capture_mode`` argument > ``POSTHOG_CAPTURE_MODE`` env - var > ``CaptureMode.V1``. An unrecognized env value logs a warning and falls - back to ``V1`` so a typo never silently flips the wire protocol. - """ - if capture_mode is not None: - return _coerce_explicit(capture_mode) - - raw = os.environ.get(CAPTURE_MODE_ENV_VAR) - if raw is None or raw.strip() == "": - return CaptureMode.V1 - - resolved = _ALIASES.get(raw.strip().lower()) - if resolved is None: - log.warning( - "Unrecognized %s=%r; falling back to %s. Expected one of %s.", - CAPTURE_MODE_ENV_VAR, - raw, - CaptureMode.V1.value, - sorted(_ALIASES), - ) - return CaptureMode.V1 - return resolved diff --git a/posthog/client.py b/posthog/client.py index d2d705d52..083a6eb03 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -31,7 +31,6 @@ CaptureCompression, _resolve_capture_compression, ) -from posthog.capture_mode import CaptureMode, _resolve_capture_mode from posthog.capture_v1 import ( _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, @@ -89,7 +88,6 @@ from posthog.poller import Poller from posthog.release_id import _resolve_release_id from posthog.request import ( - EVENTS_ENDPOINT, USER_AGENT as _USER_AGENT, APIError, QuotaLimitError, @@ -97,7 +95,6 @@ RequestsTimeout, _get as _get_with_identity, _remote_config as _remote_config_with_identity, - batch_post, determine_server_host, flags, get, @@ -377,13 +374,11 @@ def __init__( send, flush_at, flush_interval, - gzip, max_retries, timeout, historical_migration, endpoint, max_msg_size, - capture_mode, capture_compression, sdk_info, eager_start, @@ -395,13 +390,11 @@ def __init__( self.send = send self.flush_at = flush_at self.flush_interval = flush_interval - self.gzip = gzip self.max_retries = max_retries self.timeout = timeout self.historical_migration = historical_migration self.endpoint = endpoint self.max_msg_size = max_msg_size - self.capture_mode = capture_mode self.capture_compression = capture_compression self.sdk_info = sdk_info self._max_queue_size = max_queue_size @@ -430,13 +423,11 @@ def _start_locked(self) -> None: on_error=self.on_error, flush_at=self.flush_at, flush_interval=self.flush_interval, - gzip=self.gzip, retries=self.max_retries, timeout=self.timeout, historical_migration=self.historical_migration, endpoint=self.endpoint, max_msg_size=self.max_msg_size, - capture_mode=self.capture_mode, capture_compression=self.capture_compression, ) consumer._sdk_info = self.sdk_info @@ -674,13 +665,13 @@ def __init__( self, project_api_key: str, host=None, + *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, - gzip=False, max_retries=3, sync_mode=False, timeout=15, @@ -712,13 +703,10 @@ def __init__( exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, - capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, - # Appended rather than grouped with the other `capture_*` options so - # existing positional arguments keep their slots. capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, @@ -744,7 +732,6 @@ def __init__( flush_at: Number of queued events that triggers a batch upload. flush_interval: Maximum seconds a background consumer waits before flushing a partial batch. - gzip: Whether to gzip event upload payloads. max_retries: Number of upload retries. Values below 0 are treated as 0. sync_mode: If True, send each event synchronously instead of using background worker threads. This blocks the calling thread; in @@ -843,16 +830,11 @@ def __init__( interval for each exception type's bucket. exception_autocapture_refill_interval_seconds: Seconds between token refills for autocaptured exception rate limiting. - capture_mode: Capture wire protocol to use. Defaults to - ``CaptureMode.V1`` (``/i/v1/analytics/events``). Set - ``CaptureMode.V0`` (or pass the string ``"v0"``) to opt back - into the legacy ``/batch/`` endpoint. When omitted, the - ``POSTHOG_CAPTURE_MODE`` env var is consulted, then ``V1``. - capture_compression: Request-body compression for capture-v1 uploads - (ignored in V0, which uses ``gzip``). ``CaptureCompression.GZIP`` - or ``DEFLATE`` (or the strings ``"gzip"``/``"deflate"``). When - omitted, the ``POSTHOG_CAPTURE_COMPRESSION`` env var is consulted, - then the legacy ``gzip`` flag, then no compression. + capture_compression: Request-body compression for ``capture()`` + uploads. ``CaptureCompression.GZIP`` or ``DEFLATE`` (or the + strings ``"gzip"``/``"deflate"``). When omitted, the + ``POSTHOG_CAPTURE_COMPRESSION`` env var is consulted, then no + compression. Examples: ```python @@ -893,7 +875,6 @@ def __init__( self.raw_host = normalize_host(host) self.host = determine_server_host(host) self._duplicate_client_registry_key: Optional[tuple[str, str]] = None - self.gzip = gzip self.timeout = timeout self.max_retries = max(0, max_retries) self._feature_flags: Optional[list[Any]] = ( @@ -953,18 +934,10 @@ def __init__( ) self.is_server = is_server self.historical_migration = historical_migration - # Selects the capture wire protocol (V0 legacy `/batch/` vs V1 - # `/i/v1/analytics/events`). Resolved here so the env-var fallback is - # applied once; V0 is the default and keeps upgrades transparent. - self.capture_mode = _resolve_capture_mode(capture_mode) self._library_id = "posthog-python" self._library_version = VERSION self._sdk_info = f"{self._library_id}/{self._library_version}" - # v1-only request compression; falls back to the legacy `gzip` flag when - # neither the kwarg nor POSTHOG_CAPTURE_COMPRESSION is set. - self.capture_compression = _resolve_capture_compression( - capture_compression, gzip_fallback=gzip - ) + self.capture_compression = _resolve_capture_compression(capture_compression) self.super_properties = super_properties # Release id from POSTHOG_RELEASE_ID, attached to every event. Resolved # here so the env var is read once per client. @@ -1079,7 +1052,6 @@ def __init__( send=send, flush_at=flush_at, flush_interval=flush_interval, - gzip=gzip, max_retries=self.max_retries, timeout=timeout, historical_migration=historical_migration, @@ -1088,27 +1060,20 @@ def __init__( self._analytics_lane = _Lane( name="analytics", **lane_defaults, - endpoint=( - _CAPTURE_V1_PATH - if self.capture_mode == CaptureMode.V1 - else EVENTS_ENDPOINT - ), + endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, - capture_mode=self.capture_mode, capture_compression=self.capture_compression, eager_start=not sync_mode, ) - # The AI lane always posts capture v1 to the AI endpoint, whatever the - # analytics `capture_mode`, so multi-MB AI events stay off the - # analytics endpoint's smaller caps. It sends uncompressed. Lazy start, - # so the many clients that never emit AI events pay for no extra + # The AI lane posts to its own endpoint so multi-MB AI events stay off + # the analytics endpoint's smaller caps. It sends uncompressed. Lazy + # start, so the many clients that never emit AI events pay for no extra # threads. self._ai_lane = _Lane( name="ai", **lane_defaults, endpoint=_CAPTURE_AI_V1_PATH, max_msg_size=AI_MAX_MSG_SIZE, - capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.NONE, eager_start=False, ) @@ -2445,30 +2410,17 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): def send_sync() -> None: # Sync mode bypasses the lane's queue but keeps its wire config, - # so AI events post to the AI endpoint whatever `capture_mode`. - if lane.capture_mode == CaptureMode.V1: - _send_v1_batch( - self.api_key, - self.host, - [msg], - compression=lane.capture_compression, - timeout=self.timeout, - max_retries=self.max_retries, - historical_migration=self.historical_migration, - sdk_info=self._sdk_info, - path=lane.endpoint, - ) - return - - batch_post( + # so AI events still post to the AI endpoint. + _send_v1_batch( self.api_key, self.host, - gzip=self.gzip, + [msg], + compression=lane.capture_compression, timeout=self.timeout, - batch=[msg], + max_retries=self.max_retries, historical_migration=self.historical_migration, + sdk_info=self._sdk_info, path=lane.endpoint, - **self._request_identity_kwargs(), ) if lane.run_sync_if_open(send_sync): diff --git a/posthog/consumer.py b/posthog/consumer.py index 8a69844cf..0b15431a8 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -6,14 +6,10 @@ from posthog._logging import _configure_posthog_logging from posthog.capture_compression import CaptureCompression -from posthog.capture_mode import CaptureMode -from posthog.capture_v1 import _CAPTURE_V1_PATH, _backoff, _send_v1_batch +from posthog.capture_v1 import _CAPTURE_V1_PATH, _send_v1_batch from posthog.request import ( - EVENTS_ENDPOINT, USER_AGENT as _USER_AGENT, - APIError, DatetimeSerializer, - batch_post, ) from queue import Empty @@ -109,13 +105,11 @@ def __init__( host=None, on_error=None, flush_interval=5.0, - gzip=False, retries=10, timeout=15, historical_migration=False, - endpoint=None, + endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, - capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.NONE, ): """Create a consumer thread.""" @@ -128,16 +122,8 @@ def __init__( self.host = host self.on_error = on_error self.queue = queue - self.gzip = gzip - # Without an explicit endpoint, post to the analytics path of the - # selected protocol. - if endpoint is None: - endpoint = ( - _CAPTURE_V1_PATH if capture_mode == CaptureMode.V1 else EVENTS_ENDPOINT - ) self.endpoint = endpoint self.max_msg_size = max_msg_size - self.capture_mode = capture_mode self.capture_compression = capture_compression self._sdk_info = _USER_AGENT self._drain_signal: Optional[_DrainSignal] = None @@ -294,67 +280,16 @@ def next(self): return items def request(self, batch): - """Upload the batch to this consumer's `endpoint` via the wire protocol - selected by `capture_mode`. - - V1 uses the partial-retry submitter; V0 posts the whole batch. - """ - if self.capture_mode == CaptureMode.V1: - _send_v1_batch( - self.api_key, - self.host, - batch, - compression=self.capture_compression, - timeout=self.timeout, - max_retries=self.retries, - historical_migration=self.historical_migration, - sdk_info=self._sdk_info, - path=self.endpoint, - ) - return - self._send(batch, self.endpoint) - - def _send(self, batch, path): - """Attempt to upload a single batch to `path`, retrying before raising an error""" - - def is_retryable(exc): - if isinstance(exc, APIError): - # retry on server errors and client errors - # with 408 (request timeout) or 429 (rate limited), - # don't retry on other client errors - if isinstance(exc.status, int): - return not ( - (400 <= exc.status < 500) and exc.status not in (408, 429) - ) - return False - else: - # retry on all other errors (eg. network) - return True - - last_exc = None - for attempt in range(self.retries + 1): - try: - batch_post( - self.api_key, - self.host, - gzip=self.gzip, - timeout=self.timeout, - batch=batch, - historical_migration=self.historical_migration, - path=path, - **( - {"_user_agent": self._sdk_info} - if self._sdk_info != _USER_AGENT - else {} - ), - ) - return - except Exception as e: - last_exc = e - if not is_retryable(e): - raise - if attempt < self.retries: - _backoff(attempt, getattr(e, "retry_after", None)) - - if last_exc: - raise last_exc + """Upload the batch to this consumer's `endpoint` with the capture v1 + partial-retry submitter.""" + _send_v1_batch( + self.api_key, + self.host, + batch, + compression=self.capture_compression, + timeout=self.timeout, + max_retries=self.retries, + historical_migration=self.historical_migration, + sdk_info=self._sdk_info, + path=self.endpoint, + ) diff --git a/posthog/request.py b/posthog/request.py index 76df1fdc9..2f68fc38e 100644 --- a/posthog/request.py +++ b/posthog/request.py @@ -3,11 +3,8 @@ import re import socket import time -import zlib from dataclasses import dataclass from datetime import date, datetime, timezone -from gzip import GzipFile -from io import BytesIO from typing import Any, List, Optional, Tuple, Union, cast import requests @@ -216,7 +213,6 @@ def post( api_key: str, host: Optional[str] = None, path: Optional[str] = None, - gzip: bool = False, timeout: int = 15, session: Optional[requests.Session] = None, **kwargs, @@ -229,7 +225,7 @@ def post( trimmed_host = remove_trailing_slash(normalize_host(host)) url = trimmed_host + cast(str, path) body["api_key"] = api_key - data: str | bytes = json.dumps(body, cls=DatetimeSerializer) + data = json.dumps(body, cls=DatetimeSerializer) if log.isEnabledFor(logging.DEBUG): log.debug( "making request: %s to url: %s", @@ -237,17 +233,6 @@ def post( url, ) headers = {"Content-Type": "application/json", "User-Agent": user_agent} - if gzip: - try: - buf = BytesIO() - with GzipFile(fileobj=buf, mode="w") as gz: - # 'data' was produced by json.dumps(), - # whose default encoding is utf-8. - gz.write(cast(str, data).encode("utf-8")) - data = buf.getvalue() - headers["Content-Encoding"] = "gzip" - except (OSError, zlib.error) as exc: - log.warning("failed to gzip request body, sending uncompressed: %s", exc) res = (session or _get_session()).post( url, data=data, headers=headers, timeout=timeout @@ -314,7 +299,6 @@ def _feature_flags_retry_delay(failed_attempt: int) -> float: def flags( api_key: str, host: Optional[str] = None, - gzip: bool = False, timeout: int = 15, max_retries: int = 1, **kwargs, @@ -330,7 +314,6 @@ def flags( api_key, host, "/flags/?v=2", - gzip, timeout, session=_get_flags_session(), _user_agent=user_agent, @@ -382,25 +365,6 @@ def _remote_config( return response.data -EVENTS_ENDPOINT = "/batch/" -AI_EVENTS_ENDPOINT = "/i/v0/ai/batch/" - - -def batch_post( - api_key: str, - host: Optional[str] = None, - gzip: bool = False, - timeout: int = 15, - path: str = EVENTS_ENDPOINT, - **kwargs, -) -> requests.Response: - """Post the `kwargs` to the batch API endpoint for events""" - res = post(api_key, host, path, gzip, timeout, **kwargs) - return _process_response( - res, success_message="data uploaded successfully", return_json=False - ) - - def get( api_key: str, url: str, diff --git a/posthog/test/mcp/test_posthog_mcp.py b/posthog/test/mcp/test_posthog_mcp.py index 6ed662e81..8820a156a 100644 --- a/posthog/test/mcp/test_posthog_mcp.py +++ b/posthog/test/mcp/test_posthog_mcp.py @@ -9,7 +9,6 @@ import pytest from mcp.types import CallToolResult, ServerResult, TextContent, Tool -from posthog.capture_mode import CaptureMode from posthog.mcp import ( PostHogMCP, PreparedToolCall, @@ -134,20 +133,8 @@ def before_send(event): ) -def test_mcp_library_identity_reaches_capture_v0_header(): - response = mock.Mock(status_code=200) - client = PostHogMCP("phc_test", sync_mode=True) - - with mock.patch("posthog.request._session.post", return_value=response) as post: - client.capture("$mcp_custom") - - assert post.call_args.kwargs["headers"]["User-Agent"] == ( - f"posthog-python-mcp/{VERSION}" - ) - - def test_mcp_library_identity_reaches_capture_v1_header(): - client = PostHogMCP("phc_test", sync_mode=True, capture_mode=CaptureMode.V1) + client = PostHogMCP("phc_test", sync_mode=True) with mock.patch("posthog.client._send_v1_batch") as send: client.capture("$mcp_custom") diff --git a/posthog/test/snapshots/legacy_event_family.json b/posthog/test/snapshots/event_family.json similarity index 77% rename from posthog/test/snapshots/legacy_event_family.json rename to posthog/test/snapshots/event_family.json index 4e33c6bd8..0d72f4cf5 100644 --- a/posthog/test/snapshots/legacy_event_family.json +++ b/posthog/test/snapshots/event_family.json @@ -1,18 +1,16 @@ { "body": { - "api_key": "phc_snapshot_project", "batch": [ { "distinct_id": "user-123", "event": "order completed", + "options": {}, "properties": { "$geoip_disable": true, "$groups": { "company": "company-456" }, "$is_server": true, - "$lib": "posthog-python", - "$lib_version": "", "$os": "", "$os_distro": "", "$os_version": "", @@ -34,17 +32,16 @@ "uuid": "00000000-0000-4000-8000-000000000001" }, { - "$set": { - "email": "person@example.com", - "plan": "pro" - }, "distinct_id": "user-123", "event": "$set", + "options": {}, "properties": { "$geoip_disable": true, "$is_server": true, - "$lib": "posthog-python", - "$lib_version": "" + "$set": { + "email": "person@example.com", + "plan": "pro" + } }, "timestamp": "2026-01-02T03:04:05+00:00", "uuid": "00000000-0000-4000-8000-000000000002" @@ -52,11 +49,10 @@ { "distinct_id": "anonymous-789", "event": "$create_alias", + "options": {}, "properties": { "$geoip_disable": true, "$is_server": true, - "$lib": "posthog-python", - "$lib_version": "", "alias": "user-123", "distinct_id": "anonymous-789" }, @@ -66,6 +62,7 @@ { "distinct_id": "user-123", "event": "$groupidentify", + "options": {}, "properties": { "$geoip_disable": true, "$group_key": "company-456", @@ -74,21 +71,23 @@ "name": "Example Corp" }, "$group_type": "company", - "$is_server": true, - "$lib": "posthog-python", - "$lib_version": "" + "$is_server": true }, "timestamp": "2026-01-02T03:04:05+00:00", "uuid": "00000000-0000-4000-8000-000000000004" } ], - "historical_migration": false, - "sent_at": "2026-01-02T03:04:05+00:00" + "created_at": "2026-01-02T03:04:05+00:00" }, "headers": { + "Authorization": "Bearer phc_snapshot_project", "Content-Type": "application/json", + "PostHog-Attempt": "1", + "PostHog-Request-Id": "", + "PostHog-Request-Timestamp": "2026-01-02T03:04:05+00:00", + "PostHog-Sdk-Info": "posthog-python/", "User-Agent": "posthog-python/" }, "timeout": 15, - "url": "https://example.posthog.test/batch/" + "url": "https://example.posthog.test/i/v1/analytics/events" } diff --git a/posthog/test/snapshots/exception_event.json b/posthog/test/snapshots/exception_event.json index b971c7afe..59c72d814 100644 --- a/posthog/test/snapshots/exception_event.json +++ b/posthog/test/snapshots/exception_event.json @@ -1,10 +1,10 @@ { "body": { - "api_key": "phc_snapshot_project", "batch": [ { "distinct_id": "user-123", "event": "$exception", + "options": {}, "properties": { "$exception_list": [ { @@ -53,7 +53,7 @@ "", "def _exception_request():", " session = mock.MagicMock()", - " session.post.return_value = _successful_response()" + " session.post.side_effect = _capture_ok_response" ], "pre_context": [ "", @@ -113,8 +113,6 @@ "company": "company-456" }, "$is_server": true, - "$lib": "posthog-python", - "$lib_version": "", "$os": "", "$os_distro": "", "$os_version": "", @@ -127,13 +125,17 @@ "uuid": "00000000-0000-4000-8000-000000000005" } ], - "historical_migration": false, - "sent_at": "2026-01-02T03:04:05+00:00" + "created_at": "2026-01-02T03:04:05+00:00" }, "headers": { + "Authorization": "Bearer phc_snapshot_project", "Content-Type": "application/json", + "PostHog-Attempt": "1", + "PostHog-Request-Id": "", + "PostHog-Request-Timestamp": "2026-01-02T03:04:05+00:00", + "PostHog-Sdk-Info": "posthog-python/", "User-Agent": "posthog-python/" }, "timeout": 15, - "url": "https://example.posthog.test/batch/" + "url": "https://example.posthog.test/i/v1/analytics/events" } diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 91b8c7256..6b8f45348 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -7,11 +7,9 @@ from posthog.ai.utils import _capture_ai_event, finalize_ai_content, with_privacy_mode from posthog.capture_compression import CaptureCompression -from posthog.capture_mode import CaptureMode from posthog.client import Client from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE from posthog.capture_v1 import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH -from posthog.request import EVENTS_ENDPOINT from posthog.version import VERSION from posthog.test.capture_helpers import patch_capture_send, sent_batch from posthog.test.test_utils import TEST_API_KEY @@ -116,7 +114,6 @@ def test_analytics_consumers_keep_todays_parameters(self): thread=2, flush_at=7, flush_interval=0.5, - gzip=True, max_retries=4, timeout=9, historical_migration=True, @@ -129,11 +126,9 @@ def test_analytics_consumers_keep_todays_parameters(self): self.assertEqual(consumer.max_msg_size, MAX_MSG_SIZE) self.assertEqual(consumer.flush_at, 7) self.assertEqual(consumer.flush_interval, 0.5) - self.assertTrue(consumer.gzip) self.assertEqual(consumer.retries, 4) self.assertEqual(consumer.timeout, 9) self.assertTrue(consumer.historical_migration) - self.assertEqual(consumer.capture_mode, client.capture_mode) self.assertEqual(consumer.capture_compression, client.capture_compression) client.join() @@ -195,67 +190,59 @@ def test_analytics_lane_rejects_events_over_900kib(self): self.assertTrue(client.queue.empty()) -class TestAiLaneAlwaysV1(unittest.TestCase): - """The AI lane posts capture v1 to the AI endpoint whatever `capture_mode`.""" +class TestAiLaneWireConfig(unittest.TestCase): + """The AI lane posts to the AI endpoint uncompressed, whatever the + analytics `capture_compression`.""" - def test_ai_lane_consumers_use_v1_and_ai_endpoint(self): - client = Client(TEST_API_KEY, send=False, capture_mode="v0", thread=2) + def test_ai_lane_consumers_use_ai_endpoint_without_compression(self): + client = Client(TEST_API_KEY, send=False, capture_compression="gzip", thread=2) client._ai_lane.start() self.assertEqual(len(client._ai_lane.consumers), 2) for consumer in client._ai_lane.consumers: self.assertIs(consumer.queue, client._ai_lane.queue) self.assertEqual(consumer.endpoint, _CAPTURE_AI_V1_PATH) self.assertEqual(consumer.max_msg_size, AI_MAX_MSG_SIZE) - self.assertEqual(consumer.capture_mode, CaptureMode.V1) self.assertEqual(consumer.capture_compression, CaptureCompression.NONE) - def test_async_ai_events_use_v1_with_capture_mode_v0(self): - client = Client( - TEST_API_KEY, - capture_mode="v0", - capture_compression="gzip", - flush_interval=0.05, - ) - with ( - mock.patch("posthog.consumer.batch_post") as mock_post, - patch_capture_send("consumer") as mock_v1, - ): + def test_async_lanes_keep_separate_path_and_compression(self): + client = Client(TEST_API_KEY, capture_compression="gzip", flush_interval=0.05) + with patch_capture_send("consumer") as mock_send: client.capture_ai("$ai_generation", distinct_id="d") client.capture("button_clicked", distinct_id="d") client.flush() + sends = { + call.kwargs["path"]: call.kwargs["compression"] + for call in mock_send.call_args_list + } self.assertEqual( - [call.kwargs["path"] for call in mock_post.call_args_list], - [EVENTS_ENDPOINT], + sends, + { + _CAPTURE_AI_V1_PATH: CaptureCompression.NONE, + _CAPTURE_V1_PATH: CaptureCompression.GZIP, + }, ) - mock_v1.assert_called_once() - self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) self.assertEqual( - mock_v1.call_args.kwargs["compression"], CaptureCompression.NONE + _events_by_path(mock_send)[_CAPTURE_AI_V1_PATH][0]["event"], + "$ai_generation", ) - self.assertEqual([e["event"] for e in sent_batch(mock_v1)], ["$ai_generation"]) client.join() - def test_sync_ai_events_use_v1_with_capture_mode_v0(self): - client = Client( - TEST_API_KEY, - sync_mode=True, - capture_mode="v0", - capture_compression="gzip", - ) - with ( - mock.patch("posthog.client.batch_post") as mock_post, - patch_capture_send("client") as mock_v1, - ): + def test_sync_lanes_keep_separate_path_and_compression(self): + client = Client(TEST_API_KEY, sync_mode=True, capture_compression="gzip") + with patch_capture_send("client") as mock_send: client.capture_ai("$ai_generation", distinct_id="d") client.capture("button_clicked", distinct_id="d") - mock_post.assert_called_once() - self.assertEqual(mock_post.call_args.kwargs["path"], EVENTS_ENDPOINT) - mock_v1.assert_called_once() - self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) self.assertEqual( - mock_v1.call_args.kwargs["compression"], CaptureCompression.NONE + [ + (call.kwargs["path"], call.kwargs["compression"]) + for call in mock_send.call_args_list + ], + [ + (_CAPTURE_AI_V1_PATH, CaptureCompression.NONE), + (_CAPTURE_V1_PATH, CaptureCompression.GZIP), + ], ) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 9c100bc76..34c098d7d 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -10,7 +10,7 @@ import pytest -from posthog import AsyncClient, AsyncPosthog, CaptureCompression, CaptureMode +from posthog import AsyncClient, AsyncPosthog, CaptureCompression from posthog.consumer import MAX_MSG_SIZE from posthog.contexts import ( new_context, @@ -306,7 +306,6 @@ async def test_capture_immediate_uses_capture_v1_without_building_httpx_client() ): client = AsyncPosthog( "test-key", - capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.GZIP, ) event_uuid = await client.capture_immediate("event", distinct_id="user-1") @@ -319,23 +318,6 @@ async def test_capture_immediate_uses_capture_v1_without_building_httpx_client() assert send_v1.await_args.args[2][0]["uuid"] == event_uuid -@pytest.mark.asyncio -async def test_missing_async_extra_does_not_accept_undeliverable_events(): - with mock.patch( - "posthog.async_client._require_httpx", - side_effect=RuntimeError("install posthog[async]"), - ): - client = AsyncPosthog("test-key", capture_mode=CaptureMode.V0) - assert client.capture("event", distinct_id="user-1") is None - assert ( - client.set(distinct_id="user-1", properties={"email": "a@example.com"}) - is None - ) - assert client._pending_queue_items() == 0 - assert client._worker_tasks == [] - await client.shutdown() - - @pytest.mark.asyncio async def test_send_false_accepts_without_starting_workers_or_transport(): with mock.patch("posthog.async_client._build_client") as build_client: @@ -609,16 +591,17 @@ async def test_reuses_and_closes_instance_owned_http_client(): "posthog.async_client._build_client", return_value=http_client ) as build, mock.patch( - "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() - ) as batch_post, + "posthog.async_client._async_flags", + new=mock.AsyncMock(return_value={"flags": {}}), + ) as flags, ): - client = AsyncPosthog("test-key", capture_mode=CaptureMode.V0) - await client.capture_immediate("first", distinct_id="user-1") - await client.capture_immediate("second", distinct_id="user-1") + client = AsyncPosthog("test-key") + await client._get_flags_decision("user-1") + await client._get_flags_decision("user-2") await client.shutdown() build.assert_called_once_with(client.host) - assert [call.kwargs["client"] for call in batch_post.await_args_list] == [ + assert [call.kwargs["client"] for call in flags.await_args_list] == [ http_client, http_client, ] diff --git a/posthog/test/test_async_consumer.py b/posthog/test/test_async_consumer.py index 9bb3f682f..4290a9419 100644 --- a/posthog/test/test_async_consumer.py +++ b/posthog/test/test_async_consumer.py @@ -1,17 +1,12 @@ from __future__ import annotations import asyncio -import json from unittest import mock -import httpx import pytest -from freezegun import freeze_time from posthog._async_consumer import _AsyncConsumer from posthog.capture_compression import CaptureCompression -from posthog.capture_mode import CaptureMode -from posthog.request import APIError def make_consumer(*, retries: int) -> _AsyncConsumer: @@ -23,137 +18,13 @@ def make_consumer(*, retries: int) -> _AsyncConsumer: process_event=mock.AsyncMock(side_effect=lambda event: event), flush_at=100, flush_interval=1, - gzip=False, retries=retries, timeout=3, historical_migration=False, - capture_mode=CaptureMode.V0, capture_compression=CaptureCompression.NONE, - http_client=mock.Mock(), ) -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("failures", "retry_after", "expected_delays"), - [ - (1, None, [1]), - (2, None, [1, 2]), - (2, 5, [5, 5]), - ], -) -async def test_request_retries_transient_failures_until_success( - failures, retry_after, expected_delays -): - error = APIError(503, "temporary", retry_after=retry_after) - consumer = make_consumer(retries=failures) - - with ( - mock.patch( - "posthog._async_consumer.async_batch_post", - new=mock.AsyncMock(side_effect=[error] * failures + [None]), - ) as batch_post, - mock.patch( - "posthog._async_consumer.asyncio.sleep", new=mock.AsyncMock() - ) as sleep, - ): - await consumer.request([{"event": "test"}]) - - assert batch_post.await_count == failures + 1 - assert [call.args[0] for call in sleep.await_args_list] == expected_delays - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("retry_after", "expected_delay"), - [ - ("Tue, 08 Sep 2026 00:00:10 GMT", 10), - ("Tue, 08 Sep 2026 00:01:00 GMT", 30), - ("Mon, 07 Sep 2026 23:59:59 GMT", 1), - ("5", 5), - ("0", 1), - ("-1", 1), - ("invalid", 1), - (None, 1), - ], -) -async def test_request_honors_retry_after_from_http_response( - retry_after, expected_delay -): - headers = {"Retry-After": retry_after} if retry_after is not None else {} - responses = [ - httpx.Response(503, headers=headers, json={"detail": "temporary"}), - httpx.Response(200, json={"ok": True}), - ] - requests = [] - - def handle_request(request): - requests.append(request) - return responses.pop(0) - - consumer = make_consumer(retries=1) - batch = [{"event": "test", "distinct_id": "test-user"}] - async with httpx.AsyncClient( - base_url="https://example.com", transport=httpx.MockTransport(handle_request) - ) as client: - consumer.http_client = client - with ( - freeze_time("2026-09-08 00:00:00", real_asyncio=True), - mock.patch( - "posthog._async_consumer.asyncio.sleep", new=mock.AsyncMock() - ) as sleep, - ): - await consumer.request(batch) - - sleep.assert_awaited_once_with(expected_delay) - assert len(requests) == 2 - assert [json.loads(request.content)["batch"] for request in requests] == [ - batch, - batch, - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("status", [400, 401, 413]) -async def test_request_does_not_retry_terminal_client_errors(status): - consumer = make_consumer(retries=3) - - with ( - mock.patch( - "posthog._async_consumer.async_batch_post", - new=mock.AsyncMock(side_effect=APIError(status, "terminal")), - ) as batch_post, - mock.patch( - "posthog._async_consumer.asyncio.sleep", new=mock.AsyncMock() - ) as sleep, - pytest.raises(APIError), - ): - await consumer.request([{"event": "test"}]) - - batch_post.assert_awaited_once() - sleep.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_request_stops_after_configured_retry_limit(): - consumer = make_consumer(retries=2) - - with ( - mock.patch( - "posthog._async_consumer.async_batch_post", - new=mock.AsyncMock(side_effect=APIError(503, "temporary")), - ) as batch_post, - mock.patch( - "posthog._async_consumer.asyncio.sleep", new=mock.AsyncMock() - ) as sleep, - pytest.raises(APIError), - ): - await consumer.request([{"event": "test"}]) - - assert batch_post.await_count == 3 - assert [call.args[0] for call in sleep.await_args_list] == [1, 2] - - @pytest.mark.asyncio @pytest.mark.parametrize("run_worker", [False, True], ids=["wait", "worker"]) async def test_get_or_flush_cancels_waiters_on_cancellation(run_worker): diff --git a/posthog/test/test_async_request.py b/posthog/test/test_async_request.py index d95da855a..f269db258 100644 --- a/posthog/test/test_async_request.py +++ b/posthog/test/test_async_request.py @@ -1,6 +1,5 @@ from __future__ import annotations -import asyncio import json import logging import subprocess @@ -13,7 +12,6 @@ from posthog._async_request import ( _build_client, _process_response, - async_batch_post, async_flags, async_remote_config, ) @@ -84,148 +82,6 @@ def test_build_client_scopes_requests_to_host_without_following_redirects(): ) -@pytest.mark.asyncio -async def test_async_batch_post_uses_configured_host_and_sanitized_logs(caplog): - caplog.set_level(logging.DEBUG, logger="posthog") - client = FakeAsyncClient() - - await async_batch_post( - "test-secret-key", - "https://example.com", - batch=[{"properties": {"password": "super-secret"}}], - path="/batch/", - client=client, - ) - - assert client.calls[0][1] == ("https://example.com/batch/",) - assert "super-secret" not in caplog.text - assert "test-secret-key" not in caplog.text - assert "https://example.com" not in caplog.text - - -@pytest.mark.asyncio -async def test_async_batch_post_follows_same_origin_temporary_redirect(): - client = FakeAsyncClient( - [ - FakeResponse(307, headers={"Location": "/redirected-batch/"}), - FakeResponse(200), - ] - ) - - await async_batch_post( - "test-key", - "https://example.com", - batch=[{"event": "event"}], - path="/batch/", - client=client, - ) - - assert [call[1] for call in client.calls] == [ - ("https://example.com/batch/",), - ("https://example.com/redirected-batch/",), - ] - - -@pytest.mark.asyncio -@pytest.mark.parametrize("status", [307, 308]) -@pytest.mark.parametrize( - "host", ["https://example.com/proxy", "https://example.com/proxy/"] -) -@pytest.mark.parametrize( - ("location", "redirected_path"), - [ - ("/proxy/redirected-batch/", "/proxy/redirected-batch/"), - ("https://example.com/proxy/redirected-batch/", "/proxy/redirected-batch/"), - ("/redirected-batch/", "/redirected-batch/"), - ("../redirected-batch/", "/proxy/redirected-batch/"), - ("?accepted=1", "/proxy/batch/?accepted=1"), - ], -) -async def test_async_batch_post_redirects_with_host_path_prefix( - status, host, location, redirected_path -): - requests = [] - responses = [ - httpx.Response(status, headers={"Location": location}), - httpx.Response(status, headers={"Location": "?attempt=2"}), - httpx.Response(200), - ] - - def handle_request(request): - requests.append(request) - return responses.pop(0) - - batch = [{"event": "test", "distinct_id": "test-user"}] - async with httpx.AsyncClient( - base_url=host, transport=httpx.MockTransport(handle_request) - ) as client: - await async_batch_post( - "test-key", host, batch=batch, path="/batch/", client=client - ) - - assert [str(request.url) for request in requests] == [ - "https://example.com/proxy/batch/", - f"https://example.com{redirected_path}", - f"https://example.com{redirected_path.split('?')[0]}?attempt=2", - ] - assert all(request.method == "POST" for request in requests) - assert all(request.content == requests[0].content for request in requests) - assert json.loads(requests[0].content)["batch"] == batch - - -@pytest.mark.asyncio -async def test_async_batch_post_rejects_cross_origin_temporary_redirect(): - client = FakeAsyncClient( - FakeResponse( - 307, - headers={"Location": "https://attacker.example/redirected-batch/"}, - ) - ) - - with pytest.raises(APIError): - await async_batch_post( - "test-key", - "https://example.com", - batch=[{"event": "event"}], - path="/batch/", - client=client, - ) - - assert len(client.calls) == 1 - - -@pytest.mark.asyncio -async def test_async_batch_post_serializes_off_event_loop(): - client = FakeAsyncClient() - real_to_thread = asyncio.to_thread - - with mock.patch( - "posthog._async_request.asyncio.to_thread", wraps=real_to_thread - ) as to_thread: - await async_batch_post( - "test-key", - "https://example.com", - batch=[{"event": "event"}], - path="/batch/", - gzip=True, - client=client, - ) - - to_thread.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_async_batch_post_rejects_absolute_request_path(): - with pytest.raises(ValueError, match="relative"): - await async_batch_post( - "test-key", - "https://example.com", - batch=[], - path="https://attacker.example/batch/", - client=FakeAsyncClient(), - ) - - @pytest.mark.asyncio async def test_async_flags_sends_v2_request_payload(): client = FakeAsyncClient(FakeResponse(200, {"flags": {}})) diff --git a/posthog/test/test_capture_compression.py b/posthog/test/test_capture_compression.py index a4ce0154f..cb9376572 100644 --- a/posthog/test/test_capture_compression.py +++ b/posthog/test/test_capture_compression.py @@ -22,15 +22,9 @@ def setUp(self) -> None: self.addCleanup(patcher.stop) os.environ.pop(CAPTURE_COMPRESSION_ENV_VAR, None) - def test_defaults_to_none_with_no_kwarg_env_or_gzip(self) -> None: + def test_defaults_to_none_with_no_kwarg_or_env(self) -> None: self.assertIs(_resolve_capture_compression(None), CaptureCompression.NONE) - def test_gzip_fallback_used_when_nothing_else_set(self) -> None: - self.assertIs( - _resolve_capture_compression(None, gzip_fallback=True), - CaptureCompression.GZIP, - ) - @parameterized.expand( [ ("enum_gzip", CaptureCompression.GZIP, CaptureCompression.GZIP), @@ -48,12 +42,9 @@ def test_gzip_fallback_used_when_nothing_else_set(self) -> None: def test_explicit_kwarg_takes_precedence_and_coerces( self, _name, kwarg, expected ) -> None: - # Env names a different value and gzip_fallback is on, so each row proves - # the explicit kwarg wins over both lower-precedence sources. + # Env names a different value, so each row proves the explicit kwarg wins. with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: "deflate"}): - self.assertIs( - _resolve_capture_compression(kwarg, gzip_fallback=True), expected - ) + self.assertIs(_resolve_capture_compression(kwarg), expected) def test_invalid_kwarg_raises_even_with_valid_env(self) -> None: with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: "gzip"}): @@ -80,28 +71,16 @@ def test_env_var_resolution(self, _name, env_value, expected) -> None: with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: env_value}): self.assertIs(_resolve_capture_compression(None), expected) - def test_env_var_takes_precedence_over_gzip_fallback(self) -> None: - with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: "deflate"}): - self.assertIs( - _resolve_capture_compression(None, gzip_fallback=True), - CaptureCompression.DEFLATE, - ) - @parameterized.expand([("empty", ""), ("whitespace", " ")]) - def test_blank_env_var_falls_through_to_fallback(self, _name, env_value) -> None: + def test_blank_env_var_falls_through_to_none(self, _name, env_value) -> None: with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: env_value}): self.assertIs(_resolve_capture_compression(None), CaptureCompression.NONE) - self.assertIs( - _resolve_capture_compression(None, gzip_fallback=True), - CaptureCompression.GZIP, - ) - def test_unrecognized_env_var_warns_and_uses_fallback(self) -> None: + def test_unrecognized_env_var_warns_and_uses_none(self) -> None: with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: "bogus"}): with capture_message_only_logs() as stream: self.assertIs( - _resolve_capture_compression(None, gzip_fallback=True), - CaptureCompression.GZIP, + _resolve_capture_compression(None), CaptureCompression.NONE ) self.assertIn("bogus", stream.getvalue()) @@ -112,13 +91,12 @@ def test_explicit_zstd_without_package_raises(self, _name, kwarg) -> None: _resolve_capture_compression(kwarg) self.assertIn("posthog[zstd]", str(ctx.exception)) - def test_env_zstd_without_package_warns_and_uses_fallback(self) -> None: + def test_env_zstd_without_package_warns_and_uses_none(self) -> None: with mock.patch("posthog.capture_compression._zstandard", None): with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: "zstd"}): with capture_message_only_logs() as stream: self.assertIs( - _resolve_capture_compression(None, gzip_fallback=True), - CaptureCompression.GZIP, + _resolve_capture_compression(None), CaptureCompression.NONE ) self.assertIn("posthog[zstd]", stream.getvalue()) @@ -134,10 +112,6 @@ def test_client_defaults_to_none(self) -> None: client = Client(TEST_API_KEY, sync_mode=True) self.assertIs(client.capture_compression, CaptureCompression.NONE) - def test_client_gzip_flag_falls_back_to_gzip(self) -> None: - client = Client(TEST_API_KEY, sync_mode=True, gzip=True) - self.assertIs(client.capture_compression, CaptureCompression.GZIP) - @parameterized.expand( [ ("enum_deflate", CaptureCompression.DEFLATE, CaptureCompression.DEFLATE), @@ -145,11 +119,8 @@ def test_client_gzip_flag_falls_back_to_gzip(self) -> None: ("str_none", "none", CaptureCompression.NONE), ] ) - def test_client_kwarg_overrides_gzip_flag(self, _name, kwarg, expected) -> None: - # Even with the legacy gzip flag on, the explicit kwarg wins. - client = Client( - TEST_API_KEY, sync_mode=True, gzip=True, capture_compression=kwarg - ) + def test_client_kwarg_sets_compression(self, _name, kwarg, expected) -> None: + client = Client(TEST_API_KEY, sync_mode=True, capture_compression=kwarg) self.assertIs(client.capture_compression, expected) def test_client_propagates_to_consumers(self) -> None: diff --git a/posthog/test/test_capture_mode.py b/posthog/test/test_capture_mode.py deleted file mode 100644 index e63705d45..000000000 --- a/posthog/test/test_capture_mode.py +++ /dev/null @@ -1,109 +0,0 @@ -import os -import unittest -from unittest import mock - -from parameterized import parameterized - -from posthog.capture_mode import ( - CAPTURE_MODE_ENV_VAR, - CaptureMode, - _resolve_capture_mode, -) -from posthog.client import Client -from posthog.consumer import Consumer -from posthog.test.logging_helpers import capture_message_only_logs -from posthog.test.test_utils import TEST_API_KEY - - -class TestResolveCaptureMode(unittest.TestCase): - def test_defaults_to_v1_with_no_kwarg_and_no_env(self) -> None: - with mock.patch.dict(os.environ, {}, clear=False): - os.environ.pop(CAPTURE_MODE_ENV_VAR, None) - self.assertIs(_resolve_capture_mode(None), CaptureMode.V1) - - @parameterized.expand( - [ - # (name, kwarg, expected, opposite_env): the env always names the - # mode the kwarg must override, so every row proves the kwarg wins. - ("enum_v0", CaptureMode.V0, CaptureMode.V0, "v1"), - ("enum_v1", CaptureMode.V1, CaptureMode.V1, "v0"), - ("str_v0", "v0", CaptureMode.V0, "v1"), - ("str_v1", "v1", CaptureMode.V1, "v0"), - ("str_legacy_alias", "legacy", CaptureMode.V0, "v1"), - ("str_analytics_v1_alias", "analytics_v1", CaptureMode.V1, "v0"), - ("str_upper_and_padded", " V1 ", CaptureMode.V1, "v0"), - ] - ) - def test_explicit_kwarg_takes_precedence_and_coerces( - self, _name, kwarg, expected, opposite_env - ) -> None: - with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: opposite_env}): - self.assertIs(_resolve_capture_mode(kwarg), expected) - - def test_invalid_kwarg_raises_even_with_valid_env(self) -> None: - # The kwarg path is consulted before the env, so an invalid kwarg raises - # rather than silently falling back to a valid env value. - with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: "v1"}): - with self.assertRaises(ValueError): - _resolve_capture_mode("bogus") - - @parameterized.expand( - [ - ("v0", "v0", CaptureMode.V0), - ("legacy", "legacy", CaptureMode.V0), - ("v1", "v1", CaptureMode.V1), - ("analytics_v1", "analytics_v1", CaptureMode.V1), - ("uppercase", "V1", CaptureMode.V1), - ("padded", " v1 ", CaptureMode.V1), - ] - ) - def test_env_var_resolution(self, _name, env_value, expected) -> None: - with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: env_value}): - self.assertIs(_resolve_capture_mode(None), expected) - - @parameterized.expand([("empty", ""), ("whitespace", " ")]) - def test_blank_env_var_defaults_to_v1(self, _name, env_value) -> None: - with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: env_value}): - self.assertIs(_resolve_capture_mode(None), CaptureMode.V1) - - def test_unrecognized_env_var_warns_and_defaults_to_v1(self) -> None: - with mock.patch.dict(os.environ, {CAPTURE_MODE_ENV_VAR: "bogus"}): - with capture_message_only_logs() as stream: - self.assertIs(_resolve_capture_mode(None), CaptureMode.V1) - self.assertIn("bogus", stream.getvalue()) - - @parameterized.expand([("bad_str", "bogus"), ("wrong_type", 1)]) - def test_invalid_explicit_kwarg_raises(self, _name, value) -> None: - with self.assertRaises(ValueError): - _resolve_capture_mode(value) - - -class TestCaptureModePlumbing(unittest.TestCase): - def test_client_resolves_and_stores_default_v1(self) -> None: - with mock.patch.dict(os.environ, {}, clear=False): - os.environ.pop(CAPTURE_MODE_ENV_VAR, None) - client = Client(TEST_API_KEY, sync_mode=True) - self.assertIs(client.capture_mode, CaptureMode.V1) - - @parameterized.expand( - [ - ("enum_v1", CaptureMode.V1, CaptureMode.V1), - ("str_v1", "v1", CaptureMode.V1), - ("enum_v0", CaptureMode.V0, CaptureMode.V0), - ] - ) - def test_client_kwarg_sets_mode(self, _name, kwarg, expected) -> None: - client = Client(TEST_API_KEY, sync_mode=True, capture_mode=kwarg) - self.assertIs(client.capture_mode, expected) - - def test_client_propagates_mode_to_consumers(self) -> None: - # Async (non-sync) client builds Consumer threads; assert each carries - # the resolved mode. - client = Client(TEST_API_KEY, capture_mode=CaptureMode.V0, send=False, thread=2) - self.assertEqual(len(client.consumers), 2) - for consumer in client.consumers: - self.assertIs(consumer.capture_mode, CaptureMode.V0) - - def test_consumer_defaults_to_v1(self) -> None: - consumer = Consumer(None, TEST_API_KEY) - self.assertIs(consumer.capture_mode, CaptureMode.V1) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index faef807c4..10474a25c 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -3503,8 +3503,12 @@ def test_numeric_distinct_id(self): def test_debug(self): Client("bad_key", debug=True) - def test_gzip(self): - client = Client(FAKE_TEST_API_KEY, on_error=self.fail, gzip=True) + def test_gzip_compression(self): + client = Client( + FAKE_TEST_API_KEY, + on_error=self.fail, + capture_compression=CaptureCompression.GZIP, + ) for _ in range(10): client.capture( "event", distinct_id="distinct_id", properties={"trait": "value"} @@ -4660,30 +4664,18 @@ def test_debug_flag_re_raises_exceptions(self, mock_enqueue): class TestClientCaptureRetrySemantics(unittest.TestCase): - @parameterized.expand( - [ - ("v0_sync", "v0", True), - ("v1_sync", "v1", True), - ("v0_async", "v0", False), - ("v1_async", "v1", False), - ] - ) - def test_negative_max_retries_still_attempts_delivery_once( - self, _name, capture_mode, sync_mode - ): + @parameterized.expand([("sync", True), ("async", False)]) + def test_negative_max_retries_still_attempts_delivery_once(self, _name, sync_mode): response = mock.Mock(status_code=200, headers={}, text="") response.json.return_value = {"results": {}} client = None - with ( - mock.patch("posthog.client.batch_post") as sync_v0_post, - mock.patch("posthog.consumer.batch_post") as async_v0_post, - mock.patch("posthog.capture_v1._post_v1", return_value=response) as v1_post, - ): + with mock.patch( + "posthog.capture_v1._post_v1", return_value=response + ) as v1_post: try: client = Client( FAKE_TEST_API_KEY, - capture_mode=capture_mode, sync_mode=sync_mode, max_retries=-1, flush_at=1, @@ -4694,61 +4686,30 @@ def test_negative_max_retries_still_attempts_delivery_once( client.flush() self.assertEqual(client.max_retries, 0) - if capture_mode == "v1": - v1_post.assert_called_once() - sync_v0_post.assert_not_called() - async_v0_post.assert_not_called() - elif sync_mode: - sync_v0_post.assert_called_once() - async_v0_post.assert_not_called() - v1_post.assert_not_called() - else: - async_v0_post.assert_called_once() - sync_v0_post.assert_not_called() - v1_post.assert_not_called() + v1_post.assert_called_once() finally: if client is not None and not sync_mode: client.shutdown() -class TestClientSyncCaptureMode(unittest.TestCase): - """Sync-mode `_enqueue` selects the analytics submitter by `capture_mode`.""" +class TestClientSyncCapture(unittest.TestCase): + """Sync-mode `_enqueue` sends analytics events through the v1 submitter.""" def _client(self, **kwargs): return Client(FAKE_TEST_API_KEY, sync_mode=True, **kwargs) - @parameterized.expand( - [ - ("default", None, True), - ("v1", "v1", True), - ("v0", "v0", False), - ] - ) - def test_capture_mode_selects_sync_submitter(self, _name, capture_mode, expects_v1): - kwargs = {"capture_mode": capture_mode} if capture_mode else {} - with ( - mock.patch("posthog.client.batch_post") as mock_post, - mock.patch("posthog.client._send_v1_batch") as mock_v1, - ): - self._client(**kwargs).capture("evt", distinct_id="d") - if expects_v1: - mock_post.assert_not_called() - mock_v1.assert_called_once() - batch = mock_v1.call_args.args[2] - self.assertEqual(len(batch), 1) - self.assertEqual(batch[0]["event"], "evt") - self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) - else: - mock_v1.assert_not_called() - mock_post.assert_called_once() - - def test_v1_sync_forwards_config_to_submitter(self): - with ( - mock.patch("posthog.client.batch_post"), - mock.patch("posthog.client._send_v1_batch") as mock_v1, - ): + def test_sync_capture_posts_to_analytics_endpoint(self): + with mock.patch("posthog.client._send_v1_batch") as mock_v1: + self._client().capture("evt", distinct_id="d") + mock_v1.assert_called_once() + batch = mock_v1.call_args.args[2] + self.assertEqual(len(batch), 1) + self.assertEqual(batch[0]["event"], "evt") + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) + + def test_sync_forwards_config_to_submitter(self): + with mock.patch("posthog.client._send_v1_batch") as mock_v1: self._client( - capture_mode="v1", capture_compression=CaptureCompression.GZIP, max_retries=4, historical_migration=True, @@ -4758,29 +4719,15 @@ def test_v1_sync_forwards_config_to_submitter(self): self.assertEqual(kwargs["max_retries"], 4) self.assertEqual(kwargs["historical_migration"], True) - def test_v1_sync_gzip_flag_falls_back_to_gzip_compression(self): - # Legacy `gzip=True` with no explicit capture_compression -> GZIP on v1. - with ( - mock.patch("posthog.client.batch_post"), - mock.patch("posthog.client._send_v1_batch") as mock_v1, - ): - self._client(capture_mode="v1", gzip=True).capture("evt", distinct_id="d") - self.assertEqual( - mock_v1.call_args.kwargs["compression"], CaptureCompression.GZIP - ) - - def test_v1_sync_ai_named_event_through_capture_uses_v1(self): + def test_sync_ai_named_event_through_capture_uses_analytics_endpoint(self): # `capture()` never special-cases AI events: an `$ai_*`-named event - # follows `capture_mode` and rides the v1 submitter like any analytics - # event. Only `capture_ai()` reaches the AI lane. - with ( - mock.patch("posthog.client.batch_post") as mock_post, - mock.patch("posthog.client._send_v1_batch") as mock_v1, - ): - client = self._client(capture_mode="v1") + # rides the analytics endpoint like any other event. Only + # `capture_ai()` reaches the AI lane. + with mock.patch("posthog.client._send_v1_batch") as mock_v1: + client = self._client() client.capture("$ai_generation", distinct_id="d") - mock_post.assert_not_called() mock_v1.assert_called_once() batch = mock_v1.call_args.args[2] self.assertEqual(len(batch), 1) self.assertEqual(batch[0]["event"], "$ai_generation") + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index 1199d2da1..92cd9d888 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -2,9 +2,6 @@ import threading import time import unittest -from datetime import datetime, timedelta, timezone -from email.utils import format_datetime -from typing import Any from unittest import mock from parameterized import parameterized @@ -15,10 +12,8 @@ from Queue import Queue from posthog.capture_compression import CaptureCompression -from posthog.capture_mode import CaptureMode from posthog.capture_v1 import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.consumer import MAX_MSG_SIZE, Consumer, _DrainSignal -from posthog.request import AI_EVENTS_ENDPOINT, EVENTS_ENDPOINT, APIError from posthog.test.capture_helpers import patch_capture_send, sent_batch from posthog.test.logging_helpers import capture_message_only_logs from posthog.test.test_utils import TEST_API_KEY @@ -329,88 +324,15 @@ def record_batch(api_key, host, batch, **kwargs): consumer.join(15) self.assertFalse(consumer.is_alive()) - def test_request(self) -> None: - consumer = Consumer(None, TEST_API_KEY, capture_mode=CaptureMode.V0) - batch = [_track_event()] - with mock.patch("posthog.consumer.batch_post") as post: - consumer.request(batch) - post.assert_called_once_with( - TEST_API_KEY, - None, - gzip=False, - timeout=15, - batch=batch, - historical_migration=False, - path="/batch/", - ) - - def _run_retry_test( - self, - exception: Exception, - exception_count: int, - retries: int = 10, - expected_attempts: int = 3, - raises: bool = False, - ) -> None: - call_count = [0] - - def mock_post(*args: Any, **kwargs: Any) -> None: - call_count[0] += 1 - if call_count[0] <= exception_count: - raise exception - - consumer = Consumer( - None, TEST_API_KEY, retries=retries, capture_mode=CaptureMode.V0 - ) - batch = [_track_event()] - with ( - mock.patch("posthog.consumer.batch_post", side_effect=mock_post) as post, - mock.patch("posthog.consumer.time.sleep"), - ): - if raises: - with self.assertRaises(type(exception)) as raised: - consumer.request(batch) - self.assertIs(raised.exception, exception) - else: - consumer.request(batch) - self.assertEqual(post.call_count, expected_attempts) - for call in post.call_args_list: - self.assertEqual(call.kwargs["batch"], batch) - - @parameterized.expand( - [ - ("general_errors", Exception("generic exception"), 2), - ("server_errors", APIError(500, "Internal Server Error"), 2), - ("rate_limit_errors", APIError(429, "Too Many Requests"), 2), - ] - ) - def test_request_retries_on_retriable_errors( - self, _name: str, exception: Exception, exception_count: int - ) -> None: - self._run_retry_test(exception, exception_count) - - def test_request_does_not_retry_client_errors(self) -> None: - self._run_retry_test( - APIError(400, "Client Errors"), 1, expected_attempts=1, raises=True - ) - - def test_request_fails_when_exceptions_exceed_retries(self) -> None: - self._run_retry_test( - APIError(500, "Internal Server Error"), - 4, - retries=3, - expected_attempts=4, - raises=True, - ) - def test_negative_retries_still_attempts_delivery_once(self) -> None: - consumer = Consumer(None, TEST_API_KEY, retries=-1, capture_mode=CaptureMode.V0) + consumer = Consumer(None, TEST_API_KEY, retries=-1) - with mock.patch("posthog.consumer.batch_post") as mock_post: + with patch_capture_send("consumer") as mock_send: consumer.request([_track_event()]) self.assertEqual(consumer.retries, 0) - mock_post.assert_called_once() + mock_send.assert_called_once() + self.assertEqual(mock_send.call_args.kwargs["max_retries"], 0) def test_pause(self) -> None: consumer = Consumer(None, TEST_API_KEY) @@ -624,101 +546,6 @@ def test_max_batch_size(self) -> None: consumer.join(5) self.assertFalse(consumer.is_alive()) - def test_request_sleeps_with_retry_after(self) -> None: - error = APIError(429, "Too Many Requests", retry_after=5.0) - call_count = [0] - - def mock_post(*args: Any, **kwargs: Any) -> None: - call_count[0] += 1 - if call_count[0] <= 1: - raise error - - consumer = Consumer(None, TEST_API_KEY, retries=3, capture_mode=CaptureMode.V0) - with ( - mock.patch("posthog.consumer.batch_post", side_effect=mock_post), - mock.patch("posthog.consumer.time.sleep") as mock_sleep, - ): - consumer.request([_track_event()]) - mock_sleep.assert_called_once_with(5.0) - - def test_request_uses_exponential_backoff_without_retry_after(self) -> None: - error = APIError(503, "Service Unavailable") - call_count = [0] - - def mock_post(*args: Any, **kwargs: Any) -> None: - call_count[0] += 1 - if call_count[0] <= 3: - raise error - - consumer = Consumer(None, TEST_API_KEY, retries=3, capture_mode=CaptureMode.V0) - with ( - mock.patch("posthog.consumer.batch_post", side_effect=mock_post), - mock.patch("posthog.consumer.time.sleep") as mock_sleep, - ): - consumer.request([_track_event()]) - self.assertEqual( - mock_sleep.call_args_list, - [ - mock.call(1), # 2^0 - mock.call(2), # 2^1 - mock.call(4), # 2^2 - ], - ) - - @parameterized.expand( - [ - ("huge_numeric", "1000000000", [30, 30]), - ("small_numeric", "0.25", [1, 2]), - ("huge_date", "Fri, 01 Jan 2100 00:00:00 GMT", [30, 30]), - ("small_date", None, [1, 2]), - ] - ) - def test_request_bounds_retry_after_without_reducing_attempts( - self, _name: str, retry_after_header: str | None, expected_sleeps: list[int] - ) -> None: - if retry_after_header is None: - retry_after_header = format_datetime( - datetime.now(timezone.utc) + timedelta(seconds=1), usegmt=True - ) - - retry_response = mock.Mock( - status_code=503, - headers={"Retry-After": retry_after_header}, - text="Service Unavailable", - ) - retry_response.json.return_value = {"detail": "Service Unavailable"} - success_response = mock.Mock(status_code=200) - session = mock.Mock() - session.post.side_effect = [retry_response, retry_response, success_response] - - consumer = Consumer(None, TEST_API_KEY, retries=2, capture_mode=CaptureMode.V0) - with ( - mock.patch("posthog.request._get_session", return_value=session), - mock.patch("posthog.consumer.time.sleep") as mock_sleep, - ): - consumer.request([_track_event()]) - - self.assertEqual(session.post.call_count, 3) - self.assertEqual( - [call.args[0] for call in mock_sleep.call_args_list], expected_sleeps - ) - - def test_request_retries_on_408(self) -> None: - call_count = [0] - - def mock_post(*args: Any, **kwargs: Any) -> None: - call_count[0] += 1 - if call_count[0] <= 1: - raise APIError(408, "Request Timeout") - - consumer = Consumer(None, TEST_API_KEY, retries=3, capture_mode=CaptureMode.V0) - with ( - mock.patch("posthog.consumer.batch_post", side_effect=mock_post), - mock.patch("posthog.consumer.time.sleep"), - ): - consumer.request([_track_event()]) - self.assertEqual(call_count[0], 2) - @parameterized.expand( [ ("on_error_succeeds", False), @@ -755,51 +582,28 @@ def _ai_event(event_name: str = "$ai_generation") -> dict[str, str]: return {"type": "track", "event": event_name, "distinct_id": "distinct_id"} -class TestConsumerCaptureModeRouting(unittest.TestCase): - """`capture_mode` selects the submitter; both post to the consumer's `endpoint`.""" +class TestConsumerSubmitterRouting(unittest.TestCase): + """Every consumer sends through the capture v1 submitter to its `endpoint`.""" - @parameterized.expand( - [ - ("default", None, True), - ("v0", CaptureMode.V0, False), - ("v1", CaptureMode.V1, True), - ] - ) - def test_capture_mode_selects_analytics_submitter( - self, _name, mode, expects_v1 - ) -> None: - kwargs = {"capture_mode": mode} if mode else {} - consumer = Consumer(None, TEST_API_KEY, **kwargs) + def test_default_posts_to_analytics_endpoint(self) -> None: + consumer = Consumer(None, TEST_API_KEY) batch = [_track_event()] - with ( - mock.patch("posthog.consumer.batch_post") as mock_post, - mock.patch("posthog.consumer._send_v1_batch") as mock_v1, - ): + with patch_capture_send("consumer") as mock_v1: consumer.request(batch) - if expects_v1: - mock_post.assert_not_called() - mock_v1.assert_called_once() - self.assertEqual(mock_v1.call_args.args[2], batch) - self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) - else: - mock_v1.assert_not_called() - mock_post.assert_called_once() - self.assertEqual(mock_post.call_args.kwargs["path"], EVENTS_ENDPOINT) - - def test_v1_forwards_consumer_config_to_submitter(self) -> None: + mock_v1.assert_called_once() + self.assertEqual(sent_batch(mock_v1), batch) + self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_V1_PATH) + + def test_forwards_consumer_config_to_submitter(self) -> None: consumer = Consumer( None, TEST_API_KEY, - capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.DEFLATE, timeout=7, retries=4, historical_migration=True, ) - with ( - mock.patch("posthog.consumer.batch_post"), - mock.patch("posthog.consumer._send_v1_batch") as mock_v1, - ): + with patch_capture_send("consumer") as mock_v1: consumer.request([_track_event()]) kwargs = mock_v1.call_args.kwargs self.assertEqual(kwargs["compression"], CaptureCompression.DEFLATE) @@ -807,7 +611,7 @@ def test_v1_forwards_consumer_config_to_submitter(self) -> None: self.assertEqual(kwargs["max_retries"], 4) self.assertEqual(kwargs["historical_migration"], True) - def test_v1_posts_to_configured_endpoint(self) -> None: + def test_posts_to_configured_endpoint(self) -> None: consumer = Consumer(None, TEST_API_KEY, endpoint=_CAPTURE_AI_V1_PATH) batch = [_ai_event()] with patch_capture_send("consumer") as mock_v1: @@ -815,31 +619,3 @@ def test_v1_posts_to_configured_endpoint(self) -> None: mock_v1.assert_called_once() self.assertEqual(mock_v1.call_args.kwargs["path"], _CAPTURE_AI_V1_PATH) self.assertEqual(sent_batch(mock_v1), batch) - - def test_v0_posts_to_configured_endpoint(self) -> None: - consumer = Consumer( - None, - TEST_API_KEY, - endpoint=AI_EVENTS_ENDPOINT, - capture_mode=CaptureMode.V0, - ) - batch = [_ai_event()] - with mock.patch("posthog.consumer.batch_post") as mock_post: - consumer.request(batch) - mock_post.assert_called_once() - self.assertEqual(mock_post.call_args.kwargs["path"], AI_EVENTS_ENDPOINT) - self.assertEqual(mock_post.call_args.kwargs["batch"], batch) - - def test_v1_routes_whole_batch_through_v1_submitter(self) -> None: - # A consumer doesn't know AI events exist: with `capture_mode` v1, the - # whole batch (including `$ai_*`-named events) rides the v1 submitter. - consumer = Consumer(None, TEST_API_KEY, capture_mode=CaptureMode.V1) - batch = [_ai_event(), _track_event()] - with ( - mock.patch("posthog.consumer.batch_post") as mock_post, - mock.patch("posthog.consumer._send_v1_batch") as mock_v1, - ): - consumer.request(batch) - mock_v1.assert_called_once() - self.assertEqual(mock_v1.call_args.args[2], batch) - mock_post.assert_not_called() diff --git a/posthog/test/test_flag_definition_cache.py b/posthog/test/test_flag_definition_cache.py index 2d05997d4..d221245d9 100644 --- a/posthog/test/test_flag_definition_cache.py +++ b/posthog/test/test_flag_definition_cache.py @@ -17,6 +17,7 @@ ) from posthog.request import GetResponse from posthog.test.test_utils import FAKE_TEST_API_KEY +from posthog.test.capture_helpers import patch_capture_send class MockCacheProvider: @@ -94,8 +95,8 @@ class TestFlagDefinitionCacheProvider(unittest.TestCase): @classmethod def setUpClass(cls): # Prevent real HTTP requests - cls.client_post_patcher = mock.patch("posthog.client.batch_post") - cls.consumer_post_patcher = mock.patch("posthog.consumer.batch_post") + cls.client_post_patcher = patch_capture_send("client") + cls.consumer_post_patcher = patch_capture_send("consumer") cls.client_post_patcher.start() cls.consumer_post_patcher.start() diff --git a/posthog/test/test_request.py b/posthog/test/test_request.py index d4ad7c9da..c16e08f47 100644 --- a/posthog/test/test_request.py +++ b/posthog/test/test_request.py @@ -1,7 +1,5 @@ -import gzip import json import unittest -import zlib from datetime import date, datetime, timedelta from unittest import mock @@ -18,7 +16,6 @@ KEEP_ALIVE_SOCKET_OPTIONS, QuotaLimitError, _mask_tokens_in_url, - batch_post, determine_server_host, disable_connection_reuse, enable_keep_alive, @@ -132,47 +129,6 @@ def test_message_only_debug_logs_include_posthog_prefix(): class TestRequests(unittest.TestCase): - def test_valid_request(self): - response = requests.Response() - response.status_code = 200 - session = mock.Mock() - session.post.return_value = response - batch = [ - {"distinct_id": "distinct_id", "event": "python event", "type": "track"} - ] - - res = batch_post(TEST_API_KEY, batch=batch, session=session) - - self.assertIs(res, response) - session.post.assert_called_once() - self.assertTrue(session.post.call_args.args[0].endswith("/batch/")) - body = json.loads(session.post.call_args.kwargs["data"]) - self.assertEqual(body["batch"], batch) - self.assertEqual(body["api_key"], TEST_API_KEY) - - def test_invalid_request_error(self): - response = requests.Response() - response.status_code = 400 - response._content = b'{"detail": "Invalid batch"}' - session = mock.Mock() - session.post.return_value = response - - with self.assertRaises(APIError) as raised: - batch_post("testsecret", batch=[], session=session) - - self.assertEqual(raised.exception.status, 400) - self.assertEqual(raised.exception.message, "Invalid batch") - session.post.assert_called_once() - - def test_invalid_host(self): - self.assertRaises( - requests.exceptions.MissingSchema, - batch_post, - "testsecret", - "t.posthog.com/", - batch=[], - ) - def test_post_without_path_preserves_type_error(self): mock_session = mock.MagicMock() @@ -185,7 +141,7 @@ def test_post_without_path_preserves_type_error(self): mock_session.post.assert_not_called() - def test_post_sends_string_payload_without_gzip(self): + def test_post_sends_string_payload(self): mock_response = requests.Response() mock_response.status_code = 200 mock_session = mock.MagicMock() @@ -205,54 +161,6 @@ def test_post_sends_string_payload_without_gzip(self): self.assertEqual(url, "https://test.posthog.com/batch/") self.assertIsInstance(data, str) - def test_post_sends_bytes_payload_with_gzip(self): - mock_response = requests.Response() - mock_response.status_code = 200 - mock_session = mock.MagicMock() - mock_session.post.return_value = mock_response - - request_module.post( - TEST_API_KEY, - host="https://test.posthog.com", - path="/batch/", - gzip=True, - session=mock_session, - batch=[], - ) - - data = mock_session.post.call_args.kwargs["data"] - headers = mock_session.post.call_args.kwargs["headers"] - self.assertIsInstance(data, bytes) - self.assertEqual(headers["Content-Encoding"], "gzip") - body = json.loads(gzip.decompress(data)) - self.assertEqual(body["batch"], []) - self.assertEqual(body["api_key"], TEST_API_KEY) - - def test_post_falls_back_to_uncompressed_payload_when_gzip_fails(self): - for compression_error in [OSError("boom"), zlib.error("boom")]: - with self.subTest(compression_error=type(compression_error)): - mock_response = requests.Response() - mock_response.status_code = 200 - mock_session = mock.MagicMock() - mock_session.post.return_value = mock_response - - with mock.patch.object( - request_module, "GzipFile", side_effect=compression_error - ): - request_module.post( - TEST_API_KEY, - host="https://test.posthog.com", - path="/batch/", - gzip=True, - session=mock_session, - batch=[], - ) - - data = mock_session.post.call_args.kwargs["data"] - headers = mock_session.post.call_args.kwargs["headers"] - self.assertIsInstance(data, str) - self.assertNotIn("Content-Encoding", headers) - def test_datetime_serialization(self): data = {"created": datetime(2012, 3, 4, 5, 6, 7, 891011)} result = json.dumps(data, cls=DatetimeSerializer) @@ -265,30 +173,6 @@ def test_date_serialization(self): expected = '{"created": "%s"}' % today.isoformat() self.assertEqual(result, expected) - def test_should_not_timeout(self): - response = requests.Response() - response.status_code = 200 - session = mock.Mock() - session.post.return_value = response - - self.assertIs( - batch_post(TEST_API_KEY, batch=[], timeout=7, session=session), response - ) - session.post.assert_called_once() - self.assertEqual(session.post.call_args.kwargs["timeout"], 7) - - def test_should_timeout(self): - error = requests.ReadTimeout("response deadline exceeded") - session = mock.Mock() - session.post.side_effect = error - - with self.assertRaises(requests.ReadTimeout) as raised: - batch_post("key", batch=[], timeout=1, session=session) - - self.assertIs(raised.exception, error) - session.post.assert_called_once() - self.assertEqual(session.post.call_args.kwargs["timeout"], 1) - def test_quota_limited_flags_response(self): mock_response = requests.Response() mock_response.status_code = 200 diff --git a/posthog/test/test_server_payload_snapshots.py b/posthog/test/test_server_payload_snapshots.py index 2874e0761..80e03415d 100644 --- a/posthog/test/test_server_payload_snapshots.py +++ b/posthog/test/test_server_payload_snapshots.py @@ -29,6 +29,16 @@ def _successful_response() -> requests.Response: return response +def _capture_ok_response(url, data=None, **kwargs) -> requests.Response: + response = requests.Response() + response.status_code = 200 + uuids = [event["uuid"] for event in json.loads(data)["batch"]] + response._content = json.dumps( + {"results": {uuid: {"result": "ok"} for uuid in uuids}} + ).encode() + return response + + def _has_test_file_suffix(value: str) -> bool: return value.replace("\\", "/").endswith(_TEST_FILE_SUFFIX) @@ -43,8 +53,10 @@ def _normalize_snapshot_value(value): for key, item in value.items(): if key == "$lib_version": normalized[key] = "" + elif key == "PostHog-Request-Id": + normalized[key] = "" elif ( - key == "User-Agent" + key in ("User-Agent", "PostHog-Sdk-Info") and isinstance(item, str) and _USER_AGENT_PATTERN.fullmatch(item) ): @@ -99,19 +111,18 @@ def _assert_json_snapshot(name: str, value) -> None: assert actual == expected -def _legacy_event_family_request(): +def _event_family_request(): session = mock.MagicMock() - session.post.return_value = _successful_response() + session.post.side_effect = _capture_ok_response with ( freeze_time(_FIXED_TIME), - mock.patch("posthog.request._get_session", return_value=session), + mock.patch("posthog.capture_v1._get_session", return_value=session), mock.patch("posthog.client.system_context", return_value=_RUNTIME_CONTEXT), ): client = Client( "phc_snapshot_project", host="https://example.posthog.test", - capture_mode="v0", flush_at=100, flush_interval=100, ) @@ -167,18 +178,17 @@ def _raise_snapshot_exception() -> None: def _exception_request(): session = mock.MagicMock() - session.post.return_value = _successful_response() + session.post.side_effect = _capture_ok_response with ( freeze_time(_FIXED_TIME), - mock.patch("posthog.request._get_session", return_value=session), + mock.patch("posthog.capture_v1._get_session", return_value=session), mock.patch("posthog.client.system_context", return_value=_RUNTIME_CONTEXT), mock.patch("posthog.client._get_current_otel_span_properties", return_value={}), ): client = Client( "phc_snapshot_project", host="https://example.posthog.test", - capture_mode="v0", sync_mode=True, project_root=str(Path(__file__).parents[2]), capture_exception_code_variables=False, @@ -257,8 +267,8 @@ def test_does_not_normalize_unexpected_user_agent(user_agent): } -def test_legacy_capture_identify_alias_and_group_identify_request_snapshot(): - _assert_json_snapshot("legacy_event_family", _legacy_event_family_request()) +def test_capture_identify_alias_and_group_identify_request_snapshot(): + _assert_json_snapshot("event_family", _event_family_request()) def test_complete_exception_request_snapshot(): diff --git a/pyproject.toml b/pyproject.toml index 17052d31f..7ad70ea0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,6 @@ classifiers = [ ] dependencies = [ "requests>=2.7,<3.0", - "backoff>=1.10.0", "distro>=1.5.0", "typing-extensions>=4.2.0", ] diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 319451d6a..f412a2113 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -7,7 +7,6 @@ alias posthog.AsyncClient -> posthog.async_client.AsyncClient alias posthog.AsyncPosthog -> posthog.async_client.AsyncPosthog alias posthog.BeforeSendCallback -> posthog.types.BeforeSendCallback alias posthog.CaptureCompression -> posthog.capture_compression.CaptureCompression -alias posthog.CaptureMode -> posthog.capture_mode.CaptureMode alias posthog.Client -> posthog.client.Client alias posthog.DEFAULT_CODE_VARIABLES_DETECT_SECRETS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_DETECT_SECRETS alias posthog.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS @@ -235,13 +234,11 @@ alias posthog.args.SendFeatureFlagsOptions -> posthog.types.SendFeatureFlagsOpti alias posthog.client.AI_MAX_MSG_SIZE -> posthog.consumer.AI_MAX_MSG_SIZE alias posthog.client.APIError -> posthog.request.APIError alias posthog.client.CaptureCompression -> posthog.capture_compression.CaptureCompression -alias posthog.client.CaptureMode -> posthog.capture_mode.CaptureMode alias posthog.client.Consumer -> posthog.consumer.Consumer alias posthog.client.DEFAULT_CODE_VARIABLES_DETECT_SECRETS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_DETECT_SECRETS alias posthog.client.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS alias posthog.client.DEFAULT_CODE_VARIABLES_MASK_PATTERNS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_MASK_PATTERNS alias posthog.client.DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS -alias posthog.client.EVENTS_ENDPOINT -> posthog.request.EVENTS_ENDPOINT alias posthog.client.ExceptionArg -> posthog.args.ExceptionArg alias posthog.client.ExceptionCapture -> posthog.exception_capture.ExceptionCapture alias posthog.client.FeatureFlag -> posthog.types.FeatureFlag @@ -272,7 +269,6 @@ alias posthog.client.SendFeatureFlagsOptions -> posthog.types.SendFeatureFlagsOp alias posthog.client.SizeLimitedDict -> posthog.utils.SizeLimitedDict alias posthog.client.Span -> posthog.tracing.span.Span alias posthog.client.VERSION -> posthog.version.VERSION -alias posthog.client.batch_post -> posthog.request.batch_post alias posthog.client.clean -> posthog.utils.clean alias posthog.client.determine_server_host -> posthog.request.determine_server_host alias posthog.client.exc_info_from_error -> posthog.exception_utils.exc_info_from_error @@ -303,12 +299,8 @@ alias posthog.client.to_flags_and_payloads -> posthog.types.to_flags_and_payload alias posthog.client.to_payloads -> posthog.types.to_payloads alias posthog.client.to_values -> posthog.types.to_values alias posthog.client.try_attach_code_variables_to_frames -> posthog.exception_utils.try_attach_code_variables_to_frames -alias posthog.consumer.APIError -> posthog.request.APIError alias posthog.consumer.CaptureCompression -> posthog.capture_compression.CaptureCompression -alias posthog.consumer.CaptureMode -> posthog.capture_mode.CaptureMode alias posthog.consumer.DatetimeSerializer -> posthog.request.DatetimeSerializer -alias posthog.consumer.EVENTS_ENDPOINT -> posthog.request.EVENTS_ENDPOINT -alias posthog.consumer.batch_post -> posthog.request.batch_post alias posthog.contexts.Client -> posthog.client.Client alias posthog.disable_connection_reuse -> posthog.request.disable_connection_reuse alias posthog.enable_keep_alive -> posthog.request.enable_keep_alive @@ -659,9 +651,8 @@ attribute posthog.args.OptionalSetArgs.timestamp: NotRequired[Optional[Union[dat attribute posthog.args.OptionalSetArgs.uuid: NotRequired[Optional[Union[str, UUID]]] attribute posthog.async_client.AsyncClient.api_key = (project_api_key or '').strip() attribute posthog.async_client.AsyncClient.before_send = before_send -attribute posthog.async_client.AsyncClient.capture_compression = _resolve_capture_compression(capture_compression, gzip_fallback=gzip) +attribute posthog.async_client.AsyncClient.capture_compression = _resolve_capture_compression(capture_compression) attribute posthog.async_client.AsyncClient.capture_exception_code_variables = capture_exception_code_variables -attribute posthog.async_client.AsyncClient.capture_mode = _resolve_capture_mode(capture_mode) attribute posthog.async_client.AsyncClient.capture_trace_context = capture_trace_context attribute posthog.async_client.AsyncClient.code_variables_detect_secrets = code_variables_detect_secrets if code_variables_detect_secrets is not None else DEFAULT_CODE_VARIABLES_DETECT_SECRETS attribute posthog.async_client.AsyncClient.code_variables_ignore_patterns = code_variables_ignore_patterns if code_variables_ignore_patterns is not None else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS @@ -673,7 +664,6 @@ attribute posthog.async_client.AsyncClient.disabled = disabled or not self.api_k attribute posthog.async_client.AsyncClient.distinct_ids_feature_flags_reported = SizeLimitedDict(_MAX_DICT_SIZE, set) attribute posthog.async_client.AsyncClient.feature_flags_request_max_retries = max(0, feature_flags_request_max_retries) attribute posthog.async_client.AsyncClient.feature_flags_request_timeout_seconds = feature_flags_request_timeout_seconds -attribute posthog.async_client.AsyncClient.gzip = gzip attribute posthog.async_client.AsyncClient.historical_migration = historical_migration attribute posthog.async_client.AsyncClient.host = determine_server_host(host) attribute posthog.async_client.AsyncClient.in_app_modules = in_app_modules @@ -699,18 +689,14 @@ attribute posthog.capture_compression.CaptureCompression.GZIP = 'gzip' attribute posthog.capture_compression.CaptureCompression.NONE = 'none' attribute posthog.capture_compression.CaptureCompression.ZSTD = 'zstd' attribute posthog.capture_exception_code_variables = False -attribute posthog.capture_mode.CAPTURE_MODE_ENV_VAR = 'POSTHOG_CAPTURE_MODE' -attribute posthog.capture_mode.CaptureMode.V0 = 'v0' -attribute posthog.capture_mode.CaptureMode.V1 = 'v1' attribute posthog.capture_trace_context = False attribute posthog.capture_v1.CaptureV1Error.attempts = attempts attribute posthog.capture_v1.CaptureV1Error.drops = drops or [] attribute posthog.capture_v1.CaptureV1Error.request_id = request_id attribute posthog.capture_v1.CaptureV1Error.retry_exhausted = retry_exhausted or [] attribute posthog.client.Client.api_key = (project_api_key or '').strip() -attribute posthog.client.Client.capture_compression = _resolve_capture_compression(capture_compression, gzip_fallback=gzip) +attribute posthog.client.Client.capture_compression = _resolve_capture_compression(capture_compression) attribute posthog.client.Client.capture_exception_code_variables = capture_exception_code_variables -attribute posthog.client.Client.capture_mode = _resolve_capture_mode(capture_mode) attribute posthog.client.Client.capture_trace_context = capture_trace_context attribute posthog.client.Client.code_variables_detect_secrets = code_variables_detect_secrets if code_variables_detect_secrets is not None else DEFAULT_CODE_VARIABLES_DETECT_SECRETS attribute posthog.client.Client.code_variables_ignore_patterns = code_variables_ignore_patterns if code_variables_ignore_patterns is not None else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS @@ -738,7 +724,6 @@ attribute posthog.client.Client.flag_cache = self._initialize_flag_cache(flag_fa attribute posthog.client.Client.flag_definition_version = 0 attribute posthog.client.Client.flag_fallback_cache_url = flag_fallback_cache_url attribute posthog.client.Client.group_type_mapping: Optional[dict[str, str]] = None -attribute posthog.client.Client.gzip = gzip attribute posthog.client.Client.historical_migration = historical_migration attribute posthog.client.Client.host = determine_server_host(host) attribute posthog.client.Client.in_app_modules = in_app_modules @@ -769,12 +754,10 @@ attribute posthog.consumer.AI_MAX_MSG_SIZE = 8 * 1024 * 1024 attribute posthog.consumer.BATCH_SIZE_LIMIT = 5 * 1024 * 1024 attribute posthog.consumer.Consumer.api_key = api_key attribute posthog.consumer.Consumer.capture_compression = capture_compression -attribute posthog.consumer.Consumer.capture_mode = capture_mode attribute posthog.consumer.Consumer.daemon = True attribute posthog.consumer.Consumer.endpoint = endpoint attribute posthog.consumer.Consumer.flush_at = flush_at attribute posthog.consumer.Consumer.flush_interval = flush_interval -attribute posthog.consumer.Consumer.gzip = gzip attribute posthog.consumer.Consumer.historical_migration = historical_migration attribute posthog.consumer.Consumer.host = host attribute posthog.consumer.Consumer.log = logging.getLogger('posthog') @@ -1004,13 +987,11 @@ attribute posthog.privacy_mode = False attribute posthog.project_api_key = None attribute posthog.project_root = None attribute posthog.release_id.RELEASE_ID_ENV_VAR = 'POSTHOG_RELEASE_ID' -attribute posthog.request.AI_EVENTS_ENDPOINT = '/i/v0/ai/batch/' attribute posthog.request.APIError.message = message attribute posthog.request.APIError.retry_after = retry_after attribute posthog.request.APIError.status = status attribute posthog.request.DEFAULT_HOST = US_INGESTION_ENDPOINT attribute posthog.request.EU_INGESTION_ENDPOINT = 'https://eu.i.posthog.com' -attribute posthog.request.EVENTS_ENDPOINT = '/batch/' attribute posthog.request.GetResponse.data: Any attribute posthog.request.GetResponse.etag: Optional[str] = None attribute posthog.request.GetResponse.not_modified: bool = False @@ -1178,14 +1159,13 @@ class posthog.ai.types.TokenUsage class posthog.ai.types.ToolInProgress class posthog.args.OptionalCaptureArgs class posthog.args.OptionalSetArgs -class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] = None, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, gzip: bool = False, max_retries: int = 3, timeout: int = 15, thread: int = 1, disabled: bool = False, disable_geoip: bool = True, is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, capture_exception_code_variables: bool = False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, feature_flags_request_timeout_seconds: int = 3, feature_flags_request_max_retries: int = 1) +class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] = None, *, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, max_retries: int = 3, timeout: int = 15, thread: int = 1, disabled: bool = False, disable_geoip: bool = True, is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, capture_exception_code_variables: bool = False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, feature_flags_request_timeout_seconds: int = 3, feature_flags_request_max_retries: int = 1) class posthog.async_client.AsyncPosthog class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, refill_rate: Number, refill_interval_seconds: Number, on_bucket_rate_limited: Optional[Callable[[Hashable], None]] = None, clock: Callable[[], float] = time.monotonic) class posthog.capture_compression.CaptureCompression -class posthog.capture_mode.CaptureMode class posthog.capture_v1.CaptureV1Error(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) -class posthog.client.Client(project_api_key: str, host=None, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, gzip=False, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_mode: Optional[Union[CaptureMode, str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) -class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, gzip=False, retries=10, timeout=15, historical_migration=False, endpoint=None, max_msg_size=MAX_MSG_SIZE, capture_mode=CaptureMode.V1, capture_compression=CaptureCompression.NONE) +class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) +class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, retries=10, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) class posthog.exception_utils.AnnotatedValue(value, metadata) @@ -1418,14 +1398,13 @@ function posthog.mcp.session_token.encode_session_id(payload: SessionTokenPayloa function posthog.mcp.session_token.read_mcp_session_header(headers: Any) -> Optional[str] function posthog.mcp.tools.get_more_tools_result() -> Dict[str, Any] function posthog.new_context(fresh: bool = False, capture_exceptions: Optional[bool] = None, client: Optional[Client] = None) -function posthog.request.batch_post(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, path: str = EVENTS_ENDPOINT, **kwargs) -> requests.Response function posthog.request.determine_server_host(host: Optional[str]) -> str function posthog.request.disable_connection_reuse() -> None function posthog.request.enable_keep_alive() -> None -function posthog.request.flags(api_key: str, host: Optional[str] = None, gzip: bool = False, timeout: int = 15, max_retries: int = 1, **kwargs) -> Any +function posthog.request.flags(api_key: str, host: Optional[str] = None, timeout: int = 15, max_retries: int = 1, **kwargs) -> Any function posthog.request.get(api_key: str, url: str, host: Optional[str] = None, timeout: Optional[int] = None, etag: Optional[str] = None) -> GetResponse function posthog.request.normalize_host(host: Optional[str]) -> str -function posthog.request.post(api_key: str, host: Optional[str] = None, path: Optional[str] = None, gzip: bool = False, timeout: int = 15, session: Optional[requests.Session] = None, **kwargs) -> requests.Response +function posthog.request.post(api_key: str, host: Optional[str] = None, path: Optional[str] = None, timeout: int = 15, session: Optional[requests.Session] = None, **kwargs) -> requests.Response function posthog.request.remote_config(personal_api_key: str, project_api_key: str, host: Optional[str] = None, key: str = '', timeout: int = 15) -> Any function posthog.request.reset_sessions() -> None function posthog.request.set_socket_options(socket_options: Optional[SocketOptions]) -> None @@ -1762,7 +1741,6 @@ module posthog.args module posthog.async_client module posthog.bucketed_rate_limiter module posthog.capture_compression -module posthog.capture_mode module posthog.capture_v1 module posthog.client module posthog.consumer diff --git a/sdk_compliance_adapter/Dockerfile.v1 b/sdk_compliance_adapter/Dockerfile.v1 deleted file mode 100644 index 6891837a5..000000000 --- a/sdk_compliance_adapter/Dockerfile.v1 +++ /dev/null @@ -1,25 +0,0 @@ -FROM python:3.12-slim - -WORKDIR /app - -# Copy the SDK source code -COPY posthog/ /app/sdk/posthog/ -COPY setup.py pyproject.toml README.md LICENSE /app/sdk/ - -# Install the SDK from source -RUN cd /app/sdk && pip install --no-cache-dir -e . - -# Install adapter dependencies -RUN pip install --no-cache-dir flask python-dateutil - -# Copy adapter code -COPY sdk_compliance_adapter/adapter.py /app/adapter.py - -# Select the capture-v1 protocol; same adapter code, different runtime mode. -ENV CAPTURE_MODE=v1 - -# Expose port 8080 -EXPOSE 8080 - -# Run the adapter -CMD ["python", "/app/adapter.py"] diff --git a/sdk_compliance_adapter/README.md b/sdk_compliance_adapter/README.md index eae25916b..6b17bf5dd 100644 --- a/sdk_compliance_adapter/README.md +++ b/sdk_compliance_adapter/README.md @@ -30,7 +30,7 @@ The adapter implements the standard SDK adapter interface defined in the [test h ### Key Implementation Details -**Request Tracking**: The adapter monkey-patches `batch_post` to track all HTTP requests made by the SDK, including retries. +**Request Tracking**: The adapter monkey-patches the capture v1 `_post_v1` to track all HTTP requests made by the SDK, including retries. **State Management**: Thread-safe state tracking for events captured vs sent, retry attempts, and errors. diff --git a/sdk_compliance_adapter/adapter.py b/sdk_compliance_adapter/adapter.py index 2a68dbf56..b663193cb 100644 --- a/sdk_compliance_adapter/adapter.py +++ b/sdk_compliance_adapter/adapter.py @@ -17,8 +17,7 @@ from posthog.capture_compression import CaptureCompression from posthog.capture_v1 import _CAPTURE_V1_PATH from posthog.capture_v1 import _post_v1 as original_post_v1 -from posthog.request import EVENTS_ENDPOINT, USER_AGENT -from posthog.request import batch_post as original_batch_post +from posthog.request import USER_AGENT from posthog.version import VERSION # Configure logging @@ -29,16 +28,6 @@ app = Flask(__name__) -# Selects which capture protocol this adapter process speaks. Baked at build -# time via the CAPTURE_MODE env var ("v1" => capture-v1, anything else => legacy -# v0), mirroring the v0/v1 Dockerfile split. One process speaks one mode and -# advertises it via /health capabilities. -CAPTURE_MODE = os.environ.get("CAPTURE_MODE", "") - - -def is_v1() -> bool: - return CAPTURE_MODE == "v1" - class RequestInfo: """Information about an HTTP request made by the SDK""" @@ -81,7 +70,6 @@ def __init__(self): self.client: Optional[Client] = None self.remote_client: Client | None = None self.reload_thread: threading.Thread | None = None - self.retry_attempts: Dict[str, int] = {} # Track retry attempts by batch ID def reset(self): """Reset all state""" @@ -107,7 +95,6 @@ def reset(self): self.total_retries = 0 self.last_error = None self.requests_made = [] - self.retry_attempts = {} def increment_captured(self): """Increment total events captured""" @@ -115,37 +102,6 @@ def increment_captured(self): self.total_events_captured += 1 self.pending_events += 1 - def record_request(self, status_code: int, batch: List[Dict], batch_id: str): - """Record an HTTP request made by the SDK""" - with self.lock: - # Determine retry attempt for this batch - retry_attempt = self.retry_attempts.get(batch_id, 0) - - # Extract UUIDs from batch - uuid_list = [event.get("uuid", "") for event in batch] - - request_info = RequestInfo( - timestamp_ms=int(time.time() * 1000), - status_code=status_code, - retry_attempt=retry_attempt, - event_count=len(batch), - uuid_list=uuid_list, - ) - self.requests_made.append(request_info) - - # Update counters - if status_code == 200: - # Success - clear pending events - self.total_events_sent += len(batch) - self.pending_events = max(0, self.pending_events - len(batch)) - # Remove batch from retry tracking - self.retry_attempts.pop(batch_id, None) - else: - # Failure - increment retry count - self.retry_attempts[batch_id] = retry_attempt + 1 - if retry_attempt > 0: - self.total_retries += 1 - def record_request_v1( self, status_code: int, batch: List[Dict], attempt: int, terminal_count: int ): @@ -195,40 +151,6 @@ def get_state(self) -> Dict[str, Any]: state = SDKState() -def create_batch_id(batch: List[Dict]) -> str: - """Create a unique ID for a batch based on UUIDs""" - uuids = sorted([event.get("uuid", "") for event in batch]) - return "-".join(uuids[:3]) # Use first 3 UUIDs as batch ID - - -def patched_batch_post( - api_key: str, - host: Optional[str] = None, - gzip: bool = False, - timeout: int = 15, - path: str = EVENTS_ENDPOINT, - **kwargs, -): - """Patched version of batch_post that tracks requests""" - batch = kwargs.get("batch", []) - batch_id = create_batch_id(batch) - - try: - # Call original batch_post - response = original_batch_post(api_key, host, gzip, timeout, path, **kwargs) - # Record successful request - state.record_request(200, batch, batch_id) - return response - except Exception as e: - # Record failed request - status_code = ( - getattr(e, "status_code", 500) if hasattr(e, "status_code") else 500 - ) - state.record_request(status_code, batch, batch_id) - state.record_error(str(e)) - raise - - def patched_post_v1( api_key: str, host: Optional[str], @@ -244,8 +166,8 @@ def patched_post_v1( ): """Patched version of _post_v1 that records requests for /state assertions. - Mirrors the legacy `patched_batch_post`, but reads the retry attempt from the - call (1-based) and counts only terminal per-event results as sent. + Reads the retry attempt from the call (1-based) and counts only terminal + per-event results as sent. """ batch = batch_body.get("batch", []) try: @@ -287,16 +209,6 @@ def patched_post_v1( return response -# Monkey-patch the batch_post function -import posthog.request # noqa: E402 - -posthog.request.batch_post = patched_batch_post - -# Also patch in consumer module -import posthog.consumer # noqa: E402 - -posthog.consumer.batch_post = patched_batch_post - # Patch the capture-v1 submitter. `_send_v1_batch` resolves `_post_v1` as a module # global at call time, so patching it here covers both the async consumer and the # sync client paths. @@ -310,10 +222,11 @@ def health(): """Health check endpoint""" # No AI capture capability: `capture_ai` posts capture v1 to # /i/v1/ai/events, which this harness version has no suite for. - capabilities = ( - ["capture_v1", "encoding_gzip"] if is_v1() else ["capture_v0", "encoding_gzip"] - ) - capabilities.append("feature_flags_local_evaluation_v1") + capabilities = [ + "capture_v1", + "encoding_gzip", + "feature_flags_local_evaluation_v1", + ] return jsonify( { "sdk_name": "posthog-python", @@ -354,9 +267,6 @@ def init(): # Convert flush_interval from ms to seconds flush_interval = flush_interval_ms / 1000.0 - # One adapter process speaks one capture protocol, selected by CAPTURE_MODE. - capture_mode = "v1" if is_v1() else "v0" - # Explicit reloads exercise the real loader without background polling # racing the harness's per-test definition snapshots. client_options = { @@ -364,12 +274,15 @@ def init(): "host": host, "flush_at": flush_at, "flush_interval": flush_interval, - "gzip": enable_compression, + "capture_compression": ( + CaptureCompression.GZIP + if enable_compression + else CaptureCompression.NONE + ), "max_retries": max_retries, "debug": False, "disable_geoip": disable_geoip, "historical_migration": historical_migration, - "capture_mode": capture_mode, "enable_local_evaluation": False, } personal_api_key = data.get("personal_api_key") @@ -384,8 +297,8 @@ def init(): logger.info( f"Initialized SDK with api_key={api_key[:10]}..., host={host}, " f"flush_at={flush_at}, flush_interval={flush_interval}, " - f"max_retries={max_retries}, gzip={enable_compression}, " - f"capture_mode={capture_mode}, disable_geoip={disable_geoip}, " + f"max_retries={max_retries}, compression={enable_compression}, " + f"disable_geoip={disable_geoip}, " f"historical_migration={historical_migration}" ) @@ -417,9 +330,8 @@ def capture(): # Fold capture-v1 options back into the magic `$`-prefixed properties the # SDK lifts onto the wire `options` object. Renamed keys mirror the SDK's - # sentinel table; unknown keys get a bare `$` prefix. v0 has no wire - # options object, so this only applies in v1 mode. - if options and is_v1(): + # sentinel table; unknown keys get a bare `$` prefix. + if options: properties = dict(properties or {}) option_to_property = { "cookieless_mode": "$cookieless_mode", @@ -473,7 +385,7 @@ def capture_ai(): if not event: return jsonify({"error": "event is required"}), 400 - if options and is_v1(): + if options: properties = dict(properties or {}) option_to_property = { "cookieless_mode": "$cookieless_mode", diff --git a/sdk_compliance_adapter/docker-compose.yml b/sdk_compliance_adapter/docker-compose.yml index c679f5c4f..56ef72680 100644 --- a/sdk_compliance_adapter/docker-compose.yml +++ b/sdk_compliance_adapter/docker-compose.yml @@ -1,7 +1,7 @@ version: "3.8" services: - # PostHog Python SDK adapter (capture v0) + # PostHog Python SDK adapter sdk-adapter: build: context: .. @@ -11,16 +11,6 @@ services: networks: - test-network - # PostHog Python SDK adapter (capture v1) - sdk-adapter-v1: - build: - context: .. - dockerfile: sdk_compliance_adapter/Dockerfile.v1 - ports: - - "8082:8080" - networks: - - test-network - # Test harness test-harness: image: ghcr.io/posthog/sdk-test-harness:1.1.1 diff --git a/sdk_compliance_adapter/test_adapter.py b/sdk_compliance_adapter/test_adapter.py index ac1ddc294..2cd484a5d 100644 --- a/sdk_compliance_adapter/test_adapter.py +++ b/sdk_compliance_adapter/test_adapter.py @@ -20,8 +20,6 @@ def adapter(monkeypatch): # Importing the adapter installs transport instrumentation. Restore it after # every test so collecting these tests alongside SDK tests is safe. for module, name in [ - (posthog.request, "batch_post"), - (posthog.consumer, "batch_post"), (posthog.capture_v1, "_post_v1"), ]: monkeypatch.setattr(module, name, getattr(module, name)) @@ -79,14 +77,11 @@ def initialize(adapter, **overrides): return adapter.app.test_client() -@pytest.mark.parametrize("mode,capability", [("", "capture_v0"), ("v1", "capture_v1")]) -def test_health_opts_into_local_evaluation_without_losing_capture( - adapter, monkeypatch, mode, capability -): - monkeypatch.setattr(adapter, "CAPTURE_MODE", mode) +def test_health_opts_into_local_evaluation_without_losing_capture(adapter): capabilities = adapter.app.test_client().get("/health").json["capabilities"] assert "feature_flags_local_evaluation_v1" in capabilities - assert capability in capabilities + assert "capture_v1" in capabilities + assert "capture_v0" not in capabilities assert "capture_ai_v0" not in capabilities diff --git a/uv.lock b/uv.lock index 60046d858..0a2391359 100644 --- a/uv.lock +++ b/uv.lock @@ -319,15 +319,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" }, ] -[[package]] -name = "backoff" -version = "2.2.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/47/d7/5bbeb12c44d7c4f2fb5b56abce497eb5ed9f34d85701de869acedd602619/backoff-2.2.1.tar.gz", hash = "sha256:03f829f5bb1923180821643f8753b0502c3b682293992485b0eef2807afa5cba", size = 17001, upload-time = "2022-10-05T19:19:32.061Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/df/73/b6e24bd22e6720ca8ee9a85a0c4a2971af8497d8f3193fa05390cbd46e09/backoff-2.2.1-py3-none-any.whl", hash = "sha256:63579f9a0628e06278f7e47b7d7d5b6ce20dc65c5e96a6f3ca99a6adca0396e8", size = 15148, upload-time = "2022-10-05T19:19:30.546Z" }, -] - [[package]] name = "backports-asyncio-runner" version = "1.2.0" @@ -2773,7 +2764,6 @@ name = "posthog" version = "7.64.1" source = { editable = "." } dependencies = [ - { name = "backoff" }, { name = "distro" }, { name = "requests" }, { name = "typing-extensions" }, @@ -2852,7 +2842,6 @@ dev = [ [package.metadata] requires-dist = [ { name = "anthropic", marker = "extra == 'test'", specifier = ">=0.72" }, - { name = "backoff", specifier = ">=1.10.0" }, { name = "claude-agent-sdk", marker = "extra == 'test'" }, { name = "coverage", marker = "extra == 'test'" }, { name = "distro", specifier = ">=1.5.0" }, From 29225eb81a578c3ad2a7d0e491d7e2ca6c20c417 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Sat, 10 Oct 2026 16:34:22 -0700 Subject: [PATCH 04/14] refactor!: split capture_v1 into capture_event and capture_send (#1019) Rename CaptureV1Error to CaptureError and export it from posthog. Event shaping moves to posthog.capture_event and transport to posthog.capture_send. No behavior change. --- CONTRIBUTING.md | 2 +- posthog/__init__.py | 1 + posthog/_async_request.py | 2 +- posthog/ai/prompts.py | 2 +- posthog/capture_event.py | 166 ++++++++++ posthog/{capture_v1.py => capture_send.py} | 188 +---------- posthog/client.py | 2 +- posthog/consumer.py | 2 +- posthog/test/capture_helpers.py | 4 +- posthog/test/test_ai_capture_lane.py | 2 +- posthog/test/test_capture_event.py | 276 +++++++++++++++++ ...est_capture_v1.py => test_capture_send.py} | 292 +----------------- posthog/test/test_client.py | 4 +- posthog/test/test_consumer.py | 2 +- posthog/test/test_server_payload_snapshots.py | 4 +- posthog/tracing/_export.py | 2 +- references/public_api_snapshot.txt | 14 +- sdk_compliance_adapter/adapter.py | 8 +- sdk_compliance_adapter/test_adapter.py | 4 +- 19 files changed, 499 insertions(+), 478 deletions(-) create mode 100644 posthog/capture_event.py rename posthog/{capture_v1.py => capture_send.py} (71%) create mode 100644 posthog/test/test_capture_event.py rename posthog/test/{test_capture_v1.py => test_capture_send.py} (63%) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6625fec0c..bb011dfb1 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -18,7 +18,7 @@ uv sync --extra dev --extra test ## CI-aligned checks -Run the smallest relevant tests first, for example `pytest posthog/test/test_capture_v1.py --timeout=30` for v1 transport changes. Then run these core CI-aligned checks from the repository root in the activated `.venv` populated by the setup commands above: +Run the smallest relevant tests first, for example `pytest posthog/test/test_capture_send.py --timeout=30` for v1 transport changes. Then run these core CI-aligned checks from the repository root in the activated `.venv` populated by the setup commands above: ```bash ruff format --check . diff --git a/posthog/__init__.py b/posthog/__init__.py index 9e28ec5bf..5c3335e14 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -10,6 +10,7 @@ OptionalSetArgs, ) from posthog.capture_compression import CaptureCompression as CaptureCompression +from posthog.capture_send import CaptureError as CaptureError from posthog.client import Client from posthog.tracing.span import Span as Span from posthog.async_client import AsyncClient as AsyncClient diff --git a/posthog/_async_request.py b/posthog/_async_request.py index 15b218f74..9ddc7f7a2 100644 --- a/posthog/_async_request.py +++ b/posthog/_async_request.py @@ -7,7 +7,7 @@ from urllib.parse import quote from .capture_compression import CaptureCompression -from .capture_v1 import _parse_retry_after, _send_v1_batch +from .capture_send import _parse_retry_after, _send_v1_batch from .request import ( APIError, DatetimeSerializer, diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 2cbae9db4..c62841561 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Union, overload -from posthog.capture_v1 import _parse_retry_after +from posthog.capture_send import _parse_retry_after from posthog.request import USER_AGENT, _get_session from posthog.utils import remove_trailing_slash diff --git a/posthog/capture_event.py b/posthog/capture_event.py new file mode 100644 index 000000000..bc50a4628 --- /dev/null +++ b/posthog/capture_event.py @@ -0,0 +1,166 @@ +"""Event shaping for the capture v1 wire protocol. + +Transforms a legacy-shaped queued message into a v1 wire event and assembles +the batch envelope. :mod:`posthog.capture_send` posts the result. + +The v1 contract (see ``rust/capture/src/v1/analytics/types.rs``) differs from +the legacy queued-message shape in a few load-bearing ways that this module +encodes: + +- A typed ``options`` object carries a handful of sentinel properties, renamed + and strictly typed. Wrong JSON types fail deserialization of the *whole + batch*, so values are coerced to native types or omitted entirely. +- ``$set``/``$set_once`` have no top-level form in v1; the server reads them + from ``properties``. The legacy ``set()``/``set_once()`` builders emit them at + the top level, so they are relocated into ``properties`` here. +- ``$lib``/``$lib_version`` are injected server-side from the required + ``PostHog-Sdk-Info`` header and are stripped from v1 properties. +""" + +from collections.abc import Callable +from datetime import datetime, timezone +from typing import Any, Optional + +from posthog.utils import _normalize_timestamp + +# Sentinel properties lifted to top-level string fields on the event. +_TOPLEVEL_SENTINELS: tuple[tuple[str, str], ...] = ( + ("$session_id", "session_id"), + ("$window_id", "window_id"), +) + +# Top-level legacy keys relocated into properties (v1 has no top-level form). +_RELOCATE_TO_PROPERTIES = ("$set", "$set_once") + +# Properties dropped from v1 events (server injects them from PostHog-Sdk-Info). +_STRIP_FROM_PROPERTIES = ("$lib", "$lib_version") + + +def _coerce_bool(value: Any) -> Optional[bool]: + """Coerce a sentinel value to ``bool`` using the backend's truthiness rules. + + Native bool passes through; ``"true"``/``"1"`` and ``"false"``/``"0"`` + (case-insensitive, trimmed) map to the obvious bool; any other numeric value + is nonzero-truthy. Anything else returns ``None`` so the option is omitted + rather than sent with a type the strict v1 schema would reject. + """ + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in ("true", "1"): + return True + if normalized in ("false", "0"): + return False + return None + if isinstance(value, (int, float)): + return value != 0 + return None + + +def _coerce_str(value: Any) -> Optional[str]: + """Accept only ``str`` (the backend's ``product_tour_id`` is ``Option``).""" + return value if isinstance(value, str) else None + + +# Sentinel properties lifted into the typed `options` object: legacy property +# key, the backend's field name, and the coercer enforcing its strict type +# (wrong JSON types fail deserialization of the whole batch, so a value that +# won't coerce is omitted). The coercer is stored directly to keep the dispatch +# type-checked rather than keyed by a stringly-typed name. +_OPTION_SENTINELS: tuple[tuple[str, str, Callable[[Any], Any]], ...] = ( + ("$cookieless_mode", "cookieless_mode", _coerce_bool), + ("$ignore_sent_at", "disable_skew_correction", _coerce_bool), + ("$product_tour_id", "product_tour_id", _coerce_str), + ("$process_person_profile", "process_person_profile", _coerce_bool), +) + + +def _v1_timestamp(timestamp: Any) -> str: + """Return a UTC RFC3339 timestamp string. + + Messages off the queue already carry a UTC ISO-8601 string (``_enqueue`` + normalizes canonical datetimes), so that is passed through. A ``datetime`` + is normalized to UTC and serialized; a missing value defaults to now in UTC. + The v1 server parses strictly with ``DateTime::parse_from_rfc3339`` and + rejects naive timestamps. + """ + if timestamp is None: + return datetime.now(timezone.utc).isoformat() + return _normalize_timestamp(timestamp) + + +def _to_v1_event(msg: dict) -> dict: + """Transform a legacy-shaped queued message into a v1 wire event. + + Pure: the input ``msg`` is not mutated (a fresh ``properties`` dict is + built), so it remains safe to keep the original for retries or callbacks. + """ + properties = dict(msg.get("properties") or {}) + + # Relocate top-level $set/$set_once into properties; v1 has no top-level + # form. On the unusual collision where properties already carries the key, + # the properties value wins. + for key in _RELOCATE_TO_PROPERTIES: + top_val = msg.get(key) + if top_val is None: + continue + existing = properties.get(key) + if isinstance(top_val, dict) and isinstance(existing, dict): + properties[key] = {**top_val, **existing} + elif key not in properties: + properties[key] = top_val + + for key in _STRIP_FROM_PROPERTIES: + properties.pop(key, None) + + options: dict[str, Any] = {} + for prop_key, wire_key, coercer in _OPTION_SENTINELS: + if prop_key not in properties: + continue + # Always removed from properties — these sentinels must never reach v1 + # backend properties — but only emitted as an option when coercible. + coerced = coercer(properties.pop(prop_key)) + if coerced is not None: + options[wire_key] = coerced + + top_level: dict[str, str] = {} + for prop_key, field_name in _TOPLEVEL_SENTINELS: + if prop_key not in properties: + continue + coerced_str = _coerce_str(properties.pop(prop_key)) + if coerced_str is not None: + top_level[field_name] = coerced_str + + event = { + "event": msg["event"], + "uuid": msg["uuid"], + "distinct_id": msg["distinct_id"], + "timestamp": _v1_timestamp(msg.get("timestamp")), + # Always a dict so it serializes as "{}" rather than null when empty. + "options": options, + "properties": properties, + } + event.update(top_level) + return event + + +def _build_v1_batch_body( + events: list[dict], + historical_migration: bool = False, + created_at: Optional[str] = None, +) -> dict: + """Assemble the v1 batch envelope. + + Carries no ``api_key`` (Bearer auth) and no ``sent_at``. + ``historical_migration`` is omitted when False (the server defaults it). + ``created_at`` defaults to now in UTC; :func:`_send_v1_batch` passes a value + hoisted once so it stays stable across retry attempts. + """ + body: dict[str, Any] = { + "created_at": created_at or datetime.now(timezone.utc).isoformat(), + "batch": events, + } + if historical_migration: + body["historical_migration"] = True + return body diff --git a/posthog/capture_v1.py b/posthog/capture_send.py similarity index 71% rename from posthog/capture_v1.py rename to posthog/capture_send.py index 438aaf011..862eae094 100644 --- a/posthog/capture_v1.py +++ b/posthog/capture_send.py @@ -1,23 +1,9 @@ -"""Serialization and transport for the Capture V1 wire protocol. +"""Transport for the capture v1 wire protocol. -This module owns everything specific to the capture v1 endpoints +This module sends batches to the capture v1 endpoints (``POST /i/v1/analytics/events`` and ``POST /i/v1/ai/events``, which share one -wire contract): the *transform* layer (legacy-shaped queued message -> v1 wire -event + batch envelope) and the *transport* layer (a single HTTP attempt, -response parsing, and the partial-retry send loop). - -The v1 contract (see ``rust/capture/src/v1/analytics/types.rs``) differs from -the legacy ``/batch/`` shape in a few load-bearing ways that this module -encodes: - -- A typed ``options`` object carries a handful of sentinel properties, renamed - and strictly typed. Wrong JSON types fail deserialization of the *whole - batch*, so values are coerced to native types or omitted entirely. -- ``$set``/``$set_once`` have no top-level form in v1; the server reads them - from ``properties``. The legacy ``set()``/``set_once()`` builders emit them at - the top level, so they are relocated into ``properties`` here. -- ``$lib``/``$lib_version`` are injected server-side from the required - ``PostHog-Sdk-Info`` header and are stripped from v1 properties. +wire contract): a single HTTP attempt, response parsing, and the partial-retry +send loop. :mod:`posthog.capture_event` builds the events and batch envelope. The response is per-event: a 200 carries a ``results`` map keyed by event uuid, each tagged ``ok``/``warning`` (terminal-success), ``drop`` (terminal-failure), @@ -25,7 +11,7 @@ attempt, holding the ``PostHog-Request-Id`` and batch ``created_at`` stable across attempts while incrementing ``PostHog-Attempt``. ``ok``/``warning``/absent events succeed; ``drop`` and retry-exhaustion are carried on the -:class:`CaptureV1Error` raised on batch-level/terminal failure, so the consumer's +:class:`CaptureError` raised on batch-level/terminal failure, so the consumer's existing ``on_error(exc, batch)`` path surfaces them unchanged (no per-event logging of its own). @@ -37,16 +23,16 @@ import logging import time import zlib -from collections.abc import Callable from dataclasses import dataclass from datetime import datetime, timezone from email.utils import parsedate_to_datetime from gzip import GzipFile from io import BytesIO -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Optional from uuid import uuid4 from posthog.capture_compression import CaptureCompression, _zstandard +from posthog.capture_event import _build_v1_batch_body, _to_v1_event from posthog.request import ( DatetimeSerializer, USER_AGENT, @@ -54,7 +40,7 @@ _get_session, normalize_host, ) -from posthog.utils import _normalize_timestamp, remove_trailing_slash +from posthog.utils import remove_trailing_slash if TYPE_CHECKING: import requests @@ -64,7 +50,7 @@ # Only the error type is public API: it reaches user code through `on_error` # callbacks, so callers may want to catch/inspect it. Everything else is # submitter plumbing. -__all__ = ["CaptureV1Error"] +__all__ = ["CaptureError"] _CAPTURE_V1_PATH = "/i/v1/analytics/events" _CAPTURE_AI_V1_PATH = "/i/v1/ai/events" @@ -94,148 +80,6 @@ # unifies the default with posthog-go/posthog-rs (all 30s). _MAX_BACKOFF_SECONDS = 30 -# Sentinel properties lifted to top-level string fields on the event. -_TOPLEVEL_SENTINELS: tuple[tuple[str, str], ...] = ( - ("$session_id", "session_id"), - ("$window_id", "window_id"), -) - -# Top-level legacy keys relocated into properties (v1 has no top-level form). -_RELOCATE_TO_PROPERTIES = ("$set", "$set_once") - -# Properties dropped from v1 events (server injects them from PostHog-Sdk-Info). -_STRIP_FROM_PROPERTIES = ("$lib", "$lib_version") - - -def _coerce_bool(value: Any) -> Optional[bool]: - """Coerce a sentinel value to ``bool`` using the backend's truthiness rules. - - Native bool passes through; ``"true"``/``"1"`` and ``"false"``/``"0"`` - (case-insensitive, trimmed) map to the obvious bool; any other numeric value - is nonzero-truthy. Anything else returns ``None`` so the option is omitted - rather than sent with a type the strict v1 schema would reject. - """ - if isinstance(value, bool): - return value - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in ("true", "1"): - return True - if normalized in ("false", "0"): - return False - return None - if isinstance(value, (int, float)): - return value != 0 - return None - - -def _coerce_str(value: Any) -> Optional[str]: - """Accept only ``str`` (the backend's ``product_tour_id`` is ``Option``).""" - return value if isinstance(value, str) else None - - -# Sentinel properties lifted into the typed `options` object: legacy property -# key, the backend's field name, and the coercer enforcing its strict type -# (wrong JSON types fail deserialization of the whole batch, so a value that -# won't coerce is omitted). The coercer is stored directly to keep the dispatch -# type-checked rather than keyed by a stringly-typed name. -_OPTION_SENTINELS: tuple[tuple[str, str, Callable[[Any], Any]], ...] = ( - ("$cookieless_mode", "cookieless_mode", _coerce_bool), - ("$ignore_sent_at", "disable_skew_correction", _coerce_bool), - ("$product_tour_id", "product_tour_id", _coerce_str), - ("$process_person_profile", "process_person_profile", _coerce_bool), -) - - -def _v1_timestamp(timestamp: Any) -> str: - """Return a UTC RFC3339 timestamp string. - - Messages off the queue already carry a UTC ISO-8601 string (``_enqueue`` - normalizes canonical datetimes), so that is passed through. A ``datetime`` - is normalized to UTC and serialized; a missing value defaults to now in UTC. - The v1 server parses strictly with ``DateTime::parse_from_rfc3339`` and - rejects naive timestamps. - """ - if timestamp is None: - return datetime.now(timezone.utc).isoformat() - return _normalize_timestamp(timestamp) - - -def _to_v1_event(msg: dict) -> dict: - """Transform a legacy-shaped queued message into a v1 wire event. - - Pure: the input ``msg`` is not mutated (a fresh ``properties`` dict is - built), so it remains safe to keep the original for retries or callbacks. - """ - properties = dict(msg.get("properties") or {}) - - # Relocate top-level $set/$set_once into properties; v1 has no top-level - # form. On the unusual collision where properties already carries the key, - # the properties value wins. - for key in _RELOCATE_TO_PROPERTIES: - top_val = msg.get(key) - if top_val is None: - continue - existing = properties.get(key) - if isinstance(top_val, dict) and isinstance(existing, dict): - properties[key] = {**top_val, **existing} - elif key not in properties: - properties[key] = top_val - - for key in _STRIP_FROM_PROPERTIES: - properties.pop(key, None) - - options: dict[str, Any] = {} - for prop_key, wire_key, coercer in _OPTION_SENTINELS: - if prop_key not in properties: - continue - # Always removed from properties — these sentinels must never reach v1 - # backend properties — but only emitted as an option when coercible. - coerced = coercer(properties.pop(prop_key)) - if coerced is not None: - options[wire_key] = coerced - - top_level: dict[str, str] = {} - for prop_key, field_name in _TOPLEVEL_SENTINELS: - if prop_key not in properties: - continue - coerced_str = _coerce_str(properties.pop(prop_key)) - if coerced_str is not None: - top_level[field_name] = coerced_str - - event = { - "event": msg["event"], - "uuid": msg["uuid"], - "distinct_id": msg["distinct_id"], - "timestamp": _v1_timestamp(msg.get("timestamp")), - # Always a dict so it serializes as "{}" rather than null when empty. - "options": options, - "properties": properties, - } - event.update(top_level) - return event - - -def _build_v1_batch_body( - events: list[dict], - historical_migration: bool = False, - created_at: Optional[str] = None, -) -> dict: - """Assemble the v1 batch envelope. - - Carries no ``api_key`` (Bearer auth) and no ``sent_at``. - ``historical_migration`` is omitted when False (the server defaults it). - ``created_at`` defaults to now in UTC; :func:`_send_v1_batch` passes a value - hoisted once so it stays stable across retry attempts. - """ - body: dict[str, Any] = { - "created_at": created_at or datetime.now(timezone.utc).isoformat(), - "batch": events, - } - if historical_migration: - body["historical_migration"] = True - return body - @dataclass class _V1EventResult: @@ -263,7 +107,7 @@ class _V1ParsedResponse: error_message: str = "" -class CaptureV1Error(APIError): +class CaptureError(APIError): """Batch-level failure of a capture-v1 send. Subclasses :class:`APIError` so the consumer's existing ``on_error`` handling @@ -485,11 +329,11 @@ def _send_v1_batch( tagged ``retry`` after each 2xx. ``ok``/``warning``/absent events succeed. A server-chosen ``drop`` is a terminal per-event rejection. Drops are - accumulated across attempts and surfaced via :class:`CaptureV1Error` even + accumulated across attempts and surfaced via :class:`CaptureError` even when the request itself was a 2xx (a success status is not full delivery) and even when a later attempt clears the outstanding retries — matching posthog-go (per-event failure callback) and posthog-rs (``on_error`` on a - 2xx with undelivered verdicts). Raises :class:`CaptureV1Error` on any drop, + 2xx with undelivered verdicts). Raises :class:`CaptureError` on any drop, batch-level terminal failure, or retry exhaustion — carrying the accumulated ``drops`` and any exhausted uuids — so the caller's ``on_error`` fires unchanged. A transport failure re-raises the underlying exception (drops @@ -543,7 +387,7 @@ def _send_v1_batch( if parsed.is_success: if parsed.malformed: - raise CaptureV1Error( + raise CaptureError( parsed.status_code, "capture v1 returned a success status with an unparseable body", request_id=request_id, @@ -574,7 +418,7 @@ def _send_v1_batch( # surface them via on_error even though the request was a 2xx — # a success status does not mean every event was delivered. if all_drops: - raise CaptureV1Error( + raise CaptureError( parsed.status_code, f"{len(all_drops)} event(s) dropped by the server", request_id=request_id, @@ -583,7 +427,7 @@ def _send_v1_batch( ) return if last_attempt: - raise CaptureV1Error( + raise CaptureError( parsed.status_code, f"{len(retry_uuids)} event(s) still pending retry after {attempt} attempt(s)", request_id=request_id, @@ -598,7 +442,7 @@ def _send_v1_batch( # Non-2xx. Retryable transient statuses back off; everything else # (400/401/402/413/415/429/...) is terminal. Any drops collected from a # prior 2xx attempt ride along so on_error still sees them. - v1_error = CaptureV1Error( + v1_error = CaptureError( parsed.status_code, parsed.error_message, retry_after=parsed.retry_after, diff --git a/posthog/client.py b/posthog/client.py index 083a6eb03..611fda2f1 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -31,7 +31,7 @@ CaptureCompression, _resolve_capture_compression, ) -from posthog.capture_v1 import ( +from posthog.capture_send import ( _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, _send_v1_batch, diff --git a/posthog/consumer.py b/posthog/consumer.py index 0b15431a8..daad58c33 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -6,7 +6,7 @@ from posthog._logging import _configure_posthog_logging from posthog.capture_compression import CaptureCompression -from posthog.capture_v1 import _CAPTURE_V1_PATH, _send_v1_batch +from posthog.capture_send import _CAPTURE_V1_PATH, _send_v1_batch from posthog.request import ( USER_AGENT as _USER_AGENT, DatetimeSerializer, diff --git a/posthog/test/capture_helpers.py b/posthog/test/capture_helpers.py index 542d3e10f..3ead56f5e 100644 --- a/posthog/test/capture_helpers.py +++ b/posthog/test/capture_helpers.py @@ -1,8 +1,8 @@ """Intercept capture uploads at the batch submitter for client-level tests. Patching the submitter (not the HTTP layer) lets tests assert on the event -dicts the SDK built, before the wire encoding in ``capture_v1``. Wire shape is -covered by ``test_capture_v1``. +dicts the SDK built, before the wire encoding in ``capture_event``. Wire shape is +covered by ``test_capture_event``. """ import json diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 6b8f45348..568f00feb 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -9,7 +9,7 @@ from posthog.capture_compression import CaptureCompression from posthog.client import Client from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE -from posthog.capture_v1 import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.version import VERSION from posthog.test.capture_helpers import patch_capture_send, sent_batch from posthog.test.test_utils import TEST_API_KEY diff --git a/posthog/test/test_capture_event.py b/posthog/test/test_capture_event.py new file mode 100644 index 000000000..e3b48195c --- /dev/null +++ b/posthog/test/test_capture_event.py @@ -0,0 +1,276 @@ +import unittest +from datetime import datetime, timedelta, timezone + +from parameterized import parameterized + +from posthog.capture_event import ( + _build_v1_batch_body, + _coerce_bool, + _coerce_str, + _to_v1_event, +) + + +def _legacy_msg(event="my_event", properties=None, **overrides) -> dict: + """Minimal legacy-shaped message as it looks coming off the queue.""" + msg = { + "event": event, + "uuid": "0190000000007000800000000000000a", + "distinct_id": "user-1", + "timestamp": "2026-06-27T12:00:00+00:00", + "type": "capture", + "properties": {"$lib": "posthog-python", "$lib_version": "9.9.9"}, + } + if properties is not None: + msg["properties"] = properties + msg.update(overrides) + return msg + + +class TestCoercion(unittest.TestCase): + @parameterized.expand( + [ + ("bool_true", True, True), + ("bool_false", False, False), + ("str_true", "true", True), + ("str_true_upper", "TRUE", True), + ("str_true_padded", " true ", True), + ("str_one", "1", True), + ("str_false", "false", False), + ("str_zero", "0", False), + ("int_nonzero", 5, True), + ("int_zero", 0, False), + ("float_nonzero", 1.5, True), + ("float_zero", 0.0, False), + ("neg_int", -1, True), + ("str_yes_uncoercible", "yes", None), + ("str_empty_uncoercible", "", None), + ("none_uncoercible", None, None), + ("dict_uncoercible", {"a": 1}, None), + ] + ) + def test_coerce_bool(self, _name, value, expected) -> None: + self.assertIs(_coerce_bool(value), expected) + + @parameterized.expand( + [ + ("str", "tour-1", "tour-1"), + ("empty_str", "", ""), + ("int", 123, None), + ("bool", True, None), + ("none", None, None), + ] + ) + def test_coerce_str(self, _name, value, expected) -> None: + self.assertEqual(_coerce_str(value), expected) + + +class TestToV1Event(unittest.TestCase): + def test_required_fields_preserved(self) -> None: + event = _to_v1_event(_legacy_msg(event="signed_up")) + self.assertEqual(event["event"], "signed_up") + self.assertEqual(event["uuid"], "0190000000007000800000000000000a") + self.assertEqual(event["distinct_id"], "user-1") + self.assertEqual(event["timestamp"], "2026-06-27T12:00:00+00:00") + + def test_strips_lib_and_lib_version(self) -> None: + event = _to_v1_event(_legacy_msg()) + self.assertNotIn("$lib", event["properties"]) + self.assertNotIn("$lib_version", event["properties"]) + + def test_options_empty_dict_when_no_sentinels(self) -> None: + event = _to_v1_event(_legacy_msg(properties={"plain": "value"})) + self.assertEqual(event["options"], {}) + self.assertEqual(event["properties"], {"plain": "value"}) + + def test_does_not_leak_non_wire_top_level_keys(self) -> None: + event = _to_v1_event(_legacy_msg()) + # `type` is legacy-only; the v1 event carries only documented fields. + self.assertEqual( + set(event), + {"event", "uuid", "distinct_id", "timestamp", "options", "properties"}, + ) + + def test_does_not_mutate_input(self) -> None: + msg = _legacy_msg( + properties={"$cookieless_mode": True, "$session_id": "s-1"}, + **{"$set": {"name": "Max"}}, + ) + original_properties = dict(msg["properties"]) + _to_v1_event(msg) + self.assertEqual(msg["properties"], original_properties) + self.assertIn("$set", msg) # top-level $set untouched on the original + + @parameterized.expand( + [ + ("cookieless_mode", "$cookieless_mode", "cookieless_mode", True, True), + ( + "ignore_sent_at_rename", + "$ignore_sent_at", + "disable_skew_correction", + "true", + True, + ), + ( + "process_person_profile", + "$process_person_profile", + "process_person_profile", + "false", + False, + ), + ( + "product_tour_id", + "$product_tour_id", + "product_tour_id", + "tour-7", + "tour-7", + ), + ] + ) + def test_option_sentinels_lifted_renamed_and_coerced( + self, _name, prop_key, wire_key, raw, expected + ) -> None: + event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) + self.assertEqual(event["options"], {wire_key: expected}) + self.assertNotIn(prop_key, event["properties"]) + + @parameterized.expand( + [ + ("bad_bool", "$cookieless_mode", "maybe"), + ("bad_tour_id_int", "$product_tour_id", 123), + ] + ) + def test_option_sentinel_removed_but_omitted_on_bad_coercion( + self, _name, prop_key, raw + ) -> None: + event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) + # Removed from properties (sentinels must never reach v1 props) but not + # emitted as an option, so a wrong type cannot 400 the whole batch. + self.assertNotIn(prop_key, event["properties"]) + self.assertEqual(event["options"], {}) + + @parameterized.expand( + [ + ("session_id", "$session_id", "session_id", "s-123"), + ("window_id", "$window_id", "window_id", "w-456"), + ] + ) + def test_top_level_string_sentinels(self, _name, prop_key, field_name, raw) -> None: + event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) + self.assertEqual(event[field_name], raw) + self.assertNotIn(prop_key, event["properties"]) + + def test_top_level_sentinel_omitted_but_removed_when_not_string(self) -> None: + event = _to_v1_event(_legacy_msg(properties={"$session_id": 42})) + self.assertNotIn("session_id", event) + self.assertNotIn("$session_id", event["properties"]) + + def test_all_sentinels_together(self) -> None: + event = _to_v1_event( + _legacy_msg( + properties={ + "$cookieless_mode": True, + "$ignore_sent_at": "1", + "$product_tour_id": "tour-x", + "$process_person_profile": 0, + "$session_id": "s-1", + "$window_id": "w-1", + "$geoip_disable": True, + "custom": "keep", + } + ) + ) + self.assertEqual( + event["options"], + { + "cookieless_mode": True, + "disable_skew_correction": True, + "product_tour_id": "tour-x", + "process_person_profile": False, + }, + ) + self.assertEqual(event["session_id"], "s-1") + self.assertEqual(event["window_id"], "w-1") + # Non-sentinel props (including $geoip_disable) are left intact. + self.assertEqual( + event["properties"], {"$geoip_disable": True, "custom": "keep"} + ) + + @parameterized.expand([("set", "$set"), ("set_once", "$set_once")]) + def test_top_level_set_relocated_into_properties(self, _name, key) -> None: + msg = _legacy_msg(properties={}, **{key: {"email": "a@b.com"}}) + event = _to_v1_event(msg) + self.assertEqual(event["properties"][key], {"email": "a@b.com"}) + self.assertNotIn(key, event) # not a top-level v1 field + + def test_top_level_set_merges_with_existing_properties_set(self) -> None: + # properties wins on key collision. + msg = _legacy_msg( + properties={"$set": {"a": "from_props", "b": "props_only"}}, + **{"$set": {"a": "from_top", "c": "top_only"}}, + ) + event = _to_v1_event(msg) + self.assertEqual( + event["properties"]["$set"], + {"a": "from_props", "b": "props_only", "c": "top_only"}, + ) + + def test_groups_left_in_properties(self) -> None: + event = _to_v1_event(_legacy_msg(properties={"$groups": {"company": "ph"}})) + self.assertEqual(event["properties"]["$groups"], {"company": "ph"}) + + def test_timestamp_naive_datetime_made_tz_aware(self) -> None: + event = _to_v1_event(_legacy_msg(timestamp=datetime(2026, 6, 27, 12, 0, 0))) + parsed = datetime.fromisoformat(event["timestamp"]) + self.assertIsNotNone(parsed.tzinfo) + + def test_timestamp_aware_datetime_converted_to_exact_utc_instant(self) -> None: + event = _to_v1_event( + _legacy_msg( + timestamp=datetime( + 2026, + 6, + 27, + 17, + 45, + tzinfo=timezone(timedelta(hours=5, minutes=45)), + ) + ) + ) + self.assertEqual(event["timestamp"], "2026-06-27T12:00:00+00:00") + + def test_timestamp_parseable_string_converted_to_exact_utc_instant(self) -> None: + event = _to_v1_event(_legacy_msg(timestamp="2026-06-27T17:45:00+05:45")) + self.assertEqual(event["timestamp"], "2026-06-27T12:00:00+00:00") + + def test_timestamp_none_defaults_to_utc_now(self) -> None: + event = _to_v1_event(_legacy_msg(timestamp=None)) + parsed = datetime.fromisoformat(event["timestamp"]) + self.assertEqual(parsed.tzinfo, timezone.utc) + + +class TestBuildV1BatchBody(unittest.TestCase): + def test_envelope_shape_and_no_legacy_fields(self) -> None: + events = [{"event": "e"}] + body = _build_v1_batch_body(events) + self.assertEqual(body["batch"], events) + self.assertNotIn("api_key", body) + self.assertNotIn("sent_at", body) + + def test_created_at_is_tz_aware_rfc3339(self) -> None: + body = _build_v1_batch_body([]) + parsed = datetime.fromisoformat(body["created_at"]) + self.assertIsNotNone(parsed.tzinfo) + + def test_created_at_passthrough_used_verbatim(self) -> None: + # _send_v1_batch hoists created_at and passes it in so it stays stable + # across retry attempts. + body = _build_v1_batch_body([], created_at="2026-06-27T12:00:00+00:00") + self.assertEqual(body["created_at"], "2026-06-27T12:00:00+00:00") + + def test_historical_migration_omitted_when_false(self) -> None: + self.assertNotIn("historical_migration", _build_v1_batch_body([])) + + def test_historical_migration_present_when_true(self) -> None: + body = _build_v1_batch_body([], historical_migration=True) + self.assertIs(body["historical_migration"], True) diff --git a/posthog/test/test_capture_v1.py b/posthog/test/test_capture_send.py similarity index 63% rename from posthog/test/test_capture_v1.py rename to posthog/test/test_capture_send.py index e11e74c7c..637eee85f 100644 --- a/posthog/test/test_capture_v1.py +++ b/posthog/test/test_capture_send.py @@ -1,7 +1,7 @@ import json import unittest import zlib -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from unittest import mock import zstandard @@ -9,7 +9,8 @@ from parameterized import parameterized from posthog.capture_compression import CaptureCompression -from posthog.capture_v1 import ( +from posthog.capture_event import _build_v1_batch_body, _to_v1_event +from posthog.capture_send import ( _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, _HEADER_ATTEMPT, @@ -17,15 +18,11 @@ _HEADER_REQUEST_TIMESTAMP, _HEADER_SDK_INFO, _MAX_BACKOFF_SECONDS, - CaptureV1Error, - _build_v1_batch_body, + CaptureError, _parse_v1_response, _post_v1, _send_v1_batch, - _to_v1_event, _backoff, - _coerce_bool, - _coerce_str, ) from posthog.request import USER_AGENT @@ -130,271 +127,6 @@ def _results_response(directives, headers=None): return _FakeResponse(200, json_body={"results": results}, headers=headers) -def _legacy_msg(event="my_event", properties=None, **overrides) -> dict: - """Minimal legacy-shaped message as it looks coming off the queue.""" - msg = { - "event": event, - "uuid": "0190000000007000800000000000000a", - "distinct_id": "user-1", - "timestamp": "2026-06-27T12:00:00+00:00", - "type": "capture", - "properties": {"$lib": "posthog-python", "$lib_version": "9.9.9"}, - } - if properties is not None: - msg["properties"] = properties - msg.update(overrides) - return msg - - -class TestCoercion(unittest.TestCase): - @parameterized.expand( - [ - ("bool_true", True, True), - ("bool_false", False, False), - ("str_true", "true", True), - ("str_true_upper", "TRUE", True), - ("str_true_padded", " true ", True), - ("str_one", "1", True), - ("str_false", "false", False), - ("str_zero", "0", False), - ("int_nonzero", 5, True), - ("int_zero", 0, False), - ("float_nonzero", 1.5, True), - ("float_zero", 0.0, False), - ("neg_int", -1, True), - ("str_yes_uncoercible", "yes", None), - ("str_empty_uncoercible", "", None), - ("none_uncoercible", None, None), - ("dict_uncoercible", {"a": 1}, None), - ] - ) - def test_coerce_bool(self, _name, value, expected) -> None: - self.assertIs(_coerce_bool(value), expected) - - @parameterized.expand( - [ - ("str", "tour-1", "tour-1"), - ("empty_str", "", ""), - ("int", 123, None), - ("bool", True, None), - ("none", None, None), - ] - ) - def test_coerce_str(self, _name, value, expected) -> None: - self.assertEqual(_coerce_str(value), expected) - - -class TestToV1Event(unittest.TestCase): - def test_required_fields_preserved(self) -> None: - event = _to_v1_event(_legacy_msg(event="signed_up")) - self.assertEqual(event["event"], "signed_up") - self.assertEqual(event["uuid"], "0190000000007000800000000000000a") - self.assertEqual(event["distinct_id"], "user-1") - self.assertEqual(event["timestamp"], "2026-06-27T12:00:00+00:00") - - def test_strips_lib_and_lib_version(self) -> None: - event = _to_v1_event(_legacy_msg()) - self.assertNotIn("$lib", event["properties"]) - self.assertNotIn("$lib_version", event["properties"]) - - def test_options_empty_dict_when_no_sentinels(self) -> None: - event = _to_v1_event(_legacy_msg(properties={"plain": "value"})) - self.assertEqual(event["options"], {}) - self.assertEqual(event["properties"], {"plain": "value"}) - - def test_does_not_leak_non_wire_top_level_keys(self) -> None: - event = _to_v1_event(_legacy_msg()) - # `type` is legacy-only; the v1 event carries only documented fields. - self.assertEqual( - set(event), - {"event", "uuid", "distinct_id", "timestamp", "options", "properties"}, - ) - - def test_does_not_mutate_input(self) -> None: - msg = _legacy_msg( - properties={"$cookieless_mode": True, "$session_id": "s-1"}, - **{"$set": {"name": "Max"}}, - ) - original_properties = dict(msg["properties"]) - _to_v1_event(msg) - self.assertEqual(msg["properties"], original_properties) - self.assertIn("$set", msg) # top-level $set untouched on the original - - @parameterized.expand( - [ - ("cookieless_mode", "$cookieless_mode", "cookieless_mode", True, True), - ( - "ignore_sent_at_rename", - "$ignore_sent_at", - "disable_skew_correction", - "true", - True, - ), - ( - "process_person_profile", - "$process_person_profile", - "process_person_profile", - "false", - False, - ), - ( - "product_tour_id", - "$product_tour_id", - "product_tour_id", - "tour-7", - "tour-7", - ), - ] - ) - def test_option_sentinels_lifted_renamed_and_coerced( - self, _name, prop_key, wire_key, raw, expected - ) -> None: - event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) - self.assertEqual(event["options"], {wire_key: expected}) - self.assertNotIn(prop_key, event["properties"]) - - @parameterized.expand( - [ - ("bad_bool", "$cookieless_mode", "maybe"), - ("bad_tour_id_int", "$product_tour_id", 123), - ] - ) - def test_option_sentinel_removed_but_omitted_on_bad_coercion( - self, _name, prop_key, raw - ) -> None: - event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) - # Removed from properties (sentinels must never reach v1 props) but not - # emitted as an option, so a wrong type cannot 400 the whole batch. - self.assertNotIn(prop_key, event["properties"]) - self.assertEqual(event["options"], {}) - - @parameterized.expand( - [ - ("session_id", "$session_id", "session_id", "s-123"), - ("window_id", "$window_id", "window_id", "w-456"), - ] - ) - def test_top_level_string_sentinels(self, _name, prop_key, field_name, raw) -> None: - event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) - self.assertEqual(event[field_name], raw) - self.assertNotIn(prop_key, event["properties"]) - - def test_top_level_sentinel_omitted_but_removed_when_not_string(self) -> None: - event = _to_v1_event(_legacy_msg(properties={"$session_id": 42})) - self.assertNotIn("session_id", event) - self.assertNotIn("$session_id", event["properties"]) - - def test_all_sentinels_together(self) -> None: - event = _to_v1_event( - _legacy_msg( - properties={ - "$cookieless_mode": True, - "$ignore_sent_at": "1", - "$product_tour_id": "tour-x", - "$process_person_profile": 0, - "$session_id": "s-1", - "$window_id": "w-1", - "$geoip_disable": True, - "custom": "keep", - } - ) - ) - self.assertEqual( - event["options"], - { - "cookieless_mode": True, - "disable_skew_correction": True, - "product_tour_id": "tour-x", - "process_person_profile": False, - }, - ) - self.assertEqual(event["session_id"], "s-1") - self.assertEqual(event["window_id"], "w-1") - # Non-sentinel props (including $geoip_disable) are left intact. - self.assertEqual( - event["properties"], {"$geoip_disable": True, "custom": "keep"} - ) - - @parameterized.expand([("set", "$set"), ("set_once", "$set_once")]) - def test_top_level_set_relocated_into_properties(self, _name, key) -> None: - msg = _legacy_msg(properties={}, **{key: {"email": "a@b.com"}}) - event = _to_v1_event(msg) - self.assertEqual(event["properties"][key], {"email": "a@b.com"}) - self.assertNotIn(key, event) # not a top-level v1 field - - def test_top_level_set_merges_with_existing_properties_set(self) -> None: - # properties wins on key collision. - msg = _legacy_msg( - properties={"$set": {"a": "from_props", "b": "props_only"}}, - **{"$set": {"a": "from_top", "c": "top_only"}}, - ) - event = _to_v1_event(msg) - self.assertEqual( - event["properties"]["$set"], - {"a": "from_props", "b": "props_only", "c": "top_only"}, - ) - - def test_groups_left_in_properties(self) -> None: - event = _to_v1_event(_legacy_msg(properties={"$groups": {"company": "ph"}})) - self.assertEqual(event["properties"]["$groups"], {"company": "ph"}) - - def test_timestamp_naive_datetime_made_tz_aware(self) -> None: - event = _to_v1_event(_legacy_msg(timestamp=datetime(2026, 6, 27, 12, 0, 0))) - parsed = datetime.fromisoformat(event["timestamp"]) - self.assertIsNotNone(parsed.tzinfo) - - def test_timestamp_aware_datetime_converted_to_exact_utc_instant(self) -> None: - event = _to_v1_event( - _legacy_msg( - timestamp=datetime( - 2026, - 6, - 27, - 17, - 45, - tzinfo=timezone(timedelta(hours=5, minutes=45)), - ) - ) - ) - self.assertEqual(event["timestamp"], "2026-06-27T12:00:00+00:00") - - def test_timestamp_parseable_string_converted_to_exact_utc_instant(self) -> None: - event = _to_v1_event(_legacy_msg(timestamp="2026-06-27T17:45:00+05:45")) - self.assertEqual(event["timestamp"], "2026-06-27T12:00:00+00:00") - - def test_timestamp_none_defaults_to_utc_now(self) -> None: - event = _to_v1_event(_legacy_msg(timestamp=None)) - parsed = datetime.fromisoformat(event["timestamp"]) - self.assertEqual(parsed.tzinfo, timezone.utc) - - -class TestBuildV1BatchBody(unittest.TestCase): - def test_envelope_shape_and_no_legacy_fields(self) -> None: - events = [{"event": "e"}] - body = _build_v1_batch_body(events) - self.assertEqual(body["batch"], events) - self.assertNotIn("api_key", body) - self.assertNotIn("sent_at", body) - - def test_created_at_is_tz_aware_rfc3339(self) -> None: - body = _build_v1_batch_body([]) - parsed = datetime.fromisoformat(body["created_at"]) - self.assertIsNotNone(parsed.tzinfo) - - def test_created_at_passthrough_used_verbatim(self) -> None: - # _send_v1_batch hoists created_at and passes it in so it stays stable - # across retry attempts. - body = _build_v1_batch_body([], created_at="2026-06-27T12:00:00+00:00") - self.assertEqual(body["created_at"], "2026-06-27T12:00:00+00:00") - - def test_historical_migration_omitted_when_false(self) -> None: - self.assertNotIn("historical_migration", _build_v1_batch_body([])) - - def test_historical_migration_present_when_true(self) -> None: - body = _build_v1_batch_body([], historical_migration=True) - self.assertIs(body["historical_migration"], True) - - class TestPostV1(unittest.TestCase): def _post(self, response, **kwargs): session = _RecordingSession(response) @@ -472,7 +204,7 @@ def test_zstd_sets_encoding_header_and_emits_standard_frame(self) -> None: self.assertEqual(len(body["batch"]), 1) def test_zstd_without_package_raises_actionable_error(self) -> None: - with mock.patch("posthog.capture_v1._zstandard", None): + with mock.patch("posthog.capture_send._zstandard", None): with self.assertRaises(ValueError) as ctx: self._post(_results_response({}), compression=CaptureCompression.ZSTD) self.assertIn("posthog[zstd]", str(ctx.exception)) @@ -527,20 +259,20 @@ class TestSendV1Batch(unittest.TestCase): """Drives ``_send_v1_batch`` with a stubbed ``_post_v1`` and no real sleeps.""" def setUp(self) -> None: - sleep_patch = mock.patch("posthog.capture_v1.time.sleep") + sleep_patch = mock.patch("posthog.capture_send.time.sleep") self.sleep = sleep_patch.start() self.addCleanup(sleep_patch.stop) def _run(self, batch, outcomes, **kwargs): stub = _PostV1Stub(outcomes) - with mock.patch("posthog.capture_v1._post_v1", stub): + with mock.patch("posthog.capture_send._post_v1", stub): _send_v1_batch("phc_key", "https://app.posthog.com", batch, **kwargs) return stub def _run_expecting_error(self, batch, outcomes, **kwargs): stub = _PostV1Stub(outcomes) - with mock.patch("posthog.capture_v1._post_v1", stub): - with self.assertRaises(CaptureV1Error) as ctx: + with mock.patch("posthog.capture_send._post_v1", stub): + with self.assertRaises(CaptureError) as ctx: _send_v1_batch("phc_key", "https://app.posthog.com", batch, **kwargs) return stub, ctx.exception @@ -708,7 +440,7 @@ def test_transport_error_then_success(self) -> None: def test_transport_error_exhausted_reraises_original(self) -> None: stub = _PostV1Stub([ConnectionError("boom"), ConnectionError("boom")]) - with mock.patch("posthog.capture_v1._post_v1", stub): + with mock.patch("posthog.capture_send._post_v1", stub): with self.assertRaises(ConnectionError): _send_v1_batch( "phc_key", "https://app.posthog.com", [_msg("u-1")], max_retries=1 @@ -718,7 +450,7 @@ def test_transport_error_exhausted_reraises_original(self) -> None: def test_negative_max_retries_still_attempts_delivery_once(self) -> None: stub = _PostV1Stub([_results_response({"u-1": "ok"})]) - with mock.patch("posthog.capture_v1._post_v1", stub): + with mock.patch("posthog.capture_send._post_v1", stub): _send_v1_batch( "phc_key", "https://app.posthog.com", [_msg("u-1")], max_retries=-1 ) @@ -764,6 +496,6 @@ class TestBackoff(unittest.TestCase): ] ) def test_backoff(self, _name, attempt_index, retry_after, expected) -> None: - with mock.patch("posthog.capture_v1.time.sleep") as sleep: + with mock.patch("posthog.capture_send.time.sleep") as sleep: _backoff(attempt_index, retry_after) sleep.assert_called_once_with(expected) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 10474a25c..142740ad1 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -18,7 +18,7 @@ import pytest from posthog.capture_compression import CaptureCompression -from posthog.capture_v1 import _CAPTURE_V1_PATH +from posthog.capture_send import _CAPTURE_V1_PATH from posthog.client import Client from posthog.contexts import get_context_session_id, new_context, set_context_session from posthog.request import APIError, GetResponse @@ -4671,7 +4671,7 @@ def test_negative_max_retries_still_attempts_delivery_once(self, _name, sync_mod client = None with mock.patch( - "posthog.capture_v1._post_v1", return_value=response + "posthog.capture_send._post_v1", return_value=response ) as v1_post: try: client = Client( diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index 92cd9d888..2274c6906 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -12,7 +12,7 @@ from Queue import Queue from posthog.capture_compression import CaptureCompression -from posthog.capture_v1 import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.consumer import MAX_MSG_SIZE, Consumer, _DrainSignal from posthog.test.capture_helpers import patch_capture_send, sent_batch from posthog.test.logging_helpers import capture_message_only_logs diff --git a/posthog/test/test_server_payload_snapshots.py b/posthog/test/test_server_payload_snapshots.py index 80e03415d..71f437edd 100644 --- a/posthog/test/test_server_payload_snapshots.py +++ b/posthog/test/test_server_payload_snapshots.py @@ -117,7 +117,7 @@ def _event_family_request(): with ( freeze_time(_FIXED_TIME), - mock.patch("posthog.capture_v1._get_session", return_value=session), + mock.patch("posthog.capture_send._get_session", return_value=session), mock.patch("posthog.client.system_context", return_value=_RUNTIME_CONTEXT), ): client = Client( @@ -182,7 +182,7 @@ def _exception_request(): with ( freeze_time(_FIXED_TIME), - mock.patch("posthog.capture_v1._get_session", return_value=session), + mock.patch("posthog.capture_send._get_session", return_value=session), mock.patch("posthog.client.system_context", return_value=_RUNTIME_CONTEXT), mock.patch("posthog.client._get_current_otel_span_properties", return_value={}), ): diff --git a/posthog/tracing/_export.py b/posthog/tracing/_export.py index e68385f78..3a3e4b398 100644 --- a/posthog/tracing/_export.py +++ b/posthog/tracing/_export.py @@ -12,7 +12,7 @@ import time from typing import Any, Callable, List, Optional, Tuple -from ..capture_v1 import _MAX_BACKOFF_SECONDS +from ..capture_send import _MAX_BACKOFF_SECONDS from ._config import ResolvedTracesConfig from ._drops import DropLog from ._limits import truncate_attributes diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index f412a2113..c6165d97b 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -7,6 +7,7 @@ alias posthog.AsyncClient -> posthog.async_client.AsyncClient alias posthog.AsyncPosthog -> posthog.async_client.AsyncPosthog alias posthog.BeforeSendCallback -> posthog.types.BeforeSendCallback alias posthog.CaptureCompression -> posthog.capture_compression.CaptureCompression +alias posthog.CaptureError -> posthog.capture_send.CaptureError alias posthog.Client -> posthog.client.Client alias posthog.DEFAULT_CODE_VARIABLES_DETECT_SECRETS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_DETECT_SECRETS alias posthog.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS -> posthog.exception_utils.DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS @@ -689,11 +690,11 @@ attribute posthog.capture_compression.CaptureCompression.GZIP = 'gzip' attribute posthog.capture_compression.CaptureCompression.NONE = 'none' attribute posthog.capture_compression.CaptureCompression.ZSTD = 'zstd' attribute posthog.capture_exception_code_variables = False +attribute posthog.capture_send.CaptureError.attempts = attempts +attribute posthog.capture_send.CaptureError.drops = drops or [] +attribute posthog.capture_send.CaptureError.request_id = request_id +attribute posthog.capture_send.CaptureError.retry_exhausted = retry_exhausted or [] attribute posthog.capture_trace_context = False -attribute posthog.capture_v1.CaptureV1Error.attempts = attempts -attribute posthog.capture_v1.CaptureV1Error.drops = drops or [] -attribute posthog.capture_v1.CaptureV1Error.request_id = request_id -attribute posthog.capture_v1.CaptureV1Error.retry_exhausted = retry_exhausted or [] attribute posthog.client.Client.api_key = (project_api_key or '').strip() attribute posthog.client.Client.capture_compression = _resolve_capture_compression(capture_compression) attribute posthog.client.Client.capture_exception_code_variables = capture_exception_code_variables @@ -1163,7 +1164,7 @@ class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] class posthog.async_client.AsyncPosthog class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, refill_rate: Number, refill_interval_seconds: Number, on_bucket_rate_limited: Optional[Callable[[Hashable], None]] = None, clock: Callable[[], float] = time.monotonic) class posthog.capture_compression.CaptureCompression -class posthog.capture_v1.CaptureV1Error(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) +class posthog.capture_send.CaptureError(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, retries=10, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) @@ -1741,7 +1742,8 @@ module posthog.args module posthog.async_client module posthog.bucketed_rate_limiter module posthog.capture_compression -module posthog.capture_v1 +module posthog.capture_event +module posthog.capture_send module posthog.client module posthog.consumer module posthog.contexts diff --git a/sdk_compliance_adapter/adapter.py b/sdk_compliance_adapter/adapter.py index b663193cb..be2e1905b 100644 --- a/sdk_compliance_adapter/adapter.py +++ b/sdk_compliance_adapter/adapter.py @@ -15,8 +15,8 @@ from posthog import Client from posthog.capture_compression import CaptureCompression -from posthog.capture_v1 import _CAPTURE_V1_PATH -from posthog.capture_v1 import _post_v1 as original_post_v1 +from posthog.capture_send import _CAPTURE_V1_PATH +from posthog.capture_send import _post_v1 as original_post_v1 from posthog.request import USER_AGENT from posthog.version import VERSION @@ -212,9 +212,9 @@ def patched_post_v1( # Patch the capture-v1 submitter. `_send_v1_batch` resolves `_post_v1` as a module # global at call time, so patching it here covers both the async consumer and the # sync client paths. -import posthog.capture_v1 # noqa: E402 +import posthog.capture_send # noqa: E402 -posthog.capture_v1._post_v1 = patched_post_v1 +posthog.capture_send._post_v1 = patched_post_v1 @app.route("/health", methods=["GET"]) diff --git a/sdk_compliance_adapter/test_adapter.py b/sdk_compliance_adapter/test_adapter.py index 2cd484a5d..1e06c4628 100644 --- a/sdk_compliance_adapter/test_adapter.py +++ b/sdk_compliance_adapter/test_adapter.py @@ -8,7 +8,7 @@ import pytest -import posthog.capture_v1 +import posthog.capture_send import posthog.client import posthog.consumer import posthog.request @@ -20,7 +20,7 @@ def adapter(monkeypatch): # Importing the adapter installs transport instrumentation. Restore it after # every test so collecting these tests alongside SDK tests is safe. for module, name in [ - (posthog.capture_v1, "_post_v1"), + (posthog.capture_send, "_post_v1"), ]: monkeypatch.setattr(module, name, getattr(module, name)) spec = importlib.util.spec_from_file_location( From 20fdb14636f75c709cb31e6597d454d7013410e2 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Sat, 10 Oct 2026 16:34:35 -0700 Subject: [PATCH 05/14] feat!: report capture v1 results through CaptureError and on_error (#1020) * feat!: report capture v1 results through CaptureError and on_error * test: cover _uuid7 layout for the mutation gate * fix: re-raise sync capture failures in debug mode after reporting them Also document which events verdict_summary() counts. --- posthog/__init__.py | 7 +- posthog/_async_consumer.py | 41 ++++-- posthog/async_client.py | 29 ++--- posthog/capture_send.py | 202 ++++++++++++++++++++++------- posthog/client.py | 48 ++++++- posthog/consumer.py | 15 ++- posthog/test/test_async_client.py | 31 ++++- posthog/test/test_capture_send.py | 134 ++++++++++++++++--- posthog/test/test_client.py | 93 +++++++++++-- posthog/test/test_consumer.py | 6 +- posthog/test/test_utils.py | 26 ++++ posthog/utils.py | 16 +++ references/public_api_snapshot.txt | 10 +- 13 files changed, 538 insertions(+), 120 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index 5c3335e14..71e4493ba 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -313,8 +313,11 @@ def get_tags() -> Dict[str, Any]: disabled no-op global client. host: PostHog ingestion host. Defaults to the US ingestion endpoint when not set. - on_error: Optional callback invoked by background consumers when event upload - fails. Keep it short and non-blocking. Lifecycle methods can be called + on_error: Optional callback ``(error, batch)`` invoked when event upload + fails: by background consumers, or on the calling thread in ``sync_mode``. + Capture failures arrive as ``CaptureError``. Without it, each failed batch + logs one aggregate line. A capture that fails inside the callback logs that + line instead of calling it again. Keep it short and non-blocking. Lifecycle methods can be called directly and will be deferred, but the callback must not wait for another thread or task that calls ``flush()``, ``join()``, or ``shutdown()``. debug: Enable verbose SDK logging and re-raise errors from public APIs. diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 9352139ee..0ae74b995 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -11,6 +11,7 @@ from ._async_request import async_send_v1_batch from .capture_compression import CaptureCompression +from .capture_send import _CAPTURE_V1_PATH, _capture_loss_message from .consumer import BATCH_SIZE_LIMIT, MAX_MSG_SIZE from .request import DatetimeSerializer @@ -47,6 +48,33 @@ async def _invoke_callback(callback, *args): return result +# True while an SDK-invoked `on_error` runs, so a capture that fails inside the +# callback logs instead of re-entering it. +_IN_ON_ERROR: contextvars.ContextVar[bool] = contextvars.ContextVar( + "posthog_in_on_error", default=False +) + + +async def _report_capture_failure( + on_error: Optional[Callable[..., Any]], + log: logging.Logger, + error: Exception, + batch: list[dict[str, Any]], + endpoint: str, +) -> None: + """Hand a failed send to `on_error`, or log one aggregate line without one.""" + if on_error is None or _IN_ON_ERROR.get(): + log.error(_capture_loss_message(error, max(1, len(batch)), endpoint)) + return + token = _IN_ON_ERROR.set(True) + try: + await _invoke_callback(on_error, error, batch) + except Exception as callback_error: + log.error("on_error handler failed (%s)", type(callback_error).__name__) + finally: + _IN_ON_ERROR.reset(token) + + async def _serialized_event_size(event: dict[str, Any]) -> int: serialized = await asyncio.to_thread(json.dumps, event, cls=DatetimeSerializer) return len(serialized.encode()) @@ -142,18 +170,9 @@ async def upload(self, batch: list[dict[str, Any]]) -> None: try: await self.request(batch) except Exception as error: - self.log.error( - "async capture upload failed (%s, status=%s)", - type(error).__name__, - getattr(error, "status", None), + await _report_capture_failure( + self.on_error, self.log, error, batch, _CAPTURE_V1_PATH ) - if self.on_error: - try: - await _invoke_callback(self.on_error, error, batch) - except Exception as callback_error: - self.log.error( - "on_error handler failed (%s)", type(callback_error).__name__ - ) finally: for _ in batch: self.queue.task_done() diff --git a/posthog/async_client.py b/posthog/async_client.py index e46f0e494..37dfd2294 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -11,7 +11,7 @@ import weakref from datetime import datetime, timezone from typing import Any, Dict, Mapping, Optional, Union -from uuid import UUID, uuid4 +from uuid import UUID from typing_extensions import Unpack @@ -19,6 +19,7 @@ _STOP, _AsyncConsumer, _invoke_callback, + _report_capture_failure, _is_processing_event, _QueuedEvent, _run_outside_processing_event, @@ -34,6 +35,7 @@ CaptureCompression, _resolve_capture_compression, ) +from .capture_send import _CAPTURE_V1_PATH from .client import ( MAX_DICT_SIZE as _MAX_DICT_SIZE, _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, @@ -76,7 +78,13 @@ from .release_id import _resolve_release_id from .request import QuotaLimitError, determine_server_host, normalize_host from .types import FlagMetadata, FlagValue, normalize_flags_response -from .utils import SizeLimitedDict, _normalize_timestamp, clean, system_context +from .utils import ( + SizeLimitedDict, + _normalize_timestamp, + _uuid7, + clean, + system_context, +) from .version import VERSION __all__ = ["AsyncClient", "AsyncPosthog"] @@ -380,7 +388,7 @@ def _normalize_uuid(self, msg: dict[str, Any]) -> str: msg["uuid"] = normalized return normalized - normalized = str(uuid4()) + normalized = str(_uuid7()) msg["uuid"] = normalized return normalized @@ -576,20 +584,11 @@ async def capture_immediate( await consumer.request(error_batch) return sent_uuid except Exception as error: - if self.on_error: - try: - await _invoke_callback(self.on_error, error, error_batch) - except Exception as callback_error: - self.log.error( - "on_error handler failed (%s)", type(callback_error).__name__ - ) + await _report_capture_failure( + self.on_error, self.log, error, error_batch, _CAPTURE_V1_PATH + ) if self.debug: raise - self.log.error( - "Immediate async capture failed (%s, status=%s)", - type(error).__name__, - getattr(error, "status", None), - ) return None finally: remaining_calls = self._immediate_callers[current] - 1 diff --git a/posthog/capture_send.py b/posthog/capture_send.py index 862eae094..f254f165d 100644 --- a/posthog/capture_send.py +++ b/posthog/capture_send.py @@ -29,7 +29,7 @@ from gzip import GzipFile from io import BytesIO from typing import TYPE_CHECKING, Optional -from uuid import uuid4 +from uuid import UUID from posthog.capture_compression import CaptureCompression, _zstandard from posthog.capture_event import _build_v1_batch_body, _to_v1_event @@ -40,17 +40,17 @@ _get_session, normalize_host, ) -from posthog.utils import remove_trailing_slash +from posthog.utils import _uuid7, remove_trailing_slash if TYPE_CHECKING: import requests log = logging.getLogger("posthog") -# Only the error type is public API: it reaches user code through `on_error` -# callbacks, so callers may want to catch/inspect it. Everything else is +# Only the error types are public API: they reach user code through `on_error` +# callbacks, so callers may want to catch/inspect them. Everything else is # submitter plumbing. -__all__ = ["CaptureError"] +__all__ = ["CaptureError", "CaptureEventResult"] _CAPTURE_V1_PATH = "/i/v1/analytics/events" _CAPTURE_AI_V1_PATH = "/i/v1/ai/events" @@ -80,10 +80,19 @@ # unifies the default with posthog-go/posthog-rs (all 30s). _MAX_BACKOFF_SECONDS = 30 +# First retry delay; it doubles per attempt up to `_MAX_BACKOFF_SECONDS`. Matches +# posthog-go. +_BACKOFF_BASE_SECONDS = 0.1 -@dataclass -class _V1EventResult: - """A single event's directive from a 2xx ``results`` map.""" + +@dataclass(frozen=True) +class CaptureEventResult: + """One event's verdict from a 2xx capture response. + + ``result`` is ``ok``, ``warning``, ``drop`` or ``retry``, or a value newer + than this SDK, which counts as success. ``details`` is the server's reason + tag, such as ``billing_limit_exceeded``. + """ result: Optional[str] details: Optional[str] = None @@ -102,17 +111,22 @@ class _V1ParsedResponse: status_code: int is_success: bool retry_after: Optional[float] = None - results: Optional[dict[str, _V1EventResult]] = None + results: Optional[dict[str, CaptureEventResult]] = None malformed: bool = False error_message: str = "" class CaptureError(APIError): - """Batch-level failure of a capture-v1 send. + """A capture batch that was not fully delivered. - Subclasses :class:`APIError` so the consumer's existing ``on_error`` handling - (which already inspects ``status``/``retry_after``) keeps working; the extra - fields carry v1 specifics for richer logging/callbacks. + Passed to ``on_error`` with the batch. ``status`` is the HTTP status of the + last attempt, or ``0`` when no response arrived (``__cause__`` holds the + transport error). A 2xx ``status`` means the request succeeded but some + events were dropped or ran out of retries. + + ``event_results`` maps each event uuid to its last verdict from a 2xx + response. Events missing from it never got a verdict, because every + attempt that reached them failed as a whole request. """ def __init__( @@ -120,25 +134,101 @@ def __init__( status: int | str, message: str, *, + endpoint: str, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None, + event_results: Optional[dict[str, CaptureEventResult]] = None, ): super().__init__(status, message, retry_after=retry_after) + # Capture path the batch was sent to, so one on_error handler can tell + # the analytics and AI lanes apart. + self.endpoint = endpoint self.request_id = request_id self.attempts = attempts # uuids the server told us to retry but we never delivered (exhausted). self.retry_exhausted = retry_exhausted or [] # (uuid, details) pairs the server told us to drop on a 2xx response. self.drops = drops or [] + self.event_results = event_results or {} + + def verdict_summary(self) -> str: + """Count the undelivered events by verdict and reason, for example + ``drop/billing_limit_exceeded=2, retry/not_persisted=1``. + + Counts the events in ``drops`` and ``retry_exhausted``: events the + server dropped, and events still pending retry after the last 2xx + response. A retry verdict followed by a failed request, for example a + 503 on the last attempt, is not counted; ``event_results`` keeps it. + Empty when no event is counted. + """ + counts: dict[str, int] = {} + failed = [uid for uid, _ in self.drops] + self.retry_exhausted + for uid in failed: + verdict = self.event_results.get(uid) + if verdict is None or not verdict.result: + continue + tag = verdict.result + if verdict.details: + tag = f"{tag}/{verdict.details}" + counts[tag] = counts.get(tag, 0) + 1 + return ", ".join(f"{tag}={n}" for tag, n in sorted(counts.items())) def _is_success_status(status: int) -> bool: return 200 <= status < 300 +def _canonical_uuid(value: str) -> str: + """Return ``value`` in the lowercase hyphenated form capture keys results by. + + ``before_send`` can set any uuid form capture parses. A value that does not + parse is returned unchanged; capture rejects its whole batch anyway. + """ + try: + return str(UUID(value)) + except (TypeError, ValueError, AttributeError): + return value + + +def _undelivered_count(error: Exception, batch_size: int) -> int: + if not isinstance(error, CaptureError): + return batch_size + delivered = sum( + 1 + for r in error.event_results.values() + if r.result not in (_RESULT_DROP, _RESULT_RETRY) + ) + return max(0, batch_size - delivered) + + +def _capture_loss_message(error: Exception, batch_size: int, endpoint: str) -> str: + """One aggregate line for a failed batch, for callers with no ``on_error``. + + It never names individual events, includes payloads, or repeats the server's + error text: any of them may carry sensitive content, and per-event lines + scale with event volume. + """ + count = _undelivered_count(error, batch_size) + if isinstance(error, CaptureError): + endpoint = error.endpoint + if isinstance(error.status, int) and _is_success_status(error.status): + detail = f"{len(error.drops)} dropped, {len(error.retry_exhausted)} out of retries" + summary = error.verdict_summary() + if summary: + detail = f"{detail} ({summary})" + return f"{count} event(s) not persisted by {endpoint}: {detail}" + detail = type(error).__name__ + if error.__cause__ is not None: + detail = f"{detail} from {type(error.__cause__).__name__}" + status = getattr(error, "status", None) + if status is not None: + detail = f"{detail} (status={status})" + return f"{count} event(s) not persisted by {endpoint}: {detail}" + + def _parse_retry_after(header_value: Optional[str]) -> Optional[float]: """Parse a ``Retry-After`` header (delta-seconds or HTTP-date) to seconds.""" if not header_value: @@ -239,7 +329,7 @@ def _parse_v1_response(res: "requests.Response") -> _V1ParsedResponse: payload = res.json() raw_results = payload["results"] results = { - uid: _V1EventResult( + uid: CaptureEventResult( result=(r or {}).get("result"), details=(r or {}).get("details"), ) @@ -279,7 +369,7 @@ def _backoff(attempt_index: int, retry_after: Optional[float]) -> None: :data:`_MAX_BACKOFF_SECONDS`, so both sides share one ceiling and a hostile/buggy header can't park the consumer thread. """ - configured = min(2**attempt_index, _MAX_BACKOFF_SECONDS) + configured = min(_BACKOFF_BASE_SECONDS * 2**attempt_index, _MAX_BACKOFF_SECONDS) clamped_retry_after = ( min(retry_after, _MAX_BACKOFF_SECONDS) if retry_after and retry_after > 0 else 0 ) @@ -287,7 +377,7 @@ def _backoff(attempt_index: int, retry_after: Optional[float]) -> None: def _log_result_summary( - request_id: str, attempt: int, results: dict[str, _V1EventResult] + request_id: str, attempt: int, results: dict[str, CaptureEventResult] ) -> None: tally = {_RESULT_OK: 0, _RESULT_WARNING: 0, _RESULT_DROP: 0, _RESULT_RETRY: 0} other = 0 @@ -324,9 +414,11 @@ def _send_v1_batch( ) -> None: """Deliver ``batch`` to the v1 endpoint at ``path`` with partial retry. - The v1 sibling of ``Consumer._send``: it loops up to ``max_retries + 1`` - attempts, but unlike v0 it shrinks the batch to only the events the server - tagged ``retry`` after each 2xx. ``ok``/``warning``/absent events succeed. + Loops up to ``max_retries + 1`` attempts. After each 2xx it resends only + the events the server tagged ``retry``. ``ok``/``warning`` events, events + with an unrecognized verdict, and events absent from ``results`` succeed + (matching posthog-go and posthog-rs). Results are matched by canonical + uuid, because capture keys them that way whatever form was sent. A server-chosen ``drop`` is a terminal per-event rejection. Drops are accumulated across attempts and surfaced via :class:`CaptureError` even @@ -334,26 +426,47 @@ def _send_v1_batch( and even when a later attempt clears the outstanding retries — matching posthog-go (per-event failure callback) and posthog-rs (``on_error`` on a 2xx with undelivered verdicts). Raises :class:`CaptureError` on any drop, - batch-level terminal failure, or retry exhaustion — carrying the accumulated - ``drops`` and any exhausted uuids — so the caller's ``on_error`` fires - unchanged. A transport failure re-raises the underlying exception (drops - collected on an earlier attempt are still tallied in the DEBUG summary). + batch-level terminal failure, retry exhaustion, or transport failure + (``status`` 0, chained to the transport error), carrying the endpoint, the + accumulated ``drops``, any exhausted uuids, and every verdict seen so far, + so the caller's ``on_error`` gets the full picture. ``request_id`` and the batch ``created_at`` are stable across attempts; ``PostHog-Attempt`` increments. Negative ``max_retries`` values are treated as zero, so delivery is always attempted at least once. """ max_retries = max(0, max_retries) - request_id = str(uuid4()) + request_id = str(_uuid7()) # Hoisted once so the batch envelope is byte-identical across retry attempts # (only the events list shrinks and the attempt header increments). created_at = datetime.now(timezone.utc).isoformat() pending_events = [_to_v1_event(m) for m in batch] - pending_uuids = [e["uuid"] for e in pending_events] - last_exc: Optional[Exception] = None + pending_uuids = [_canonical_uuid(e["uuid"]) for e in pending_events] # (uuid, details) for every event the server dropped, across all attempts. # Accumulated (not per-attempt) so a drop seen early is not lost when a # later attempt succeeds or clears the outstanding retries. all_drops: list[tuple[str, Optional[str]]] = [] + # Latest 2xx verdict per uuid, across all attempts. + event_results: dict[str, CaptureEventResult] = {} + + def capture_error( + status: int, + message: str, + *, + attempts: int, + retry_after: Optional[float] = None, + retry_exhausted: Optional[list[str]] = None, + ) -> CaptureError: + return CaptureError( + status, + message, + endpoint=path, + retry_after=retry_after, + request_id=request_id, + attempts=attempts, + retry_exhausted=retry_exhausted, + drops=all_drops, + event_results=event_results, + ) for attempt_index in range(max_retries + 1): attempt = attempt_index + 1 @@ -376,10 +489,11 @@ def _send_v1_batch( path=path, ) except Exception as e: - # Transport-level failure (connection/timeout): retry like v0 does. - last_exc = e + # Transport-level failure (connection/timeout): retryable. if last_attempt: - raise + raise capture_error( + 0, f"{type(e).__name__}: {e}", attempts=attempt + ) from e _backoff(attempt_index, None) continue @@ -387,12 +501,10 @@ def _send_v1_batch( if parsed.is_success: if parsed.malformed: - raise CaptureError( + raise capture_error( parsed.status_code, "capture v1 returned a success status with an unparseable body", - request_id=request_id, attempts=attempt, - drops=all_drops, ) results = parsed.results or {} _log_result_summary(request_id, attempt, results) @@ -402,8 +514,11 @@ def _send_v1_batch( for event, uid in zip(pending_events, pending_uuids): directive = results.get(uid) if directive is None: - # Absent from the map: treated as accepted (matches posthog-rs). + # Absent from the map: treated as accepted (matches + # posthog-go and posthog-rs). Capture answers every event + # of a 2xx, so this only happens if it breaks that contract. continue + event_results[uid] = directive if directive.result == _RESULT_RETRY: retry_events.append(event) retry_uuids.append(uid) @@ -418,22 +533,18 @@ def _send_v1_batch( # surface them via on_error even though the request was a 2xx — # a success status does not mean every event was delivered. if all_drops: - raise CaptureError( + raise capture_error( parsed.status_code, f"{len(all_drops)} event(s) dropped by the server", - request_id=request_id, attempts=attempt, - drops=all_drops, ) return if last_attempt: - raise CaptureError( + raise capture_error( parsed.status_code, f"{len(retry_uuids)} event(s) still pending retry after {attempt} attempt(s)", - request_id=request_id, attempts=attempt, retry_exhausted=retry_uuids, - drops=all_drops, ) pending_events, pending_uuids = retry_events, retry_uuids _backoff(attempt_index, parsed.retry_after) @@ -442,22 +553,13 @@ def _send_v1_batch( # Non-2xx. Retryable transient statuses back off; everything else # (400/401/402/413/415/429/...) is terminal. Any drops collected from a # prior 2xx attempt ride along so on_error still sees them. - v1_error = CaptureError( + v1_error = capture_error( parsed.status_code, parsed.error_message, retry_after=parsed.retry_after, - request_id=request_id, attempts=attempt, - drops=all_drops, ) - if parsed.status_code in _RETRYABLE_STATUSES: - last_exc = v1_error - if last_attempt: - raise v1_error + if parsed.status_code in _RETRYABLE_STATUSES and not last_attempt: _backoff(attempt_index, parsed.retry_after) continue raise v1_error - - # Unreachable in practice (every branch returns or continues). - if last_exc: - raise last_exc diff --git a/posthog/client.py b/posthog/client.py index 611fda2f1..a5e23fe2d 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -34,6 +34,7 @@ from posthog.capture_send import ( _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, + _capture_loss_message, _send_v1_batch, ) from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE, Consumer, _DrainSignal @@ -124,6 +125,7 @@ SizeLimitedDict, clean, _normalize_timestamp, + _uuid7, guess_timezone as guess_timezone, system_context, ) @@ -141,6 +143,9 @@ # overrun it by up to the client's request timeout. _TRACES_SHUTDOWN_FLUSH_SECONDS = 30.0 _atexit_deadline: Optional[float] = None +# Marks a thread that is running a sync-mode `on_error`, so a capture that fails +# inside the callback logs instead of re-entering it. +_on_error_state = threading.local() _atexit_deadline_lock = threading.Lock() @@ -246,14 +251,13 @@ def _stringify_event_uuid(value) -> str: ) try: - UUID(stringified) + # Canonical form, because capture keys per-event results by it. + return str(UUID(stringified)) except ValueError: raise ValueError( f"Invalid event uuid {value!r}. Expected a valid UUID string or uuid.UUID instance." ) from None - return stringified - def add_context_tags(properties): properties = properties or {} @@ -724,8 +728,12 @@ def __init__( API methods. max_queue_size: Maximum number of events buffered before upload. send: If False, queueing succeeds but events are not sent. - on_error: Optional callback invoked by background consumers when an - upload fails. Keep it short and non-blocking. Calling lifecycle + on_error: Optional callback ``(error, batch)`` invoked when an upload + fails: by background consumers, or on the calling thread in + ``sync_mode``. Capture failures arrive as ``CaptureError``. + Without it, each failed batch logs one aggregate line. A capture + that fails inside the callback logs that line instead of calling + it again. Keep it short and non-blocking. Calling lifecycle methods directly is safe and deferred, but do not start another thread or task that calls ``flush()``, ``join()``, or ``shutdown()`` and then wait for it from the callback. @@ -2310,7 +2318,26 @@ def _normalize_event_uuid(self, msg): if "uuid" not in msg: # Always send a uuid, so we can always return one - msg["uuid"] = stringify_id(uuid4()) + msg["uuid"] = str(_uuid7()) + + def _report_capture_failure( + self, error: Exception, batch: list[dict], endpoint: str + ) -> None: + """Hand a failed sync-mode send to `on_error`, like the queued path does. + + Without a handler, or when the failed capture ran inside `on_error`, + log one aggregate line instead, so the loss stays visible without recursion. + """ + if not self.on_error or getattr(_on_error_state, "active", False): + self.log.error(_capture_loss_message(error, len(batch), endpoint)) + return + _on_error_state.active = True + try: + self.on_error(error, batch) + except Exception as callback_error: + self.log.error("on_error handler failed: %s", callback_error) + finally: + _on_error_state.active = False def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): # type: (...) -> Optional[str] @@ -2423,7 +2450,14 @@ def send_sync() -> None: path=lane.endpoint, ) - if lane.run_sync_if_open(send_sync): + try: + admitted = lane.run_sync_if_open(send_sync) + except Exception as e: + self._report_capture_failure(e, [msg], lane.endpoint) + if self.debug: + raise + return None + if admitted: return sent_uuid self.log.warning( "%s lane received event %s after shutdown, dropping it", diff --git a/posthog/consumer.py b/posthog/consumer.py index daad58c33..22944064d 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -6,7 +6,11 @@ from posthog._logging import _configure_posthog_logging from posthog.capture_compression import CaptureCompression -from posthog.capture_send import _CAPTURE_V1_PATH, _send_v1_batch +from posthog.capture_send import ( + _CAPTURE_V1_PATH, + _capture_loss_message, + _send_v1_batch, +) from posthog.request import ( USER_AGENT as _USER_AGENT, DatetimeSerializer, @@ -105,7 +109,7 @@ def __init__( host=None, on_error=None, flush_interval=5.0, - retries=10, + retries=3, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, @@ -172,13 +176,14 @@ def upload(self): self.request(batch) success = True except Exception as e: - self.log.error("error uploading: %s", e) success = False if self.on_error: try: self.on_error(e, batch) - except Exception as e: - self.log.error("on_error handler failed: %s", e) + except Exception as callback_error: + self.log.error("on_error handler failed: %s", callback_error) + else: + self.log.error(_capture_loss_message(e, len(batch), self.endpoint)) finally: # mark items as acknowledged from queue for item in batch: diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 34c098d7d..02ce2ce4c 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -7,6 +7,7 @@ import sys import threading from unittest import mock +from uuid import UUID import pytest @@ -56,7 +57,7 @@ async def send_batch(api_key, host, batch, **kwargs): distinct_id="user-1", properties={"plan": "pro"}, ) - assert isinstance(event_uuid, str) + assert UUID(event_uuid).version == 7 await client.flush(timeout_seconds=1) assert len(batches) == 1 @@ -706,6 +707,34 @@ def on_error(error, batch): assert callback_thread != threading.get_ident() +@pytest.mark.asyncio +@pytest.mark.parametrize("immediate", [False, True]) +async def test_capture_failure_inside_on_error_logs_instead_of_recursing( + caplog, immediate +): + calls = [] + client: AsyncPosthog + + async def on_error(error, batch): + calls.append(batch[0]["event"]) + await client.capture_immediate("from_callback", distinct_id="user-1") + + with patch_async_capture_send(side_effect=APIError(503, "unavailable")): + client = AsyncPosthog("test-key", on_error=on_error, flush_at=1, max_retries=0) + if immediate: + await client.capture_immediate("event", distinct_id="user-1") + else: + client.capture("event", distinct_id="user-1") + await client.flush(timeout_seconds=1) + await client.shutdown() + + assert calls == ["event"] + assert ( + "1 event(s) not persisted by /i/v1/analytics/events: APIError (status=503)" + in caplog.text + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("immediate", [False, True]) async def test_failed_capture_does_not_log_server_response_detail(caplog, immediate): diff --git a/posthog/test/test_capture_send.py b/posthog/test/test_capture_send.py index 637eee85f..1f8c95062 100644 --- a/posthog/test/test_capture_send.py +++ b/posthog/test/test_capture_send.py @@ -3,6 +3,7 @@ import zlib from datetime import datetime, timedelta from unittest import mock +from uuid import UUID import zstandard @@ -19,6 +20,8 @@ _HEADER_SDK_INFO, _MAX_BACKOFF_SECONDS, CaptureError, + CaptureEventResult, + _capture_loss_message, _parse_v1_response, _post_v1, _send_v1_batch, @@ -317,6 +320,7 @@ def test_request_id_and_created_at_stable_attempt_increments(self) -> None: ], ) self.assertEqual(stub.calls[0]["request_id"], stub.calls[1]["request_id"]) + self.assertEqual(UUID(stub.calls[0]["request_id"]).version, 7) # created_at is hoisted once, so the envelope timestamp is identical # across retry attempts (only the attempt header increments). self.assertEqual(stub.calls[0]["created_at"], stub.calls[1]["created_at"]) @@ -383,18 +387,56 @@ def test_retry_exhausted_raises_with_uuids(self) -> None: def test_retry_exhausted_carries_earlier_drops(self) -> None: # A drop seen on attempt 1 rides along on the retry-exhaustion error. - batch = [_msg("u-drop"), _msg("u-retry")] + batch = [_msg("u-ok"), _msg("u-drop"), _msg("u-retry")] stub, exc = self._run_expecting_error( batch, [ - _results_response({"u-drop": ("drop", "billing"), "u-retry": "retry"}), - _results_response({"u-retry": "retry"}), + _results_response( + { + "u-ok": "ok", + "u-drop": ("drop", "billing"), + "u-retry": "retry", + } + ), + _results_response({"u-retry": ("retry", "not_persisted")}), ], max_retries=1, ) self.assertEqual(len(stub.calls), 2) + self.assertEqual(exc.endpoint, _CAPTURE_V1_PATH) self.assertEqual(exc.retry_exhausted, ["u-retry"]) self.assertEqual(exc.drops, [("u-drop", "billing")]) + self.assertEqual( + exc.event_results, + { + "u-ok": CaptureEventResult("ok"), + "u-drop": CaptureEventResult("drop", "billing"), + "u-retry": CaptureEventResult("retry", "not_persisted"), + }, + ) + self.assertEqual(exc.verdict_summary(), "drop/billing=1, retry/not_persisted=1") + + @parameterized.expand( + [ + ("uppercase", "AAAAAAAA-AAAA-4AAA-8AAA-AAAAAAAAAAAA"), + ("no_hyphens", "aaaaaaaaaaaa4aaa8aaaaaaaaaaaaaaa"), + ("braced", "{aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa}"), + ("urn", "urn:uuid:aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"), + ] + ) + def test_results_match_non_canonical_uuid(self, _name, sent_uuid) -> None: + # Capture parses any of these forms but keys results canonically, so a + # verdict for a uuid `before_send` rewrote must still reach its event. + canonical = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" + stub, exc = self._run_expecting_error( + [_msg(sent_uuid)], + [ + _results_response({canonical: "retry"}), + _results_response({canonical: ("drop", "billing")}), + ], + ) + self.assertEqual(len(stub.calls), 2) + self.assertEqual(exc.drops, [(canonical, "billing")]) def test_malformed_2xx_is_terminal(self) -> None: stub, exc = self._run_expecting_error( @@ -438,14 +480,22 @@ def test_transport_error_then_success(self) -> None: ) self.assertEqual(len(stub.calls), 2) - def test_transport_error_exhausted_reraises_original(self) -> None: + def test_transport_error_exhausted_raises_capture_error(self) -> None: stub = _PostV1Stub([ConnectionError("boom"), ConnectionError("boom")]) with mock.patch("posthog.capture_send._post_v1", stub): - with self.assertRaises(ConnectionError): + with self.assertRaises(CaptureError) as ctx: _send_v1_batch( - "phc_key", "https://app.posthog.com", [_msg("u-1")], max_retries=1 + "phc_key", + "https://app.posthog.com", + [_msg("u-1")], + max_retries=1, + path=_CAPTURE_AI_V1_PATH, ) self.assertEqual(len(stub.calls), 2) + self.assertEqual(ctx.exception.status, 0) + self.assertEqual(ctx.exception.endpoint, _CAPTURE_AI_V1_PATH) + self.assertEqual(ctx.exception.attempts, 2) + self.assertIsInstance(ctx.exception.__cause__, ConnectionError) def test_negative_max_retries_still_attempts_delivery_once(self) -> None: stub = _PostV1Stub([_results_response({"u-1": "ok"})]) @@ -461,20 +511,69 @@ def test_negative_max_retries_still_attempts_delivery_once(self) -> None: def test_small_retry_after_does_not_shorten_backoff(self) -> None: # A Retry-After smaller than the configured backoff must not make the # client retry earlier than its own schedule (Retry-After is a minimum). - # attempt_index=1 -> configured backoff 2s; Retry-After 0.5s is ignored. + # attempt_index=1 -> configured backoff 0.2s; Retry-After 0.1s is ignored. stub = self._run( [_msg("u-1")], [ _results_response({"u-1": "retry"}), - _results_response({"u-1": "retry"}, headers={"Retry-After": "0.5"}), + _results_response({"u-1": "retry"}, headers={"Retry-After": "0.1"}), _results_response({"u-1": "ok"}), ], max_retries=3, ) self.assertEqual(len(stub.calls), 3) - # First backoff (attempt_index 0) waits 1s; second (attempt_index 1) - # keeps the 2s configured backoff rather than the smaller 0.5s header. - self.assertEqual([c.args[0] for c in self.sleep.call_args_list], [1, 2]) + # First backoff (attempt_index 0) waits 0.1s; second (attempt_index 1) + # keeps the 0.2s configured backoff rather than the smaller 0.1s header. + self.assertEqual([c.args[0] for c in self.sleep.call_args_list], [0.1, 0.2]) + + +class TestCaptureLossMessage(unittest.TestCase): + @parameterized.expand( + [ + ( + "partial_2xx", + CaptureError( + 200, + "2 event(s) not delivered", + endpoint=_CAPTURE_V1_PATH, + drops=[("u-drop", "billing")], + retry_exhausted=["u-retry"], + event_results={ + "u-ok": CaptureEventResult("ok"), + "u-drop": CaptureEventResult("drop", "billing"), + "u-retry": CaptureEventResult("retry", "not_persisted"), + }, + ), + 3, + "2 event(s) not persisted by /i/v1/analytics/events: 1 dropped, " + "1 out of retries (drop/billing=1, retry/not_persisted=1)", + ), + ( + "request_failure_after_partial_success", + CaptureError( + 503, + "unavailable", + endpoint=_CAPTURE_AI_V1_PATH, + event_results={ + "u-ok": CaptureEventResult("warning"), + "u-retry": CaptureEventResult("retry"), + }, + ), + 2, + "1 event(s) not persisted by /i/v1/ai/events: CaptureError (status=503)", + ), + ( + "other_exception_uses_caller_endpoint", + ValueError("bad payload"), + 4, + "4 event(s) not persisted by /i/v1/analytics/events: ValueError", + ), + ] + ) + def test_message(self, _name, error, batch_size, expected) -> None: + self.assertEqual( + _capture_loss_message(error, batch_size, _CAPTURE_V1_PATH), expected + ) class TestBackoff(unittest.TestCase): @@ -483,13 +582,13 @@ class TestBackoff(unittest.TestCase): @parameterized.expand( [ # (attempt_index, retry_after, expected sleep seconds) - ("first_no_header", 0, None, 1), - ("second_no_header", 1, None, 2), + ("first_no_header", 0, None, 0.1), + ("second_no_header", 1, None, 0.2), ("exp_capped_at_30", 10, None, 30), - ("zero_header_uses_backoff", 0, 0, 1), + ("zero_header_uses_backoff", 0, 0, 0.1), ("larger_header_wins", 0, 5.0, 5.0), - ("smaller_header_ignored", 3, 2.0, 8), # configured 8 > 2.0 - ("equal_header_and_backoff", 0, 1.0, 1), + ("smaller_header_ignored", 3, 0.5, 0.8), # configured 0.8 > 0.5 + ("equal_header_and_backoff", 1, 0.2, 0.2), ("header_at_ceiling", 0, 30.0, 30), ("header_above_ceiling_clamped", 0, 120.0, _MAX_BACKOFF_SECONDS), ("absurd_header_clamped", 0, 10**9, _MAX_BACKOFF_SECONDS), @@ -498,4 +597,5 @@ class TestBackoff(unittest.TestCase): def test_backoff(self, _name, attempt_index, retry_after, expected) -> None: with mock.patch("posthog.capture_send.time.sleep") as sleep: _backoff(attempt_index, retry_after) - sleep.assert_called_once_with(expected) + sleep.assert_called_once() + self.assertAlmostEqual(sleep.call_args.args[0], expected) diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 142740ad1..37d862b5c 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -470,6 +470,24 @@ def test_basic_capture_with_uuid_object(self): msg = sent_batch(mock_post)[0] self.assertEqual(msg["uuid"], str(uuid)) + @parameterized.expand( + [ + ("uppercase", "0190A3B4-C5D6-7E8F-9A0B-1C2D3E4F5A6B"), + ("no hyphens", "0190a3b4c5d67e8f9a0b1c2d3e4f5a6b"), + ("braced", "{0190a3b4-c5d6-7e8f-9a0b-1c2d3e4f5a6b}"), + ] + ) + def test_capture_sends_and_returns_canonical_uuid(self, _name, supplied): + with patch_capture_send("client") as mock_post: + client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) + msg_uuid = client.capture( + "python test event", distinct_id="distinct_id", uuid=supplied + ) + + canonical = "0190a3b4-c5d6-7e8f-9a0b-1c2d3e4f5a6b" + self.assertEqual(msg_uuid, canonical) + self.assertEqual(sent_batch(mock_post)[0]["uuid"], canonical) + @parameterized.expand( [ ("empty string", ""), @@ -489,7 +507,7 @@ def test_capture_with_invalid_uuid_logs_and_falls_back_to_generated_uuid( ) self.assertIsNotNone(msg_uuid) - UUID(msg_uuid) + self.assertEqual(UUID(msg_uuid).version, 7) mock_post.assert_called_once() msg = sent_batch(mock_post)[0] self.assertEqual(msg["uuid"], msg_uuid) @@ -2707,20 +2725,79 @@ def flush(timeout_seconds): self.assertEqual(timeouts, [0, None]) - def test_sync_send_failure_does_not_invoke_async_on_error(self): - on_error = mock.Mock() + def test_sync_send_failure_invokes_on_error_after_releasing_send_slot(self): + error = Exception("upload failed") + active_sends_in_callback = [] + client: Client + + def on_error(err, batch): + active_sends_in_callback.append(client._analytics_lane._active_sync_sends) + self.assertIs(err, error) + self.assertEqual([event["event"] for event in batch], ["event"]) + client.shutdown() + + client = Client(FAKE_TEST_API_KEY, sync_mode=True, on_error=on_error) + + with patch_capture_send("client", side_effect=error): + result = client.capture("event", distinct_id="distinct_id") + + self.assertIsNone(result) + self.assertEqual(active_sends_in_callback, [0]) + + def test_sync_send_failure_inside_on_error_logs_instead_of_recursing(self): + calls = [] + client: Client + + def on_error(err, batch): + calls.append(batch[0]["event"]) + client.capture("from_callback", distinct_id="distinct_id") + + client = Client(FAKE_TEST_API_KEY, sync_mode=True, on_error=on_error) + + with patch_capture_send("client", side_effect=Exception("upload failed")): + with self.assertLogs("posthog", level="ERROR") as logs: + client.capture("event", distinct_id="distinct_id") + + self.assertEqual(calls, ["event"]) + self.assertIn( + "1 event(s) not persisted by /i/v1/analytics/events: Exception", + "\n".join(logs.output), + ) + + @parameterized.expand([("capture",), ("capture_ai",)]) + def test_sync_send_failure_in_debug_reports_once_then_raises(self, method_name): + errors = [] client = Client( FAKE_TEST_API_KEY, sync_mode=True, - on_error=on_error, + debug=True, + on_error=lambda err, batch: errors.append(err), ) + failure = APIError(503, "unavailable") - with patch_capture_send("client", side_effect=Exception("upload failed")): - result = client.capture("event", distinct_id="distinct_id") + with patch_capture_send("client", side_effect=failure): + with self.assertRaises(APIError) as raised: + getattr(client, method_name)("event", distinct_id="distinct_id") + + self.assertIs(raised.exception, failure) + self.assertEqual(errors, [failure]) + + def test_sync_send_failure_without_on_error_logs_aggregate_line(self): + client = Client(FAKE_TEST_API_KEY, sync_mode=True) + + with patch_capture_send( + "client", side_effect=APIError(400, "password=server-secret") + ): + with self.assertLogs("posthog", level="ERROR") as logs: + result = client.capture("event", distinct_id="distinct_id") self.assertIsNone(result) - on_error.assert_not_called() - self.assertEqual(client._analytics_lane._active_sync_sends, 0) + output = "\n".join(logs.output) + self.assertIn( + "1 event(s) not persisted by /i/v1/analytics/events: APIError (status=400)", + output, + ) + self.assertNotIn("server-secret", output) def test_on_error_can_request_shutdown_with_pending_work(self): first_send_started = threading.Event() diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index 2274c6906..ca4748829 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -232,9 +232,11 @@ def test_message_only_error_logs_include_posthog_prefix(self) -> None: # other tests can log into the same stream. Assert on the line under # test rather than on the entire capture. upload_logs = [ - line for line in logs.getvalue().splitlines() if "error uploading" in line + line for line in logs.getvalue().splitlines() if "not persisted" in line ] - expected_log = "[PostHog] error uploading: boom" + expected_log = ( + "[PostHog] 1 event(s) not persisted by /i/v1/analytics/events: Exception" + ) self.assertEqual( [line for line in upload_logs if line == expected_log], [expected_log] ) diff --git a/posthog/test/test_utils.py b/posthog/test/test_utils.py index e655c8cd5..c2dd1a0c5 100644 --- a/posthog/test/test_utils.py +++ b/posthog/test/test_utils.py @@ -261,6 +261,32 @@ def test_clean_fn(self): def test_remove_slash(self, input_url, expected_url): assert expected_url == utils.remove_trailing_slash(input_url) + @parameterized.expand( + [ + ("random_ones", b"\xff" * 10, "01234567-89ab-7fff-bfff-ffffffffffff"), + ("random_zeros", b"\x00" * 10, "01234567-89ab-7000-8000-000000000000"), + ( + "random_mixed", + bytes(range(0x10, 0x1A)), + "01234567-89ab-7011-9213-141516171819", + ), + ] + ) + def test_uuid7_layout(self, _name, random_bytes, expected): + # Bit 48 of the millisecond clock is set to check that it is masked off. + time_ns = ((1 << 48) + 0x0123456789AB) * 1_000_000 + 999_999 + with ( + mock.patch.object(utils.time, "time_ns", return_value=time_ns), + mock.patch.object( + utils.os, + "urandom", + side_effect=lambda n: (random_bytes + b"\xaa" * n)[:n], + ), + ): + value = utils._uuid7() + self.assertEqual(str(value), expected) + self.assertEqual(value.version, 7) + def test_clean_pydantic(self): class ModelV2(BaseModel): foo: str diff --git a/posthog/utils.py b/posthog/utils.py index 05608983f..840478411 100644 --- a/posthog/utils.py +++ b/posthog/utils.py @@ -1,6 +1,7 @@ import json import logging import numbers +import os import re import threading import time @@ -51,6 +52,21 @@ def guess_timezone(dt: datetime) -> datetime: return dt.astimezone(timezone.utc) +def _uuid7() -> UUID: + """Return a time-ordered version 7 UUID (RFC 9562). + + Capture generates v7 server-side when a client omits the uuid, so + SDK-generated event uuids sort the same way. + """ + unix_ms = time.time_ns() // 1_000_000 + # Python 3.10 requires the byte order; from 3.11 "big" is the default. + random_bits = int.from_bytes(os.urandom(10), "big") # pragma: no mutate + value = ((unix_ms & ((1 << 48) - 1)) << 80) | random_bits + value = (value & ~(0xF << 76)) | (0x7 << 76) + value = (value & ~(0x3 << 62)) | (0x2 << 62) + return UUID(int=value) + + def _normalize_timestamp(timestamp: Union[datetime, str]) -> str: """Normalize a datetime or parseable ISO datetime string to UTC.""" parsed_timestamp: datetime diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index c6165d97b..0cb0bc898 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -692,8 +692,12 @@ attribute posthog.capture_compression.CaptureCompression.ZSTD = 'zstd' attribute posthog.capture_exception_code_variables = False attribute posthog.capture_send.CaptureError.attempts = attempts attribute posthog.capture_send.CaptureError.drops = drops or [] +attribute posthog.capture_send.CaptureError.endpoint = endpoint +attribute posthog.capture_send.CaptureError.event_results = event_results or {} attribute posthog.capture_send.CaptureError.request_id = request_id attribute posthog.capture_send.CaptureError.retry_exhausted = retry_exhausted or [] +attribute posthog.capture_send.CaptureEventResult.details: Optional[str] = None +attribute posthog.capture_send.CaptureEventResult.result: Optional[str] attribute posthog.capture_trace_context = False attribute posthog.client.Client.api_key = (project_api_key or '').strip() attribute posthog.client.Client.capture_compression = _resolve_capture_compression(capture_compression) @@ -1164,9 +1168,10 @@ class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] class posthog.async_client.AsyncPosthog class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, refill_rate: Number, refill_interval_seconds: Number, on_bucket_rate_limited: Optional[Callable[[Hashable], None]] = None, clock: Callable[[], float] = time.monotonic) class posthog.capture_compression.CaptureCompression -class posthog.capture_send.CaptureError(status: int | str, message: str, *, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None) +class posthog.capture_send.CaptureError(status: int | str, message: str, *, endpoint: str, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None, event_results: Optional[dict[str, CaptureEventResult]] = None) +class posthog.capture_send.CaptureEventResult(result: Optional[str], details: Optional[str] = None) class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) -class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, retries=10, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=CaptureCompression.NONE) +class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, retries=3, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) class posthog.exception_utils.AnnotatedValue(value, metadata) @@ -1559,6 +1564,7 @@ method posthog.async_client.AsyncClient.set_once(**kwargs: Unpack[OptionalSetArg method posthog.async_client.AsyncClient.shutdown() -> None method posthog.bucketed_rate_limiter.BucketedRateLimiter.consume_rate_limit(key: Hashable) -> bool method posthog.bucketed_rate_limiter.BucketedRateLimiter.stop() -> None +method posthog.capture_send.CaptureError.verdict_summary() -> str method posthog.client.Client.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None) -> Optional[str] method posthog.client.Client.capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.client.Client.capture_ai(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] From a15eff43a73e1ec17f538e0dcbe95d1ef7be7ae2 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Sat, 10 Oct 2026 16:34:47 -0700 Subject: [PATCH 06/14] feat: AI lane config and size-aware batching (#1022) * feat: AI lane config and size-aware batching * chore: tighten AI lane comments * fix: reject non-finite AI lane config values --- AGENTS.md | 2 +- posthog/capture_compression.py | 22 +++- posthog/client.py | 142 +++++++++++++++++++++--- posthog/consumer.py | 36 ++++-- posthog/test/test_ai_capture_lane.py | 160 ++++++++++++++++++++++++--- posthog/test/test_consumer.py | 21 ++++ references/public_api_snapshot.txt | 7 +- 7 files changed, 345 insertions(+), 45 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 6b7a8d74f..4d67e9efc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -23,7 +23,7 @@ Follow [Public API changes](./CONTRIBUTING.md#public-api-changes). As an agent, Before changing capture configuration, serialization, routing, or retries, read the relevant implementation and tests. -Capture v1 is the only capture protocol (`capture` posts to `/i/v1/analytics/events`, `capture_ai` to `/i/v1/ai/events`); strictly typed v1 options and `$set`/`$set_once` relocation; compression (gzip, zlib-wrapped deflate, optional zstd, default none); partial-only per-event retries with stable identity; accumulated drop reporting even on 2xx; terminal v1 `429`; `Retry-After` as a minimum bounded by the shared 30s ceiling; and inline blocking retries with `sync_mode=True`. +Capture v1 is the only capture protocol (`capture` posts to `/i/v1/analytics/events`, `capture_ai` to `/i/v1/ai/events`); strictly typed v1 options and `$set`/`$set_once` relocation; compression (gzip, zlib-wrapped deflate, optional zstd, default none), set per lane by `capture_compression` and `capture_ai_compression`; partial-only per-event retries with stable identity; accumulated drop reporting even on 2xx; terminal v1 `429`; `Retry-After` as a minimum bounded by the shared 30s ceiling; and inline blocking retries with `sync_mode=True`. ## Mirror and build safety diff --git a/posthog/capture_compression.py b/posthog/capture_compression.py index 52a1d4c2d..257c80c88 100644 --- a/posthog/capture_compression.py +++ b/posthog/capture_compression.py @@ -54,6 +54,7 @@ def _zstd_available() -> bool: def _coerce_explicit( value: Union[CaptureCompression, str], + name: str = "capture_compression", ) -> CaptureCompression: """Normalize an explicitly-supplied compression to a ``CaptureCompression``. @@ -68,7 +69,7 @@ def _coerce_explicit( if resolved is not None: return resolved raise ValueError( - f"invalid capture_compression {value!r}; expected a CaptureCompression " + f"invalid {name} {value!r}; expected a CaptureCompression " f"or one of {sorted(_ALIASES)}" ) @@ -122,3 +123,22 @@ def _resolve_capture_compression( ) return fallback return env_resolved + + +def _resolve_capture_ai_compression( + capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, +) -> CaptureCompression: + """Resolve the AI lane's request-body compression. + + Explicit argument only, defaulting to ``NONE``. ``POSTHOG_CAPTURE_COMPRESSION`` + does not apply, so changing analytics compression never changes AI uploads. + """ + if capture_ai_compression is None: + return CaptureCompression.NONE + resolved = _coerce_explicit(capture_ai_compression, "capture_ai_compression") + if resolved is CaptureCompression.ZSTD and not _zstd_available(): + raise ValueError( + "capture_ai_compression 'zstd' requires the zstandard package; " + "install posthog[zstd]" + ) + return resolved diff --git a/posthog/client.py b/posthog/client.py index a5e23fe2d..a0af85319 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -2,7 +2,9 @@ import hashlib as _hashlib import inspect import json +from contextlib import contextmanager import logging +import math import os import sys import threading @@ -29,6 +31,7 @@ from posthog.tracing._span import inert_span as _inert_span from posthog.capture_compression import ( CaptureCompression, + _resolve_capture_ai_compression, _resolve_capture_compression, ) from posthog.capture_send import ( @@ -91,6 +94,7 @@ from posthog.request import ( USER_AGENT as _USER_AGENT, APIError, + DatetimeSerializer as _DatetimeSerializer, QuotaLimitError, RequestsConnectionError, RequestsTimeout, @@ -227,6 +231,28 @@ def _get_atexit_deadline() -> float: return _atexit_deadline +def _positive_config_value( + name: str, value, *, integer: bool = False, maximum: Optional[int] = None +): + """Return ``value`` if it is finite, positive and no larger than ``maximum``. + + Bad lane config is a programming error, so it raises instead of falling + back to a default. + """ + allowed = (int,) if integer else (int, float) + if ( + isinstance(value, bool) + or not isinstance(value, allowed) + or not math.isfinite(value) + or value <= 0 + ): + kind = "integer" if integer else "number" + raise ValueError(f"{name} must be a positive {kind}, got {value!r}") + if maximum is not None and value > maximum: + raise ValueError(f"{name} must be at most {maximum}, got {value!r}") + return value + + def get_identity_state(passed) -> tuple[str, bool]: """Returns the distinct id to use, and whether this is a personless event or not""" stringified = stringify_id(passed) @@ -708,6 +734,10 @@ def __init__( exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, + capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, + capture_ai_max_queue_size: int = 1000, + capture_ai_timeout: float = 30, + capture_ai_max_event_bytes: int = AI_MAX_MSG_SIZE, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, @@ -726,7 +756,8 @@ def __init__( the corresponding ingestion host. debug: Enable verbose SDK logging and re-raise errors from public API methods. - max_queue_size: Maximum number of events buffered before upload. + max_queue_size: Maximum number of analytics events buffered before + upload. AI events use ``capture_ai_max_queue_size``. send: If False, queueing succeeds but events are not sent. on_error: Optional callback ``(error, batch)`` invoked when an upload fails: by background consumers, or on the calling thread in @@ -843,6 +874,21 @@ def __init__( strings ``"gzip"``/``"deflate"``). When omitted, the ``POSTHOG_CAPTURE_COMPRESSION`` env var is consulted, then no compression. + capture_ai_compression: Request-body compression for + ``capture_ai()`` uploads, set independently of + ``capture_compression``. Defaults to no compression, and the + env var does not apply. ``CaptureCompression.ZSTD`` suits large + AI payloads. + capture_ai_max_queue_size: Maximum number of AI events buffered + before upload. Defaults to 1000, lower than ``max_queue_size`` + because AI events are much larger. + capture_ai_timeout: Seconds allowed for one AI upload request. + Defaults to 30, longer than ``timeout`` because AI batches are + much larger. + capture_ai_max_event_bytes: Largest serialized AI event the SDK + sends; a larger one is dropped with an error log. Defaults to + the AI endpoint's ceiling plus envelope headroom, and may only + be lowered. Examples: ```python @@ -946,6 +992,21 @@ def __init__( self._library_version = VERSION self._sdk_info = f"{self._library_id}/{self._library_version}" self.capture_compression = _resolve_capture_compression(capture_compression) + self.capture_ai_compression = _resolve_capture_ai_compression( + capture_ai_compression + ) + capture_ai_max_queue_size = _positive_config_value( + "capture_ai_max_queue_size", capture_ai_max_queue_size, integer=True + ) + capture_ai_timeout = _positive_config_value( + "capture_ai_timeout", capture_ai_timeout + ) + capture_ai_max_event_bytes = _positive_config_value( + "capture_ai_max_event_bytes", + capture_ai_max_event_bytes, + integer=True, + maximum=AI_MAX_MSG_SIZE, + ) self.super_properties = super_properties # Release id from POSTHOG_RELEASE_ID, attached to every event. Resolved # here so the env var is read once per client. @@ -1055,34 +1116,34 @@ def __init__( api_key=self.api_key, host=self.host, on_error=on_error, - max_queue_size=max_queue_size, thread_count=thread, send=send, flush_at=flush_at, flush_interval=flush_interval, max_retries=self.max_retries, - timeout=timeout, historical_migration=historical_migration, sdk_info=self._sdk_info, ) self._analytics_lane = _Lane( name="analytics", **lane_defaults, + max_queue_size=max_queue_size, + timeout=timeout, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=self.capture_compression, eager_start=not sync_mode, ) - # The AI lane posts to its own endpoint so multi-MB AI events stay off - # the analytics endpoint's smaller caps. It sends uncompressed. Lazy - # start, so the many clients that never emit AI events pay for no extra - # threads. + # A separate endpoint keeps multi-MB AI events off the analytics caps. + # Lazy start, so clients that never send AI events pay for no threads. self._ai_lane = _Lane( name="ai", **lane_defaults, + max_queue_size=capture_ai_max_queue_size, + timeout=capture_ai_timeout, endpoint=_CAPTURE_AI_V1_PATH, - max_msg_size=AI_MAX_MSG_SIZE, - capture_compression=CaptureCompression.NONE, + max_msg_size=capture_ai_max_event_bytes, + capture_compression=self.capture_ai_compression, eager_start=False, ) self._lanes = [self._analytics_lane, self._ai_lane] @@ -2435,6 +2496,23 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): if self.sync_mode: self.log.debug("enqueued with blocking %s.", msg["event"]) + try: + event_size = len(json.dumps(msg, cls=_DatetimeSerializer).encode()) + except Exception: + self.log.error("Unable to serialize event for sizing, dropping.") + return None + if event_size > lane.max_msg_size: + # Log only name and size: AI events may carry unredacted + # multimodal payloads that must not leak into logs. + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", + msg["event"], + event_size, + lane.max_msg_size // 1024, + lane.endpoint, + ) + return None + def send_sync() -> None: # Sync mode bypasses the lane's queue but keeps its wire config, # so AI events still post to the AI endpoint. @@ -2443,7 +2521,7 @@ def send_sync() -> None: self.host, [msg], compression=lane.capture_compression, - timeout=self.timeout, + timeout=lane.timeout, max_retries=self.max_retries, historical_migration=self.historical_migration, sdk_info=self._sdk_info, @@ -2681,13 +2759,14 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: # Spans drain with events: serverless handlers call flush(), not # shutdown(), and leaving spans on their own timer would lose them. span_flush = self._start_span_flush(timeout_seconds) - if timeout_seconds is None: - for lane in self._lanes: - lane.flush(None) - else: - deadline = time.monotonic() + timeout_seconds - for lane in self._lanes: - lane.flush(max(0.0, deadline - time.monotonic())) + with self._drain_lanes_together(): + if timeout_seconds is None: + for lane in self._lanes: + lane.flush(None) + else: + deadline = time.monotonic() + timeout_seconds + for lane in self._lanes: + lane.flush(max(0.0, deadline - time.monotonic())) if span_flush is not None: # The last span request is bounded only by the request # timeout, so the wait is not. @@ -2824,7 +2903,36 @@ def _run_lifecycle_cleanup( self.log.exception(log_message) errors.append(error) + @contextmanager + def _drain_lanes_together(self): + """Signal every lane to drain before waiting on any of them. + + Lanes then drain in parallel under one budget. Otherwise a lane keeps + batching on its normal cadence while the client waits on the lane before it. + """ + signals: list[_DrainSignal] = [] + try: + for lane in self._lanes: + signal = lane._drain_signal + signal.request() + signals.append(signal) + yield + finally: + for signal in signals: + signal.complete() + def _flush_or_discard_queues(self, errors: list[Exception]) -> None: + try: + with self._drain_lanes_together(): + self._flush_or_discard_each_lane(errors) + return + except Exception as error: + self.log.exception("Failed to signal lane drains during lifecycle cleanup") + errors.append(error) + # Each lane's flush signals its own drain, so lanes still drain one by one. + self._flush_or_discard_each_lane(errors) + + def _flush_or_discard_each_lane(self, errors: list[Exception]) -> None: for lane in self._lanes: try: if any(consumer.is_alive() for consumer in lane.consumers): diff --git a/posthog/consumer.py b/posthog/consumer.py index 22944064d..46c8c6abb 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -21,13 +21,18 @@ MAX_MSG_SIZE = 900 * 1024 # 900KiB per event -# AI events carry LLM inputs/outputs and post to a dedicated endpoint whose -# pipeline accepts larger messages than analytics ingestion, so the AI lane -# grants a higher per-event ceiling. `next()` appends an item before checking -# BATCH_SIZE_LIMIT, so worst-case request body is BATCH_SIZE_LIMIT + -# AI_MAX_MSG_SIZE (~13MiB) — keep that sum under the 20MiB server body cap. -AI_MAX_MSG_SIZE = 8 * 1024 * 1024 # 8MiB per event - +# The AI endpoint's per-event ceiling. The endpoint applies it to the +# serialized properties alone. +AI_MAX_PROPERTIES_SIZE = 8 * 1024 * 1024 +# The local guard measures the whole event, not only its properties, so it +# allows this much more to keep events at the endpoint's ceiling. +AI_ENVELOPE_HEADROOM = 64 * 1024 +# The AI lane's per-event guard, and the upper bound for +# `capture_ai_max_event_bytes`. +AI_MAX_MSG_SIZE = AI_MAX_PROPERTIES_SIZE + AI_ENVELOPE_HEADROOM + +# A batch closes before it appends an event that would take it past this, so +# a request carries at most this much event data, or one larger event alone. # The maximum request body size is currently 20MiB, let's be conservative # in case we want to lower it in the future. BATCH_SIZE_LIMIT = 5 * 1024 * 1024 @@ -265,6 +270,11 @@ def next(self): queue.task_done() pending_items -= 1 continue + if items and total_size + item_size > BATCH_SIZE_LIMIT: + self._return_to_queue_head(item) + pending_items -= 1 + self.log.debug("hit batch size limit (size: %d)", total_size) + break items.append(item) total_size += item_size if total_size >= BATCH_SIZE_LIMIT: @@ -284,6 +294,18 @@ def next(self): return items + def _return_to_queue_head(self, item) -> None: + """Put a dequeued event back at the head of the queue for the next batch. + + The event stays counted in ``unfinished_tasks``, because it was never + marked done. Keeping it in the queue, not in the consumer, means a + stop, a discard or a fork accounts for it like any other queued event. + """ + queue = self.queue + with queue.not_empty: + queue.queue.appendleft(item) + queue.not_empty.notify() + def request(self, batch): """Upload the batch to this consumer's `endpoint` with the capture v1 partial-retry submitter.""" diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 568f00feb..2234de1b3 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -1,14 +1,17 @@ +import os import threading import unittest import uuid from unittest import mock +from parameterized import parameterized + import posthog from posthog.ai.utils import _capture_ai_event, finalize_ai_content, with_privacy_mode -from posthog.capture_compression import CaptureCompression +from posthog.capture_compression import CAPTURE_COMPRESSION_ENV_VAR, CaptureCompression from posthog.client import Client -from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE +from posthog.consumer import AI_MAX_MSG_SIZE, AI_MAX_PROPERTIES_SIZE, MAX_MSG_SIZE from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.version import VERSION from posthog.test.capture_helpers import patch_capture_send, sent_batch @@ -174,14 +177,45 @@ def test_ai_lane_accepts_multi_megabyte_events(self): batch = consumer.next() self.assertEqual([e["event"] for e in batch], ["$ai_generation"]) - def test_ai_lane_drops_events_over_its_cap(self): - client = self._client() + @parameterized.expand( + [ + ("properties_at_endpoint_ceiling", {}, AI_MAX_PROPERTIES_SIZE, True), + ("over_guard", {}, AI_MAX_MSG_SIZE, False), + ( + "over_lowered_cap", + {"capture_ai_max_event_bytes": 1024 * 1024}, + 2 * 1024 * 1024, + False, + ), + ] + ) + def test_ai_lane_size_guard(self, _name, config, payload_bytes, accepted): + client = Client(TEST_API_KEY, send=False, flush_interval=0.05, **config) client._ai_lane.start() consumer = client._ai_lane.consumers[0] - client._ai_lane.queue.put(self._sized_event("$ai_generation", AI_MAX_MSG_SIZE)) - self.assertEqual(consumer.next(), []) + client._ai_lane.queue.put(self._sized_event("$ai_generation", payload_bytes)) + self.assertEqual( + [e["event"] for e in consumer.next()], + ["$ai_generation"] if accepted else [], + ) self.assertTrue(client._ai_lane.queue.empty()) + def test_sync_mode_ai_event_over_cap_is_not_sent(self): + client = Client( + TEST_API_KEY, sync_mode=True, capture_ai_max_event_bytes=1024 * 1024 + ) + with patch_capture_send("client") as mock_send: + with self.assertLogs("posthog", level="ERROR") as logs: + result = client.capture_ai( + "$ai_generation", + distinct_id="d", + properties={"p": "x" * (2 * 1024 * 1024)}, + ) + + self.assertIsNone(result) + mock_send.assert_not_called() + self.assertIn("exceeds the 1024KiB limit", "\n".join(logs.output)) + def test_analytics_lane_rejects_events_over_900kib(self): client = self._client() consumer = client.consumers[0] @@ -191,18 +225,76 @@ def test_analytics_lane_rejects_events_over_900kib(self): class TestAiLaneWireConfig(unittest.TestCase): - """The AI lane posts to the AI endpoint uncompressed, whatever the - analytics `capture_compression`.""" + """The AI lane has its own endpoint, compression, timeout, queue and size + guard, independent of the analytics lane's settings.""" - def test_ai_lane_consumers_use_ai_endpoint_without_compression(self): - client = Client(TEST_API_KEY, send=False, capture_compression="gzip", thread=2) + @parameterized.expand( + [ + ("defaults", {}, CaptureCompression.NONE, 30, 1000, AI_MAX_MSG_SIZE), + ( + "configured", + { + "capture_ai_compression": "zstd", + "capture_ai_timeout": 45, + "capture_ai_max_queue_size": 50, + "capture_ai_max_event_bytes": 1024 * 1024, + }, + CaptureCompression.ZSTD, + 45, + 50, + 1024 * 1024, + ), + ] + ) + def test_ai_lane_consumers_use_ai_config( + self, _name, config, compression, timeout, queue_size, max_event_bytes + ): + with mock.patch.dict(os.environ, {CAPTURE_COMPRESSION_ENV_VAR: "deflate"}): + client = Client( + TEST_API_KEY, + send=False, + capture_compression="gzip", + timeout=9, + max_queue_size=77, + thread=2, + **config, + ) client._ai_lane.start() + self.assertEqual(client._ai_lane.queue.maxsize, queue_size) self.assertEqual(len(client._ai_lane.consumers), 2) for consumer in client._ai_lane.consumers: self.assertIs(consumer.queue, client._ai_lane.queue) self.assertEqual(consumer.endpoint, _CAPTURE_AI_V1_PATH) - self.assertEqual(consumer.max_msg_size, AI_MAX_MSG_SIZE) - self.assertEqual(consumer.capture_compression, CaptureCompression.NONE) + self.assertEqual(consumer.max_msg_size, max_event_bytes) + self.assertEqual(consumer.capture_compression, compression) + self.assertEqual(consumer.timeout, timeout) + self.assertEqual(client.queue.maxsize, 77) + for consumer in client._analytics_lane.consumers: + self.assertEqual(consumer.capture_compression, CaptureCompression.GZIP) + self.assertEqual(consumer.timeout, 9) + client.join() + + @parameterized.expand( + [ + ( + "event_bytes_over_ceiling", + "capture_ai_max_event_bytes", + AI_MAX_MSG_SIZE + 1, + ), + ("event_bytes_zero", "capture_ai_max_event_bytes", 0), + ("event_bytes_float", "capture_ai_max_event_bytes", 1024.5), + ("event_bytes_bool", "capture_ai_max_event_bytes", True), + ("queue_size_negative", "capture_ai_max_queue_size", -1), + ("queue_size_string", "capture_ai_max_queue_size", "100"), + ("timeout_zero", "capture_ai_timeout", 0), + ("timeout_nan", "capture_ai_timeout", float("nan")), + ("timeout_inf", "capture_ai_timeout", float("inf")), + ("compression_unknown", "capture_ai_compression", "br"), + ] + ) + def test_invalid_ai_config_raises(self, _name, field, value): + with self.assertRaisesRegex(ValueError, field): + Client(TEST_API_KEY, send=False, **{field: value}) def test_async_lanes_keep_separate_path_and_compression(self): client = Client(TEST_API_KEY, capture_compression="gzip", flush_interval=0.05) @@ -228,23 +320,57 @@ def test_async_lanes_keep_separate_path_and_compression(self): ) client.join() - def test_sync_lanes_keep_separate_path_and_compression(self): - client = Client(TEST_API_KEY, sync_mode=True, capture_compression="gzip") + def test_sync_lanes_keep_separate_path_compression_and_timeout(self): + client = Client( + TEST_API_KEY, + sync_mode=True, + capture_compression="gzip", + capture_ai_compression="deflate", + timeout=9, + capture_ai_timeout=45, + ) with patch_capture_send("client") as mock_send: client.capture_ai("$ai_generation", distinct_id="d") client.capture("button_clicked", distinct_id="d") self.assertEqual( [ - (call.kwargs["path"], call.kwargs["compression"]) + ( + call.kwargs["path"], + call.kwargs["compression"], + call.kwargs["timeout"], + ) for call in mock_send.call_args_list ], [ - (_CAPTURE_AI_V1_PATH, CaptureCompression.NONE), - (_CAPTURE_V1_PATH, CaptureCompression.GZIP), + (_CAPTURE_AI_V1_PATH, CaptureCompression.DEFLATE, 45), + (_CAPTURE_V1_PATH, CaptureCompression.GZIP, 9), ], ) + def test_flush_drains_ai_lane_while_waiting_on_analytics(self): + release_analytics = threading.Event() + ai_sent = threading.Event() + + def send(api_key, host, batch, **kwargs): + if kwargs["path"] == _CAPTURE_AI_V1_PATH: + ai_sent.set() + else: + release_analytics.wait(5) + + client = Client(TEST_API_KEY, flush_interval=30) + with patch_capture_send("consumer", side_effect=send): + client.capture("button_clicked", distinct_id="d") + client.capture_ai("$ai_generation", distinct_id="d") + flusher = threading.Thread(target=client.flush, args=(10,)) + flusher.start() + try: + self.assertTrue(ai_sent.wait(2)) + finally: + release_analytics.set() + flusher.join(5) + client.join() + class TestAiLaneLazyStart(unittest.TestCase): def test_no_ai_consumers_until_first_capture_ai(self): diff --git a/posthog/test/test_consumer.py b/posthog/test/test_consumer.py index ca4748829..f4cf3188d 100644 --- a/posthog/test/test_consumer.py +++ b/posthog/test/test_consumer.py @@ -204,6 +204,27 @@ def test_max_msg_size_param_raises_per_event_ceiling(self) -> None: q.put(big_msg) self.assertEqual(consumer.next(), [big_msg]) + @parameterized.expand( + [ + # A small event serializes to 18 bytes, the "big" one to 81. + ("closes_before_overflow", 43, [0, 1, 2, 3], [[0, 1], [2, 3]]), + ("event_over_limit_goes_alone", 30, [0, "big", 1], [[0], ["big"], [1]]), + ] + ) + def test_batch_byte_limit_is_checked_before_appending( + self, _name, limit, ids, expected + ) -> None: + q = Queue() + consumer = Consumer(q, "", flush_at=10, flush_interval=0.01) + for i in ids: + q.put({"m": "x" * (60 if i == "big" else 1), "i": i}) + + with mock.patch("posthog.consumer.BATCH_SIZE_LIMIT", limit): + batches = [[e["i"] for e in consumer.next()] for _ in expected] + + self.assertEqual(batches, expected) + self.assertEqual(q.unfinished_tasks, len(ids)) + def test_upload(self) -> None: q = Queue() consumer = Consumer(q, TEST_API_KEY, flush_at=1) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 0cb0bc898..282645df7 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -700,6 +700,7 @@ attribute posthog.capture_send.CaptureEventResult.details: Optional[str] = None attribute posthog.capture_send.CaptureEventResult.result: Optional[str] attribute posthog.capture_trace_context = False attribute posthog.client.Client.api_key = (project_api_key or '').strip() +attribute posthog.client.Client.capture_ai_compression = _resolve_capture_ai_compression(capture_ai_compression) attribute posthog.client.Client.capture_compression = _resolve_capture_compression(capture_compression) attribute posthog.client.Client.capture_exception_code_variables = capture_exception_code_variables attribute posthog.client.Client.capture_trace_context = capture_trace_context @@ -755,7 +756,9 @@ attribute posthog.code_variables_detect_secrets = DEFAULT_CODE_VARIABLES_DETECT_ attribute posthog.code_variables_ignore_patterns = DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS attribute posthog.code_variables_mask_patterns = DEFAULT_CODE_VARIABLES_MASK_PATTERNS attribute posthog.code_variables_mask_url_credentials = DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS -attribute posthog.consumer.AI_MAX_MSG_SIZE = 8 * 1024 * 1024 +attribute posthog.consumer.AI_ENVELOPE_HEADROOM = 64 * 1024 +attribute posthog.consumer.AI_MAX_MSG_SIZE = AI_MAX_PROPERTIES_SIZE + AI_ENVELOPE_HEADROOM +attribute posthog.consumer.AI_MAX_PROPERTIES_SIZE = 8 * 1024 * 1024 attribute posthog.consumer.BATCH_SIZE_LIMIT = 5 * 1024 * 1024 attribute posthog.consumer.Consumer.api_key = api_key attribute posthog.consumer.Consumer.capture_compression = capture_compression @@ -1170,7 +1173,7 @@ class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, ref class posthog.capture_compression.CaptureCompression class posthog.capture_send.CaptureError(status: int | str, message: str, *, endpoint: str, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None, event_results: Optional[dict[str, CaptureEventResult]] = None) class posthog.capture_send.CaptureEventResult(result: Optional[str], details: Optional[str] = None) -class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) +class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_max_queue_size: int = 1000, capture_ai_timeout: float = 30, capture_ai_max_event_bytes: int = AI_MAX_MSG_SIZE, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, retries=3, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) From 252886c0ea4ebc61aef8647ef3c8a12d4b10da17 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:42:19 -0700 Subject: [PATCH 07/14] test(mcp): fix get_tool checks for fastmcp 4.0.11 (#1050) test(mcp): drop get_tool checks that fastmcp 4.0.11 invalidated fastmcp 4.0.11 resolves middleware-injected tools through the server's tool lookup (PrefectHQ/fastmcp#5431), so get_tool() now returns them. The two setup assertions that expected None fail on any fresh install. The SDK behavior the tests check is unchanged and passes on 4.0.10 and 4.0.11. (cherry picked from commit 897cd43d322156d9903d4950ee4caa06fcaa7d3f) Co-authored-by: Claude Opus 5.5 --- posthog/test/mcp/test_fastmcp_v4.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/posthog/test/mcp/test_fastmcp_v4.py b/posthog/test/mcp/test_fastmcp_v4.py index ac29c3a9f..1d5afe45c 100644 --- a/posthog/test/mcp/test_fastmcp_v4.py +++ b/posthog/test/mcp/test_fastmcp_v4.py @@ -152,7 +152,6 @@ def echo(context: str, conversation_id: str, llm_model: str) -> str: tool = Tool.from_function(echo) if source == "middleware": child.add_middleware(ToolInjectionMiddleware(tools=[tool])) - assert await child.get_tool("echo") is None else: child.add_tool(tool) server = FastMCP("example-app") if source == "mounted" else child @@ -225,7 +224,6 @@ def echo(text: str) -> str: server = FastMCP("example-middleware-tools") server.add_middleware(ToolInjectionMiddleware(tools=[Tool.from_function(echo)])) - assert await server.get_tool("echo") is None sink = FakeClient() instrument( server, From 305e74e756e45dbb88cb9f90115492d4b6654874 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:45:58 -0700 Subject: [PATCH 08/14] feat!: add per-event capture options (#1023) feat!: per-event capture options and legacy property hoisting --- posthog/__init__.py | 8 +- posthog/args.py | 6 + posthog/async_client.py | 13 +- posthog/capture_event.py | 108 +++++++++------- posthog/client.py | 38 ++++-- posthog/test/test_capture_event.py | 120 +++++++++--------- posthog/test/test_event_options.py | 111 ++++++++++++++++ .../test_feature_flag_called_minimization.py | 29 +++++ posthog/test/test_module.py | 2 + references/public_api_snapshot.txt | 15 ++- 10 files changed, 318 insertions(+), 132 deletions(-) create mode 100644 posthog/test/test_event_options.py diff --git a/posthog/__init__.py b/posthog/__init__.py index 71e4493ba..f7ed8e522 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -638,6 +638,7 @@ def group_identify( uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None, + options: Optional[Dict[str, Any]] = None, ) -> Optional[str]: """ Set properties on a group. @@ -653,6 +654,7 @@ def group_identify( uuid: Optional UUID for the event disable_geoip: Whether to disable GeoIP lookup distinct_id: Optional distinct ID of the user performing the action + options: Optional capture options for the event, sent as given Examples: ```python @@ -676,6 +678,7 @@ def group_identify( uuid=uuid, disable_geoip=disable_geoip, distinct_id=distinct_id, + options=options, ) @@ -685,6 +688,7 @@ def alias( timestamp: Optional[Union[datetime.datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, + options: Optional[Dict[str, Any]] = None, ) -> Optional[str]: """ Associate user behaviour before and after they e.g. register, login, or perform some other identifying action. @@ -696,6 +700,7 @@ def alias( datetimes and parseable ISO timestamp strings are converted to UTC. uuid: Optional UUID for the event disable_geoip: Whether to disable GeoIP lookup + options: Optional capture options for the event, sent as given Details: To marry up whatever a user does before they sign up or log in with what they do after you need to make an alias call. This will allow you to answer questions like "Which marketing channels leads to users churning after a month?" or "What do users do on our website before signing up?". Particularly useful for associating user behaviour before and after they e.g. register, login, or perform some other identifying action. @@ -717,6 +722,7 @@ def alias( timestamp=timestamp, uuid=uuid, disable_geoip=disable_geoip, + options=options, ) @@ -730,7 +736,7 @@ def capture_exception( Args: exception: The exception to capture. If not provided, the current exception is captured via `sys.exc_info()` **kwargs: Optional capture arguments including distinct_id, properties, - timestamp, uuid, groups, flags, send_feature_flags, and disable_geoip. + timestamp, uuid, groups, flags, send_feature_flags, disable_geoip, and options. Details: Capture exception is idempotent - if it is called twice with the same exception instance, only a occurrence will be tracked in posthog. This is because, generally, contexts will cause exceptions to be captured automatically. However, to ensure you track an exception, if you catch and do not re-raise it, capturing it manually is recommended, unless you are certain it will have crossed a context boundary (e.g. by existing a `with posthog.new_context():` block already). If the passed exception was raised and caught, the captured stack trace will consist of every frame between where the exception was raised and the point at which it is captured (the "traceback"). If the passed exception was never raised, e.g. if you call `posthog.capture_exception(ValueError("Some Error"))`, the stack trace captured will be the full stack trace at the moment the exception was captured. Note that heavy use of contexts will lead to truncated stack traces, as the exception will be captured by the context entered most recently, which may not be the point you catch the exception for the final time in your code. It's recommended to use contexts sparingly, for this reason. `capture_exception` takes the same set of optional arguments as `capture`. diff --git a/posthog/args.py b/posthog/args.py index 42083b923..1b2aa4118 100644 --- a/posthog/args.py +++ b/posthog/args.py @@ -48,12 +48,16 @@ class OptionalCaptureArgs(TypedDict): hidden ``/flags`` request on capture and may return different values than the ones the code branched on. disable_geoip: Whether to disable GeoIP lookup for this event. Defaults to False. + options: Capture options for this event, such as ``{"process_person_profile": False}``. + Sent as given, for PostHog to validate. An option wins over its legacy ``$`` property, + such as ``$process_person_profile``. A value that is not a dict is logged and ignored. """ distinct_id: NotRequired[Optional[ID_TYPES]] properties: NotRequired[Optional[Dict[str, Any]]] timestamp: NotRequired[Optional[Union[datetime, str]]] uuid: NotRequired[Optional[Union[str, UUID]]] + options: NotRequired[Optional[Dict[str, Any]]] groups: NotRequired[Optional[Dict[str, str]]] flags: NotRequired[Optional["FeatureFlagEvaluations"]] send_feature_flags: NotRequired[ @@ -83,6 +87,7 @@ class OptionalSetArgs(TypedDict): it must be a valid UUID string or uuid.UUID instance; invalid values are ignored and replaced with a newly generated UUID. disable_geoip: Whether to disable GeoIP lookup for this operation. Defaults to False. + options: Capture options for this event, sent as given. See ``OptionalCaptureArgs``. """ distinct_id: NotRequired[Optional[ID_TYPES]] @@ -90,6 +95,7 @@ class OptionalSetArgs(TypedDict): timestamp: NotRequired[Optional[Union[datetime, str]]] uuid: NotRequired[Optional[Union[str, UUID]]] disable_geoip: NotRequired[Optional[bool]] + options: NotRequired[Optional[Dict[str, Any]]] ExcInfo = Union[ diff --git a/posthog/async_client.py b/posthog/async_client.py index 37dfd2294..5597d79c2 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -35,6 +35,7 @@ CaptureCompression, _resolve_capture_compression, ) +from .capture_event import _canonical_event_uuid, _event_options from .capture_send import _CAPTURE_V1_PATH from .client import ( MAX_DICT_SIZE as _MAX_DICT_SIZE, @@ -378,9 +379,8 @@ def enqueue_on_bound_loop() -> None: def _normalize_uuid(self, msg: dict[str, Any]) -> str: raw_uuid = msg.pop("uuid", None) if raw_uuid is not None: - try: - normalized = str(UUID(str(raw_uuid))) - except (TypeError, ValueError, AttributeError): + normalized = _canonical_event_uuid(raw_uuid) + if normalized is None: self.log.error( "Invalid UUID %r. Falling back to a generated UUID.", raw_uuid ) @@ -493,6 +493,7 @@ def _build_capture_event( "distinct_id": distinct_id, "event": event, "uuid": kwargs.get("uuid"), + "options": _event_options(kwargs.get("options")), }, kwargs.get("disable_geoip"), kwargs.get("_property_allowlist"), @@ -611,6 +612,7 @@ def _build_person_properties_event( property_key: properties, "event": event, "uuid": kwargs.get("uuid"), + "options": _event_options(kwargs.get("options")), } def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: @@ -646,6 +648,7 @@ def group_identify( uuid: Optional[Union[str, UUID]] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None, + options: Optional[Dict[str, Any]] = None, ) -> Optional[str]: try: if not _stringify_id(group_type): @@ -670,6 +673,7 @@ def group_identify( "distinct_id": resolved_distinct_id, "timestamp": timestamp, "uuid": uuid, + "options": _event_options(options), } session_id = _get_context_session_id() if session_id: @@ -688,6 +692,7 @@ def alias( timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, + options: Optional[Dict[str, Any]] = None, ) -> Optional[str]: try: resolved_previous_id = _stringify_id(previous_id) @@ -707,6 +712,7 @@ def alias( "event": "$create_alias", "distinct_id": resolved_previous_id, "uuid": uuid, + "options": _event_options(options), } session_id = _get_context_session_id() if session_id: @@ -809,6 +815,7 @@ def capture_exception( groups=kwargs.get("groups"), flags=kwargs.get("flags"), disable_geoip=kwargs.get("disable_geoip"), + options=kwargs.get("options"), ) if exception is not None and result is not None: mark_exception_as_captured(exception, result) diff --git a/posthog/capture_event.py b/posthog/capture_event.py index bc50a4628..5afb31563 100644 --- a/posthog/capture_event.py +++ b/posthog/capture_event.py @@ -7,9 +7,10 @@ the legacy queued-message shape in a few load-bearing ways that this module encodes: -- A typed ``options`` object carries a handful of sentinel properties, renamed - and strictly typed. Wrong JSON types fail deserialization of the *whole - batch*, so values are coerced to native types or omitted entirely. +- An ``options`` object carries per-event processing options. Options the + caller sets are sent as given, for PostHog to validate. Four legacy ``$`` + properties fill the matching option when the caller left it unset, and are + always removed from ``properties``. - ``$set``/``$set_once`` have no top-level form in v1; the server reads them from ``properties``. The legacy ``set()``/``set_once()`` builders emit them at the top level, so they are relocated into ``properties`` here. @@ -17,12 +18,16 @@ ``PostHog-Sdk-Info`` header and are stripped from v1 properties. """ -from collections.abc import Callable +import logging +import re from datetime import datetime, timezone from typing import Any, Optional +from uuid import UUID from posthog.utils import _normalize_timestamp +log = logging.getLogger("posthog") + # Sentinel properties lifted to top-level string fields on the event. _TOPLEVEL_SENTINELS: tuple[tuple[str, str], ...] = ( ("$session_id", "session_id"), @@ -35,45 +40,49 @@ # Properties dropped from v1 events (server injects them from PostHog-Sdk-Info). _STRIP_FROM_PROPERTIES = ("$lib", "$lib_version") +# Legacy properties and the option each one fills. The order matches posthog-rs +# and posthog-go. +_LEGACY_OPTION_PROPERTIES: tuple[tuple[str, str], ...] = ( + ("$cookieless_mode", "cookieless_mode"), + ("$ignore_sent_at", "disable_skew_correction"), + ("$product_tour_id", "product_tour_id"), + ("$process_person_profile", "process_person_profile"), +) + +# The uuid forms Go's uuid.Validate accepts. Python's UUID() also accepts +# misplaced hyphens and a bare "uuid:" prefix, which other SDKs reject. +_EVENT_UUID_PATTERN = re.compile( + r"(?:urn:uuid:)?[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}" + r"|\{[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\}" + r"|[0-9a-f]{32}", + re.IGNORECASE, +) + -def _coerce_bool(value: Any) -> Optional[bool]: - """Coerce a sentinel value to ``bool`` using the backend's truthiness rules. +def _canonical_event_uuid(value: Any) -> Optional[str]: + """Return the canonical form of a caller's event uuid, or None if invalid. - Native bool passes through; ``"true"``/``"1"`` and ``"false"``/``"0"`` - (case-insensitive, trimmed) map to the obvious bool; any other numeric value - is nonzero-truthy. Anything else returns ``None`` so the option is omitted - rather than sent with a type the strict v1 schema would reject. + Capture keys per-event results by the canonical lowercase hyphenated form, + so a uuid sent in any other form would never match its result. """ - if isinstance(value, bool): - return value - if isinstance(value, str): - normalized = value.strip().lower() - if normalized in ("true", "1"): - return True - if normalized in ("false", "0"): - return False + if isinstance(value, UUID): + return str(value) + if not isinstance(value, str) or not _EVENT_UUID_PATTERN.fullmatch(value): return None - if isinstance(value, (int, float)): - return value != 0 - return None - - -def _coerce_str(value: Any) -> Optional[str]: - """Accept only ``str`` (the backend's ``product_tour_id`` is ``Option``).""" - return value if isinstance(value, str) else None - - -# Sentinel properties lifted into the typed `options` object: legacy property -# key, the backend's field name, and the coercer enforcing its strict type -# (wrong JSON types fail deserialization of the whole batch, so a value that -# won't coerce is omitted). The coercer is stored directly to keep the dispatch -# type-checked rather than keyed by a stringly-typed name. -_OPTION_SENTINELS: tuple[tuple[str, str, Callable[[Any], Any]], ...] = ( - ("$cookieless_mode", "cookieless_mode", _coerce_bool), - ("$ignore_sent_at", "disable_skew_correction", _coerce_bool), - ("$product_tour_id", "product_tour_id", _coerce_str), - ("$process_person_profile", "process_person_profile", _coerce_bool), -) + return str(UUID(value.lower())) + + +def _event_options(value: Any) -> dict[str, Any]: + """Return a copy of a caller's ``options``, or ``{}`` when it is not a dict.""" + if value is None: + return {} + if not isinstance(value, dict): + log.error( + "options must be a dict, got %s. Sending the event without them.", + type(value).__name__, + ) + return {} + return dict(value) def _v1_timestamp(timestamp: Any) -> str: @@ -114,23 +123,24 @@ def _to_v1_event(msg: dict) -> dict: for key in _STRIP_FROM_PROPERTIES: properties.pop(key, None) - options: dict[str, Any] = {} - for prop_key, wire_key, coercer in _OPTION_SENTINELS: + options = _event_options(msg.get("options")) + for prop_key, option_key in _LEGACY_OPTION_PROPERTIES: if prop_key not in properties: continue - # Always removed from properties — these sentinels must never reach v1 - # backend properties — but only emitted as an option when coercible. - coerced = coercer(properties.pop(prop_key)) - if coerced is not None: - options[wire_key] = coerced + legacy = properties.pop(prop_key) + # A null option counts as unset, so the legacy value fills it. + if options.get(option_key) is None: + options[option_key] = legacy top_level: dict[str, str] = {} for prop_key, field_name in _TOPLEVEL_SENTINELS: if prop_key not in properties: continue - coerced_str = _coerce_str(properties.pop(prop_key)) - if coerced_str is not None: - top_level[field_name] = coerced_str + # Always removed. A non-string value would fail the whole batch, so it + # is dropped. + value = properties.pop(prop_key) + if isinstance(value, str): + top_level[field_name] = value event = { "event": msg["event"], diff --git a/posthog/client.py b/posthog/client.py index a0af85319..2825f4bfe 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -34,6 +34,7 @@ _resolve_capture_ai_compression, _resolve_capture_compression, ) +from posthog.capture_event import _canonical_event_uuid, _event_options from posthog.capture_send import ( _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, @@ -267,22 +268,12 @@ def get_identity_state(passed) -> tuple[str, bool]: def _stringify_event_uuid(value) -> str: - if isinstance(value, UUID): - return str(value) - - stringified = stringify_id(value) - if not stringified: + canonical = _canonical_event_uuid(value) + if canonical is None: raise ValueError( f"Invalid event uuid {value!r}. Expected a valid UUID string or uuid.UUID instance." ) - - try: - # Canonical form, because capture keys per-event results by it. - return str(UUID(stringified)) - except ValueError: - raise ValueError( - f"Invalid event uuid {value!r}. Expected a valid UUID string or uuid.UUID instance." - ) from None + return canonical def add_context_tags(properties): @@ -352,8 +343,14 @@ def wrapper(self, *args, **kwargs): # Correctness-required "$groups", "$process_person_profile", + # Option sources: legacy properties move into options after this + # allowlist runs, so dropping them would drop the option. + "$cookieless_mode", + "$ignore_sent_at", + "$product_tour_id", # Linkage / SDK identity "$session_id", + "$window_id", "$lib", "$lib_version", "$is_server", @@ -1636,6 +1633,10 @@ def capture( evaluate_flags(). When truthy, evaluates flags during capture and attaches them to the event. disable_geoip: Whether to disable GeoIP for this event. + options: Capture options for this event, such as + ``{"process_person_profile": False}``. Sent as given, for + PostHog to validate. An option wins over its legacy ``$`` + property, such as ``$process_person_profile``. Examples: ```python @@ -1710,6 +1711,7 @@ def _capture( flags_snapshot = kwargs.get("flags", None) send_feature_flags = kwargs.get("send_feature_flags", False) disable_geoip = kwargs.get("disable_geoip", None) + options = _event_options(kwargs.get("options", None)) # Internal, set for minimal $feature_flag_called events: a strict allowlist # applied to the fully-enriched properties dict just before enqueueing. property_allowlist = kwargs.get("_property_allowlist", None) @@ -1733,6 +1735,7 @@ def _capture( "distinct_id": distinct_id, "event": event, "uuid": uuid, + "options": options, } if groups: @@ -1918,6 +1921,7 @@ def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: "$set": properties, "event": "$set", "uuid": uuid, + "options": _event_options(kwargs.get("options", None)), } return self._enqueue(msg, disable_geoip) @@ -1967,6 +1971,7 @@ def set_once(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: "$set_once": properties, "event": "$set_once", "uuid": uuid, + "options": _event_options(kwargs.get("options", None)), } return self._enqueue(msg, disable_geoip) @@ -1981,6 +1986,7 @@ def group_identify( uuid: Optional[Union[str, UUID]] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None, + options: Optional[Dict[str, Any]] = None, ) -> Optional[str]: """ Identify a group and set its properties. @@ -1998,6 +2004,7 @@ def group_identify( ignored and replaced with a newly generated UUID. disable_geoip: Whether to disable GeoIP for this event. distinct_id: The distinct ID of the user performing the action. + options: Capture options for this event, sent as given. Examples: ```python @@ -2039,6 +2046,7 @@ def group_identify( "distinct_id": distinct_id, "timestamp": timestamp, "uuid": uuid, + "options": _event_options(options), } # NOTE - group_identify doesn't generally use context properties - should it? @@ -2055,6 +2063,7 @@ def alias( timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, + options: Optional[Dict[str, Any]] = None, ) -> Optional[str]: """ Create an alias between two distinct IDs. @@ -2071,6 +2080,7 @@ def alias( valid UUID string or uuid.UUID instance; invalid values are ignored and replaced with a newly generated UUID. disable_geoip: Whether to disable GeoIP for this event. + options: Capture options for this event, sent as given. Examples: ```python @@ -2107,6 +2117,7 @@ def alias( "event": "$create_alias", "distinct_id": previous_id, "uuid": uuid, + "options": _event_options(options), } if get_context_session_id(): @@ -2251,6 +2262,7 @@ def capture_exception( flags=flags_snapshot, send_feature_flags=send_feature_flags, disable_geoip=disable_geoip, + options=kwargs.get("options", None), ) # Mark the exception as captured to prevent duplicate captures diff --git a/posthog/test/test_capture_event.py b/posthog/test/test_capture_event.py index e3b48195c..1e17e76a6 100644 --- a/posthog/test/test_capture_event.py +++ b/posthog/test/test_capture_event.py @@ -3,10 +3,11 @@ from parameterized import parameterized +from uuid import UUID + from posthog.capture_event import ( _build_v1_batch_body, - _coerce_bool, - _coerce_str, + _canonical_event_uuid, _to_v1_event, ) @@ -27,42 +28,37 @@ def _legacy_msg(event="my_event", properties=None, **overrides) -> dict: return msg -class TestCoercion(unittest.TestCase): +_CANONICAL_UUID = "0190a8f3-1b2c-7d4e-8f90-123456789abc" + + +class TestCanonicalEventUuid(unittest.TestCase): @parameterized.expand( [ - ("bool_true", True, True), - ("bool_false", False, False), - ("str_true", "true", True), - ("str_true_upper", "TRUE", True), - ("str_true_padded", " true ", True), - ("str_one", "1", True), - ("str_false", "false", False), - ("str_zero", "0", False), - ("int_nonzero", 5, True), - ("int_zero", 0, False), - ("float_nonzero", 1.5, True), - ("float_zero", 0.0, False), - ("neg_int", -1, True), - ("str_yes_uncoercible", "yes", None), - ("str_empty_uncoercible", "", None), - ("none_uncoercible", None, None), - ("dict_uncoercible", {"a": 1}, None), + ("canonical", _CANONICAL_UUID), + ("uppercase", _CANONICAL_UUID.upper()), + ("unhyphenated", _CANONICAL_UUID.replace("-", "")), + ("braced", "{" + _CANONICAL_UUID + "}"), + ("urn", "urn:uuid:" + _CANONICAL_UUID), + ("urn_uppercase", "URN:UUID:" + _CANONICAL_UUID.upper()), + ("uuid_instance", UUID(_CANONICAL_UUID)), ] ) - def test_coerce_bool(self, _name, value, expected) -> None: - self.assertIs(_coerce_bool(value), expected) + def test_accepted_forms_are_canonicalized(self, _name, value) -> None: + self.assertEqual(_canonical_event_uuid(value), _CANONICAL_UUID) @parameterized.expand( [ - ("str", "tour-1", "tour-1"), - ("empty_str", "", ""), - ("int", 123, None), - ("bool", True, None), - ("none", None, None), + ("bare_uuid_prefix", "uuid:" + _CANONICAL_UUID), + ("misplaced_hyphen", "0190a8f31-b2c-7d4e-8f90-123456789abc"), + ("surrounding_whitespace", " " + _CANONICAL_UUID), + ("braced_unhyphenated", "{" + _CANONICAL_UUID.replace("-", "") + "}"), + ("too_short", _CANONICAL_UUID[:-1]), + ("empty", ""), + ("int", 123), ] ) - def test_coerce_str(self, _name, value, expected) -> None: - self.assertEqual(_coerce_str(value), expected) + def test_other_forms_are_rejected(self, _name, value) -> None: + self.assertIsNone(_canonical_event_uuid(value)) class TestToV1Event(unittest.TestCase): @@ -94,60 +90,64 @@ def test_does_not_leak_non_wire_top_level_keys(self) -> None: def test_does_not_mutate_input(self) -> None: msg = _legacy_msg( properties={"$cookieless_mode": True, "$session_id": "s-1"}, + options={"product_tour_id": "tour-1"}, **{"$set": {"name": "Max"}}, ) original_properties = dict(msg["properties"]) _to_v1_event(msg) self.assertEqual(msg["properties"], original_properties) + self.assertEqual(msg["options"], {"product_tour_id": "tour-1"}) self.assertIn("$set", msg) # top-level $set untouched on the original @parameterized.expand( [ - ("cookieless_mode", "$cookieless_mode", "cookieless_mode", True, True), - ( - "ignore_sent_at_rename", - "$ignore_sent_at", - "disable_skew_correction", - "true", - True, - ), + ("cookieless_mode", "$cookieless_mode", "cookieless_mode", True), + ("ignore_sent_at", "$ignore_sent_at", "disable_skew_correction", "true"), ( "process_person_profile", "$process_person_profile", "process_person_profile", - "false", - False, - ), - ( - "product_tour_id", - "$product_tour_id", - "product_tour_id", - "tour-7", - "tour-7", + 0, ), + ("product_tour_id", "$product_tour_id", "product_tour_id", 123), ] ) - def test_option_sentinels_lifted_renamed_and_coerced( - self, _name, prop_key, wire_key, raw, expected + def test_legacy_property_fills_option_unchanged( + self, _name, prop_key, option_key, raw ) -> None: event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) - self.assertEqual(event["options"], {wire_key: expected}) + self.assertEqual(event["options"], {option_key: raw}) self.assertNotIn(prop_key, event["properties"]) @parameterized.expand( [ - ("bad_bool", "$cookieless_mode", "maybe"), - ("bad_tour_id_int", "$product_tour_id", 123), + ("option_set", {"cookieless_mode": False}, False), + ("option_null", {"cookieless_mode": None}, True), + ("option_missing", {}, True), ] ) - def test_option_sentinel_removed_but_omitted_on_bad_coercion( - self, _name, prop_key, raw + def test_caller_option_wins_over_legacy_property( + self, _name, options, expected ) -> None: - event = _to_v1_event(_legacy_msg(properties={prop_key: raw})) - # Removed from properties (sentinels must never reach v1 props) but not - # emitted as an option, so a wrong type cannot 400 the whole batch. - self.assertNotIn(prop_key, event["properties"]) - self.assertEqual(event["options"], {}) + event = _to_v1_event( + _legacy_msg(properties={"$cookieless_mode": True}, options=options) + ) + self.assertEqual(event["options"], {"cookieless_mode": expected}) + self.assertNotIn("$cookieless_mode", event["properties"]) + + def test_caller_options_pass_through_unchanged(self) -> None: + options = {"process_person_profile": "false", "future_option": {"a": [1]}} + event = _to_v1_event(_legacy_msg(options=options)) + self.assertEqual(event["options"], options) + + @parameterized.expand([("list", ["x"]), ("string", "x"), ("int", 1)]) + def test_non_dict_options_are_logged_and_ignored(self, _name, options) -> None: + with self.assertLogs("posthog", level="ERROR") as logs: + event = _to_v1_event( + _legacy_msg(properties={"$cookieless_mode": True}, options=options) + ) + self.assertEqual(event["options"], {"cookieless_mode": True}) + self.assertIn("options must be a dict", logs.output[0]) @parameterized.expand( [ @@ -170,9 +170,9 @@ def test_all_sentinels_together(self) -> None: _legacy_msg( properties={ "$cookieless_mode": True, - "$ignore_sent_at": "1", + "$ignore_sent_at": True, "$product_tour_id": "tour-x", - "$process_person_profile": 0, + "$process_person_profile": False, "$session_id": "s-1", "$window_id": "w-1", "$geoip_disable": True, diff --git a/posthog/test/test_event_options.py b/posthog/test/test_event_options.py new file mode 100644 index 000000000..33e6c3a50 --- /dev/null +++ b/posthog/test/test_event_options.py @@ -0,0 +1,111 @@ +import logging + +import pytest + +from posthog import AsyncPosthog +from posthog.capture_event import _to_v1_event +from posthog.client import Client +from posthog.test.capture_helpers import ( + patch_async_capture_send, + patch_capture_send, + sent_events, +) +from posthog.test.test_utils import FAKE_TEST_API_KEY + +OPTIONS = {"cookieless_mode": True, "future_option": "kept"} + +CAPTURE_CALLS = { + "capture": lambda c, o: c.capture("e", distinct_id="u", options=o), + "capture_ai": lambda c, o: c.capture_ai( + "$ai_generation", distinct_id="u", options=o + ), + "set": lambda c, o: c.set(distinct_id="u", properties={"a": 1}, options=o), + "set_once": lambda c, o: c.set_once( + distinct_id="u", properties={"a": 1}, options=o + ), + "group_identify": lambda c, o: c.group_identify( + "company", "acme", distinct_id="u", options=o + ), + "alias": lambda c, o: c.alias("previous", "u", options=o), + "capture_exception": lambda c, o: c.capture_exception( + ValueError("boom"), distinct_id="u", options=o + ), +} +ASYNC_CAPTURE_CALLS = {k: v for k, v in CAPTURE_CALLS.items() if k != "capture_ai"} + + +def _sync_wire_events(call, before_send=None) -> list[dict]: + with patch_capture_send("client") as send: + client = Client(FAKE_TEST_API_KEY, sync_mode=True, before_send=before_send) + assert call(client) is not None + return [_to_v1_event(msg) for msg in sent_events(send)] + + +async def _async_wire_events(call, before_send=None) -> list[dict]: + batches: list[list[dict]] = [] + + async def send_batch(api_key, host, batch, **kwargs): + batches.append(batch) + + with patch_async_capture_send(side_effect=send_batch): + async with AsyncPosthog("test-key", before_send=before_send) as client: + assert call(client) is not None + await client.flush(timeout_seconds=1) + return [_to_v1_event(msg) for batch in batches for msg in batch] + + +@pytest.mark.parametrize("method", list(CAPTURE_CALLS)) +def test_sync_options_reach_the_wire(method): + events = _sync_wire_events(lambda c: CAPTURE_CALLS[method](c, OPTIONS)) + assert [e["options"] for e in events] == [OPTIONS] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", list(ASYNC_CAPTURE_CALLS)) +async def test_async_options_reach_the_wire(method): + events = await _async_wire_events(lambda c: ASYNC_CAPTURE_CALLS[method](c, OPTIONS)) + assert [e["options"] for e in events] == [OPTIONS] + + +def _hook(msg): + properties = {k: v for k, v in msg["properties"].items() if k != "$product_tour_id"} + return {**msg, "options": {"cookieless_mode": False}, "properties": properties} + + +def _hooked_capture(client): + return client.capture( + "e", + distinct_id="u", + properties={"$product_tour_id": "tour-1"}, + options={"cookieless_mode": True}, + ) + + +def test_sync_before_send_changes_options_and_legacy_properties(): + events = _sync_wire_events(_hooked_capture, before_send=_hook) + assert events[0]["options"] == {"cookieless_mode": False} + + +@pytest.mark.asyncio +async def test_async_before_send_changes_options_and_legacy_properties(): + events = await _async_wire_events(_hooked_capture, before_send=_hook) + assert events[0]["options"] == {"cookieless_mode": False} + + +def test_sync_non_dict_options_are_logged_and_event_is_sent(caplog): + with caplog.at_level(logging.ERROR, logger="posthog"): + events = _sync_wire_events( + lambda c: c.capture("e", distinct_id="u", options=["cookieless_mode"]) + ) + assert events[0]["options"] == {} + assert "options must be a dict" in caplog.text + + +@pytest.mark.asyncio +async def test_async_non_dict_options_are_logged_and_event_is_sent(caplog): + with caplog.at_level(logging.ERROR, logger="posthog"): + events = await _async_wire_events( + lambda c: c.capture("e", distinct_id="u", options=["cookieless_mode"]) + ) + assert events[0]["options"] == {} + assert "options must be a dict" in caplog.text diff --git a/posthog/test/test_feature_flag_called_minimization.py b/posthog/test/test_feature_flag_called_minimization.py index c32b4ea27..b355a68d6 100644 --- a/posthog/test/test_feature_flag_called_minimization.py +++ b/posthog/test/test_feature_flag_called_minimization.py @@ -13,6 +13,7 @@ from parameterized import parameterized +from posthog.capture_event import _to_v1_event from posthog.client import _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, Client from posthog.request import GetResponse from posthog.test.test_utils import FAKE_TEST_API_KEY @@ -153,6 +154,34 @@ def test_gated_non_experiment_flag_with_groups_keeps_group_context( self.assertEqual(properties["$groups"], {"organization": "org-1"}) self.assertLessEqual(set(properties), _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES) + @mock.patch("posthog.client.flags") + def test_gated_event_keeps_option_sources(self, patch_flags): + patch_flags.return_value = _flags_response(has_experiment=False, gate=True) + client, captured = self._make_client() + client.super_properties = { + "$cookieless_mode": True, + "$ignore_sent_at": True, + "$product_tour_id": "tour-1", + "$window_id": "w-1", + "app_version": "1.2.3", + } + + client.get_feature_flag_result("person-flag", "some-distinct-id") + + event = _to_v1_event( + next(m for m in captured if m["event"] == "$feature_flag_called") + ) + self.assertEqual( + event["options"], + { + "cookieless_mode": True, + "disable_skew_correction": True, + "product_tour_id": "tour-1", + }, + ) + self.assertEqual(event["window_id"], "w-1") + self.assertNotIn("app_version", event["properties"]) + @mock.patch("posthog.client.flags") def test_remote_evaluation_uses_this_responses_gate_not_a_concurrent_update( self, patch_flags diff --git a/posthog/test/test_module.py b/posthog/test/test_module.py index 6ebd8fe4e..655b1d2cc 100644 --- a/posthog/test/test_module.py +++ b/posthog/test/test_module.py @@ -169,6 +169,7 @@ def test_group_identify_propagates_distinct_id(self): "company_123", {"name": "Awesome Inc."}, distinct_id="user_456", + options={"cookieless_mode": True}, ) self.mock_client.group_identify.assert_called_once_with( group_type="company", @@ -178,6 +179,7 @@ def test_group_identify_propagates_distinct_id(self): uuid=None, disable_geoip=None, distinct_id="user_456", + options={"cookieless_mode": True}, ) def test_group_identify_distinct_id_defaults_to_none(self): diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 282645df7..07db58cf9 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -641,12 +641,14 @@ attribute posthog.args.OptionalCaptureArgs.disable_geoip: NotRequired[Optional[b attribute posthog.args.OptionalCaptureArgs.distinct_id: NotRequired[Optional[ID_TYPES]] attribute posthog.args.OptionalCaptureArgs.flags: NotRequired[Optional[FeatureFlagEvaluations]] attribute posthog.args.OptionalCaptureArgs.groups: NotRequired[Optional[Dict[str, str]]] +attribute posthog.args.OptionalCaptureArgs.options: NotRequired[Optional[Dict[str, Any]]] attribute posthog.args.OptionalCaptureArgs.properties: NotRequired[Optional[Dict[str, Any]]] attribute posthog.args.OptionalCaptureArgs.send_feature_flags: NotRequired[Optional[Union[bool, SendFeatureFlagsOptions]]] attribute posthog.args.OptionalCaptureArgs.timestamp: NotRequired[Optional[Union[datetime, str]]] attribute posthog.args.OptionalCaptureArgs.uuid: NotRequired[Optional[Union[str, UUID]]] attribute posthog.args.OptionalSetArgs.disable_geoip: NotRequired[Optional[bool]] attribute posthog.args.OptionalSetArgs.distinct_id: NotRequired[Optional[ID_TYPES]] +attribute posthog.args.OptionalSetArgs.options: NotRequired[Optional[Dict[str, Any]]] attribute posthog.args.OptionalSetArgs.properties: NotRequired[Optional[Dict[str, Any]]] attribute posthog.args.OptionalSetArgs.timestamp: NotRequired[Optional[Union[datetime, str]]] attribute posthog.args.OptionalSetArgs.uuid: NotRequired[Optional[Union[str, UUID]]] @@ -689,6 +691,7 @@ attribute posthog.capture_compression.CaptureCompression.DEFLATE = 'deflate' attribute posthog.capture_compression.CaptureCompression.GZIP = 'gzip' attribute posthog.capture_compression.CaptureCompression.NONE = 'none' attribute posthog.capture_compression.CaptureCompression.ZSTD = 'zstd' +attribute posthog.capture_event.log = logging.getLogger('posthog') attribute posthog.capture_exception_code_variables = False attribute posthog.capture_send.CaptureError.attempts = attempts attribute posthog.capture_send.CaptureError.drops = drops or [] @@ -1301,7 +1304,7 @@ function posthog.ai.utils.merge_system_prompt(kwargs: Dict[str, Any], provider: function posthog.ai.utils.merge_usage_stats(target: TokenUsage, source: TokenUsage, mode: str = 'incremental') -> None function posthog.ai.utils.serialize_raw_usage(raw_usage: Any) -> Optional[Dict[str, Any]] function posthog.ai.utils.with_privacy_mode(ph_client: PostHogClient, privacy_mode: bool, value: Any) -function posthog.alias(previous_id: ID_TYPES, distinct_id: str, timestamp: Optional[Union[datetime.datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None) -> Optional[str] +function posthog.alias(previous_id: ID_TYPES, distinct_id: str, timestamp: Optional[Union[datetime.datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] function posthog.capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] function posthog.capture_ai(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] function posthog.capture_exception(exception: Optional[ExceptionArg] = None, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] @@ -1388,7 +1391,7 @@ function posthog.get_feature_flag_payload(key: str, distinct_id: ID_TYPES, match function posthog.get_feature_flag_result(key: str, distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, send_feature_flag_events: bool = True, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None) -> Optional[FeatureFlagResult] function posthog.get_remote_config_payload(key: str) function posthog.get_tags() -> Dict[str, Any] -function posthog.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime.datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None) -> Optional[str] +function posthog.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime.datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] function posthog.identify_context(distinct_id: str) function posthog.integrations.django.markcoroutinefunction(func) function posthog.join() -> None @@ -1553,14 +1556,14 @@ method posthog.ai.prompts.Prompts.get(name: str, *, with_metadata: Optional[bool method posthog.ai.prompts.Prompts.get_all(*, label: Optional[str] = None) -> Dict[str, PromptResult] method posthog.ai.stream.AsyncStreamWrapper.aclose() -> None method posthog.ai.stream.AsyncStreamWrapper.close() -> None -method posthog.async_client.AsyncClient.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None) -> Optional[str] +method posthog.async_client.AsyncClient.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] method posthog.async_client.AsyncClient.capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.async_client.AsyncClient.capture_exception(exception: Optional[ExceptionArg] = None, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.async_client.AsyncClient.capture_immediate(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.async_client.AsyncClient.evaluate_flags(distinct_id: Optional[ID_TYPES] = None, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, disable_geoip: Optional[bool] = None, flag_keys: Optional[list[str]] = None, device_id: Optional[str] = None) -> FeatureFlagEvaluations method posthog.async_client.AsyncClient.flush(timeout_seconds: Optional[float] = 10) -> None method posthog.async_client.AsyncClient.get_remote_config_payload(key: str) -> Optional[Any] -method posthog.async_client.AsyncClient.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[Union[str, UUID]] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None) -> Optional[str] +method posthog.async_client.AsyncClient.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[Union[str, UUID]] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] method posthog.async_client.AsyncClient.join() -> None method posthog.async_client.AsyncClient.set(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] method posthog.async_client.AsyncClient.set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] @@ -1568,7 +1571,7 @@ method posthog.async_client.AsyncClient.shutdown() -> None method posthog.bucketed_rate_limiter.BucketedRateLimiter.consume_rate_limit(key: Hashable) -> bool method posthog.bucketed_rate_limiter.BucketedRateLimiter.stop() -> None method posthog.capture_send.CaptureError.verdict_summary() -> str -method posthog.client.Client.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None) -> Optional[str] +method posthog.client.Client.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] method posthog.client.Client.capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.client.Client.capture_ai(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.client.Client.capture_exception(exception: Optional[ExceptionArg], **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] @@ -1590,7 +1593,7 @@ method posthog.client.Client.get_feature_variants(distinct_id: ID_TYPES, groups: method posthog.client.Client.get_flags_decision(distinct_id: Optional[ID_TYPES] = None, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, disable_geoip: Optional[bool] = None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> FlagsResponse method posthog.client.Client.get_remote_config_payload(key: str) method posthog.client.Client.get_tags() -> Dict[str, Any] -method posthog.client.Client.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[Union[str, UUID]] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None) -> Optional[str] +method posthog.client.Client.group_identify(group_type: str, group_key: str, properties: Optional[Dict[str, Any]] = None, timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[Union[str, UUID]] = None, disable_geoip: Optional[bool] = None, distinct_id: Optional[ID_TYPES] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] method posthog.client.Client.identify_context(distinct_id: str) -> None method posthog.client.Client.join() -> None method posthog.client.Client.load_feature_flags() From 4cc19c173dea471e279e4763bd5525cd666c700e Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:46:16 -0700 Subject: [PATCH 09/14] feat!: add global and context option layers (#1024) * feat!: super_options, context options and personless as the lowest option layer * docs: state the option layer rules in capture docstrings * feat!: fill context and super values after before_send Context tags, context options, super_properties, super_options, the derived personless option and the environment $release_id now fill only keys an event leaves unset, after before_send runs. before_send sees only the event's own values and SDK enrichment, and its changes beat every default. A None option counts as unset and is filled. Legacy properties still fill only unset options and are always removed. * feat!: fill context and super values before before_send Context tags, context options, super_properties, super_options, the derived personless option and the environment $release_id now fill before before_send runs, still only into keys the event leaves unset. before_send sees every value and has the final say, including removing a super property. $set, $set_once, $groups and $group_set fill one level deep when both values are dicts. The set() and set_once() values now win key by key over a $set or $set_once in properties, as ingestion merges them. Hoisting still runs once, after the hook. * feat!: SDK values are defaults that caller values override $is_server, $geoip_disable and system context now fill last, only keys the event, context tags and super_properties left unset, before before_send. A super property $geoip_disable: False now wins over disable_geoip=True. The groups argument merges into a $groups property key by key and wins, in Client, AsyncClient and posthog.mcp events. * fix: copy async default layers at admission; typed groups win on MCP exception events AsyncPosthog filled defaults in the worker from references to the live super_properties and options, so a change after capture could reach a queued event. Each layer is now shallow-copied when the event is admitted, as 7.x merged super_properties into a new dict at capture. The MCP exception event now merges typed groups after custom properties, like the primary event, so both events of a failing call get the same groups. Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- posthog/__init__.py | 63 ++- posthog/_async_consumer.py | 13 +- posthog/async_client.py | 132 +++++-- posthog/capture_event.py | 113 +++++- posthog/client.py | 178 +++++++-- posthog/contexts.py | 56 +++ posthog/mcp/_posthog_events.py | 9 +- posthog/test/capture_helpers.py | 14 + posthog/test/mcp/test_pipeline.py | 25 +- posthog/test/test_async_consumer.py | 2 +- posthog/test/test_capture_event.py | 3 +- posthog/test/test_client.py | 18 +- posthog/test/test_contexts.py | 22 ++ posthog/test/test_event_options.py | 374 +++++++++++++++++- .../test_feature_flag_called_minimization.py | 17 +- posthog/test/test_module.py | 14 +- posthog/test/test_release_id.py | 18 +- references/public_api_snapshot.txt | 18 +- 18 files changed, 942 insertions(+), 147 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index f7ed8e522..d228d811e 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -52,6 +52,12 @@ from posthog.contexts import ( get_tags as inner_get_tags, ) +from posthog.contexts import ( + set_context_option as inner_set_context_option, +) +from posthog.contexts import ( + get_context_options as inner_get_context_options, +) from posthog.exception_utils import ( DEFAULT_CODE_VARIABLES_DETECT_SECRETS, DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS, @@ -299,6 +305,44 @@ def get_tags() -> Dict[str, Any]: return inner_get_tags() +def set_context_option(key: str, value: Any) -> None: + """ + Set a capture option for every event captured within the current context. + + Context options fill options an event leaves unset, before ``before_send`` + runs, so the hook sees them. They override ``super_options``. An event's + own ``options`` override them, and ``before_send`` can change them. + + Args: + key: The option name, such as ``"process_person_profile"`` + value: The option value, sent as given + + Examples: + ```python + from posthog import new_context, set_context_option + with new_context(): + set_context_option("process_person_profile", False) + ``` + + Category: + Contexts + """ + return inner_set_context_option(key, value) + + +def get_context_options() -> Dict[str, Any]: + """ + Get all capture options from the current context. + + Returns: + Dict of all capture options in the current context + + Category: + Contexts + """ + return inner_get_context_options() + + """Settings. These module-level settings configure the legacy global PostHog client used by @@ -340,7 +384,18 @@ def get_tags() -> Dict[str, Any]: feature_flags_request_max_retries: Number of retries for feature flag requests after network, transport, or timeout failures. Defaults to 1. Set to 0 to disable retries. - super_properties: Properties merged into every captured event. + super_properties: Properties for every captured event. They fill only + keys the event leaves unset, before ``before_send`` runs. ``$set``, + ``$set_once``, ``$groups`` and ``$group_set`` fill one level deep. An + event's own properties and context tags override them, and + ``before_send`` can change or remove them. + super_options: Capture options for every captured event, such as + ``{"cookieless_mode": True}``. They fill only options the event leaves + unset, before ``before_send`` runs. An event's own ``options`` and + context options override them, and ``before_send`` can change them. They also + win over an event's legacy property for the same key, such as + ``$cookieless_mode``, so pass per-event overrides of that key as + ``options``. metrics: Config dict for the ``client.metrics`` API (``service_name``, ``service_version``, ``environment``, ``flush_interval``, ...). Applied when ``setup()`` builds the global client, or on a later ``setup()`` @@ -364,7 +419,9 @@ def get_tags() -> Dict[str, Any]: project_root: Root path used to determine in-app exception stack frames. privacy_mode: Capture AI usage metadata without prompt inputs or outputs. before_send: Optional callback that can modify or drop events before upload. - Return ``None`` to drop an event. + Return ``None`` to drop an event. Context tags, context options, + ``super_properties`` and ``super_options`` fill in before it runs, so + it can change or remove them. enable_local_evaluation: Whether to poll feature flag definitions for local evaluation when a personal API key is configured. flag_definition_cache_provider: Optional external cache provider for sharing @@ -414,6 +471,7 @@ def get_tags() -> Dict[str, Any]: feature_flags_request_timeout_seconds = 3 # type: int feature_flags_request_max_retries = 1 # type: int super_properties = None # type: Optional[Dict] +super_options = None # type: Optional[Dict] metrics = None # type: Optional[Dict] traces = None # type: Optional[Dict] enable_exception_autocapture = False # type: bool @@ -1406,6 +1464,7 @@ def setup() -> Client: feature_flags_request_timeout_seconds=feature_flags_request_timeout_seconds, feature_flags_request_max_retries=feature_flags_request_max_retries, super_properties=super_properties, + super_options=super_options, metrics=metrics, traces=traces, # TODO: Currently this monitoring begins only when the Client is initialised (which happens when you do something with the SDK) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 0ae74b995..ce0b209c2 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -11,6 +11,7 @@ from ._async_request import async_send_v1_batch from .capture_compression import CaptureCompression +from .capture_event import _EventDefaults from .capture_send import _CAPTURE_V1_PATH, _capture_loss_message from .consumer import BATCH_SIZE_LIMIT, MAX_MSG_SIZE from .request import DatetimeSerializer @@ -37,6 +38,7 @@ async def _run_outside_processing_event(awaitable): class _QueuedEvent: event: dict[str, Any] context: contextvars.Context + defaults: Optional[_EventDefaults] = None async def _invoke_callback(callback, *args): @@ -92,7 +94,10 @@ def __init__( *, host: Optional[str], on_error: Optional[Callable[[Exception, list[dict[str, Any]]], Any]], - process_event: Callable[[dict[str, Any]], Awaitable[Optional[dict[str, Any]]]], + process_event: Callable[ + [dict[str, Any], Optional[_EventDefaults]], + Awaitable[Optional[dict[str, Any]]], + ], flush_at: int, flush_interval: float, retries: int, @@ -158,11 +163,11 @@ async def _get_or_flush(self, timeout: float) -> tuple[Any, bool]: return None, False async def _process_queued_event( - self, event: dict[str, Any] + self, queued: _QueuedEvent ) -> Optional[dict[str, Any]]: token = _PROCESSING_EVENT.set(True) try: - return await self.process_event(event) + return await self.process_event(queued.event, queued.defaults) finally: _PROCESSING_EVENT.reset(token) @@ -207,7 +212,7 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: try: process_task = queued.context.run( - asyncio.create_task, self._process_queued_event(queued.event) + asyncio.create_task, self._process_queued_event(queued) ) item = await process_task except Exception as error: diff --git a/posthog/async_client.py b/posthog/async_client.py index 5597d79c2..00732d399 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -9,6 +9,7 @@ import threading import warnings import weakref +from dataclasses import replace from datetime import datetime, timezone from typing import Any, Dict, Mapping, Optional, Union from uuid import UUID @@ -35,7 +36,14 @@ CaptureCompression, _resolve_capture_compression, ) -from .capture_event import _canonical_event_uuid, _event_options +from .capture_event import ( + _build_event_defaults, + _canonical_event_uuid, + _event_options, + _EventDefaults, + _fill_event_defaults, + _merge_groups, +) from .capture_send import _CAPTURE_V1_PATH from .client import ( MAX_DICT_SIZE as _MAX_DICT_SIZE, @@ -43,6 +51,9 @@ Client as _SyncClient, _metadata_has_experiment, _parse_flag_payload, + _add_context_session_id, + _context_tag_defaults, + _personless_options, add_context_tags as _add_context_tags, get_identity_state as _get_identity_state, stringify_id as _stringify_id, @@ -55,6 +66,7 @@ get_code_variables_mask_url_credentials_context, get_context_device_id as _get_context_device_id, get_context_distinct_id as _get_context_distinct_id, + get_context_options as _get_context_options, get_context_session_id as _get_context_session_id, ) from .exception_utils import ( @@ -120,6 +132,7 @@ def __init__( is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, + super_options: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, @@ -154,6 +167,7 @@ def __init__( self.is_server = is_server self.historical_migration = historical_migration self.super_properties = super_properties + self.super_options = super_options self._release_id = _resolve_release_id() self.capture_compression = _resolve_capture_compression(capture_compression) self.capture_trace_context = capture_trace_context @@ -335,10 +349,18 @@ def _ensure_workers_started(self) -> None: self._consumers.append(consumer) self._worker_tasks.append(asyncio.create_task(consumer.run())) - def _enqueue_prepared_event(self, prepared: dict[str, Any]) -> bool: + def _enqueue_prepared_event( + self, prepared: dict[str, Any], defaults: _EventDefaults + ) -> bool: if not self._accepting or self._closed: return False - queued_event = _QueuedEvent(prepared, contextvars.copy_context()) + # Shallow copy, as 7.x merged super_properties into a new dict at capture. + defaults = replace( + defaults, + property_layers=tuple(dict(layer) for layer in defaults.property_layers), + option_layers=tuple(dict(layer) for layer in defaults.option_layers), + ) + queued_event = _QueuedEvent(prepared, contextvars.copy_context(), defaults) try: running_loop = asyncio.get_running_loop() except RuntimeError: @@ -395,7 +417,6 @@ def _normalize_uuid(self, msg: dict[str, Any]) -> str: def _prepare_event( self, msg: dict[str, Any], - disable_geoip: Optional[bool], property_allowlist=None, ) -> tuple[Optional[dict[str, Any]], Optional[str]]: if self.disabled or not self._accepting: @@ -418,16 +439,7 @@ def _prepare_event( properties["$lib"] = "posthog-python" properties["$lib_version"] = VERSION - if disable_geoip is None: - disable_geoip = self.disable_geoip - if disable_geoip: - properties["$geoip_disable"] = True - if self.super_properties: - msg["properties"] = {**properties, **self.super_properties} - if self._release_id is not None: - msg["properties"].setdefault("$release_id", self._release_id) - if self.is_server: - msg["properties"]["$is_server"] = True + msg["options"] = msg.get("options") or {} if property_allowlist is not None: msg["properties"] = { key: value @@ -439,7 +451,13 @@ def _prepare_event( cleaned = clean(msg) return cleaned, sent_uuid - async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: + async def _process_event( + self, msg: dict[str, Any], defaults: Optional[_EventDefaults] + ) -> Optional[dict[str, Any]]: + _fill_event_defaults(msg, defaults) + return await self._run_before_send(msg) + + async def _run_before_send(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: if self.before_send is None: return msg @@ -458,21 +476,42 @@ async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: self.log.error("Error in before_send callback (%s)", type(error).__name__) return None + def _event_defaults( + self, + context_properties: Optional[dict[str, Any]] = None, + context_options: Optional[dict[str, Any]] = None, + derived_options: Optional[dict[str, Any]] = None, + property_allowlist=None, + disable_geoip: Optional[bool] = None, + system_properties: Optional[dict[str, Any]] = None, + ) -> _EventDefaults: + if disable_geoip is None: + disable_geoip = self.disable_geoip + return _build_event_defaults( + super_properties=self.super_properties, + super_options=self.super_options, + release_id=self._release_id, + context_properties=context_properties, + context_options=context_options, + derived_options=derived_options, + property_allowlist=property_allowlist, + is_server=self.is_server, + disable_geoip=disable_geoip, + system_properties=system_properties, + ) + def _build_capture_event( self, event: str, kwargs: OptionalCaptureArgs - ) -> tuple[dict[str, Any], Optional[bool], Any]: - properties = {**(kwargs.get("properties") or {}), **system_context()} + ) -> tuple[dict[str, Any], Any, _EventDefaults]: + properties = dict(kwargs.get("properties") or {}) if self.capture_trace_context: properties = {**_get_current_otel_span_properties(), **properties} - properties = _add_context_tags(properties) - assert properties is not None + properties = _add_context_session_id(properties) distinct_id, personless = _get_identity_state(kwargs.get("distinct_id")) - if personless and "$process_person_profile" not in properties: - properties["$process_person_profile"] = False groups = kwargs.get("groups") if groups: - properties["$groups"] = groups + _merge_groups(properties, groups) flags_snapshot = kwargs.get("flags") send_feature_flags = kwargs.get("send_feature_flags") @@ -495,8 +534,15 @@ def _build_capture_event( "uuid": kwargs.get("uuid"), "options": _event_options(kwargs.get("options")), }, - kwargs.get("disable_geoip"), kwargs.get("_property_allowlist"), + self._event_defaults( + context_properties=_context_tag_defaults(), + context_options=_get_context_options(), + derived_options=_personless_options(personless), + property_allowlist=kwargs.get("_property_allowlist"), + disable_geoip=kwargs.get("disable_geoip"), + system_properties=system_context(), + ), ) def capture( @@ -504,18 +550,14 @@ def capture( ) -> Optional[str]: """Queue an event without blocking for network delivery.""" try: - msg, disable_geoip, property_allowlist = self._build_capture_event( - event, kwargs - ) - prepared, sent_uuid = self._prepare_event( - msg, disable_geoip, property_allowlist - ) + msg, property_allowlist, defaults = self._build_capture_event(event, kwargs) + prepared, sent_uuid = self._prepare_event(msg, property_allowlist) if prepared is None or sent_uuid is None: return None if not self.send: return sent_uuid - if not self._enqueue_prepared_event(prepared): + if not self._enqueue_prepared_event(prepared, defaults): return None self.log.debug("queued async event %s", event) return sent_uuid @@ -551,15 +593,11 @@ async def capture_immediate( self._immediate_callers[current] = self._immediate_callers.get(current, 0) + 1 error_batch: list[dict[str, Any]] = [] try: - msg, disable_geoip, property_allowlist = self._build_capture_event( - event, kwargs - ) - prepared, sent_uuid = self._prepare_event( - msg, disable_geoip, property_allowlist - ) + msg, property_allowlist, defaults = self._build_capture_event(event, kwargs) + prepared, sent_uuid = self._prepare_event(msg, property_allowlist) if prepared is None or sent_uuid is None: return None - processed = await self._process_event(prepared) + processed = await self._process_event(prepared, defaults) if processed is None: return None error_batch = [processed] @@ -620,7 +658,9 @@ def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: msg = self._build_person_properties_event("$set", "$set", kwargs) if msg is None: return None - return self._enqueue_built_event(msg, kwargs.get("disable_geoip")) + return self._enqueue_built_event( + msg, kwargs.get("disable_geoip"), _get_context_options() + ) except Exception as error: if self.debug: raise @@ -632,7 +672,9 @@ def set_once(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: msg = self._build_person_properties_event("$set_once", "$set_once", kwargs) if msg is None: return None - return self._enqueue_built_event(msg, kwargs.get("disable_geoip")) + return self._enqueue_built_event( + msg, kwargs.get("disable_geoip"), _get_context_options() + ) except Exception as error: if self.debug: raise @@ -725,14 +767,20 @@ def alias( return None def _enqueue_built_event( - self, msg: dict[str, Any], disable_geoip: Optional[bool] + self, + msg: dict[str, Any], + disable_geoip: Optional[bool], + context_options: Optional[dict[str, Any]] = None, ) -> Optional[str]: - prepared, sent_uuid = self._prepare_event(msg, disable_geoip) + prepared, sent_uuid = self._prepare_event(msg) if prepared is None or sent_uuid is None: return None if not self.send: return sent_uuid - if not self._enqueue_prepared_event(prepared): + defaults = self._event_defaults( + context_options=context_options, disable_geoip=disable_geoip + ) + if not self._enqueue_prepared_event(prepared, defaults): return None return sent_uuid diff --git a/posthog/capture_event.py b/posthog/capture_event.py index 5afb31563..a7acabd7b 100644 --- a/posthog/capture_event.py +++ b/posthog/capture_event.py @@ -20,11 +20,14 @@ import logging import re +from collections.abc import Collection, Mapping +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any, Optional from uuid import UUID from posthog.utils import _normalize_timestamp +from posthog.utils import clean as _clean log = logging.getLogger("posthog") @@ -37,6 +40,9 @@ # Top-level legacy keys relocated into properties (v1 has no top-level form). _RELOCATE_TO_PROPERTIES = ("$set", "$set_once") +# Properties that defaults fill one level deep when both values are dicts. +_NESTED_FILL_PROPERTIES = frozenset({"$set", "$set_once", "$groups", "$group_set"}) + # Properties dropped from v1 events (server injects them from PostHog-Sdk-Info). _STRIP_FROM_PROPERTIES = ("$lib", "$lib_version") @@ -85,6 +91,105 @@ def _event_options(value: Any) -> dict[str, Any]: return dict(value) +@dataclass(frozen=True) +class _EventDefaults: + """Context, global and SDK-derived values for one event, highest layer first. + + They fill in before ``before_send``, so the hook sees them and can change + or remove them. The event's own values win over every default. + """ + + property_layers: tuple[Mapping[str, Any], ...] = () + option_layers: tuple[Mapping[str, Any], ...] = () + property_allowlist: Optional[Collection[str]] = None + + +def _build_event_defaults( + *, + super_properties: Optional[Mapping[str, Any]], + super_options: Any, + release_id: Optional[str], + context_properties: Optional[Mapping[str, Any]] = None, + context_options: Optional[Mapping[str, Any]] = None, + derived_options: Optional[Mapping[str, Any]] = None, + property_allowlist: Optional[Collection[str]] = None, + is_server: bool = False, + disable_geoip: bool = False, + system_properties: Optional[Mapping[str, Any]] = None, +) -> _EventDefaults: + """Order the layers: context, then global, then values the SDK derives.""" + # A value in the event, the context or super properties wins over every + # value the SDK adds, including `$is_server` and `$geoip_disable`. + sdk_properties = dict(system_properties or {}) + if release_id is not None: + sdk_properties["$release_id"] = release_id + if is_server: + sdk_properties["$is_server"] = True + if disable_geoip: + sdk_properties["$geoip_disable"] = True + return _EventDefaults( + property_layers=( + context_properties or {}, + super_properties or {}, + sdk_properties, + ), + option_layers=( + context_options or {}, + _event_options(super_options), + derived_options or {}, + ), + property_allowlist=property_allowlist, + ) + + +def _fill_event_defaults( + msg: dict[str, Any], defaults: Optional[_EventDefaults] +) -> None: + """Fill the keys an event left unset from ``defaults``, layer by layer. + + A property is unset only when its key is missing. An option is unset when + it is missing or ``None``, the same rule hoisting uses. ``$set``, + ``$set_once``, ``$groups`` and ``$group_set`` fill one level deep when the + event's value and the default are both dicts. + """ + if defaults is None: + return + allowlist = defaults.property_allowlist + properties = msg.get("properties") + if not isinstance(properties, dict): + properties = {} + msg["properties"] = properties + for property_layer in defaults.property_layers: + for key, value in property_layer.items(): + if allowlist is not None and key not in allowlist: + continue + if key not in properties: + properties[key] = _clean(value) + continue + existing = properties[key] + if ( + key in _NESTED_FILL_PROPERTIES + and isinstance(existing, dict) + and isinstance(value, Mapping) + ): + properties[key] = {**_clean(dict(value)), **existing} + options = _event_options(msg.get("options")) + for option_layer in defaults.option_layers: + for key, value in option_layer.items(): + if options.get(key) is None: + options[key] = _clean(value) + msg["options"] = options + + +def _merge_groups(properties: dict[str, Any], groups: Mapping[str, Any]) -> None: + """Merge typed ``groups`` into the ``$groups`` property; ``groups`` wins key by key.""" + existing = properties.get("$groups") + if isinstance(existing, dict): + properties["$groups"] = {**existing, **groups} + else: + properties["$groups"] = dict(groups) + + def _v1_timestamp(timestamp: Any) -> str: """Return a UTC RFC3339 timestamp string. @@ -108,16 +213,16 @@ def _to_v1_event(msg: dict) -> dict: properties = dict(msg.get("properties") or {}) # Relocate top-level $set/$set_once into properties; v1 has no top-level - # form. On the unusual collision where properties already carries the key, - # the properties value wins. + # form. The top-level value comes from the set() or set_once() call, so it + # wins key by key over a $set in properties, as ingestion merges them. for key in _RELOCATE_TO_PROPERTIES: top_val = msg.get(key) if top_val is None: continue existing = properties.get(key) if isinstance(top_val, dict) and isinstance(existing, dict): - properties[key] = {**top_val, **existing} - elif key not in properties: + properties[key] = {**existing, **top_val} + else: properties[key] = top_val for key in _STRIP_FROM_PROPERTIES: diff --git a/posthog/client.py b/posthog/client.py index 2825f4bfe..1ce5108ba 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -34,7 +34,13 @@ _resolve_capture_ai_compression, _resolve_capture_compression, ) -from posthog.capture_event import _canonical_event_uuid, _event_options +from posthog.capture_event import ( + _build_event_defaults, + _canonical_event_uuid, + _event_options, + _fill_event_defaults, + _merge_groups, +) from posthog.capture_send import ( _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH, @@ -52,11 +58,13 @@ get_context_device_id, get_context_distinct_id, get_context_session_id, + get_context_options as _context_get_context_options, get_tags as _context_get_tags, identify_context as _context_identify_context, _scoped as _context_scoped, new_context, set_context_device_id as _context_set_context_device_id, + set_context_option as _context_set_context_option, set_context_session as _context_set_context_session, tag as _context_tag, ) @@ -276,6 +284,11 @@ def _stringify_event_uuid(value) -> str: return canonical +def _personless_options(personless: bool) -> dict[str, Any]: + """The lowest option layer: no person profile for a generated distinct ID.""" + return {"process_person_profile": False} if personless else {} + + def add_context_tags(properties): properties = properties or {} current_context = _get_current_context() @@ -286,12 +299,24 @@ def add_context_tags(properties): context_tags.update(properties) properties = context_tags + return _add_context_session_id(properties) + + +def _add_context_session_id(properties): if "$session_id" not in properties and get_context_session_id(): properties["$session_id"] = get_context_session_id() - return properties +def _context_tag_defaults() -> dict[str, Any]: + """The current context's tags, as event property defaults.""" + current_context = _get_current_context() + if not current_context: + return {} + context_tags = current_context.collect_tags() + return {**context_tags, "$context_tags": set(context_tags.keys())} + + def no_throw(default_return=None): """ Decorator to prevent raising exceptions from public API methods. @@ -712,6 +737,7 @@ def __init__( feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, + super_options=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, @@ -787,17 +813,31 @@ def __init__( accepts a Project Secret API Key. disabled: If True, disable captures and API requests. Useful in tests. disable_geoip: Whether to disable server-side GeoIP enrichment. - Defaults to True. + Defaults to True. Events get ``$geoip_disable`` only when the + event, context tags and ``super_properties`` leave it unset. is_server: Whether events are emitted from a server-side runtime. Defaults to True; set to False when using the SDK as a client/CLI - so the device OS is attributed to the person normally. + so the device OS is attributed to the person normally. Events + get ``$is_server`` only when the event, context tags and + ``super_properties`` leave it unset. historical_migration: Mark events as historical migration imports. feature_flags_request_timeout_seconds: Timeout in seconds for feature flag and remote config requests. feature_flags_request_max_retries: Number of retries for feature flag requests after network, transport, or timeout failures. Defaults to 1. Set to 0 to disable retries. - super_properties: Properties merged into every captured event. + super_properties: Properties for every captured event. They fill + only keys the event leaves unset, before ``before_send`` runs. + ``$set``, ``$set_once``, ``$groups`` and ``$group_set`` fill + one level deep. An event's own properties and context tags + override them, and ``before_send`` can change or remove them. + super_options: Capture options for every captured event, such as + ``{"cookieless_mode": True}``. They fill only options the event + leaves unset, before ``before_send`` runs. An event's own + ``options`` and context options override them, and + ``before_send`` can change them. They also win over an event's legacy property + for the same key, such as ``$cookieless_mode``, so pass + per-event overrides of that key as ``options``. enable_exception_autocapture: Automatically capture uncaught exceptions. log_captured_exceptions: Also log exceptions captured by error @@ -812,7 +852,11 @@ def __init__( through unredacted. ``privacy_mode`` always wins. Defaults to False. before_send: Optional callback that can modify or drop events before - upload. Return ``None`` to drop an event. + upload. Return ``None`` to drop an event. Context tags, + context options, ``super_properties``, ``super_options`` and + the values the SDK adds, such as ``$is_server``, + ``$geoip_disable`` and ``$os``, fill in before it runs, so it + can change or remove them. flag_fallback_cache_url: Optional feature flag fallback cache URL, such as ``memory://local/?ttl=300&size=10000`` or a Redis URL. enable_local_evaluation: Whether to poll feature flag definitions for @@ -1005,6 +1049,7 @@ def __init__( maximum=AI_MAX_MSG_SIZE, ) self.super_properties = super_properties + self.super_options = super_options # Release id from POSTHOG_RELEASE_ID, attached to every event. Resolved # here so the env var is read once per client. self._release_id = _resolve_release_id() @@ -1315,6 +1360,36 @@ def get_tags(self) -> Dict[str, Any]: """ return _context_get_tags() + def set_context_option(self, key: str, value: Any) -> None: + """ + Set a capture option for every event captured within the current context. + + Context options fill options an event leaves unset, before + ``before_send`` runs, so the hook sees them. They override + ``super_options``. An event's own ``options`` override them, and + ``before_send`` can change them. + + Args: + key: The option name, such as ``"process_person_profile"``. + value: The option value, sent as given. + + Category: + Contexts + """ + _context_set_context_option(key, value) + + def get_context_options(self) -> Dict[str, Any]: + """ + Get all capture options from the current context. + + Returns: + Dict of all capture options in the current context. + + Category: + Contexts + """ + return _context_get_context_options() + def identify_context(self, distinct_id: str) -> None: """ Identify the current context with a distinct ID. @@ -1625,18 +1700,24 @@ def capture( uuid: A unique identifier for the event. If provided, it must be a valid UUID string or uuid.UUID instance; invalid values are ignored and replaced with a newly generated UUID. - groups: A dictionary of group information. + groups: A dictionary of group information. It merges into a + ``$groups`` property and wins key by key. flags: A FeatureFlagEvaluations snapshot from evaluate_flags(). The exact values from the snapshot are attached with no extra /flags request. send_feature_flags: Deprecated. Prefer flags=... from evaluate_flags(). When truthy, evaluates flags during capture and attaches them to the event. - disable_geoip: Whether to disable GeoIP for this event. + disable_geoip: Whether to disable GeoIP for this event. A + ``$geoip_disable`` property in the event, context tags or + ``super_properties`` wins. options: Capture options for this event, such as ``{"process_person_profile": False}``. Sent as given, for - PostHog to validate. An option wins over its legacy ``$`` - property, such as ``$process_person_profile``. + PostHog to validate. They override context options and + ``super_options``, which fill in before ``before_send`` runs. + An option set at any layer wins over its + legacy ``$`` property, such as ``$process_person_profile``, + set at any layer. A ``None`` option counts as unset. Examples: ```python @@ -1716,19 +1797,15 @@ def _capture( # applied to the fully-enriched properties dict just before enqueueing. property_allowlist = kwargs.get("_property_allowlist", None) - properties = {**(properties or {}), **system_context()} + properties = dict(properties or {}) if self.capture_trace_context: properties = {**_get_current_otel_span_properties(), **properties} - properties = add_context_tags(properties) - assert properties is not None # Type hint for mypy + properties = _add_context_session_id(properties) (distinct_id, personless) = get_identity_state(distinct_id) - if personless and "$process_person_profile" not in properties: - properties["$process_person_profile"] = False - msg = { "properties": properties, "timestamp": timestamp, @@ -1739,7 +1816,7 @@ def _capture( } if groups: - properties["$groups"] = groups + _merge_groups(properties, groups) extra_properties: dict[str, Any] = {} @@ -1833,7 +1910,14 @@ def _capture( msg["properties"] = properties return self._enqueue( - msg, disable_geoip, lane, property_allowlist=property_allowlist + msg, + disable_geoip, + lane, + property_allowlist=property_allowlist, + context_properties=_context_tag_defaults(), + context_options=_context_get_context_options(), + derived_options=_personless_options(personless), + system_properties=system_context(), ) def _parse_send_feature_flags(self, send_feature_flags) -> SendFeatureFlagsOptions: @@ -1924,7 +2008,9 @@ def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: "options": _event_options(kwargs.get("options", None)), } - return self._enqueue(msg, disable_geoip) + return self._enqueue( + msg, disable_geoip, context_options=_context_get_context_options() + ) @no_throw() def set_once(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: @@ -1974,7 +2060,9 @@ def set_once(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: "options": _event_options(kwargs.get("options", None)), } - return self._enqueue(msg, disable_geoip) + return self._enqueue( + msg, disable_geoip, context_options=_context_get_context_options() + ) @no_throw() def group_identify( @@ -2412,7 +2500,17 @@ def _report_capture_failure( finally: _on_error_state.active = False - def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): + def _enqueue( + self, + msg, + disable_geoip, + lane=None, + property_allowlist=None, + context_properties=None, + context_options=None, + derived_options=None, + system_properties=None, + ): # type: (...) -> Optional[str] """Push a new `msg` onto a lane's queue (analytics when unspecified), return the event uuid or None.""" @@ -2445,25 +2543,27 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None): if disable_geoip is None: disable_geoip = self.disable_geoip - if disable_geoip: - msg["properties"]["$geoip_disable"] = True - - if self.super_properties: - msg["properties"] = {**msg["properties"], **self.super_properties} - - # Set after the super_properties merge so an explicit `$release_id` from - # the caller's properties or the super properties wins over the env var. - if self._release_id is not None: - msg["properties"].setdefault("$release_id", self._release_id) - - # Set after the super_properties merge so this SDK's server classification - # can't be silently overridden by a user-provided super property. - if self.is_server: - msg["properties"]["$is_server"] = True + msg["options"] = msg.get("options") or {} + + _fill_event_defaults( + msg, + _build_event_defaults( + super_properties=self.super_properties, + super_options=self.super_options, + release_id=self._release_id, + context_properties=context_properties, + context_options=context_options, + derived_options=derived_options, + property_allowlist=property_allowlist, + is_server=self.is_server, + disable_geoip=disable_geoip, + system_properties=system_properties, + ), + ) - # Applied after every enrichment step (system context, context tags, super - # properties, $lib/$lib_version) so the final event shape is exactly the - # allowlist regardless of where a property came from. + # Applied after every enrichment step (system context, $lib/$lib_version, + # filled defaults), so the final event shape is exactly the allowlist + # regardless of where a property came from. if property_allowlist is not None: msg["properties"] = { k: v for k, v in msg["properties"].items() if k in property_allowlist diff --git a/posthog/contexts.py b/posthog/contexts.py index f897d1d96..a6af01635 100644 --- a/posthog/contexts.py +++ b/posthog/contexts.py @@ -28,6 +28,7 @@ def __init__( self.distinct_id: Optional[str] = None self.device_id: Optional[str] = None self.tags: Dict[str, Any] = {} + self.options: Dict[str, Any] = {} self.capture_exception_code_variables: Optional[bool] = None self.code_variables_mask_patterns: Optional[list] = None self.code_variables_ignore_patterns: Optional[list] = None @@ -46,6 +47,9 @@ def set_device_id(self, device_id: str): def add_tag(self, key: str, value: Any): self.tags[key] = value + def add_option(self, key: str, value: Any): + self.options[key] = value + def set_capture_exception_code_variables(self, enabled: bool): self.capture_exception_code_variables = enabled @@ -94,6 +98,13 @@ def collect_tags(self) -> Dict[str, Any]: return tags return self.tags.copy() + def collect_options(self) -> Dict[str, Any]: + if self.parent and not self.fresh: + options = self.parent.collect_options() + options.update(self.options) + return options + return self.options.copy() + def get_capture_exception_code_variables(self) -> Optional[bool]: if self.capture_exception_code_variables is not None: return self.capture_exception_code_variables @@ -288,6 +299,51 @@ def get_tags() -> Dict[str, Any]: return {} +def set_context_option(key: str, value: Any) -> None: + """ + Set a capture option for every event captured within the current context. + + Context options fill options an event leaves unset, before ``before_send`` + runs, so the hook sees them. They override the client's ``super_options``. + An event's own ``options`` override them, and ``before_send`` can change + them. Child contexts inherit them unless they are fresh. + + Args: + key: The option name, such as ``"process_person_profile"`` + value: The option value, sent as given + + Example: + ```python + with posthog.new_context(): + posthog.set_context_option("process_person_profile", False) + posthog.capture("health_check") + ``` + + Category: + Contexts + """ + current_context = _get_current_context() + if current_context: + current_context.add_option(key, value) + + +def get_context_options() -> Dict[str, Any]: + """ + Get all capture options from the current context. Note, modifying + the returned dictionary will not affect the current context. + + Returns: + Dict of all capture options in the current context + + Category: + Contexts + """ + current_context = _get_current_context() + if current_context: + return current_context.collect_options() + return {} + + def identify_context(distinct_id: str) -> None: """ Identify the current context with a distinct ID, associating all events captured in this or diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index 90040011a..512d9a37f 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -10,6 +10,8 @@ from datetime import datetime, timezone from typing import Any, Dict, List +from posthog.capture_event import _merge_groups + from .constants import ( POSTHOG_MCP_ANALYTICS_SOURCE, PostHogMCPAnalyticsEvent, @@ -63,9 +65,9 @@ def _build_capture_event(event: Event) -> PostHogCaptureEvent: _add_session_id(event, properties) _add_conversation_id(event, properties) _add_person_processing(event, properties) - _add_groups(event, properties) _add_common_properties(event, properties) _add_custom_properties(event, properties) + _add_groups(event, properties) _add_server_build(event, properties) event_name = ( @@ -92,9 +94,10 @@ def _add_conversation_id(event: Event, properties: Dict[str, Any]) -> None: def _add_groups(event: Event, properties: Dict[str, Any]) -> None: + """Merge the typed groups over a custom ``$groups`` property, key by key.""" groups = event.get("groups") if groups: - properties["$groups"] = groups + _merge_groups(properties, groups) def _add_person_processing(event: Event, properties: Dict[str, Any]) -> None: @@ -264,7 +267,6 @@ def _build_exception_event(event: Event) -> PostHogCaptureEvent: _add_session_id(event, properties) _add_conversation_id(event, properties) _add_person_processing(event, properties) - _add_groups(event, properties) error = event.get("error") if error: @@ -292,6 +294,7 @@ def _build_exception_event(event: Event) -> PostHogCaptureEvent: properties[_P.PROTOCOL_VERSION] = event["protocol_version"] _add_custom_properties(event, properties) + _add_groups(event, properties) _add_server_build(event, properties) return { diff --git a/posthog/test/capture_helpers.py b/posthog/test/capture_helpers.py index 3ead56f5e..d32f4741f 100644 --- a/posthog/test/capture_helpers.py +++ b/posthog/test/capture_helpers.py @@ -38,6 +38,20 @@ def patch_capture_send(site: str = "client", **kwargs) -> "mock._patch": return mock.patch(f"posthog.{site}.{_SUBMITTER}", **kwargs) +def record_sync_capture_sends(test_case) -> list[dict]: + """Record every event a ``sync_mode`` client uploads until ``test_case`` ends.""" + events: list[dict] = [] + + def record(*args, **kwargs): + events.extend(args[2] if len(args) > 2 else kwargs["batch"]) + return mock.DEFAULT + + patcher = patch_capture_send("client", side_effect=record) + patcher.start() + test_case.addCleanup(patcher.stop) + return events + + def patch_async_capture_send(**kwargs) -> "mock._patch": """Patch the submitter the ``AsyncPosthog`` consumer awaits.""" return mock.patch("posthog._async_consumer.async_send_v1_batch", **kwargs) diff --git a/posthog/test/mcp/test_pipeline.py b/posthog/test/mcp/test_pipeline.py index 81b6db644..413941c28 100644 --- a/posthog/test/mcp/test_pipeline.py +++ b/posthog/test/mcp/test_pipeline.py @@ -925,6 +925,7 @@ def test_identity_enables_person_processing_and_set(): "identify_actor_given_id": "user_1", "identify_actor_data": {"email": "a@b.com"}, "groups": {"organization": "org_1"}, + "properties": {"$groups": {"organization": "custom", "project": "p1"}}, "timestamp": datetime.now(timezone.utc), } [capture] = build_posthog_capture_events(event) @@ -932,7 +933,29 @@ def test_identity_enables_person_processing_and_set(): assert capture["distinct_id"] == "user_1" assert "$process_person_profile" not in props assert props["$set"] == {"email": "a@b.com"} - assert props["$groups"] == {"organization": "org_1"} + assert props["$groups"] == {"organization": "org_1", "project": "p1"} + + +def test_typed_groups_win_on_primary_and_exception_events(): + event = { + "event_type": MCPAnalyticsEventType.MCP_TOOLS_CALL, + "session_id": "ses_abc", + "is_error": True, + "error": {"$exception_list": [{"type": "ValueError", "value": "boom"}]}, + "groups": {"organization": "org_1"}, + "properties": {"$groups": {"organization": "custom", "project": "p1"}}, + "timestamp": datetime.now(timezone.utc), + } + captures = build_posthog_capture_events(event) + assert [c["event"] for c in captures] == [ + PostHogMCPAnalyticsEvent.TOOL_CALL, + PostHogMCPAnalyticsEvent.EXCEPTION, + ] + for capture in captures: + assert capture["properties"]["$groups"] == { + "organization": "org_1", + "project": "p1", + } def test_listed_tool_names_only_on_tools_list(): diff --git a/posthog/test/test_async_consumer.py b/posthog/test/test_async_consumer.py index 4290a9419..db3e17feb 100644 --- a/posthog/test/test_async_consumer.py +++ b/posthog/test/test_async_consumer.py @@ -15,7 +15,7 @@ def make_consumer(*, retries: int) -> _AsyncConsumer: "test-key", host="https://example.com", on_error=None, - process_event=mock.AsyncMock(side_effect=lambda event: event), + process_event=mock.AsyncMock(side_effect=lambda event, defaults: event), flush_at=100, flush_interval=1, retries=retries, diff --git a/posthog/test/test_capture_event.py b/posthog/test/test_capture_event.py index 1e17e76a6..21380fb79 100644 --- a/posthog/test/test_capture_event.py +++ b/posthog/test/test_capture_event.py @@ -204,7 +204,6 @@ def test_top_level_set_relocated_into_properties(self, _name, key) -> None: self.assertNotIn(key, event) # not a top-level v1 field def test_top_level_set_merges_with_existing_properties_set(self) -> None: - # properties wins on key collision. msg = _legacy_msg( properties={"$set": {"a": "from_props", "b": "props_only"}}, **{"$set": {"a": "from_top", "c": "top_only"}}, @@ -212,7 +211,7 @@ def test_top_level_set_merges_with_existing_properties_set(self) -> None: event = _to_v1_event(msg) self.assertEqual( event["properties"]["$set"], - {"a": "from_props", "b": "props_only", "c": "top_only"}, + {"a": "from_top", "b": "props_only", "c": "top_only"}, ) def test_groups_left_in_properties(self) -> None: diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index 37d862b5c..c2c3e9ad9 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -420,20 +420,6 @@ def test_capture_omits_is_server_when_disabled(self): self.assertEqual(msg["properties"]["$lib"], "posthog-python") self.assertNotIn("$is_server", msg["properties"]) - def test_is_server_not_overridden_by_super_properties(self): - with patch_capture_send("client") as mock_post: - client = Client( - FAKE_TEST_API_KEY, - on_error=self.set_fail, - sync_mode=True, - super_properties={"$is_server": False}, - ) - client.capture("python test event", distinct_id="distinct_id") - self.assertFalse(self.failed) - - msg = sent_batch(mock_post)[0] - self.assertEqual(msg["properties"]["$is_server"], True) - def test_basic_capture_with_uuid(self): with patch_capture_send("client") as mock_post: client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, sync_mode=True) @@ -2390,10 +2376,10 @@ def test_session_id_with_different_event_types( [ # test_name, super_properties, event_session_id, expected_session_id, expected_super_props ( - "super_properties_override_session_id", + "event_session_id_overrides_super_properties", {"$session_id": "super-session", "source": "test"}, "event-session-808", - "super-session", + "event-session-808", {"source": "test"}, ), ( diff --git a/posthog/test/test_contexts.py b/posthog/test/test_contexts.py index e75a6610d..73a1ab466 100644 --- a/posthog/test/test_contexts.py +++ b/posthog/test/test_contexts.py @@ -7,7 +7,9 @@ import posthog from posthog.client import Client from posthog.contexts import ( + get_context_options, get_tags, + set_context_option, new_context, scoped, tag, @@ -384,6 +386,26 @@ def context_state(): ), ) + def test_context_options_inherit_like_tags(self): + with new_context(fresh=True): + set_context_option("cookieless_mode", True) + set_context_option("process_person_profile", False) + + with new_context(fresh=False): + set_context_option("process_person_profile", True) + assert get_context_options() == { + "cookieless_mode": True, + "process_person_profile": True, + } + + with new_context(fresh=True): + assert get_context_options() == {} + + assert get_context_options() == { + "cookieless_mode": True, + "process_person_profile": False, + } + def test_child_tags_override_parent_tags_in_non_fresh_context(self): with new_context(fresh=True): tag("shared_key", "parent_value") diff --git a/posthog/test/test_event_options.py b/posthog/test/test_event_options.py index 33e6c3a50..7e9f18fda 100644 --- a/posthog/test/test_event_options.py +++ b/posthog/test/test_event_options.py @@ -1,8 +1,9 @@ +import inspect import logging import pytest -from posthog import AsyncPosthog +from posthog import AsyncPosthog, new_context, set_context_option, tag from posthog.capture_event import _to_v1_event from posthog.client import Client from posthog.test.capture_helpers import ( @@ -34,22 +35,29 @@ ASYNC_CAPTURE_CALLS = {k: v for k, v in CAPTURE_CALLS.items() if k != "capture_ai"} -def _sync_wire_events(call, before_send=None) -> list[dict]: +def _sync_wire_events(call, before_send=None, **config) -> list[dict]: with patch_capture_send("client") as send: - client = Client(FAKE_TEST_API_KEY, sync_mode=True, before_send=before_send) + client = Client( + FAKE_TEST_API_KEY, sync_mode=True, before_send=before_send, **config + ) assert call(client) is not None return [_to_v1_event(msg) for msg in sent_events(send)] -async def _async_wire_events(call, before_send=None) -> list[dict]: +async def _async_wire_events(call, before_send=None, **config) -> list[dict]: batches: list[list[dict]] = [] async def send_batch(api_key, host, batch, **kwargs): batches.append(batch) with patch_async_capture_send(side_effect=send_batch): - async with AsyncPosthog("test-key", before_send=before_send) as client: - assert call(client) is not None + async with AsyncPosthog( + "test-key", before_send=before_send, **config + ) as client: + result = call(client) + if inspect.isawaitable(result): + result = await result + assert result is not None await client.flush(timeout_seconds=1) return [_to_v1_event(msg) for batch in batches for msg in batch] @@ -109,3 +117,357 @@ async def test_async_non_dict_options_are_logged_and_event_is_sent(caplog): ) assert events[0]["options"] == {} assert "options must be a dict" in caplog.text + + +PP = "process_person_profile" + +# Layers, lowest first: personless (no distinct_id), super, context, event. +LAYER_CASES = { + "personless_alone": ({}, {}, {}, None, {PP: False}), + "identified_sends_no_option": ({}, {}, {}, "u", {}), + "super_beats_personless": ({PP: True}, {}, {}, None, {PP: True}), + "context_beats_super": ({PP: True}, {PP: False}, {}, "u", {PP: False}), + "event_beats_context": ({}, {PP: False}, {PP: True}, "u", {PP: True}), + "event_beats_all": ({PP: False}, {PP: False}, {PP: True}, None, {PP: True}), + "null_event_option_is_filled": ({PP: True}, {}, {PP: None}, "u", {PP: True}), + "layers_merge_by_key": ( + {"cookieless_mode": True}, + {"product_tour_id": "t"}, + {}, + "u", + {"cookieless_mode": True, "product_tour_id": "t"}, + ), +} + + +def _layered_capture(context_options, event_options, distinct_id): + def call(client): + with new_context(fresh=True): + for key, value in context_options.items(): + set_context_option(key, value) + return client.capture("e", distinct_id=distinct_id, options=event_options) + + return call + + +@pytest.mark.parametrize("case", list(LAYER_CASES)) +def test_sync_option_layers(case): + super_options, context, event, distinct_id, expected = LAYER_CASES[case] + events = _sync_wire_events( + _layered_capture(context, event, distinct_id), super_options=super_options + ) + assert events[0]["options"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", list(LAYER_CASES)) +async def test_async_option_layers(case): + super_options, context, event, distinct_id, expected = LAYER_CASES[case] + events = await _async_wire_events( + _layered_capture(context, event, distinct_id), super_options=super_options + ) + assert events[0]["options"] == expected + + +@pytest.mark.parametrize( + "method", ["capture", "set", "set_once", "group_identify", "alias"] +) +def test_sync_super_options_reach_every_path(method): + events = _sync_wire_events( + lambda c: CAPTURE_CALLS[method](c, None), super_options=OPTIONS + ) + assert [e["options"] for e in events] == [OPTIONS] + + +@pytest.mark.parametrize("method", ["capture", "set", "set_once"]) +def test_sync_context_options_reach_user_paths(method): + def call(client): + with new_context(fresh=True): + set_context_option("cookieless_mode", True) + return CAPTURE_CALLS[method](client, None) + + events = _sync_wire_events(call) + assert [e["options"] for e in events] == [{"cookieless_mode": True}] + + +def test_personless_option_beats_legacy_property(): + events = _sync_wire_events( + lambda c: c.capture("e", properties={"$process_person_profile": True}) + ) + assert events[0]["options"] == {PP: False} + assert "$process_person_profile" not in events[0]["properties"] + + +def test_legacy_super_property_fills_option_when_unset(): + events = _sync_wire_events( + lambda c: c.capture("e", distinct_id="u"), + super_properties={"$process_person_profile": False}, + ) + assert events[0]["options"] == {PP: False} + + +def _layered_properties(client): + with new_context(fresh=True): + tag("from_context", "context") + tag("shared", "context") + return client.capture( + "e", distinct_id="u", properties={"shared": "event", "only_event": 1} + ) + + +SUPER_PROPERTIES = {"shared": "super", "from_context": "super", "only_super": 1} + + +def test_sync_event_and_context_properties_beat_super_properties(): + events = _sync_wire_events(_layered_properties, super_properties=SUPER_PROPERTIES) + properties = events[0]["properties"] + assert ( + properties["shared"], + properties["from_context"], + properties["only_super"], + properties["only_event"], + ) == ("event", "context", 1, 1) + + +@pytest.mark.asyncio +async def test_async_event_and_context_properties_beat_super_properties(): + events = await _async_wire_events( + _layered_properties, super_properties=SUPER_PROPERTIES + ) + properties = events[0]["properties"] + assert ( + properties["shared"], + properties["from_context"], + properties["only_super"], + properties["only_event"], + ) == ("event", "context", 1, 1) + + +DEFAULTS_CONFIG = { + "super_properties": {"shared": "super", "only_super": 1}, + "super_options": {"cookieless_mode": True}, +} + + +def _recording_hook(seen): + def hook(msg): + seen.append((dict(msg["properties"]), dict(msg["options"]))) + removed = {"only_super", "$is_server"} + properties = { + **{k: v for k, v in msg["properties"].items() if k not in removed}, + "shared": "hook", + } + return {**msg, "properties": properties, "options": {"cookieless_mode": False}} + + return hook + + +def _set_context_values(): + tag("from_context", "context") + set_context_option("product_tour_id", "context-tour") + + +def _context_capture(method): + def call(client): + if method == "capture": + with new_context(fresh=True): + _set_context_values() + return client.capture("e", distinct_id="u") + + async def immediate(): + with new_context(fresh=True): + _set_context_values() + return await client.capture_immediate("e", distinct_id="u") + + return immediate() + + return call + + +def _assert_hook_sees_defaults_and_has_final_say(seen, events): + hook_properties, hook_options = seen[0] + assert ( + hook_properties["shared"], + hook_properties["only_super"], + hook_properties["from_context"], + hook_properties["$is_server"], + hook_properties["$geoip_disable"], + ) == ("super", 1, "context", True, True) + assert "$os" in hook_properties + assert hook_options == {"cookieless_mode": True, "product_tour_id": "context-tour"} + properties = events[0]["properties"] + assert (properties["shared"], properties["from_context"]) == ("hook", "context") + assert not {"only_super", "$is_server"} & properties.keys() + assert events[0]["options"] == {"cookieless_mode": False} + + +def test_sync_defaults_fill_before_before_send(): + seen: list = [] + events = _sync_wire_events( + _context_capture("capture"), + before_send=_recording_hook(seen), + **DEFAULTS_CONFIG, + ) + _assert_hook_sees_defaults_and_has_final_say(seen, events) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["capture", "capture_immediate"]) +async def test_async_defaults_fill_before_before_send(method): + seen: list = [] + events = await _async_wire_events( + _context_capture(method), before_send=_recording_hook(seen), **DEFAULTS_CONFIG + ) + _assert_hook_sees_defaults_and_has_final_say(seen, events) + + +@pytest.mark.asyncio +async def test_async_queued_event_keeps_capture_time_defaults(): + def capture_then_change_defaults(client): + result = client.capture("e", distinct_id="d") + client.super_properties["plan"] = "pro" + client.super_options["cookieless_mode"] = True + return result + + events = await _async_wire_events( + capture_then_change_defaults, + super_properties={"plan": "free"}, + super_options={"cookieless_mode": False}, + ) + assert events[0]["properties"]["plan"] == "free" + assert events[0]["options"]["cookieless_mode"] is False + + +NESTED_SUPER_PROPERTIES = { + "$set": {"plan": "free", "source": "super"}, + "$groups": {"company": "acme", "team": "core", "region": "us"}, + "$unset": ["stale"], +} + + +def _nested_capture(client): + return client.capture( + "e", + distinct_id="u", + properties={ + "$set": {"plan": "pro"}, + "$unset": ["other"], + "$groups": {"company": "event", "team": "event"}, + }, + groups={"company": "posthog"}, + ) + + +def _assert_nested_properties(events): + properties = events[0]["properties"] + assert properties["$set"] == {"plan": "pro", "source": "super"} + assert properties["$groups"] == { + "company": "posthog", + "team": "event", + "region": "us", + } + assert properties["$unset"] == ["other"] + + +def test_nested_properties_fill_one_level_deep(): + events = _sync_wire_events( + _nested_capture, super_properties=NESTED_SUPER_PROPERTIES + ) + _assert_nested_properties(events) + + +@pytest.mark.asyncio +async def test_async_nested_properties_fill_one_level_deep(): + events = await _async_wire_events( + _nested_capture, super_properties=NESTED_SUPER_PROPERTIES + ) + _assert_nested_properties(events) + + +CALLER_SDK_VALUES = {"$is_server": False, "$geoip_disable": False, "$os": "caller-os"} +SDK_FLAGS = {"$is_server": False, "$geoip_disable": False} + + +def _capture_with_caller_sdk_values(source): + def call(client): + with new_context(fresh=True): + if source == "context": + for key, value in CALLER_SDK_VALUES.items(): + tag(key, value) + properties = CALLER_SDK_VALUES if source == "event" else None + return client.capture("e", distinct_id="u", properties=properties) + + return call + + +def _sdk_value_config(source): + return {"super_properties": CALLER_SDK_VALUES} if source == "super" else {} + + +def _assert_caller_values(events, expected): + properties = events[0]["properties"] + assert {key: properties.get(key) for key in expected} == expected + + +@pytest.mark.parametrize("source", ["event", "context", "super"]) +def test_sync_caller_values_beat_sdk_values(source): + events = _sync_wire_events( + _capture_with_caller_sdk_values(source), **_sdk_value_config(source) + ) + _assert_caller_values(events, CALLER_SDK_VALUES) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["event", "context", "super"]) +async def test_async_caller_values_beat_sdk_values(source): + events = await _async_wire_events( + _capture_with_caller_sdk_values(source), **_sdk_value_config(source) + ) + _assert_caller_values(events, CALLER_SDK_VALUES) + + +@pytest.mark.parametrize("method", list(CAPTURE_CALLS)) +def test_sync_super_properties_beat_sdk_values_on_every_path(method): + events = _sync_wire_events( + lambda c: CAPTURE_CALLS[method](c, None), super_properties=SDK_FLAGS + ) + _assert_caller_values(events, SDK_FLAGS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", list(ASYNC_CAPTURE_CALLS)) +async def test_async_super_properties_beat_sdk_values_on_every_path(method): + events = await _async_wire_events( + lambda c: ASYNC_CAPTURE_CALLS[method](c, None), super_properties=SDK_FLAGS + ) + _assert_caller_values(events, SDK_FLAGS) + + +@pytest.mark.parametrize("method", ["set", "set_once"]) +def test_set_call_wins_over_super_person_properties(method): + key = f"${method}" + events = _sync_wire_events( + lambda c: getattr(c, method)(distinct_id="u", properties={"plan": "pro"}), + super_properties={key: {"plan": "free", "source": "super"}}, + ) + assert events[0]["properties"][key] == {"plan": "pro", "source": "super"} + + +def test_late_options_replace_and_remove_legacy_properties(): + def call(client): + with new_context(fresh=True): + set_context_option("product_tour_id", "context-tour") + return client.capture( + "e", distinct_id="u", properties={"$product_tour_id": "event-tour"} + ) + + events = _sync_wire_events( + call, + super_properties={"$cookieless_mode": False}, + super_options={"cookieless_mode": True}, + ) + assert events[0]["options"] == { + "cookieless_mode": True, + "product_tour_id": "context-tour", + } + assert not {"$cookieless_mode", "$product_tour_id"} & events[0]["properties"].keys() diff --git a/posthog/test/test_feature_flag_called_minimization.py b/posthog/test/test_feature_flag_called_minimization.py index b355a68d6..81cf76a1f 100644 --- a/posthog/test/test_feature_flag_called_minimization.py +++ b/posthog/test/test_feature_flag_called_minimization.py @@ -16,6 +16,7 @@ from posthog.capture_event import _to_v1_event from posthog.client import _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, Client from posthog.request import GetResponse +from posthog.test.capture_helpers import record_sync_capture_sends from posthog.test.test_utils import FAKE_TEST_API_KEY from posthog.utils import system_context @@ -65,21 +66,15 @@ def _local_flag_definition(has_experiment): class _CapturedEventsMixin: - """Builds a non-sending client whose fully-enriched events are captured via - ``before_send``, so tests assert the exact wire shape after every enrichment - step (system context, super properties, $lib, ...).""" + """Builds a sync client whose fully-enriched events are recorded at upload, + so tests assert the exact wire shape after every enrichment step (system + context, super properties, $lib, ...).""" def _make_client(self, **kwargs): - captured = [] - - def before_send(msg): - captured.append(msg) - return msg - + captured = record_sync_capture_sends(self) client = Client( FAKE_TEST_API_KEY, - send=False, - before_send=before_send, + sync_mode=True, super_properties={"app_version": "1.2.3"}, **kwargs, ) diff --git a/posthog/test/test_module.py b/posthog/test/test_module.py index 655b1d2cc..09d6ea82c 100644 --- a/posthog/test/test_module.py +++ b/posthog/test/test_module.py @@ -65,6 +65,7 @@ def setUp(self): self._original_api_key = posthog.api_key self._original_project_api_key = posthog.project_api_key self._original_project_root = posthog.project_root + self._original_super_options = posthog.super_options self._original_privacy_mode = posthog.privacy_mode self._original_disabled = posthog.disabled self._original_send = posthog.send @@ -81,6 +82,7 @@ def tearDown(self): posthog.api_key = self._original_api_key posthog.project_api_key = self._original_project_api_key posthog.project_root = self._original_project_root + posthog.super_options = self._original_super_options posthog.privacy_mode = self._original_privacy_mode posthog.disabled = self._original_disabled posthog.send = self._original_send @@ -129,13 +131,19 @@ def test_setup_uses_api_key_when_project_api_key_is_blank(self): self.assertEqual(client.api_key, "phc_api_key") self.assertFalse(client.disabled) - def test_setup_propagates_project_root(self): + @parameterized.expand( + [ + ("project_root", "/path/to/project"), + ("super_options", {"cookieless_mode": True}), + ] + ) + def test_setup_propagates_config(self, setting, value): posthog.api_key = "phc_test" - posthog.project_root = "/path/to/project" + setattr(posthog, setting, value) client = posthog.setup() - self.assertEqual(client.project_root, "/path/to/project") + self.assertEqual(getattr(client, setting), value) def test_setup_propagates_and_updates_privacy_mode(self): posthog.api_key = "phc_test" diff --git a/posthog/test/test_release_id.py b/posthog/test/test_release_id.py index 109f79f44..8d72a12bf 100644 --- a/posthog/test/test_release_id.py +++ b/posthog/test/test_release_id.py @@ -10,7 +10,10 @@ from posthog.client import _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, Client from posthog.release_id import RELEASE_ID_ENV_VAR, _resolve_release_id from posthog.test.test_utils import FAKE_TEST_API_KEY -from posthog.test.capture_helpers import patch_async_capture_send +from posthog.test.capture_helpers import ( + patch_async_capture_send, + record_sync_capture_sends, +) # (name, call, expected event): one row per public event-producing method, shared # by the sync and async clients. Each call builds its own arguments, because a @@ -80,17 +83,10 @@ def test_env_var_resolution(self, _name, env_value, expected) -> None: class TestClientReleaseId(unittest.TestCase): def _client(self, env_value, **kwargs): - """Build a client under `env_value` and collect the events it would send.""" - events = [] - - def before_send(msg): - events.append(msg) - return msg - + """Build a client under `env_value` and collect the events it sends.""" + events = record_sync_capture_sends(self) with _release_id_env(env_value): - client = Client( - FAKE_TEST_API_KEY, send=False, before_send=before_send, **kwargs - ) + client = Client(FAKE_TEST_API_KEY, sync_mode=True, **kwargs) return client, events @parameterized.expand(EVENT_CALLS) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 07db58cf9..a0214f6f5 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -314,6 +314,7 @@ alias posthog.feature_flags.FlagValue -> posthog.types.FlagValue alias posthog.feature_flags.convert_to_datetime_aware -> posthog.utils.convert_to_datetime_aware alias posthog.feature_flags.is_valid_regex -> posthog.utils.is_valid_regex alias posthog.feature_flags.utils -> posthog.utils +alias posthog.inner_get_context_options -> posthog.contexts.get_context_options alias posthog.inner_get_tags -> posthog.contexts.get_tags alias posthog.inner_identify_context -> posthog.contexts.identify_context alias posthog.inner_new_context -> posthog.contexts.new_context @@ -324,6 +325,7 @@ alias posthog.inner_set_code_variables_ignore_patterns_context -> posthog.contex alias posthog.inner_set_code_variables_mask_patterns_context -> posthog.contexts.set_code_variables_mask_patterns_context alias posthog.inner_set_code_variables_mask_url_credentials_context -> posthog.contexts.set_code_variables_mask_url_credentials_context alias posthog.inner_set_context_device_id -> posthog.contexts.set_context_device_id +alias posthog.inner_set_context_option -> posthog.contexts.set_context_option alias posthog.inner_set_context_session -> posthog.contexts.set_context_session alias posthog.inner_tag -> posthog.contexts.tag alias posthog.integrations.django.Client -> posthog.client.Client @@ -680,6 +682,7 @@ attribute posthog.async_client.AsyncClient.project_root = os.getcwd() attribute posthog.async_client.AsyncClient.raw_host = normalize_host(host) attribute posthog.async_client.AsyncClient.secret_key = (resolved_secret_key.strip() if isinstance(resolved_secret_key, str) else resolved_secret_key) or None attribute posthog.async_client.AsyncClient.send = send +attribute posthog.async_client.AsyncClient.super_options = super_options attribute posthog.async_client.AsyncClient.super_properties = super_properties attribute posthog.async_client.AsyncClient.timeout = timeout attribute posthog.before_send = None @@ -751,6 +754,7 @@ attribute posthog.client.Client.queue: Queue attribute posthog.client.Client.raw_host = normalize_host(host) attribute posthog.client.Client.secret_key = (resolved_secret_key.strip() if isinstance(resolved_secret_key, str) else resolved_secret_key) or None attribute posthog.client.Client.send = send +attribute posthog.client.Client.super_options = super_options attribute posthog.client.Client.super_properties = super_properties attribute posthog.client.Client.sync_mode = sync_mode attribute posthog.client.Client.timeout = timeout @@ -789,6 +793,7 @@ attribute posthog.contexts.ContextScope.code_variables_mask_url_credentials: Opt attribute posthog.contexts.ContextScope.device_id: Optional[str] = None attribute posthog.contexts.ContextScope.distinct_id: Optional[str] = None attribute posthog.contexts.ContextScope.fresh = fresh +attribute posthog.contexts.ContextScope.options: Dict[str, Any] = {} attribute posthog.contexts.ContextScope.parent = parent attribute posthog.contexts.ContextScope.session_id: Optional[str] = None attribute posthog.contexts.ContextScope.tags: Dict[str, Any] = {} @@ -1018,6 +1023,7 @@ attribute posthog.request.USER_AGENT = 'posthog-python/' + VERSION attribute posthog.request.US_INGESTION_ENDPOINT = 'https://us.i.posthog.com' attribute posthog.secret_key = None attribute posthog.send = True +attribute posthog.super_options = None attribute posthog.super_properties = None attribute posthog.sync_mode = False attribute posthog.traces = None @@ -1170,13 +1176,13 @@ class posthog.ai.types.TokenUsage class posthog.ai.types.ToolInProgress class posthog.args.OptionalCaptureArgs class posthog.args.OptionalSetArgs -class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] = None, *, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, max_retries: int = 3, timeout: int = 15, thread: int = 1, disabled: bool = False, disable_geoip: bool = True, is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, capture_exception_code_variables: bool = False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, feature_flags_request_timeout_seconds: int = 3, feature_flags_request_max_retries: int = 1) +class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] = None, *, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, max_retries: int = 3, timeout: int = 15, thread: int = 1, disabled: bool = False, disable_geoip: bool = True, is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, super_options: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, capture_exception_code_variables: bool = False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, feature_flags_request_timeout_seconds: int = 3, feature_flags_request_max_retries: int = 1) class posthog.async_client.AsyncPosthog class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, refill_rate: Number, refill_interval_seconds: Number, on_bucket_rate_limited: Optional[Callable[[Hashable], None]] = None, clock: Callable[[], float] = time.monotonic) class posthog.capture_compression.CaptureCompression class posthog.capture_send.CaptureError(status: int | str, message: str, *, endpoint: str, retry_after: Optional[float] = None, request_id: Optional[str] = None, attempts: Optional[int] = None, retry_exhausted: Optional[list[str]] = None, drops: Optional[list[tuple[str, Optional[str]]]] = None, event_results: Optional[dict[str, CaptureEventResult]] = None) class posthog.capture_send.CaptureEventResult(result: Optional[str], details: Optional[str] = None) -class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_max_queue_size: int = 1000, capture_ai_timeout: float = 30, capture_ai_max_event_bytes: int = AI_MAX_MSG_SIZE, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) +class posthog.client.Client(project_api_key: str, host=None, *, debug=False, max_queue_size=10000, send=True, on_error=None, flush_at=100, flush_interval=5.0, max_retries=3, sync_mode=False, timeout=15, thread=1, poll_interval=30, personal_api_key=None, disabled=False, disable_geoip=True, is_server=True, historical_migration=False, feature_flags_request_timeout_seconds=3, feature_flags_request_max_retries=1, super_properties=None, super_options=None, enable_exception_autocapture=False, log_captured_exceptions=False, project_root=None, privacy_mode=False, before_send=None, flag_fallback_cache_url=None, enable_local_evaluation=True, flag_definition_cache_provider: Optional[FlagDefinitionCacheProvider] = None, capture_exception_code_variables=False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: list[str] | None = None, enable_exception_autocapture_rate_limiting=False, exception_autocapture_bucket_size=ExceptionCapture.DEFAULT_BUCKET_SIZE, exception_autocapture_refill_rate=ExceptionCapture.DEFAULT_REFILL_RATE, exception_autocapture_refill_interval_seconds=ExceptionCapture.DEFAULT_REFILL_INTERVAL_SECONDS, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_max_queue_size: int = 1000, capture_ai_timeout: float = 30, capture_ai_max_event_bytes: int = AI_MAX_MSG_SIZE, secret_key=None, metrics: Optional[dict] = None, enable_full_ai_capture=False, capture_trace_context=False, _use_ai_lane=False, _enable_multimodal_capture=False, traces: Optional[dict] = None) class posthog.consumer.Consumer(queue, api_key, flush_at=100, host=None, on_error=None, flush_interval=5.0, retries=3, timeout=15, historical_migration=False, endpoint=_CAPTURE_V1_PATH, max_msg_size=MAX_MSG_SIZE, capture_compression=CaptureCompression.NONE) class posthog.contexts.ContextScope(parent=None, fresh: bool = False, capture_exceptions: bool = True, client: Optional[Client] = None) class posthog.exception_capture.ExceptionCapture(client: Client, rate_limiting_enabled=False, bucket_size=DEFAULT_BUCKET_SIZE, refill_rate=DEFAULT_REFILL_RATE, refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS) @@ -1319,6 +1325,7 @@ function posthog.contexts.get_code_variables_mask_patterns_context() -> Optional function posthog.contexts.get_code_variables_mask_url_credentials_context() -> Optional[bool] function posthog.contexts.get_context_device_id() -> Optional[str] function posthog.contexts.get_context_distinct_id() -> Optional[str] +function posthog.contexts.get_context_options() -> Dict[str, Any] function posthog.contexts.get_context_session_id() -> Optional[str] function posthog.contexts.get_tags() -> Dict[str, Any] function posthog.contexts.identify_context(distinct_id: str) -> None @@ -1330,6 +1337,7 @@ function posthog.contexts.set_code_variables_ignore_patterns_context(ignore_patt function posthog.contexts.set_code_variables_mask_patterns_context(mask_patterns: list) -> None function posthog.contexts.set_code_variables_mask_url_credentials_context(enabled: bool) -> None function posthog.contexts.set_context_device_id(device_id: str) -> None +function posthog.contexts.set_context_option(key: str, value: Any) -> None function posthog.contexts.set_context_session(session_id: str) -> None function posthog.contexts.tag(key: str, value: Any) -> None function posthog.evaluate_flags(distinct_id: Optional[ID_TYPES] = None, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, flag_keys: Optional[list[str]] = None, device_id: Optional[str] = None) -> FeatureFlagEvaluations @@ -1384,6 +1392,7 @@ function posthog.flush(timeout_seconds: Optional[float] = 10) -> None function posthog.get_active_span() -> Optional[Span] function posthog.get_all_flags(distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None, flag_keys_to_evaluate: Optional[list[str]] = None) -> Optional[dict[str, FlagValue]] function posthog.get_all_flags_and_payloads(distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None, flag_keys_to_evaluate: Optional[list[str]] = None) -> FlagsAndPayloads +function posthog.get_context_options() -> Dict[str, Any] function posthog.get_feature_flag(key: str, distinct_id: ID_TYPES, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, send_feature_flag_events: bool = True, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None) -> Optional[FlagValue] function posthog.get_feature_flag_evaluation_runtime(key: str) -> Optional[FeatureFlagEvaluationRuntime] function posthog.get_feature_flag_keys_by_evaluation_runtime(evaluation_runtime: Union[FeatureFlagEvaluationRuntime, str]) -> list[str] @@ -1428,6 +1437,7 @@ function posthog.set_code_variables_ignore_patterns_context(ignore_patterns: lis function posthog.set_code_variables_mask_patterns_context(mask_patterns: list[str]) function posthog.set_code_variables_mask_url_credentials_context(enabled: bool) function posthog.set_context_device_id(device_id: str) +function posthog.set_context_option(key: str, value: Any) -> None function posthog.set_context_session(session_id: str) function posthog.set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] function posthog.setup() -> Client @@ -1582,6 +1592,7 @@ method posthog.client.Client.flush(timeout_seconds: Optional[float] = 10) -> Non method posthog.client.Client.get_active_span() -> Optional[Span] method posthog.client.Client.get_all_flags(distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> Optional[dict[str, Union[bool, str]]] method posthog.client.Client.get_all_flags_and_payloads(distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, disable_geoip: Optional[bool] = None, flag_keys_to_evaluate: Optional[list[str]] = None, device_id: Optional[str] = None) -> FlagsAndPayloads +method posthog.client.Client.get_context_options() -> Dict[str, Any] method posthog.client.Client.get_feature_flag(key: str, distinct_id: ID_TYPES, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, only_evaluate_locally: bool = False, send_feature_flag_events: bool = True, disable_geoip: Optional[bool] = None, device_id: Optional[str] = None) -> Optional[FlagValue] method posthog.client.Client.get_feature_flag_evaluation_runtime(key: str) -> Optional[FeatureFlagEvaluationRuntime] method posthog.client.Client.get_feature_flag_keys_by_evaluation_runtime(evaluation_runtime: Union[FeatureFlagEvaluationRuntime, str]) -> list[str] @@ -1601,6 +1612,7 @@ method posthog.client.Client.new_context(fresh=False, capture_exceptions: Option method posthog.client.Client.scoped(fresh=False, capture_exceptions: Optional[bool] = None) method posthog.client.Client.set(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] method posthog.client.Client.set_context_device_id(device_id: str) -> None +method posthog.client.Client.set_context_option(key: str, value: Any) -> None method posthog.client.Client.set_context_session(session_id: str) -> None method posthog.client.Client.set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str] method posthog.client.Client.shutdown() -> None @@ -1611,7 +1623,9 @@ method posthog.consumer.Consumer.pause() method posthog.consumer.Consumer.request(batch) method posthog.consumer.Consumer.run() method posthog.consumer.Consumer.upload() +method posthog.contexts.ContextScope.add_option(key: str, value: Any) method posthog.contexts.ContextScope.add_tag(key: str, value: Any) +method posthog.contexts.ContextScope.collect_options() -> Dict[str, Any] method posthog.contexts.ContextScope.collect_tags() -> Dict[str, Any] method posthog.contexts.ContextScope.get_capture_exception_code_variables() -> Optional[bool] method posthog.contexts.ContextScope.get_code_variables_detect_secrets() -> Optional[bool] From f4b67b54280e02b808e9051d1821c334d36f2433 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:46:33 -0700 Subject: [PATCH 10/14] feat!: AI and MCP send personless as an event option (#1025) feat!: AI wrappers and MCP send personless as a per-event option --- posthog/ai/claude_agent_sdk/client.py | 4 +- posthog/ai/claude_agent_sdk/processor.py | 42 ++++++----- posthog/ai/gemini/_shared.py | 4 +- posthog/ai/langchain/callbacks.py | 8 +-- posthog/ai/openai/_embeddings.py | 4 +- posthog/ai/openai_agents/processor.py | 42 +++++------ posthog/ai/utils.py | 30 +++++--- posthog/mcp/_posthog_events.py | 10 +-- posthog/mcp/_sink.py | 1 + posthog/mcp/types.py | 1 + posthog/test/ai/anthropic/test_anthropic.py | 69 ++++++++++++------- .../ai/claude_agent_sdk/test_processor.py | 5 +- posthog/test/ai/gemini/test_gemini.py | 3 +- posthog/test/ai/gemini/test_gemini_async.py | 3 +- posthog/test/ai/langchain/test_callbacks.py | 24 +++---- posthog/test/ai/openai/test_async_parity.py | 3 +- .../test/ai/openai_agents/test_processor.py | 8 +-- posthog/test/mcp/_helpers.py | 7 +- posthog/test/mcp/test_fastmcp.py | 2 +- posthog/test/mcp/test_pipeline.py | 39 ++++++++++- posthog/test/mcp/test_v2_mcpserver.py | 3 +- posthog/test/test_event_options.py | 29 ++++++++ 22 files changed, 221 insertions(+), 120 deletions(-) diff --git a/posthog/ai/claude_agent_sdk/client.py b/posthog/ai/claude_agent_sdk/client.py index f9a718d80..2766ec925 100644 --- a/posthog/ai/claude_agent_sdk/client.py +++ b/posthog/ai/claude_agent_sdk/client.py @@ -242,14 +242,12 @@ async def disconnect(self) -> None: **self._extra_props, } - if resolved_id is None: - properties["$process_person_profile"] = False - self._processor._capture_event( "$ai_trace", properties, resolved_id or self._trace_id, self._groups, + personless=resolved_id is None, ) try: diff --git a/posthog/ai/claude_agent_sdk/processor.py b/posthog/ai/claude_agent_sdk/processor.py index 1d45b61c2..15a830556 100644 --- a/posthog/ai/claude_agent_sdk/processor.py +++ b/posthog/ai/claude_agent_sdk/processor.py @@ -192,6 +192,7 @@ def _capture_event( properties: Dict[str, Any], distinct_id: Optional[str] = None, groups: Optional[Dict[str, Any]] = None, + personless: bool = False, ) -> None: _capture_processor_event( self._client, @@ -200,6 +201,7 @@ def _capture_event( default_properties=self._properties, distinct_id=distinct_id, groups=groups if groups is not None else self._groups, + personless=personless, ) async def query( @@ -449,11 +451,12 @@ def _emit_generation( if gen.stop_reason is not None: properties["$ai_stop_reason"] = gen.stop_reason - if resolved_id is None: - properties["$process_person_profile"] = False - self._capture_event( - "$ai_generation", properties, resolved_id or trace_id, groups + "$ai_generation", + properties, + resolved_id or trace_id, + groups, + personless=resolved_id is None, ) def _emit_generation_from_result( @@ -522,11 +525,12 @@ def _emit_generation_from_result( if result.total_cost_usd is not None: properties["$ai_total_cost_usd"] = result.total_cost_usd - if resolved_id is None: - properties["$process_person_profile"] = False - self._capture_event( - "$ai_generation", properties, resolved_id or trace_id, groups + "$ai_generation", + properties, + resolved_id or trace_id, + groups, + personless=resolved_id is None, ) def _emit_tool_span( @@ -557,10 +561,13 @@ def _emit_tool_span( _ensure_serializable(block.input), self._client ) - if resolved_id is None: - properties["$process_person_profile"] = False - - self._capture_event("$ai_span", properties, resolved_id or trace_id, groups) + self._capture_event( + "$ai_span", + properties, + resolved_id or trace_id, + groups, + personless=resolved_id is None, + ) def _emit_trace( self, @@ -592,10 +599,13 @@ def _emit_trace( if result.total_cost_usd is not None: properties["$ai_total_cost_usd"] = result.total_cost_usd - if resolved_id is None: - properties["$process_person_profile"] = False - - self._capture_event("$ai_trace", properties, resolved_id or trace_id, groups) + self._capture_event( + "$ai_trace", + properties, + resolved_id or trace_id, + groups, + personless=resolved_id is None, + ) # Flush to ensure events are sent before process exits try: diff --git a/posthog/ai/gemini/_shared.py b/posthog/ai/gemini/_shared.py index fc3635589..633c0ac86 100644 --- a/posthog/ai/gemini/_shared.py +++ b/posthog/ai/gemini/_shared.py @@ -266,13 +266,11 @@ def _capture_embedding_outcome( event_properties["$ai_is_error"] = True event_properties["$ai_error"] = str(error) - if distinct_id is None: - event_properties["$process_person_profile"] = False - _capture_ai_event( self._ph_client, "$ai_embedding", distinct_id=distinct_id or trace_id, properties=event_properties, groups=groups, + personless=distinct_id is None, ) diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index 3b9e188b0..a80cd1a1d 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -577,15 +577,13 @@ def _capture_trace_or_span( finalize_ai_content(outputs, self._ph_client), ) - if self._distinct_id is None: - event_properties["$process_person_profile"] = False - _capture_ai_event( self._ph_client, event_name, distinct_id=self._distinct_id or run_id, properties=event_properties, groups=self._groups, + personless=self._distinct_id is None, ) def _pop_run_and_capture_generation( @@ -662,9 +660,6 @@ def _capture_generation( if self._properties: event_properties.update(self._properties) - if self._distinct_id is None: - event_properties["$process_person_profile"] = False - if isinstance(output, BaseException): event_properties["$ai_http_status"] = _get_http_status(output) event_properties["$ai_error"] = _stringify_exception(output) @@ -737,6 +732,7 @@ def _capture_generation( distinct_id=self._distinct_id or trace_id, properties=event_properties, groups=self._groups, + personless=self._distinct_id is None, ) def _log_debug_event( diff --git a/posthog/ai/openai/_embeddings.py b/posthog/ai/openai/_embeddings.py index 8cdb04d33..49dc0f5b5 100644 --- a/posthog/ai/openai/_embeddings.py +++ b/posthog/ai/openai/_embeddings.py @@ -38,9 +38,6 @@ def _capture_embedding_event( **(properties or {}), } - if distinct_id is None: - event_properties["$process_person_profile"] = False - if hasattr(posthog_client, "capture"): _capture_ai_event( posthog_client, @@ -48,4 +45,5 @@ def _capture_embedding_event( distinct_id=distinct_id or trace_id, properties=event_properties, groups=groups, + personless=distinct_id is None, ) diff --git a/posthog/ai/openai_agents/processor.py b/posthog/ai/openai_agents/processor.py index 8bdea6c7e..e90432c2f 100644 --- a/posthog/ai/openai_agents/processor.py +++ b/posthog/ai/openai_agents/processor.py @@ -159,14 +159,18 @@ def _capture_event( properties: Dict[str, Any], distinct_id: Optional[str] = None, ) -> None: - """Capture an event without allowing telemetry failures to escape.""" + """Capture an event without allowing telemetry failures to escape. + + Without a distinct ID, the event is personless and uses its trace ID. + """ _capture_processor_event( self._client, event, properties, default_properties=self._properties, - distinct_id=distinct_id, + distinct_id=distinct_id or properties.get("$ai_trace_id"), groups=self._groups, + personless=distinct_id is None, ) def on_trace_start(self, trace: Trace) -> None: @@ -228,12 +232,9 @@ def on_trace_end(self, trace: Trace) -> None: if metadata: properties["$ai_trace_metadata"] = _ensure_serializable(metadata) - if distinct_id is None: - properties["$process_person_profile"] = False - self._capture_event( event="$ai_trace", - distinct_id=distinct_id or trace_id, + distinct_id=distinct_id, properties=properties, ) except Exception as e: @@ -267,8 +268,8 @@ def on_span_end(self, span: Span[Any]) -> None: latency = (ended - started) if (started and ended) else 0 # Get user-provided distinct ID from trace metadata (resolved at trace start). - # None means no user-provided ID — use trace_id as fallback in personless mode, - # matching the langchain/openai pattern: `distinct_id or trace_id`. + # None means no user-provided ID: _capture_event falls back to the + # trace ID and makes the event personless. trace_info = self._trace_metadata.get(trace_id, {}) distinct_id = trace_info.get("distinct_id") or self._get_distinct_id(None) @@ -317,11 +318,6 @@ def on_span_end(self, span: Span[Any]) -> None: "$ai_error_type": error_type, } - # Personless mode: no user-provided distinct_id, fallback to trace_id - if distinct_id is None: - error_properties["$process_person_profile"] = False - distinct_id = trace_id - # Dispatch based on span data type if isinstance(span_data, GenerationSpanData): self._handle_generation_span( @@ -471,7 +467,7 @@ def _handle_generation_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -554,7 +550,7 @@ def _handle_function_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -587,7 +583,7 @@ def _handle_agent_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -616,7 +612,7 @@ def _handle_handoff_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -640,7 +636,7 @@ def _handle_guardrail_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -663,7 +659,7 @@ def _handle_response_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -732,7 +728,7 @@ def _handle_custom_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -757,7 +753,7 @@ def _handle_audio_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -832,7 +828,7 @@ def _handle_mcp_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: @@ -856,7 +852,7 @@ def _handle_generic_span( span_id: str, parent_id: Optional[str], latency: float, - distinct_id: str, + distinct_id: Optional[str], group_id: Optional[str], error_properties: Dict[str, Any], ) -> None: diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index ffc2c864c..ae652e80a 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -69,12 +69,24 @@ def _ai_lane_enabled(ph_client) -> bool: return _full_ai_capture_enabled(ph_client) -def _capture_ai_event(ph_client, event: str, **kwargs): - """Capture a wrapper-emitted AI event with the PostHog AI library identity.""" +def _capture_ai_event(ph_client, event: str, *, personless: bool = False, **kwargs): + """Capture a wrapper-emitted AI event with the PostHog AI library identity. + + ``personless`` marks an event whose distinct ID is a fallback, such as a + trace ID, not a real user. + """ kwargs["properties"] = { **_AI_LIB_PROPERTIES, **(kwargs.get("properties") or {}), } + if personless: + # A per-event option beats legacy properties and global or context + # options, because a person profile for a fallback ID is always wrong. + # before_send can still change it. + kwargs["options"] = { + **(kwargs.get("options") or {}), + "process_person_profile": False, + } if _ai_lane_enabled(ph_client): capture_ai = getattr(ph_client, "capture_ai", None) if callable(capture_ai): @@ -90,6 +102,7 @@ def _capture_processor_event( default_properties: Optional[Dict[str, Any]] = None, distinct_id: Optional[str] = None, groups: Optional[Dict[str, Any]] = None, + personless: bool = False, ) -> None: """Apply the shared capture policy used by AI SDK processors.""" try: @@ -103,6 +116,7 @@ def _capture_processor_event( distinct_id=distinct_id or "unknown", properties={**properties, **(default_properties or {})}, groups=groups, + personless=personless, ) except Exception as exc: logging.getLogger("posthog").debug("Failed to capture PostHog event: %s", exc) @@ -586,9 +600,6 @@ def call_llm_and_track_usage( if stop_reason is not None: tag("$ai_stop_reason", stop_reason) - if not has_person_distinct_id: - tag("$process_person_profile", False) - # Process instructions for Responses API if provider == "openai" and kwargs.get("instructions") is not None: tag( @@ -615,6 +626,7 @@ def call_llm_and_track_usage( distinct_id=contexts.get_context_distinct_id(), properties=merged_properties, groups=posthog_groups, + personless=not has_person_distinct_id, ) if error: @@ -753,9 +765,6 @@ async def call_llm_and_track_usage_async( if stop_reason is not None: tag("$ai_stop_reason", stop_reason) - if not has_person_distinct_id: - tag("$process_person_profile", False) - # Process instructions for Responses API if provider == "openai" and kwargs.get("instructions") is not None: tag( @@ -782,6 +791,7 @@ async def call_llm_and_track_usage_async( distinct_id=contexts.get_context_distinct_id(), properties=merged_properties, groups=posthog_groups, + personless=not has_person_distinct_id, ) if error: @@ -964,9 +974,6 @@ def capture_streaming_event( event_data["kwargs"]["instructions"], ) - if event_data.get("distinct_id") is None: - event_properties["$process_person_profile"] = False - # Send event to PostHog if hasattr(ph_client, "capture"): _capture_ai_event( @@ -975,4 +982,5 @@ def capture_streaming_event( distinct_id=event_data.get("distinct_id") or trace_id, properties=event_properties, groups=event_data.get("groups"), + personless=event_data.get("distinct_id") is None, ) diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index 512d9a37f..0e6a37e54 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -64,7 +64,6 @@ def _build_capture_event(event: Event) -> PostHogCaptureEvent: properties: Dict[str, Any] = {_P.SOURCE: POSTHOG_MCP_ANALYTICS_SOURCE} _add_session_id(event, properties) _add_conversation_id(event, properties) - _add_person_processing(event, properties) _add_common_properties(event, properties) _add_custom_properties(event, properties) _add_groups(event, properties) @@ -77,6 +76,7 @@ def _build_capture_event(event: Event) -> PostHogCaptureEvent: "event": event_name, "distinct_id": _get_distinct_id(event), "properties": properties, + "options": _get_options(event), "timestamp": _get_timestamp(event), } @@ -100,11 +100,13 @@ def _add_groups(event: Event, properties: Dict[str, Any]) -> None: _merge_groups(properties, groups) -def _add_person_processing(event: Event, properties: Dict[str, Any]) -> None: +def _get_options(event: Event) -> Dict[str, Any]: # Without a resolved identity the distinct id is just the session id, so # processing a person profile would mint one anonymous person per session. + # As an option it beats any custom $process_person_profile property. if not event.get("identify_actor_given_id"): - properties["$process_person_profile"] = False + return {"process_person_profile": False} + return {} def _is_tool_call(event: Event) -> bool: @@ -266,7 +268,6 @@ def _build_exception_event(event: Event) -> PostHogCaptureEvent: properties: Dict[str, Any] = {} _add_session_id(event, properties) _add_conversation_id(event, properties) - _add_person_processing(event, properties) error = event.get("error") if error: @@ -301,5 +302,6 @@ def _build_exception_event(event: Event) -> PostHogCaptureEvent: "event": PostHogMCPAnalyticsEvent.EXCEPTION, "distinct_id": _get_distinct_id(event), "properties": properties, + "options": _get_options(event), "timestamp": _get_timestamp(event), } diff --git a/posthog/mcp/_sink.py b/posthog/mcp/_sink.py index 7d1db967f..15a5e789a 100644 --- a/posthog/mcp/_sink.py +++ b/posthog/mcp/_sink.py @@ -106,6 +106,7 @@ async def capture(self, event: McpEvent, options: McpCaptureOptions) -> None: capture_event["event"], distinct_id=capture_event["distinct_id"], properties=capture_event["properties"], + options=capture_event.get("options"), timestamp=capture_event.get("timestamp"), uuid=_uuid7(), ) diff --git a/posthog/mcp/types.py b/posthog/mcp/types.py index f65ae0d26..5b0efbc1a 100644 --- a/posthog/mcp/types.py +++ b/posthog/mcp/types.py @@ -78,6 +78,7 @@ class PostHogCaptureEvent(TypedDict, total=False): distinct_id: str event: str properties: Dict[str, Any] + options: Dict[str, Any] timestamp: datetime diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index 84275e9ad..a4a3e9bc6 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -1735,25 +1735,52 @@ async def run_test(): # ======================= -def test_no_distinct_id_uses_trace_id_and_personless( - mock_client, mock_anthropic_response +@pytest.mark.parametrize("stream", [False, True], ids=["response", "stream"]) +@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"]) +async def test_no_distinct_id_uses_trace_id_and_personless( + mock_client, + mock_anthropic_response, + mock_anthropic_stream_with_tools, + is_async, + stream, ): """When no distinct_id is provided and no outer context, trace_id is used and event is personless.""" - with patch( - "anthropic.resources.Messages.create", return_value=mock_anthropic_response - ): - client = Anthropic(api_key="test-key", posthog_client=mock_client) - client.messages.create( - model="claude-3-opus-20240229", - messages=[{"role": "user", "content": "Hello"}], - posthog_trace_id="trace-123", + request = { + "model": "claude-3-opus-20240229", + "messages": [{"role": "user", "content": "Hello"}], + "posthog_trace_id": "trace-123", + "posthog_properties": {"$process_person_profile": True}, + **({"stream": True} if stream else {}), + } + if is_async: + + async def stream_events(): + for event in mock_anthropic_stream_with_tools: + yield event + + async def create(**kwargs): + return stream_events() if stream else mock_anthropic_response + + with patch("anthropic.resources.AsyncMessages.create", side_effect=create): + client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) + response = await client.messages.create(**request) + if stream: + [event async for event in response] + else: + provider_response = ( + mock_anthropic_stream_with_tools if stream else mock_anthropic_response ) + with patch( + "anthropic.resources.Messages.create", return_value=provider_response + ): + client = Anthropic(api_key="test-key", posthog_client=mock_client) + response = client.messages.create(**request) + if stream: + list(response) - call_args = mock_client.capture.call_args[1] - props = call_args["properties"] - - assert call_args["distinct_id"] == "trace-123" - assert props["$process_person_profile"] is False + call_args = mock_client.capture.call_args[1] + assert call_args["distinct_id"] == "trace-123" + assert call_args["options"] == {"process_person_profile": False} def test_explicit_distinct_id_creates_person_profile( @@ -1775,10 +1802,8 @@ def test_explicit_distinct_id_creates_person_profile( props = call_args["properties"] assert call_args["distinct_id"] == "user-123" - assert ( - "$process_person_profile" not in props - or props["$process_person_profile"] is not False - ) + assert "options" not in call_args + assert "$process_person_profile" not in props def test_outer_context_distinct_id_is_used(mock_client, mock_anthropic_response): @@ -1799,10 +1824,8 @@ def test_outer_context_distinct_id_is_used(mock_client, mock_anthropic_response) props = call_args["properties"] assert call_args["distinct_id"] == "outer-user-456" - assert ( - "$process_person_profile" not in props - or props["$process_person_profile"] is not False - ) + assert "options" not in call_args + assert "$process_person_profile" not in props def test_explicit_distinct_id_overrides_outer_context( diff --git a/posthog/test/ai/claude_agent_sdk/test_processor.py b/posthog/test/ai/claude_agent_sdk/test_processor.py index 3dc776002..1fb49c664 100644 --- a/posthog/test/ai/claude_agent_sdk/test_processor.py +++ b/posthog/test/ai/claude_agent_sdk/test_processor.py @@ -644,7 +644,7 @@ async def test_data_url_prompt_and_tool_use_base64_are_redacted( class TestPersonlessMode: @pytest.mark.asyncio - async def test_no_distinct_id_sets_process_person_profile_false(self, mock_client): + async def test_no_distinct_id_sets_personless_option(self, mock_client): proc = PostHogClaudeAgentProcessor( client=mock_client, distinct_id=None, @@ -668,8 +668,7 @@ async def test_no_distinct_id_sets_process_person_profile_false(self, mock_clien pass for call in mock_client.capture.call_args_list: - props = call.kwargs.get("properties") or call[1].get("properties") - assert props.get("$process_person_profile") is False + assert call.kwargs["options"] == {"process_person_profile": False} assert call.kwargs["distinct_id"] == "trace-fallback" diff --git a/posthog/test/ai/gemini/test_gemini.py b/posthog/test/ai/gemini/test_gemini.py index f2a801e68..0079e73db 100644 --- a/posthog/test/ai/gemini/test_gemini.py +++ b/posthog/test/ai/gemini/test_gemini.py @@ -1376,7 +1376,8 @@ def test_embed_content_no_distinct_id( # Should fall back to trace_id as distinct_id assert call_args["distinct_id"] == props["$ai_trace_id"] - assert props["$process_person_profile"] is False + assert call_args["options"] == {"process_person_profile": False} + assert "$process_person_profile" not in props def test_embed_content_default_params( diff --git a/posthog/test/ai/gemini/test_gemini_async.py b/posthog/test/ai/gemini/test_gemini_async.py index 8dc51ca34..1373a331e 100644 --- a/posthog/test/ai/gemini/test_gemini_async.py +++ b/posthog/test/ai/gemini/test_gemini_async.py @@ -995,7 +995,8 @@ async def test_async_embed_content_no_distinct_id( # Should fall back to trace_id as distinct_id assert call_args["distinct_id"] == props["$ai_trace_id"] - assert props["$process_person_profile"] is False + assert call_args["options"] == {"process_person_profile": False} + assert "$process_person_profile" not in props async def test_async_embed_content_default_params( diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 2556294ae..7bc1b4166 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -543,13 +543,13 @@ def test_personless_mode(mock_client): # span assert span_args["event"] == "$ai_span" - assert span_args["properties"]["$process_person_profile"] is False + assert span_args["options"] == {"process_person_profile": False} # generation assert generation_args["event"] == "$ai_generation" - assert generation_args["properties"]["$process_person_profile"] is False + assert generation_args["options"] == {"process_person_profile": False} # trace assert trace_args["event"] == "$ai_trace" - assert trace_args["properties"]["$process_person_profile"] is False + assert trace_args["options"] == {"process_person_profile": False} id = uuid.uuid4() chain.invoke( @@ -561,13 +561,13 @@ def test_personless_mode(mock_client): trace_args = mock_client.capture.call_args_list[5][1] # span - assert "$process_person_profile" not in span_args["properties"] + assert "options" not in span_args assert span_args["distinct_id"] == id # generation - assert "$process_person_profile" not in generation_args["properties"] + assert "options" not in generation_args assert generation_args["distinct_id"] == id # trace - assert "$process_person_profile" not in trace_args["properties"] + assert "options" not in trace_args assert trace_args["distinct_id"] == id @@ -602,13 +602,13 @@ def test_personless_mode_exception(mock_client, unauthorized_http_client): # span assert span_args["event"] == "$ai_span" - assert span_args["properties"]["$process_person_profile"] is False + assert span_args["options"] == {"process_person_profile": False} # generation assert generation_args["event"] == "$ai_generation" - assert generation_args["properties"]["$process_person_profile"] is False + assert generation_args["options"] == {"process_person_profile": False} # trace assert trace_args["event"] == "$ai_trace" - assert trace_args["properties"]["$process_person_profile"] is False + assert trace_args["options"] == {"process_person_profile": False} id = uuid.uuid4() with pytest.raises(AuthenticationError): @@ -622,17 +622,17 @@ def test_personless_mode_exception(mock_client, unauthorized_http_client): # span assert span_args["event"] == "$ai_span" - assert "$process_person_profile" not in span_args["properties"] + assert "options" not in span_args assert span_args["distinct_id"] == id # generation assert generation_args["event"] == "$ai_generation" - assert "$process_person_profile" not in generation_args["properties"] + assert "options" not in generation_args assert generation_args["distinct_id"] == id # trace assert trace_args["event"] == "$ai_trace" - assert "$process_person_profile" not in trace_args["properties"] + assert "options" not in trace_args assert trace_args["distinct_id"] == id diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index 850b2a885..75ab907c5 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -186,7 +186,8 @@ async def test_embedding_telemetry_has_sync_async_parity(mock_client): assert async_props["$ai_input"] is None assert async_props["$ai_input_tokens"] == 12 assert async_props["$ai_provider"] == "azure" - assert async_props["$process_person_profile"] is False + assert async_capture.kwargs["options"] == sync_capture.kwargs["options"] + assert async_capture.kwargs["options"] == {"process_person_profile": False} assert async_capture.kwargs["distinct_id"] == "shared-trace" assert async_capture.kwargs["groups"] == {"company": "test-company"} diff --git a/posthog/test/ai/openai_agents/test_processor.py b/posthog/test/ai/openai_agents/test_processor.py index 3cbc6d276..69745413b 100644 --- a/posthog/test/ai/openai_agents/test_processor.py +++ b/posthog/test/ai/openai_agents/test_processor.py @@ -157,7 +157,7 @@ def test_personless_mode_when_no_distinct_id(self, mock_client, mock_trace): processor.on_trace_end(mock_trace) call_kwargs = mock_client.capture.call_args[1] - assert call_kwargs["properties"]["$process_person_profile"] is False + assert call_kwargs["options"] == {"process_person_profile": False} # Should fallback to trace_id as the distinct_id assert call_kwargs["distinct_id"] == mock_trace.trace_id @@ -179,7 +179,7 @@ def test_personless_mode_for_spans_when_no_distinct_id( processor.on_span_end(mock_span) call_kwargs = mock_client.capture.call_args[1] - assert call_kwargs["properties"]["$process_person_profile"] is False + assert call_kwargs["options"] == {"process_person_profile": False} assert call_kwargs["distinct_id"] == mock_span.trace_id def test_personless_mode_when_callable_returns_none( @@ -205,7 +205,7 @@ def resolver(trace): processor.on_span_end(mock_span) call_kwargs = mock_client.capture.call_args[1] - assert call_kwargs["properties"]["$process_person_profile"] is False + assert call_kwargs["options"] == {"process_person_profile": False} assert call_kwargs["distinct_id"] == mock_span.trace_id def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace): @@ -219,7 +219,7 @@ def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace) processor.on_trace_end(mock_trace) call_kwargs = mock_client.capture.call_args[1] - assert "$process_person_profile" not in call_kwargs["properties"] + assert "options" not in call_kwargs def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace): """Test that on_trace_end clears stored trace metadata.""" diff --git a/posthog/test/mcp/_helpers.py b/posthog/test/mcp/_helpers.py index 836ef8ad6..cd4227c8e 100644 --- a/posthog/test/mcp/_helpers.py +++ b/posthog/test/mcp/_helpers.py @@ -29,7 +29,12 @@ def capture( **kwargs, ): self.events.append( - {"event": event, "distinct_id": distinct_id, "properties": properties or {}} + { + "event": event, + "distinct_id": distinct_id, + "properties": properties or {}, + "options": kwargs.get("options") or {}, + } ) return None diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py index df725009e..f3e904982 100644 --- a/posthog/test/mcp/test_fastmcp.py +++ b/posthog/test/mcp/test_fastmcp.py @@ -255,7 +255,7 @@ async def test_identify_sets_distinct_id_and_groups(): calls = _events(client, "$mcp_tool_call") assert calls[0]["distinct_id"] == "user_42" assert calls[0]["properties"]["$groups"] == {"organization": "org_7"} - assert "$process_person_profile" not in calls[0]["properties"] + assert calls[0]["options"] == {} # an $identify event was emitted assert _events(client, "$identify") diff --git a/posthog/test/mcp/test_pipeline.py b/posthog/test/mcp/test_pipeline.py index 413941c28..729a8332d 100644 --- a/posthog/test/mcp/test_pipeline.py +++ b/posthog/test/mcp/test_pipeline.py @@ -19,7 +19,10 @@ sanitize_captured_value, sanitize_event, ) -from posthog.mcp._sink import McpCaptureOptions, process_mcp_event +from posthog.capture_event import _to_v1_event +from posthog.client import Client +from posthog.mcp._sink import McpCaptureOptions, McpEventSink, process_mcp_event +from posthog.test.capture_helpers import patch_capture_send, sent_events from posthog.mcp._truncation import MAX_EVENT_BYTES, normalize, truncate_event # --- ids --------------------------------------------------------------------- @@ -895,7 +898,8 @@ def test_build_tool_call_event_properties(): assert props[PostHogMCPAnalyticsProperty.LLM_MODEL_SOURCE] == "client_metadata" assert props[PostHogMCPAnalyticsProperty.SESSION_ID] == "ses_abc" # anonymous (no identity) => person processing disabled - assert props["$process_person_profile"] is False + assert capture["options"] == {"process_person_profile": False} + assert "$process_person_profile" not in props def test_protocol_version_on_primary_and_exception_events(): @@ -931,7 +935,7 @@ def test_identity_enables_person_processing_and_set(): [capture] = build_posthog_capture_events(event) props = capture["properties"] assert capture["distinct_id"] == "user_1" - assert "$process_person_profile" not in props + assert capture["options"] == {} assert props["$set"] == {"email": "a@b.com"} assert props["$groups"] == {"organization": "org_1", "project": "p1"} @@ -1131,3 +1135,32 @@ async def before_send(e): result = await process_mcp_event(event, McpCaptureOptions(before_send=before_send)) _, captures = result assert captures[0]["properties"]["added"] is True + + +def _opt_in(e): + return {**e, "options": {"process_person_profile": True}} + + +@pytest.mark.parametrize( + "before_send, expected", + [(None, False), (_opt_in, True)], + ids=["sdk_option_wins", "before_send_can_change_it"], +) +async def test_anonymous_event_sends_personless_option(before_send, expected): + event = { + "event_type": MCPAnalyticsEventType.MCP_TOOLS_CALL, + "session_id": "ses_abc", + "properties": {"$process_person_profile": True}, + "timestamp": datetime.now(timezone.utc), + } + with patch_capture_send("client") as send: + client = Client( + "test-key", + sync_mode=True, + super_options={"process_person_profile": True}, + ) + await McpEventSink(client).capture( + event, McpCaptureOptions(before_send=before_send) + ) + [wire] = [_to_v1_event(msg) for msg in sent_events(send)] + assert wire["options"] == {"process_person_profile": expected} diff --git a/posthog/test/mcp/test_v2_mcpserver.py b/posthog/test/mcp/test_v2_mcpserver.py index c17e6de15..7fbf0919d 100644 --- a/posthog/test/mcp/test_v2_mcpserver.py +++ b/posthog/test/mcp/test_v2_mcpserver.py @@ -544,7 +544,8 @@ async def test_anonymous_events_do_not_create_person_profiles(): await _flush() calls = _events(client, "$mcp_tool_call") - assert calls[0]["properties"]["$process_person_profile"] is False + assert calls[0]["options"] == {"process_person_profile": False} + assert "$process_person_profile" not in calls[0]["properties"] async def test_a_failed_registry_lookup_delegates_instead_of_swallowing(): diff --git a/posthog/test/test_event_options.py b/posthog/test/test_event_options.py index 7e9f18fda..b94f70143 100644 --- a/posthog/test/test_event_options.py +++ b/posthog/test/test_event_options.py @@ -4,6 +4,7 @@ import pytest from posthog import AsyncPosthog, new_context, set_context_option, tag +from posthog.ai.utils import _capture_ai_event from posthog.capture_event import _to_v1_event from posthog.client import Client from posthog.test.capture_helpers import ( @@ -471,3 +472,31 @@ def call(client): "product_tour_id": "context-tour", } assert not {"$cookieless_mode", "$product_tour_id"} & events[0]["properties"].keys() + + +def _personless_ai_capture(client): + with new_context(fresh=True): + set_context_option(PP, True) + return _capture_ai_event( + client, + "$ai_generation", + distinct_id="trace-1", + properties={"$process_person_profile": True}, + options={PP: True, "cookieless_mode": True}, + personless=True, + ) + + +@pytest.mark.parametrize( + "before_send, expected", + [ + (None, {PP: False, "cookieless_mode": True}), + (lambda m: {**m, "options": {PP: True}}, {PP: True}), + ], + ids=["sdk_option_wins", "before_send_can_change_it"], +) +def test_personless_ai_event_overrides_every_user_layer(before_send, expected): + events = _sync_wire_events( + _personless_ai_capture, before_send=before_send, super_options={PP: True} + ) + assert events[0]["options"] == expected From 49399edce37f1cd8e62cecc6735b434796b40ca0 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:46:50 -0700 Subject: [PATCH 11/14] chore(ci): run compliance harness 1.13.1 once per codec (#1026) * chore(ci): run compliance harness 1.13.1 once per codec The adapter now advertises capture_ai_v1 and event_options, passes options and uuids through unchanged, and takes its codec from COMPRESSION. The workflow runs the harness containers directly in a gzip/deflate/zstd matrix, so it no longer needs pull-requests: write. * chore(ci): install zstd extra for adapter protocol tests * chore: describe the single capture compliance job and harness 1.13.1 The standalone adapter-test setup installs the zstd extra, matching CI. --- .github/workflows/sdk-compliance.yml | 85 +++++++++-- sdk_compliance_adapter/CONTRIBUTING.md | 6 +- sdk_compliance_adapter/Dockerfile | 4 +- sdk_compliance_adapter/README.md | 4 +- sdk_compliance_adapter/adapter.py | 166 ++++++++++------------ sdk_compliance_adapter/docker-compose.yml | 5 +- sdk_compliance_adapter/test_adapter.py | 52 ++++++- 7 files changed, 207 insertions(+), 115 deletions(-) diff --git a/.github/workflows/sdk-compliance.yml b/.github/workflows/sdk-compliance.yml index 55cfc8ebe..b77681cad 100644 --- a/.github/workflows/sdk-compliance.yml +++ b/.github/workflows/sdk-compliance.yml @@ -2,8 +2,6 @@ name: SDK Compliance Tests permissions: contents: read - packages: read - pull-requests: write on: pull_request: @@ -23,16 +21,81 @@ jobs: with: python-version: "3.12" - name: Install adapter and test dependencies - run: python -m pip install -e . -r sdk_compliance_adapter/requirements.txt pytest pytest-timeout pytest-asyncio + run: python -m pip install -e ".[zstd]" -r sdk_compliance_adapter/requirements.txt pytest pytest-timeout pytest-asyncio - name: Test adapter protocol run: python -m pytest sdk_compliance_adapter/test_adapter.py --timeout=30 compliance: - name: PostHog SDK compliance tests - uses: PostHog/posthog-sdk-test-harness/.github/workflows/test-sdk-action.yml@4593de8b423f61fa222115da592e5c18dc82ad3c # 1.11.0 - with: - adapter-dockerfile: "sdk_compliance_adapter/Dockerfile" - adapter-context: "." - test-harness-version: "1.1.1" - continue-on-error: false - report-name: "sdk-compliance-report" + # The harness has one enable_compression flag but a test per codec, so each + # job runs the adapter with one codec. The reusable harness workflow cannot + # pass the COMPRESSION env var, so this job runs the containers directly. + # The SDK has no brotli support, so there is no br job. + name: PostHog SDK compliance tests (${{ matrix.compression }}) + runs-on: ubuntu-latest + timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + compression: [gzip, deflate, zstd] + concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ matrix.compression }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + env: + HARNESS_IMAGE: ghcr.io/posthog/sdk-test-harness:1.13.1 + steps: + - name: Check out source code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - name: Pull test harness + run: docker pull "$HARNESS_IMAGE" + + - name: Build SDK adapter + run: docker build -f sdk_compliance_adapter/Dockerfile -t sdk-adapter:test . + + - name: Start SDK adapter + run: | + docker network create test-network + docker run -d \ + --name sdk-adapter \ + --network test-network \ + -e COMPRESSION=${{ matrix.compression }} \ + sdk-adapter:test + + - name: Run compliance suite + run: | + mkdir -p report + docker run --rm \ + --name test-harness \ + --network test-network \ + -v "${GITHUB_WORKSPACE}/report:/report" \ + "$HARNESS_IMAGE" \ + run \ + --adapter-url http://sdk-adapter:8080 \ + --mock-url http://test-harness:8081 \ + --output text \ + --report /report/sdk-compliance-report.md \ + --sdk-type server + + - name: Publish report to the job summary + if: always() + run: | + if [ -f report/sdk-compliance-report.md ]; then + cat report/sdk-compliance-report.md >> "$GITHUB_STEP_SUMMARY" + fi + + - name: Upload compliance report + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: sdk-compliance-report-${{ matrix.compression }} + path: report/sdk-compliance-report.md + + - name: Show adapter logs on failure + if: failure() + run: docker logs sdk-adapter || true + + - name: Stop containers + if: always() + run: | + docker rm -f sdk-adapter || true + docker network rm test-network || true diff --git a/sdk_compliance_adapter/CONTRIBUTING.md b/sdk_compliance_adapter/CONTRIBUTING.md index 3033609b3..ffddda4b3 100644 --- a/sdk_compliance_adapter/CONTRIBUTING.md +++ b/sdk_compliance_adapter/CONTRIBUTING.md @@ -4,12 +4,12 @@ This package contains the PostHog Python SDK compliance adapter used with the Po ## Running tests -Tests run automatically in CI via GitHub Actions against harness **1.1.1**, for both capture protocols. Both jobs opt into the local-evaluation suite and fail on compliance regressions. +Tests run automatically in CI via GitHub Actions against harness **1.13.1**. The compliance job runs once for each capture compression (gzip, deflate and zstd). Each run opts into the local-evaluation suite and fails on compliance regressions. Run adapter protocol tests from the repository root in an activated virtual environment: ```bash -python -m pip install -e . -r sdk_compliance_adapter/requirements.txt pytest pytest-timeout pytest-asyncio +python -m pip install -e ".[zstd]" -r sdk_compliance_adapter/requirements.txt pytest pytest-timeout pytest-asyncio python -m pytest sdk_compliance_adapter/test_adapter.py --timeout=30 ``` @@ -44,7 +44,7 @@ docker run -d --name sdk-adapter --network test-network -p 8080:8080 posthog-pyt docker run --rm \ --name test-harness \ --network test-network \ - ghcr.io/posthog/sdk-test-harness:1.1.1 \ + ghcr.io/posthog/sdk-test-harness:1.13.1 \ run --adapter-url http://sdk-adapter:8080 --mock-url http://test-harness:8081 # Cleanup diff --git a/sdk_compliance_adapter/Dockerfile b/sdk_compliance_adapter/Dockerfile index c347f3266..bfc691dd3 100644 --- a/sdk_compliance_adapter/Dockerfile +++ b/sdk_compliance_adapter/Dockerfile @@ -6,8 +6,8 @@ WORKDIR /app COPY posthog/ /app/sdk/posthog/ COPY setup.py pyproject.toml README.md LICENSE /app/sdk/ -# Install the SDK from source -RUN cd /app/sdk && pip install --no-cache-dir -e . +# Install the SDK from source, with zstd for the zstd compliance job +RUN cd /app/sdk && pip install --no-cache-dir -e ".[zstd]" # Install adapter dependencies RUN pip install --no-cache-dir flask python-dateutil diff --git a/sdk_compliance_adapter/README.md b/sdk_compliance_adapter/README.md index 6b17bf5dd..9efbee7ce 100644 --- a/sdk_compliance_adapter/README.md +++ b/sdk_compliance_adapter/README.md @@ -38,8 +38,8 @@ The adapter implements the standard SDK adapter interface defined in the [test h ### Local feature flag evaluation -Both capture adapters advertise `feature_flags_local_evaluation_v1` for harness -**1.1.1**. The capability versions the adapter protocol and tests both legacy and +The adapter advertises `feature_flags_local_evaluation_v1` for harness +**1.13.1**. The capability versions the adapter protocol and tests both legacy and explicit property matching; it does not change the SDK's default matching mode. - `/init` maps optional `personal_api_key` to the SDK's `secret_key`. Ordinary diff --git a/sdk_compliance_adapter/adapter.py b/sdk_compliance_adapter/adapter.py index be2e1905b..19068152b 100644 --- a/sdk_compliance_adapter/adapter.py +++ b/sdk_compliance_adapter/adapter.py @@ -28,6 +28,19 @@ app = Flask(__name__) +# The harness has one enable_compression flag but a test per codec, so each +# adapter instance takes its codec from COMPRESSION and advertises only that +# codec's encoding_ capability. The SDK has no brotli support. +COMPRESSION_CODECS = { + "gzip": CaptureCompression.GZIP, + "deflate": CaptureCompression.DEFLATE, + "zstd": CaptureCompression.ZSTD, +} + + +def compression_name() -> str: + return os.environ.get("COMPRESSION") or "gzip" + class RequestInfo: """Information about an HTTP request made by the SDK""" @@ -220,11 +233,11 @@ def patched_post_v1( @app.route("/health", methods=["GET"]) def health(): """Health check endpoint""" - # No AI capture capability: `capture_ai` posts capture v1 to - # /i/v1/ai/events, which this harness version has no suite for. capabilities = [ "capture_v1", - "encoding_gzip", + "capture_ai_v1", + "event_options", + "encoding_" + compression_name(), "feature_flags_local_evaluation_v1", ] return jsonify( @@ -267,6 +280,11 @@ def init(): # Convert flush_interval from ms to seconds flush_interval = flush_interval_ms / 1000.0 + compression = ( + COMPRESSION_CODECS[compression_name()] + if enable_compression + else CaptureCompression.NONE + ) # Explicit reloads exercise the real loader without background polling # racing the harness's per-test definition snapshots. client_options = { @@ -274,11 +292,8 @@ def init(): "host": host, "flush_at": flush_at, "flush_interval": flush_interval, - "capture_compression": ( - CaptureCompression.GZIP - if enable_compression - else CaptureCompression.NONE - ), + "capture_compression": compression, + "capture_ai_compression": compression, "max_retries": max_retries, "debug": False, "disable_geoip": disable_geoip, @@ -308,56 +323,55 @@ def init(): return jsonify({"error": str(e)}), 500 -@app.route("/capture", methods=["POST"]) -def capture(): - """Capture a single event""" - try: - if not state.client: - return jsonify({"error": "SDK not initialized"}), 400 +def _capture_with(capture_method_name: str): + """Capture one event through the named SDK method. - data = request.json or {} + Options pass through unchanged, and a supplied uuid stays on the wire. The + response carries the uuid the SDK actually sent, which differs from the + supplied one when the SDK canonicalizes or replaces it. + """ + if not state.client: + return jsonify({"error": "SDK not initialized"}), 400 - distinct_id = data.get("distinct_id") - event = data.get("event") - properties = data.get("properties") - timestamp = data.get("timestamp") - options = data.get("options") + data = request.json or {} - if not distinct_id: - return jsonify({"error": "distinct_id is required"}), 400 - if not event: - return jsonify({"error": "event is required"}), 400 - - # Fold capture-v1 options back into the magic `$`-prefixed properties the - # SDK lifts onto the wire `options` object. Renamed keys mirror the SDK's - # sentinel table; unknown keys get a bare `$` prefix. - if options: - properties = dict(properties or {}) - option_to_property = { - "cookieless_mode": "$cookieless_mode", - "disable_skew_correction": "$ignore_sent_at", - "process_person_profile": "$process_person_profile", - "product_tour_id": "$product_tour_id", - } - for key, value in options.items(): - properties[option_to_property.get(key, "$" + key)] = value + distinct_id = data.get("distinct_id") + event = data.get("event") + timestamp = data.get("timestamp") + + if not distinct_id: + return jsonify({"error": "distinct_id is required"}), 400 + if not event: + return jsonify({"error": "event is required"}), 400 + + kwargs = { + "distinct_id": distinct_id, + "properties": data.get("properties"), + "options": data.get("options"), + } + if timestamp: + # Parse ISO8601 timestamp + from dateutil.parser import parse - # Capture event - kwargs = {"distinct_id": distinct_id, "properties": properties} - if timestamp: - # Parse ISO8601 timestamp - from dateutil.parser import parse + kwargs["timestamp"] = parse(timestamp) + if data.get("uuid"): + kwargs["uuid"] = data["uuid"] - kwargs["timestamp"] = parse(timestamp) + uuid = getattr(state.client, capture_method_name)(event, **kwargs) - uuid = state.client.capture(event, **kwargs) + # Track that we captured an event + state.increment_captured() - # Track that we captured an event - state.increment_captured() + logger.info(f"Captured event: {event} for {distinct_id}, uuid={uuid}") - logger.info(f"Captured event: {event} for {distinct_id}, uuid={uuid}") + return jsonify({"success": True, "uuid": uuid}) - return jsonify({"success": True, "uuid": uuid}) + +@app.route("/capture", methods=["POST"]) +def capture(): + """Capture a single event""" + try: + return _capture_with("capture") except Exception as e: logger.exception("Error capturing event") state.record_error(str(e)) @@ -368,50 +382,7 @@ def capture(): def capture_ai(): """Capture a single AI event on the dedicated AI capture endpoint""" try: - if not state.client: - return jsonify({"error": "SDK not initialized"}), 400 - - data = request.json or {} - - distinct_id = data.get("distinct_id") - event = data.get("event") - properties = data.get("properties") - timestamp = data.get("timestamp") - options = data.get("options") - supplied_uuid = data.get("uuid") - - if not distinct_id: - return jsonify({"error": "distinct_id is required"}), 400 - if not event: - return jsonify({"error": "event is required"}), 400 - - if options: - properties = dict(properties or {}) - option_to_property = { - "cookieless_mode": "$cookieless_mode", - "disable_skew_correction": "$ignore_sent_at", - "process_person_profile": "$process_person_profile", - "product_tour_id": "$product_tour_id", - } - for key, value in options.items(): - properties[option_to_property.get(key, "$" + key)] = value - - kwargs = {"distinct_id": distinct_id, "properties": properties} - if timestamp: - from dateutil.parser import parse - - kwargs["timestamp"] = parse(timestamp) - # Unlike /capture, forward a supplied uuid so it's echoed back to the caller. - if supplied_uuid: - kwargs["uuid"] = supplied_uuid - - uuid = state.client.capture_ai(event, **kwargs) - - state.increment_captured() - - logger.info(f"Captured AI event: {event} for {distinct_id}, uuid={uuid}") - - return jsonify({"success": True, "uuid": uuid}) + return _capture_with("capture_ai") except Exception as e: logger.exception("Error capturing AI event") state.record_error(str(e)) @@ -638,7 +609,14 @@ def reset(): def main(): """Main entry point""" port = int(os.environ.get("PORT", 8080)) - logger.info(f"Starting SDK Test Adapter on port {port}") + if compression_name() not in COMPRESSION_CODECS: + raise SystemExit( + f"unsupported COMPRESSION {compression_name()!r}: " + f"want one of {', '.join(COMPRESSION_CODECS)}" + ) + logger.info( + f"Starting SDK Test Adapter on port {port} (compression={compression_name()})" + ) app.run(host="0.0.0.0", port=port, debug=False) diff --git a/sdk_compliance_adapter/docker-compose.yml b/sdk_compliance_adapter/docker-compose.yml index 56ef72680..9281fa8e3 100644 --- a/sdk_compliance_adapter/docker-compose.yml +++ b/sdk_compliance_adapter/docker-compose.yml @@ -6,6 +6,9 @@ services: build: context: .. dockerfile: sdk_compliance_adapter/Dockerfile + environment: + # One of gzip, deflate or zstd; selects the encoding_ test. + COMPRESSION: ${COMPRESSION:-gzip} ports: - "8080:8080" networks: @@ -13,7 +16,7 @@ services: # Test harness test-harness: - image: ghcr.io/posthog/sdk-test-harness:1.1.1 + image: ghcr.io/posthog/sdk-test-harness:1.13.1 command: ["run", "--adapter-url", "http://sdk-adapter:8080", "--mock-url", "http://test-harness:8081"] networks: - test-network diff --git a/sdk_compliance_adapter/test_adapter.py b/sdk_compliance_adapter/test_adapter.py index 1e06c4628..02eeac261 100644 --- a/sdk_compliance_adapter/test_adapter.py +++ b/sdk_compliance_adapter/test_adapter.py @@ -77,12 +77,60 @@ def initialize(adapter, **overrides): return adapter.app.test_client() -def test_health_opts_into_local_evaluation_without_losing_capture(adapter): +@pytest.mark.parametrize( + "env_value,codec", + [(None, "gzip"), ("gzip", "gzip"), ("deflate", "deflate"), ("zstd", "zstd")], +) +def test_health_and_init_use_one_codec(adapter, monkeypatch, env_value, codec): + if env_value is None: + monkeypatch.delenv("COMPRESSION", raising=False) + else: + monkeypatch.setenv("COMPRESSION", env_value) capabilities = adapter.app.test_client().get("/health").json["capabilities"] assert "feature_flags_local_evaluation_v1" in capabilities - assert "capture_v1" in capabilities + assert {"capture_v1", "capture_ai_v1", "event_options"} <= set(capabilities) assert "capture_v0" not in capabilities assert "capture_ai_v0" not in capabilities + assert [c for c in capabilities if c.startswith("encoding_")] == [ + f"encoding_{codec}" + ] + + initialize(adapter, enable_compression=True) + assert adapter.state.client.capture_compression.value == codec + assert adapter.state.client.capture_ai_compression.value == codec + + +@pytest.mark.parametrize( + "route,method", [("/capture", "capture"), ("/capture_ai", "capture_ai")] +) +def test_capture_passes_options_and_uuid_unchanged(adapter, monkeypatch, route, method): + client = initialize(adapter, personal_api_key=None) + sent = Mock(return_value="0190a5a8-0000-7000-8000-000000000001") + monkeypatch.setattr(adapter.state.client, method, sent) + options = {"cookieless_mode": "no", "process_person_profile": 0.0, "custom": [1]} + + response = client.post( + route, + json={ + "distinct_id": "user", + "event": "$ai_generation", + "properties": {"$ignore_sent_at": "maybe"}, + "options": options, + "uuid": "0190A5A8-0000-7000-8000-000000000001", + }, + ) + + assert response.json == { + "success": True, + "uuid": "0190a5a8-0000-7000-8000-000000000001", + } + sent.assert_called_once_with( + "$ai_generation", + distinct_id="user", + properties={"$ignore_sent_at": "maybe"}, + options=options, + uuid="0190A5A8-0000-7000-8000-000000000001", + ) def test_init_enables_explicit_definitions_loading_without_polling(adapter): From dd9c12296b03ae7029e99392417e785a489f119c Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:47:06 -0700 Subject: [PATCH 12/14] feat!: AI wrappers always send through capture_ai (#1027) enable_full_ai_capture now controls only truncation and media redaction. A client object without capture_ai still gets capture calls. --- posthog/__init__.py | 4 +- posthog/ai/utils.py | 15 +- posthog/client.py | 10 +- posthog/test/ai/anthropic/test_anthropic.py | 149 +++++----- .../ai/claude_agent_sdk/test_processor.py | 50 ++-- posthog/test/ai/gemini/test_gemini.py | 124 ++++---- posthog/test/ai/gemini/test_gemini_async.py | 80 +++--- posthog/test/ai/gemini/test_gemini_parity.py | 4 +- posthog/test/ai/langchain/test_callbacks.py | 270 +++++++++--------- posthog/test/ai/langchain/test_middleware.py | 6 +- posthog/test/ai/openai/test_async_parity.py | 22 +- posthog/test/ai/openai/test_openai.py | 200 ++++++------- .../test/ai/openai_agents/test_processor.py | 105 +++---- posthog/test/ai/test_token_reporting.py | 4 +- posthog/test/test_ai_capture_lane.py | 83 +----- 15 files changed, 486 insertions(+), 640 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index d228d811e..027d6425b 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -485,8 +485,8 @@ def get_context_options() -> Dict[str, Any]: # We recommend setting this to False if you are only using the personalApiKey for evaluating remote config payloads via `get_remote_config_payload` and not using local evaluation. enable_local_evaluation = True # type: bool flag_definition_cache_provider = None # type: Optional[FlagDefinitionCacheProvider] -# Routes AI SDK wrapper events through the dedicated AI capture lane, skips -# truncation, and passes media unredacted. `privacy_mode` always wins. +# Captures full AI content in AI SDK wrapper events: skips truncation and +# passes media unredacted. `privacy_mode` always wins. enable_full_ai_capture = False # type: bool # Deprecated aliases for `enable_full_ai_capture`. _use_ai_lane = False # type: bool diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index ae652e80a..2a6d20cce 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -5,7 +5,7 @@ from posthog import get_tags, identify_context, new_context, tag, contexts from posthog.ai.gateway import warn_if_posthog_ai_gateway -from posthog.ai.sanitization import _full_ai_capture_enabled, redact_media +from posthog.ai.sanitization import redact_media from posthog.ai.sanitization import sanitize_messages # noqa: F401 -- re-exported for back-compat from posthog.ai.types import FormattedMessage, StreamingEventData, TokenUsage from posthog.client import Client as PostHogClient @@ -64,11 +64,6 @@ def _get_tokens_source( return "sdk" -def _ai_lane_enabled(ph_client) -> bool: - """The client's full-AI-capture opt-in routes wrapper events onto the AI lane.""" - return _full_ai_capture_enabled(ph_client) - - def _capture_ai_event(ph_client, event: str, *, personless: bool = False, **kwargs): """Capture a wrapper-emitted AI event with the PostHog AI library identity. @@ -87,10 +82,10 @@ def _capture_ai_event(ph_client, event: str, *, personless: bool = False, **kwar **(kwargs.get("options") or {}), "process_person_profile": False, } - if _ai_lane_enabled(ph_client): - capture_ai = getattr(ph_client, "capture_ai", None) - if callable(capture_ai): - return capture_ai(event=event, **kwargs) + # Clients without capture_ai, such as custom stand-ins, keep the analytics lane. + capture_ai = getattr(ph_client, "capture_ai", None) + if callable(capture_ai): + return capture_ai(event=event, **kwargs) return ph_client.capture(event=event, **kwargs) diff --git a/posthog/client.py b/posthog/client.py index 1ce5108ba..0f7b415ef 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -846,11 +846,11 @@ def __init__( captured exceptions. Defaults to the current working directory. privacy_mode: For AI observability, capture usage metadata without prompt inputs or outputs. - enable_full_ai_capture: Route PostHog AI wrapper events through - the dedicated AI capture endpoint and capture full AI content: - skips string truncation and passes media (base64/data URIs) - through unredacted. ``privacy_mode`` always wins. Defaults to - False. + enable_full_ai_capture: Capture full AI content in PostHog AI + wrapper events: skips string truncation and passes media + (base64/data URIs) through unredacted. ``privacy_mode`` always + wins. Wrapper events use the AI capture endpoint whether or not + this is set. Defaults to False. before_send: Optional callback that can modify or drop events before upload. Return ``None`` to drop an event. Context tags, context options, ``super_properties``, ``super_options`` and diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index a4a3e9bc6..2774933aa 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -430,9 +430,9 @@ def test_basic_completion(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -476,7 +476,7 @@ def test_tokens_source_passthrough(mock_client, mock_anthropic_response): posthog_properties={"$ai_input_tokens": 99999}, ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_tokens_source"] == "passthrough" assert props["$ai_input_tokens"] == 99999 @@ -494,9 +494,9 @@ def test_groups(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "test_company"} @@ -513,9 +513,9 @@ def test_privacy_mode_local(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -535,9 +535,9 @@ def test_privacy_mode_global(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -561,9 +561,9 @@ def test_basic_integration(mock_client): system="You must always answer with 'Bar'.", ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" assert call_args["event"] == "$ai_generation" @@ -607,9 +607,9 @@ async def mock_async_create(**kwargs): posthog_properties={"foo": "bar"}, ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -669,9 +669,9 @@ async def async_create_wrapper(**kwargs): [c async for c in response] # Capture happens in the async finally block before generator completes - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ @@ -691,9 +691,9 @@ def test_error(mock_client, mock_anthropic_response): messages=[{"role": "user", "content": "Hello"}], ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "Test error" @@ -713,9 +713,9 @@ def test_cached_tokens(mock_client, mock_anthropic_response_with_cached_tokens): ) assert response == mock_anthropic_response_with_cached_tokens - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -754,7 +754,7 @@ def test_preserves_cache_creation_ttl_breakdown_non_streaming( ) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -776,7 +776,7 @@ async def mock_async_create(**kwargs): ) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -794,7 +794,7 @@ def test_preserves_cache_creation_ttl_breakdown_streaming(mock_client): list(response) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -820,7 +820,7 @@ async def mock_async_create(**kwargs): [event async for event in response] assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -843,7 +843,7 @@ def test_anthropic_bedrock_preserves_cache_creation_ttl_breakdown( ) assert_cache_creation_ttl_breakdown_preserved( - mock_client.capture.call_args.kwargs["properties"] + mock_client.capture_ai.call_args.kwargs["properties"] ) @@ -882,9 +882,9 @@ def test_tool_definition(mock_client, mock_anthropic_response): ) assert response == mock_anthropic_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -936,9 +936,9 @@ def test_tool_calls_in_output_choices( ) assert response == mock_anthropic_response_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -999,9 +999,9 @@ def test_tool_calls_only_no_content( ) assert response == mock_anthropic_response_tool_calls_only - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1068,9 +1068,9 @@ async def run_test(): response = asyncio.run(run_test()) assert response == mock_anthropic_response_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1136,9 +1136,9 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools list(response) # Capture happens synchronously when generator is exhausted - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1278,7 +1278,7 @@ def stream_generator(): ) list(response) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_output_choices"] == [ { "role": "assistant", @@ -1345,9 +1345,9 @@ async def run_test(): asyncio.run(run_test()) # Capture completes before asyncio.run() returns - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1427,9 +1427,9 @@ def test_streaming_with_thinking(mock_client, mock_anthropic_stream_with_thinkin # Consume the stream - this triggers the finally block synchronously list(response) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] output_choices = props["$ai_output_choices"] @@ -1479,9 +1479,9 @@ async def run_test(): asyncio.run(run_test()) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] output_choices = props["$ai_output_choices"] @@ -1533,9 +1533,9 @@ def __init__(self): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1615,9 +1615,9 @@ def test_streaming_with_web_search(mock_client, mock_anthropic_stream_with_web_s list(response) # Capture happens synchronously when generator is exhausted - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1672,9 +1672,9 @@ async def run_test(): response = asyncio.run(run_test()) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1719,9 +1719,9 @@ async def run_test(): asyncio.run(run_test()) # Capture completes before asyncio.run() returns - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured @@ -1778,7 +1778,7 @@ async def create(**kwargs): if stream: list(response) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["distinct_id"] == "trace-123" assert call_args["options"] == {"process_person_profile": False} @@ -1798,7 +1798,7 @@ def test_explicit_distinct_id_creates_person_profile( posthog_trace_id="trace-123", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "user-123" @@ -1820,7 +1820,7 @@ def test_outer_context_distinct_id_is_used(mock_client, mock_anthropic_response) posthog_trace_id="trace-123", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "outer-user-456" @@ -1845,7 +1845,7 @@ def test_explicit_distinct_id_overrides_outer_context( posthog_trace_id="trace-123", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["distinct_id"] == "explicit-user-789" @@ -1862,7 +1862,7 @@ def test_integration_stop_reason(mock_client): posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] in ("end_turn", "max_tokens") assert props["$ai_provider"] == "anthropic" assert props["$ai_input_tokens"] > 0 @@ -1963,9 +1963,9 @@ async def test_streaming_sync_async_accumulation_parity(mock_client): 0, ) list(sync_response) - sync_properties = mock_client.capture.call_args.kwargs["properties"] + sync_properties = mock_client.capture_ai.call_args.kwargs["properties"] - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) async_response = async_client.messages._track_streaming_response( @@ -1979,7 +1979,7 @@ async def test_streaming_sync_async_accumulation_parity(mock_client): 0, ) [event async for event in async_response] - async_properties = mock_client.capture.call_args.kwargs["properties"] + async_properties = mock_client.capture_ai.call_args.kwargs["properties"] parity_keys = ( "$ai_input", @@ -2019,9 +2019,9 @@ async def test_streaming_sync_async_exceptions_capture_and_close(mock_client): list(stream) assert sync_source.closed is True - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() async_source = FailingRecordingAsyncStream([]) async_client = AsyncAnthropic(api_key="test-key", posthog_client=mock_client) @@ -2040,7 +2040,7 @@ async def test_streaming_sync_async_exceptions_capture_and_close(mock_client): [event async for event in stream] assert async_source.closed is True - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 def test_messages_stream_preserves_native_manager_helpers_close_and_tracking( @@ -2064,8 +2064,8 @@ def test_messages_stream_preserves_native_manager_helpers_close_and_tracking( assert text == ["Hi"] assert source.closed is True - assert mock_client.capture.call_count == 1 - assert mock_client.capture.call_args.kwargs["distinct_id"] == "test-user" + assert mock_client.capture_ai.call_count == 1 + assert mock_client.capture_ai.call_args.kwargs["distinct_id"] == "test-user" assert "posthog_distinct_id" not in client.post.call_args.kwargs @@ -2106,8 +2106,8 @@ async def test_async_messages_stream_preserves_provider_contract_and_manager( assert text == ["Hi"] assert source.closed is True - assert mock_client.capture.call_count == 1 - assert mock_client.capture.call_args.kwargs["distinct_id"] == "test-user" + assert mock_client.capture_ai.call_count == 1 + assert mock_client.capture_ai.call_args.kwargs["distinct_id"] == "test-user" assert "posthog_distinct_id" not in client.post.call_args.kwargs @@ -2153,7 +2153,7 @@ async def mock_async_create(**kwargs): events = [event async for event in stream] assert len(events) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 def test_messages_streaming_early_exit_closes_provider_stream(mock_client): @@ -2176,7 +2176,7 @@ def test_messages_streaming_early_exit_closes_provider_stream(mock_client): break assert source.closed is True - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 @pytest.mark.asyncio @@ -2205,24 +2205,7 @@ async def mock_async_create(**kwargs): break assert source.closed is True - assert mock_client.capture.call_count == 1 - - -def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_anthropic_response): - mock_client.enable_full_ai_capture = True - with patch( - "anthropic.resources.Messages.create", return_value=mock_anthropic_response - ): - client = Anthropic(api_key="test-key", posthog_client=mock_client) - client.messages.create( - model="claude-3-opus-20240229", - messages=[{"role": "user", "content": "Hello"}], - posthog_distinct_id="test-id", - ) - - mock_client.capture.assert_not_called() assert mock_client.capture_ai.call_count == 1 - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation" def test_multimodal_client_skips_media_redaction(mock_client, mock_anthropic_response): diff --git a/posthog/test/ai/claude_agent_sdk/test_processor.py b/posthog/test/ai/claude_agent_sdk/test_processor.py index 1fb49c664..2b888b128 100644 --- a/posthog/test/ai/claude_agent_sdk/test_processor.py +++ b/posthog/test/ai/claude_agent_sdk/test_processor.py @@ -202,7 +202,7 @@ async def test_emits_generation_from_stream_events(self, processor, mock_client) collected.append(msg) # Should have captured $ai_generation + $ai_trace - calls = mock_client.capture.call_args_list + calls = mock_client.capture_ai.call_args_list events = [c.kwargs.get("event") or c[1].get("event") for c in calls] assert "$ai_generation" in events assert "$ai_trace" in events @@ -247,7 +247,7 @@ async def test_emits_cache_creation_ttl_breakdown_from_stream_events( generation_call = next( call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" ) props = generation_call.kwargs.get("properties") or generation_call[1].get( @@ -286,7 +286,7 @@ async def test_emits_multiple_generations_for_multi_turn( gen_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation" ] assert len(gen_calls) == 2 @@ -318,7 +318,7 @@ async def test_fallback_generation_from_result_when_no_stream_events( gen_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation" ] assert len(gen_calls) == 1 @@ -349,7 +349,7 @@ async def test_fallback_copies_result_usage_before_capture( generation_call = next( call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" ) props = generation_call.kwargs.get("properties") or generation_call[1].get( @@ -391,7 +391,7 @@ async def test_emits_span_for_tool_use(self, processor, mock_client): span_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_span" ] assert len(span_calls) == 1 @@ -422,7 +422,7 @@ async def test_emits_trace_on_result(self, processor, mock_client): trace_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_trace" ] assert len(trace_calls) == 1 @@ -449,7 +449,7 @@ async def test_trace_emits_error_status(self, processor, mock_client): trace_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_trace" ] assert len(trace_calls) == 1 @@ -473,7 +473,7 @@ async def test_client_privacy_mode_redacts_generation_input(self, mock_client): async for _ in proc.query(prompt="secret", options=ClaudeAgentOptions()): pass - properties = mock_client.capture.call_args.kwargs["properties"] + properties = mock_client.capture_ai.call_args.kwargs["properties"] assert properties["$ai_input"] is None @pytest.mark.asyncio @@ -502,7 +502,7 @@ async def test_per_call_privacy_mode_cannot_disable_processor_mode( ): pass - properties = mock_client.capture.call_args.kwargs["properties"] + properties = mock_client.capture_ai.call_args.kwargs["properties"] assert properties["$ai_input"] is None @pytest.mark.asyncio @@ -559,7 +559,7 @@ async def test_privacy_mode_redacts_tool_input(self, mock_client): span_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_span" ] assert len(span_calls) == 1 @@ -570,7 +570,7 @@ async def test_privacy_mode_redacts_tool_input(self, mock_client): generation_call = next( call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if (call.kwargs.get("event") or call[1].get("event")) == "$ai_generation" ) generation_props = generation_call.kwargs.get("properties") or generation_call[ @@ -619,7 +619,7 @@ async def test_data_url_prompt_and_tool_use_base64_are_redacted( gen_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_generation" ] assert len(gen_calls) == 2 @@ -667,7 +667,7 @@ async def test_no_distinct_id_sets_personless_option(self, mock_client): ): pass - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: assert call.kwargs["options"] == {"process_person_profile": False} assert call.kwargs["distinct_id"] == "trace-fallback" @@ -681,7 +681,7 @@ def test_explicit_groups_override_preserves_empty_groups(self, mock_client): proc._capture_event("$ai_trace", {}, groups={}) - assert mock_client.capture.call_args.kwargs["groups"] == {} + assert mock_client.capture_ai.call_args.kwargs["groups"] == {} def test_default_properties_keep_existing_precedence(self, mock_client): proc = PostHogClaudeAgentProcessor( @@ -694,7 +694,7 @@ def test_default_properties_keep_existing_precedence(self, mock_client): {"environment": "event", "$ai_trace_id": "trace-id"}, ) - assert mock_client.capture.call_args.kwargs["properties"] == { + assert mock_client.capture_ai.call_args.kwargs["properties"] == { "$ai_lib": "posthog-ai", "$ai_lib_version": VERSION, "environment": "processor", @@ -707,7 +707,7 @@ def test_client_without_capture_capability_is_ignored(self): proc._capture_event("$ai_trace", {}) def test_capture_errors_are_logged_and_suppressed(self, mock_client, caplog): - mock_client.capture.side_effect = RuntimeError("capture failed") + mock_client.capture_ai.side_effect = RuntimeError("capture failed") proc = PostHogClaudeAgentProcessor(client=mock_client) with caplog.at_level(logging.DEBUG, logger="posthog"): @@ -738,7 +738,7 @@ async def test_instance_properties_merged(self, mock_client): async for _ in proc.query(prompt="Hi", options=ClaudeAgentOptions()): pass - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: props = call.kwargs.get("properties") or call[1].get("properties") assert props.get("app") == "stamphog" assert props.get("version") == "1.0" @@ -763,7 +763,7 @@ async def test_per_call_properties_merged(self, processor, mock_client): ): pass - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: props = call.kwargs.get("properties") or call[1].get("properties") assert props.get("pr_number") == 42 @@ -791,7 +791,7 @@ def resolver(result): trace_calls = [ c - for c in mock_client.capture.call_args_list + for c in mock_client.capture_ai.call_args_list if (c.kwargs.get("event") or c[1].get("event")) == "$ai_trace" ] assert len(trace_calls) == 1 @@ -894,13 +894,3 @@ async def test_non_config_errors_propagate(self): prompt="Hello", options=ClaudeAgentOptions() ): pass - - -def test_ai_lane_client_routes_through_capture_ai(mock_client): - mock_client.enable_full_ai_capture = True - processor = PostHogClaudeAgentProcessor(client=mock_client, distinct_id="test-user") - processor._capture_event(event="$ai_trace", properties={}, distinct_id="d") - - mock_client.capture.assert_not_called() - mock_client.capture_ai.assert_called_once() - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_trace" diff --git a/posthog/test/ai/gemini/test_gemini.py b/posthog/test/ai/gemini/test_gemini.py index 0079e73db..48d5cbd08 100644 --- a/posthog/test/ai/gemini/test_gemini.py +++ b/posthog/test/ai/gemini/test_gemini.py @@ -188,9 +188,9 @@ def test_new_client_basic_generation( ) assert response == mock_gemini_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -261,8 +261,8 @@ def mock_streaming_response(): assert chunks[1].text == "world!" # Check that the streaming event was captured - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -343,8 +343,8 @@ def mock_streaming_response(): assert chunks[1].text == "the weather" # Check that the streaming event was captured with tools - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -373,7 +373,7 @@ def test_new_client_groups(mock_client, mock_google_genai_client, mock_gemini_re posthog_groups={"company": "company_123"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "company_123"} @@ -392,7 +392,7 @@ def test_new_client_privacy_mode_local( posthog_privacy_mode=True, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -414,7 +414,7 @@ def test_new_client_privacy_mode_global( posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -432,7 +432,7 @@ def test_new_client_different_input_formats( client.models.generate_content( model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] @@ -443,7 +443,7 @@ def test_new_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "hey"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ {"role": "user", "content": [{"type": "text", "text": "hey"}]} @@ -456,7 +456,7 @@ def test_new_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ { @@ -469,11 +469,11 @@ def test_new_client_different_input_formats( ] # Test list input with string - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() client.models.generate_content( model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "List item"}] @@ -494,7 +494,7 @@ def test_new_client_model_parameters( max_tokens=100, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model_parameters"]["temperature"] == 0.7 assert props["$ai_model_parameters"]["max_tokens"] == 100 @@ -518,7 +518,7 @@ def test_new_client_default_settings( # Call without overriding defaults client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"]) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default_user" @@ -551,7 +551,7 @@ def test_new_client_override_defaults( posthog_groups={"organization": "special_org"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check overrides @@ -676,9 +676,9 @@ def test_tool_use_response(mock_client, mock_google_genai_client, mock_gemini_re ) assert response == mock_gemini_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -718,9 +718,9 @@ def test_function_calls_in_output_choices( ) assert response == mock_gemini_response_with_function_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -767,9 +767,9 @@ def test_function_calls_only_no_content( ) assert response == mock_gemini_response_function_calls_only - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -826,9 +826,9 @@ def test_cache_and_reasoning_tokens(mock_client, mock_google_genai_client): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present @@ -890,9 +890,9 @@ def test_streaming_cache_and_reasoning_tokens(mock_client, mock_google_genai_cli assert len(result) == 2 # Check PostHog capture was called - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present (should use final chunk's usage) @@ -963,9 +963,9 @@ def test_web_search_grounding(mock_client, mock_google_genai_client): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -1034,9 +1034,9 @@ def mock_streaming_response(): chunks = list(response) assert len(chunks) == 2 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -1098,9 +1098,9 @@ def test_empty_grounding_metadata_no_web_search(mock_client, mock_google_genai_c ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is 0 (not present in properties when 0) @@ -1161,9 +1161,9 @@ def test_empty_array_grounding_metadata_no_web_search( ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is 0 (not present in properties when 0) @@ -1203,9 +1203,9 @@ def test_falsy_candidates_no_web_search( ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should not crash and web search count should not be present @@ -1257,9 +1257,9 @@ def test_embed_content_basic( ) assert response == mock_embed_content_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1293,7 +1293,7 @@ def test_embed_content_with_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input_tokens"] == 13 # 5 + 8 @@ -1311,7 +1311,7 @@ def test_embed_content_without_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] # No embedding carried a token count, so the property is omitted, not 0. assert "$ai_input_tokens" not in props @@ -1334,7 +1334,7 @@ def test_streaming_without_usage_omits_token_counts( ) list(response) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert "$ai_input_tokens" not in props assert "$ai_output_tokens" not in props @@ -1354,7 +1354,7 @@ def test_embed_content_privacy_mode( posthog_privacy_mode=True, ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input"] is None @@ -1371,7 +1371,7 @@ def test_embed_content_no_distinct_id( contents="Hello", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should fall back to trace_id as distinct_id @@ -1400,7 +1400,7 @@ def test_embed_content_default_params( posthog_properties={"extra": "prop"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default-id" @@ -1422,8 +1422,8 @@ def test_embed_content_error_handling(mock_client, mock_google_genai_client): ) # Event should still be captured - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "API error" assert props["$ai_http_status"] == 0 @@ -1467,8 +1467,8 @@ def test_embed_content_integration(mock_client): assert all(isinstance(v, float) for v in response.embeddings[0].values) # Verify event captured - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_embedding" assert props["$ai_provider"] == "gemini" @@ -1531,7 +1531,7 @@ def test_stop_reason_captured( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] == finish_reason_name @@ -1580,8 +1580,8 @@ def mock_streaming_response(): list(response) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] == "STOP" @@ -1597,26 +1597,8 @@ def test_integration_stop_reason(mock_client): posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] is not None assert isinstance(props["$ai_stop_reason"], str) assert props["$ai_provider"] == "gemini" assert props["$ai_input_tokens"] > 0 - - -def test_ai_lane_client_routes_through_capture_ai( - mock_client, mock_google_genai_client, mock_gemini_response -): - mock_google_genai_client.models.generate_content.return_value = mock_gemini_response - - mock_client.enable_full_ai_capture = True - client = Client(api_key="test-key", posthog_client=mock_client) - client.models.generate_content( - model="gemini-2.0-flash", - contents=["Tell me a fun fact about hedgehogs"], - posthog_distinct_id="test-id", - ) - - mock_client.capture.assert_not_called() - assert mock_client.capture_ai.call_count == 1 - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation" diff --git a/posthog/test/ai/gemini/test_gemini_async.py b/posthog/test/ai/gemini/test_gemini_async.py index 1373a331e..701b7ebdc 100644 --- a/posthog/test/ai/gemini/test_gemini_async.py +++ b/posthog/test/ai/gemini/test_gemini_async.py @@ -134,9 +134,9 @@ async def test_async_client_basic_generation( ) assert response == mock_gemini_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -199,8 +199,8 @@ async def mock_streaming_response(): assert chunks[1].text == "world!" # Check that the streaming event was captured - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -283,8 +283,8 @@ async def mock_streaming_response(): assert chunks[1].text == "the weather" # Check that the streaming event was captured with tools - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -317,7 +317,7 @@ async def test_async_client_groups( posthog_groups={"company": "company_123"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "company_123"} @@ -338,7 +338,7 @@ async def test_async_client_privacy_mode_local( posthog_privacy_mode=True, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -362,7 +362,7 @@ async def test_async_client_privacy_mode_global( posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -382,7 +382,7 @@ async def test_async_client_different_input_formats( await client.models.generate_content( model="gemini-2.0-flash", contents="Hello", posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "Hello"}] @@ -393,7 +393,7 @@ async def test_async_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "hey"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ {"role": "user", "content": [{"type": "text", "text": "hey"}]} @@ -406,7 +406,7 @@ async def test_async_client_different_input_formats( contents=[{"role": "user", "parts": [{"text": "Hello "}, {"text": "world"}]}], posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [ { @@ -419,11 +419,11 @@ async def test_async_client_different_input_formats( ] # Test list input with string - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() await client.models.generate_content( model="gemini-2.0-flash", contents=["List item"], posthog_distinct_id="test-id" ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] == [{"role": "user", "content": "List item"}] @@ -446,7 +446,7 @@ async def test_async_client_model_parameters( max_tokens=100, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model_parameters"]["temperature"] == 0.7 assert props["$ai_model_parameters"]["max_tokens"] == 100 @@ -472,7 +472,7 @@ async def test_async_client_default_settings( # Call without overriding defaults await client.models.generate_content(model="gemini-2.0-flash", contents=["Hello"]) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default_user" @@ -507,7 +507,7 @@ async def test_async_client_override_defaults( posthog_groups={"organization": "special_org"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check overrides @@ -586,9 +586,9 @@ async def test_async_function_calls_in_output_choices( ) assert response == mock_gemini_response_with_function_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -649,9 +649,9 @@ async def test_async_cache_and_reasoning_tokens(mock_client, mock_google_genai_c ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present @@ -708,9 +708,9 @@ async def mock_streaming_response(): assert len(result) == 2 # Check PostHog capture was called - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Check that all token types are present (should use final chunk's usage) @@ -772,9 +772,9 @@ async def test_async_web_search_grounding(mock_client, mock_google_genai_client) ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -845,9 +845,9 @@ async def mock_streaming_response(): chunks.append(chunk) assert len(chunks) == 2 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary for grounding) @@ -899,9 +899,9 @@ async def test_async_embed_content_basic( ) assert response == mock_embed_content_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -935,7 +935,7 @@ async def test_async_embed_content_with_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input_tokens"] == 13 # 5 + 8 @@ -953,7 +953,7 @@ async def test_async_embed_content_without_token_counts( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] # No embedding carried a token count, so the property is omitted, not 0. assert "$ai_input_tokens" not in props @@ -973,7 +973,7 @@ async def test_async_embed_content_privacy_mode( posthog_privacy_mode=True, ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_input"] is None @@ -990,7 +990,7 @@ async def test_async_embed_content_no_distinct_id( contents="Hello", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should fall back to trace_id as distinct_id @@ -1019,7 +1019,7 @@ async def test_async_embed_content_default_params( posthog_properties={"extra": "prop"}, ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "default-id" @@ -1045,8 +1045,8 @@ async def test_async_embed_content_error_handling( ) # Event should still be captured - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "API error" assert props["$ai_http_status"] == 0 @@ -1090,8 +1090,8 @@ async def test_async_embed_content_integration(mock_client): assert all(isinstance(v, float) for v in response.embeddings[0].values) # Verify event captured - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_embedding" assert props["$ai_provider"] == "gemini" @@ -1148,7 +1148,7 @@ async def mock_streaming_response(): chunks.append(chunk) assert len(chunks) == 1 - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] assert call_args["event"] == "$ai_generation" assert call_args["properties"]["$ai_provider"] == "gemini" diff --git a/posthog/test/ai/gemini/test_gemini_parity.py b/posthog/test/ai/gemini/test_gemini_parity.py index de9f9b82f..a31299032 100644 --- a/posthog/test/ai/gemini/test_gemini_parity.py +++ b/posthog/test/ai/gemini/test_gemini_parity.py @@ -178,8 +178,8 @@ async def test_sync_client_aio_models_tracks_generations( assert response is gemini_response provider_client.aio.models.generate_content.assert_awaited_once() - assert posthog_client.capture.call_count == 1 - call_args = posthog_client.capture.call_args[1] + assert posthog_client.capture_ai.call_count == 1 + call_args = posthog_client.capture_ai.call_args[1] assert call_args["distinct_id"] == "test-id" assert call_args["event"] == "$ai_generation" assert call_args["properties"]["$ai_model"] == "gemini-2.0-flash" diff --git a/posthog/test/ai/langchain/test_callbacks.py b/posthog/test/ai/langchain/test_callbacks.py index 7bc1b4166..326eb0ddb 100644 --- a/posthog/test/ai/langchain/test_callbacks.py +++ b/posthog/test/ai/langchain/test_callbacks.py @@ -185,15 +185,15 @@ def test_basic_chat_chain(mock_client, stream): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] + span_args = mock_client.capture_ai.call_args_list[0][1] span_props = span_args["properties"] - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] # Span is first @@ -262,13 +262,13 @@ async def test_async_basic_chat_chain(mock_client, stream): else: result = await chain.ainvoke({}, config={"callbacks": callbacks}) assert result.content == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] + span_args = mock_client.capture_ai.call_args_list[0][1] span_props = span_args["properties"] - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] # Span is first @@ -349,7 +349,7 @@ def _capture_anthropic_usage( ) cb._pop_run_and_capture_generation(run_id, None, response) - return mock_client.capture.call_args.kwargs["properties"] + return mock_client.capture_ai.call_args.kwargs["properties"] @pytest.mark.parametrize( @@ -380,8 +380,8 @@ def test_basic_llm_chain(mock_client, Model, stream): ) assert result == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 1 - args = mock_client.capture.call_args_list[0][1] + assert mock_client.capture_ai.call_count == 1 + args = mock_client.capture_ai.call_args_list[0][1] props = args["properties"] assert args["event"] == "$ai_generation" @@ -425,8 +425,8 @@ async def test_async_basic_llm_chain(mock_client, Model, stream): ) assert result == "The Los Angeles Dodgers won the World Series in 2020." - assert mock_client.capture.call_count == 1 - args = mock_client.capture.call_args_list[0][1] + assert mock_client.capture_ai.call_count == 1 + args = mock_client.capture_ai.call_args_list[0][1] props = args["properties"] assert args["event"] == "$ai_generation" @@ -455,21 +455,21 @@ def test_trace_id_and_inputs_for_multiple_chains(mock_client): assert result.content == "Bar" # span, generation, span, generation, trace - assert mock_client.capture.call_count == 5 + assert mock_client.capture_ai.call_count == 5 - first_span_args = mock_client.capture.call_args_list[0][1] + first_span_args = mock_client.capture_ai.call_args_list[0][1] first_span_props = first_span_args["properties"] - first_generation_args = mock_client.capture.call_args_list[1][1] + first_generation_args = mock_client.capture_ai.call_args_list[1][1] first_generation_props = first_generation_args["properties"] - second_span_args = mock_client.capture.call_args_list[2][1] + second_span_args = mock_client.capture_ai.call_args_list[2][1] second_span_props = second_span_args["properties"] - second_generation_args = mock_client.capture.call_args_list[3][1] + second_generation_args = mock_client.capture_ai.call_args_list[3][1] second_generation_props = second_generation_args["properties"] - trace_args = mock_client.capture.call_args_list[4][1] + trace_args = mock_client.capture_ai.call_args_list[4][1] trace_props = trace_args["properties"] # Prompt span @@ -536,10 +536,10 @@ def test_personless_mode(mock_client): prompt = ChatPromptTemplate.from_messages([("user", "Foo")]) chain = prompt | FakeMessagesListChatModel(responses=[AIMessage(content="Bar")]) chain.invoke({}, config={"callbacks": [CallbackHandler(mock_client)]}) - assert mock_client.capture.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] - generation_args = mock_client.capture.call_args_list[1][1] - trace_args = mock_client.capture.call_args_list[2][1] + assert mock_client.capture_ai.call_count == 3 + span_args = mock_client.capture_ai.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] # span assert span_args["event"] == "$ai_span" @@ -555,10 +555,10 @@ def test_personless_mode(mock_client): chain.invoke( {}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]} ) - assert mock_client.capture.call_count == 6 - span_args = mock_client.capture.call_args_list[3][1] - generation_args = mock_client.capture.call_args_list[4][1] - trace_args = mock_client.capture.call_args_list[5][1] + assert mock_client.capture_ai.call_count == 6 + span_args = mock_client.capture_ai.call_args_list[3][1] + generation_args = mock_client.capture_ai.call_args_list[4][1] + trace_args = mock_client.capture_ai.call_args_list[5][1] # span assert "options" not in span_args @@ -595,10 +595,10 @@ def test_personless_mode_exception(mock_client, unauthorized_http_client): callbacks = CallbackHandler(mock_client) with pytest.raises(AuthenticationError): chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - span_args = mock_client.capture.call_args_list[0][1] - generation_args = mock_client.capture.call_args_list[1][1] - trace_args = mock_client.capture.call_args_list[2][1] + assert mock_client.capture_ai.call_count == 3 + span_args = mock_client.capture_ai.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] # span assert span_args["event"] == "$ai_span" @@ -615,10 +615,10 @@ def test_personless_mode_exception(mock_client, unauthorized_http_client): chain.invoke( {}, config={"callbacks": [CallbackHandler(mock_client, distinct_id=id)]} ) - assert mock_client.capture.call_count == 6 - span_args = mock_client.capture.call_args_list[3][1] - generation_args = mock_client.capture.call_args_list[4][1] - trace_args = mock_client.capture.call_args_list[5][1] + assert mock_client.capture_ai.call_count == 6 + span_args = mock_client.capture_ai.call_args_list[3][1] + generation_args = mock_client.capture_ai.call_args_list[4][1] + trace_args = mock_client.capture_ai.call_args_list[5][1] # span assert span_args["event"] == "$ai_span" @@ -655,9 +655,9 @@ def test_metadata(mock_client): result = chain.invoke({"plan": None}, config={"callbacks": callbacks}) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_call_args = mock_client.capture.call_args_list[0][1] + span_call_args = mock_client.capture_ai.call_args_list[0][1] span_call_props = span_call_args["properties"] assert span_call_args["distinct_id"] == "test_id" assert span_call_args["event"] == "$ai_span" @@ -666,7 +666,7 @@ def test_metadata(mock_client): assert "$ai_parent_id" in span_call_props assert "$ai_span_id" in span_call_props - generation_call_args = mock_client.capture.call_args_list[1][1] + generation_call_args = mock_client.capture_ai.call_args_list[1][1] generation_call_props = generation_call_args["properties"] assert generation_call_args["distinct_id"] == "test_id" assert generation_call_args["event"] == "$ai_generation" @@ -679,7 +679,7 @@ def test_metadata(mock_client): assert generation_call_props["$ai_http_status"] == 200 assert isinstance(generation_call_props["$ai_latency"], float) - trace_call_args = mock_client.capture.call_args_list[2][1] + trace_call_args = mock_client.capture_ai.call_args_list[2][1] trace_call_props = trace_call_args["properties"] assert trace_call_args["distinct_id"] == "test_id" assert trace_call_args["event"] == "$ai_trace" @@ -742,8 +742,8 @@ def test_graph_state(mock_client): assert isinstance(result["messages"][2], AIMessage) assert result["messages"][2].content == "It's a type of greeble." - assert mock_client.capture.call_count == 6 - calls = [call[1] for call in mock_client.capture.call_args_list] + assert mock_client.capture_ai.call_count == 6 + calls = [call[1] for call in mock_client.capture_ai.call_args_list] # The trace event is captured at the end trace_args = calls[-1] @@ -867,8 +867,8 @@ def runnable(_): assert callbacks._runs == {} assert callbacks._parent_tree == {} - assert mock_client.capture.call_count == 1 - trace_call_args = mock_client.capture.call_args_list[0][1] + assert mock_client.capture_ai.call_count == 1 + trace_call_args = mock_client.capture_ai.call_args_list[0][1] assert trace_call_args["event"] == "$ai_trace" assert trace_call_args["properties"]["$ai_span_name"] == "runnable" @@ -895,8 +895,8 @@ def unauthorized(request): assert callbacks._runs == {} assert callbacks._parent_tree == {} - assert mock_client.capture.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_args = mock_client.capture_ai.call_args_list[1][1] props = generation_args["properties"] assert props["$ai_http_status"] == 401 assert props["$ai_input"] == [{"role": "user", "content": "Foo"}] @@ -928,9 +928,9 @@ def test_openai_chain(mock_client): approximate_latency = math.floor(time.time() - start_time) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] assert gen_args["event"] == "$ai_generation" assert gen_props["$ai_trace_id"] == "test-trace-id" @@ -990,11 +990,11 @@ def test_openai_captures_multiple_generations(mock_client): result = chain.invoke({}, config={"callbacks": [callbacks]}) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1053,11 +1053,11 @@ def test_openai_streaming(mock_client): result = sum(result[1:], result[0]) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1097,11 +1097,11 @@ async def test_async_openai_streaming(mock_client): result = sum(result[1:], result[0]) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1135,8 +1135,8 @@ def test_base_url_retrieval(mock_client, unauthorized_http_client): with pytest.raises(AuthenticationError): chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_base_url"] == "https://test.posthog.com" @@ -1152,8 +1152,8 @@ def test_groups(mock_client): callbacks = CallbackHandler(mock_client, groups={"company": "test_company"}) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["groups"] == {"company": "test_company"} @@ -1169,8 +1169,8 @@ def test_privacy_mode_local(mock_client): callbacks = CallbackHandler(mock_client, privacy_mode=True) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_input"] is None assert generation_call["properties"]["$ai_output_choices"] is None @@ -1188,8 +1188,8 @@ def test_privacy_mode_global(mock_client): callbacks = CallbackHandler(mock_client) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_input"] is None assert generation_call["properties"]["$ai_output_choices"] is None @@ -1219,11 +1219,11 @@ def test_anthropic_chain(mock_client): approximate_latency = math.floor(time.time() - start_time) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1279,11 +1279,11 @@ async def test_async_anthropic_streaming(mock_client): result = sum(result[1:], result[0]) assert result.content == "Bar" - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - gen_args = mock_client.capture.call_args_list[1][1] + gen_args = mock_client.capture_ai.call_args_list[1][1] gen_props = gen_args["properties"] - trace_args = mock_client.capture.call_args_list[2][1] + trace_args = mock_client.capture_ai.call_args_list[2][1] trace_props = trace_args["properties"] assert gen_args["event"] == "$ai_generation" @@ -1385,8 +1385,8 @@ def test_tool_calls(mock_client): callbacks = CallbackHandler(mock_client) chain.invoke({}, config={"callbacks": [callbacks]}) - assert mock_client.capture.call_count == 3 - generation_call = mock_client.capture.call_args_list[1][1] + assert mock_client.capture_ai.call_count == 3 + generation_call = mock_client.capture_ai.call_args_list[1][1] assert generation_call["properties"]["$ai_output_choices"][0]["tool_calls"] == [ { "type": "function", @@ -1420,10 +1420,10 @@ async def sleep(x): # -> Any: chain2.ainvoke({}, config={"callbacks": [cb]}), ) approximate_latency = math.floor(time.time() - start_time) - assert mock_client.capture.call_count == 4 + assert mock_client.capture_ai.call_count == 4 first_call, second_call, third_call, fourth_call = ( - mock_client.capture.call_args_list + mock_client.capture_ai.call_args_list ) assert first_call[1]["event"] == "$ai_span" assert second_call[1]["event"] == "$ai_generation" @@ -1460,7 +1460,7 @@ def get_weather(city: Literal["nyc", "sf"]): mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id" ) graph.invoke(inputs, config={"callbacks": [cb]}) - calls = [call[1] for call in mock_client.capture.call_args_list] + calls = [call[1] for call in mock_client.capture_ai.call_args_list] assert len(calls) == 15 for call in calls: assert call["properties"]["$ai_trace_id"] == "test-trace-id" @@ -1486,15 +1486,15 @@ def test_span_set_parent_ids(mock_client, trace_id): chain = prompt | model chain.invoke({}, config={"callbacks": callbacks}) - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - span_props = mock_client.capture.call_args_list[0][1] + span_props = mock_client.capture_ai.call_args_list[0][1] assert ( span_props["properties"]["$ai_trace_id"] == span_props["properties"]["$ai_parent_id"] ) - generation_props = mock_client.capture.call_args_list[1][1] + generation_props = mock_client.capture_ai.call_args_list[1][1] assert ( generation_props["properties"]["$ai_trace_id"] == generation_props["properties"]["$ai_parent_id"] @@ -1516,9 +1516,9 @@ def span_3(_): chain = RunnableLambda(span_1) chain.invoke({}, config={"callbacks": callbacks}) - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - calls = mock_client.capture.call_args_list + calls = mock_client.capture_ai.call_args_list span_props_2 = calls[0][1]["properties"] span_props_1 = calls[1][1]["properties"] trace_props = calls[2][1]["properties"] @@ -1538,12 +1538,12 @@ def span(_): except ValueError: pass - assert mock_client.capture.call_count == 2 + assert mock_client.capture_ai.call_count == 2 assert ( - mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"] + mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_error"] == "ValueError: test" ) - assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"] + assert mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_is_error"] def test_captures_error_without_details_in_span(mock_client): @@ -1557,12 +1557,12 @@ def span(_): except ValueError: pass - assert mock_client.capture.call_count == 2 + assert mock_client.capture_ai.call_count == 2 assert ( - mock_client.capture.call_args_list[1][1]["properties"]["$ai_error"] + mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_error"] == "ValueError" ) - assert mock_client.capture.call_args_list[1][1]["properties"]["$ai_is_error"] + assert mock_client.capture_ai.call_args_list[1][1]["properties"]["$ai_is_error"] def test_openai_reasoning_tokens(mock_client): @@ -1591,9 +1591,9 @@ def test_openai_reasoning_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Let me think through this step by step..." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -1626,9 +1626,9 @@ def test_anthropic_cache_write_and_read_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "I've analyzed the document and cached the context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -1661,9 +1661,9 @@ def test_anthropic_cache_write_and_read_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Using cached analysis to provide quick response." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -1714,7 +1714,7 @@ def test_anthropic_provider_subtracts_cache_tokens(mock_client): cb._pop_run_and_capture_generation(run_id, None, response) - generation_args = mock_client.capture.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[0][1] assert generation_args["properties"]["$ai_input_tokens"] == 400 # 1200 - 800 assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800 @@ -1757,7 +1757,7 @@ def test_anthropic_provider_subtracts_cache_write_tokens(mock_client): cb._pop_run_and_capture_generation(run_id, None, response) - generation_args = mock_client.capture.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[0][1] assert generation_args["properties"]["$ai_input_tokens"] == 200 # 1000 - 800 assert generation_args["properties"]["$ai_cache_creation_input_tokens"] == 800 @@ -1930,7 +1930,7 @@ def test_anthropic_provider_subtracts_both_cache_read_and_write_tokens(mock_clie cb._pop_run_and_capture_generation(run_id, None, response) - generation_args = mock_client.capture.call_args_list[0][1] + generation_args = mock_client.capture_ai.call_args_list[0][1] # 2000 - 800 (read) - 500 (write) = 700 assert generation_args["properties"]["$ai_input_tokens"] == 700 assert generation_args["properties"]["$ai_cache_read_input_tokens"] == 800 @@ -1966,9 +1966,9 @@ def test_openai_cache_read_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response using cached prompt context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2010,9 +2010,9 @@ def test_openai_cache_creation_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Created cache for the prompt context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2052,9 +2052,9 @@ def test_combined_reasoning_and_cache_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Let me reason through this using cached context..." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2074,7 +2074,7 @@ def test_openai_reasoning_tokens_o4_mini(mock_client): mock_client, trace_id="test-trace-id", distinct_id="test-distinct-id" ) model.invoke("what is the weather in sf", config={"callbacks": [cb]}) - call = mock_client.capture.call_args_list[0][1] + call = mock_client.capture_ai.call_args_list[0][1] assert call["properties"]["$ai_reasoning_tokens"] is not None assert call["properties"]["$ai_input_tokens"] is not None assert call["properties"]["$ai_output_tokens"] is not None @@ -2103,7 +2103,7 @@ def test_callback_handler_without_client(): assert result.content == "Bar" # Verify that the mock client was used for capturing events - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 def test_convert_message_to_dict_tool_calls(): @@ -2210,8 +2210,8 @@ def test_tool_definition(mock_client): parent_run_id=None, ) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == run_id @@ -2261,9 +2261,9 @@ def test_cache_read_tokens_subtraction_from_input_tokens(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response using cached prompt context." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2303,9 +2303,9 @@ def test_cache_read_tokens_subtraction_prevents_negative(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response with edge case token counts." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2344,9 +2344,9 @@ def test_no_cache_read_tokens_no_subtraction(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response without cache." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2383,9 +2383,9 @@ def test_zero_input_tokens_with_cache_read(mock_client): result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Response." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2424,9 +2424,9 @@ def test_non_anthropic_cache_write_tokens_not_subtracted_from_input(mock_client) result = chain.invoke({}, config={"callbacks": callbacks}) assert result.content == "Creating cache." - assert mock_client.capture.call_count == 3 + assert mock_client.capture_ai.call_count == 3 - generation_args = mock_client.capture.call_args_list[1][1] + generation_args = mock_client.capture_ai.call_args_list[1][1] generation_props = generation_args["properties"] assert generation_args["event"] == "$ai_generation" @@ -2475,8 +2475,8 @@ def test_agent_action_and_finish_imports(): callbacks.on_agent_finish(finish, run_id=run_id, parent_run_id=parent_run_id) # Verify capture was called - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] assert call_args["event"] == "$ai_span" @@ -2583,8 +2583,8 @@ def test_billable_property_in_generation_event(mock_client): parent_run_id=None, ) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_generation" @@ -2603,7 +2603,7 @@ def test_billable_defaults_to_false_in_event(mock_client): chain.invoke({}, config={"callbacks": callbacks}) generation_call = None - for call in mock_client.capture.call_args_list: + for call in mock_client.capture_ai.call_args_list: if call[1]["event"] == "$ai_generation": generation_call = call break @@ -2647,8 +2647,8 @@ def test_billable_with_real_chain(mock_client): parent_run_id=None, ) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_generation" @@ -2706,7 +2706,7 @@ def failing_span(_): # Find the span event (should have $ai_is_error=True) span_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("properties", {}).get("$ai_is_error") is True ] assert len(span_calls) >= 1 @@ -2739,7 +2739,7 @@ def failing_span(_): # But the span event should still have error info span_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("properties", {}).get("$ai_is_error") is True ] assert len(span_calls) >= 1 @@ -2777,7 +2777,7 @@ def test_exception_autocapture_on_llm_generation_error(mock_client): # Verify the generation event has $exception_event_id generation_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("event") == "$ai_generation" ] assert len(generation_calls) == 1 @@ -2851,7 +2851,7 @@ def failing_span(_): # Span event should NOT have $exception_event_id span_calls = [ call - for call in mock_client.capture.call_args_list + for call in mock_client.capture_ai.call_args_list if call[1].get("properties", {}).get("$ai_is_error") is True ] assert len(span_calls) >= 1 @@ -2860,20 +2860,6 @@ def failing_span(_): assert "$exception_event_id" not in span_props -def test_ai_lane_client_routes_through_capture_ai(mock_client): - prompt = ChatPromptTemplate.from_messages([("user", "Who won the world series?")]) - model = FakeMessagesListChatModel(responses=[AIMessage(content="The Dodgers.")]) - mock_client.enable_full_ai_capture = True - callbacks = [CallbackHandler(mock_client)] - - (prompt | model).invoke({}, config={"callbacks": callbacks}) - - mock_client.capture.assert_not_called() - events = [c[1]["event"] for c in mock_client.capture_ai.call_args_list] - assert "$ai_generation" in events - assert "$ai_trace" in events - - def test_served_service_tier_merges_into_model_parameters(mock_client): from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, LLMResult @@ -2894,7 +2880,7 @@ def test_served_service_tier_merges_into_model_parameters(mock_client): cb._pop_run_and_capture_generation(run_id, None, response) - props = mock_client.capture.call_args.kwargs["properties"] + props = mock_client.capture_ai.call_args.kwargs["properties"] assert props["$ai_model_parameters"]["service_tier"] == "flex" assert props["$ai_model_parameters"]["temperature"] == 0.5 assert props["$ai_service_tier"] == "flex" @@ -2951,5 +2937,5 @@ def test_stop_reason_resolution( cb._pop_run_and_capture_generation(run_id, None, response) - props = mock_client.capture.call_args.kwargs["properties"] + props = mock_client.capture_ai.call_args.kwargs["properties"] assert props.get("$ai_stop_reason") == expected diff --git a/posthog/test/ai/langchain/test_middleware.py b/posthog/test/ai/langchain/test_middleware.py index ed68bca85..441cdc229 100644 --- a/posthog/test/ai/langchain/test_middleware.py +++ b/posthog/test/ai/langchain/test_middleware.py @@ -121,7 +121,7 @@ def _client() -> MagicMock: def _events(client: MagicMock) -> list[dict[str, Any]]: - return [call.kwargs for call in client.capture.call_args_list] + return [call.kwargs for call in client.capture_ai.call_args_list] def _event(client: MagicMock, event_name: str) -> dict[str, Any]: @@ -353,7 +353,7 @@ async def test_async_hooks_do_not_capture_on_the_event_loop_thread() -> None: client.sync_mode = True event_loop_thread = threading.get_ident() capture_threads: list[int] = [] - client.capture.side_effect = lambda **_: capture_threads.append( + client.capture_ai.side_effect = lambda **_: capture_threads.append( threading.get_ident() ) middleware = PostHogMiddleware(client) @@ -841,7 +841,7 @@ def test_privacy_mode_redacts_agent_and_model_content() -> None: def test_capture_failure_never_changes_agent_result() -> None: client = _client() - client.capture.side_effect = RuntimeError("telemetry unavailable") + client.capture_ai.side_effect = RuntimeError("telemetry unavailable") middleware = PostHogMiddleware(client) state = _middleware_state(middleware) request = _model_request(StubAgentModel(responses=[]), state) diff --git a/posthog/test/ai/openai/test_async_parity.py b/posthog/test/ai/openai/test_async_parity.py index 75ab907c5..8a96bfec8 100644 --- a/posthog/test/ai/openai/test_async_parity.py +++ b/posthog/test/ai/openai/test_async_parity.py @@ -46,7 +46,7 @@ def _sync_props(mock_client, chunks): posthog_distinct_id="test-id", ) ) - return mock_client.capture.call_args[1]["properties"] + return mock_client.capture_ai.call_args[1]["properties"] async def _async_props(mock_client, chunks): @@ -68,7 +68,7 @@ async def it(): ) async for _ in stream: pass - return mock_client.capture.call_args[1]["properties"] + return mock_client.capture_ai.call_args[1]["properties"] @pytest.mark.asyncio @@ -76,7 +76,7 @@ async def test_async_streaming_emits_the_same_properties_as_sync( mock_client, streaming_tool_call_chunks ): sync_props = _sync_props(mock_client, streaming_tool_call_chunks) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() async_props = await _async_props(mock_client, streaming_tool_call_chunks) # Guard: if the sync side stopped emitting these, the comparison below is vacuous. @@ -120,7 +120,7 @@ async def test_responses_streaming_properties_have_sync_async_parity(mock_client ): client = OpenAI(api_key="test-key", posthog_client=mock_client) list(client.responses.create(**request)) - sync_props = mock_client.capture.call_args.kwargs["properties"] + sync_props = mock_client.capture_ai.call_args.kwargs["properties"] async def create(self, **kwargs): async def chunks(): @@ -128,13 +128,13 @@ async def chunks(): return chunks() - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() with patch("openai.resources.responses.AsyncResponses.create", new=create): client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) stream = await client.responses.create(**request) async for _ in stream: pass - async_props = mock_client.capture.call_args.kwargs["properties"] + async_props = mock_client.capture_ai.call_args.kwargs["properties"] sync_without_latency = {k: v for k, v in sync_props.items() if k != "$ai_latency"} async_without_latency = {k: v for k, v in async_props.items() if k != "$ai_latency"} @@ -167,15 +167,15 @@ async def test_embedding_telemetry_has_sync_async_parity(mock_client): client = OpenAI(api_key="test-key", posthog_client=mock_client) assert client.embeddings.create(**request) is response sync_create.assert_called_once_with(**provider_request) - sync_capture = mock_client.capture.call_args + sync_capture = mock_client.capture_ai.call_args async_create = AsyncMock(return_value=response) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() with patch("openai.resources.embeddings.AsyncEmbeddings.create", new=async_create): client = AsyncOpenAI(api_key="test-key", posthog_client=mock_client) assert await client.embeddings.create(**request) is response async_create.assert_awaited_once_with(**provider_request) - async_capture = mock_client.capture.call_args + async_capture = mock_client.capture_ai.call_args sync_props = sync_capture.kwargs["properties"] async_props = async_capture.kwargs["properties"] @@ -209,5 +209,5 @@ def test_sync_stream_close_after_early_exit_captures_partial_state( assert next(stream) == streaming_tool_call_chunks[0] stream.close() - assert mock_client.capture.call_count == 1 - assert mock_client.capture.call_args.kwargs["properties"]["$ai_model"] == "gpt-4" + assert mock_client.capture_ai.call_count == 1 + assert mock_client.capture_ai.call_args.kwargs["properties"]["$ai_model"] == "gpt-4" diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index 0d22f393a..01f6f842d 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -393,9 +393,9 @@ def test_basic_completion(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -440,9 +440,9 @@ def test_embeddings(mock_client, mock_embedding_response): ) assert response == mock_embedding_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -470,9 +470,9 @@ def test_groups(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] assert call_args["groups"] == {"company": "test_company"} @@ -491,9 +491,9 @@ def test_privacy_mode_local(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -514,9 +514,9 @@ def test_privacy_mode_global(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input"] is None assert props["$ai_output_choices"] is None @@ -533,9 +533,9 @@ def test_error(mock_client, mock_openai_response): model="gpt-4", messages=[{"role": "user", "content": "Hello"}] ) - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_is_error"] is True assert props["$ai_error"] == "Test error" @@ -555,9 +555,9 @@ def test_cached_tokens(mock_client, mock_openai_response_with_cached_tokens): ) assert response == mock_openai_response_with_cached_tokens - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -594,9 +594,9 @@ def test_null_token_details_do_not_crash( ) assert response == mock_openai_response_with_null_token_details - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_input_tokens"] == 20 @@ -630,9 +630,9 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls): ) assert response == mock_openai_response_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -701,9 +701,9 @@ def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls ) assert response == mock_openai_response_tool_calls_only - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -756,9 +756,9 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call ) assert response == mock_responses_api_with_tool_calls - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -825,9 +825,9 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks): assert chunks == streaming_tool_call_chunks # Verify the capture was called with the right arguments - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -901,9 +901,9 @@ def test_responses_api(mock_client, mock_openai_response_with_responses_api): posthog_properties={"foo": "bar"}, ) assert response == mock_openai_response_with_responses_api - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -966,9 +966,9 @@ def test_responses_parse(mock_client, mock_parsed_response): ) assert response == mock_parsed_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1018,9 +1018,9 @@ def test_chat_completions_parse(mock_client, mock_openai_response): assert response == mock_openai_response assert mock_parse.call_count == 1 assert "posthog_distinct_id" not in mock_parse.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1059,9 +1059,9 @@ async def test_async_chat_completions_parse(mock_client, mock_openai_response): assert response == mock_openai_response mock_parse.assert_awaited_once() assert "posthog_distinct_id" not in mock_parse.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1184,9 +1184,9 @@ def mock_streaming_response(**kwargs): assert "stream_options" not in captured_kwargs # Verify capture was called - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify tokens are captured correctly from response.usage (not 0) @@ -1249,8 +1249,8 @@ async def chunk_iterable(): assert len(chunks) == len(streaming_tool_call_chunks) assert chunks == streaming_tool_call_chunks - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1317,8 +1317,8 @@ async def chunk_iterable(): kwargs = captured_kwargs["kwargs"] assert "stream_options" not in kwargs - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1347,9 +1347,9 @@ async def test_async_embeddings_create(mock_client, mock_embedding_response): assert response == mock_embedding_response assert mock_create.await_count == 1 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1398,9 +1398,9 @@ def test_tool_definition(mock_client, mock_openai_response): ) assert response == mock_openai_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1462,9 +1462,9 @@ def __init__(self): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is detected (binary detection) @@ -1521,9 +1521,9 @@ def mock_create_call(**kwargs): ) assert result == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify exact web search count @@ -1608,9 +1608,9 @@ def test_streaming_with_web_search(mock_client, streaming_web_search_chunks): # Verify the chunks were returned correctly assert len(chunks) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured (binary detection = 1) @@ -1644,9 +1644,9 @@ def test_streaming_with_web_search_on_non_usage_chunk( # Verify the chunks were returned correctly assert len(chunks) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured even though citations were on first chunk @@ -1700,9 +1700,9 @@ async def mock_create(self, **kwargs): ) assert response == mock_response - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured (binary detection = 1) @@ -1749,9 +1749,9 @@ async def chunk_iterable(): # Verify the chunks were returned correctly assert len(chunks) == 3 - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Verify web search count is captured (binary detection = 1) @@ -1816,8 +1816,8 @@ def test_streaming_chat_extracts_model_from_chunk_when_not_in_kwargs(mock_client # Consume the generator list(response_generator) - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from chunk, not kwargs @@ -1861,7 +1861,7 @@ def test_streaming_chat_prefers_kwargs_model_over_chunk_model(mock_client): list(response_generator) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # kwargs model should take precedence @@ -1909,7 +1909,7 @@ def test_streaming_responses_api_extracts_model_from_response_object(mock_client list(response_generator) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from chunk.response.model @@ -1983,7 +1983,7 @@ def test_streaming_responses_api_captures_usage_and_output_when_incomplete( ) list(response_generator) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] == "max_output_tokens" assert props["$ai_input_tokens"] == 20 assert props["$ai_output_tokens"] == 16 @@ -2033,7 +2033,7 @@ def test_non_streaming_extracts_model_from_response(mock_client): ) assert response == mock_response - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from response.model @@ -2093,7 +2093,7 @@ def test_non_streaming_responses_api_extracts_model_from_response(mock_client): ) assert response == mock_response - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Model should be extracted from response.model @@ -2139,7 +2139,7 @@ def test_non_streaming_returns_none_when_no_model(mock_client): posthog_distinct_id="test-id", ) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should be None, NOT "unknown" (to avoid incorrect cost matching) @@ -2178,7 +2178,7 @@ def test_streaming_falls_back_to_unknown_when_no_model(mock_client): list(response_generator) - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] # Should fall back to "unknown" @@ -2231,7 +2231,7 @@ async def chunk_iterable(): async for _ in response_stream: pass - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model"] == "gpt-4o-async-stored" @@ -2281,7 +2281,7 @@ async def chunk_iterable(): async for _ in response_stream: pass - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_model"] == "gpt-4o-mini-async-stored" @@ -2300,7 +2300,7 @@ def test_integration_stop_reason(mock_client): posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_stop_reason"] in ("stop", "length") assert props["$ai_provider"] == "openai" assert props["$ai_input_tokens"] > 0 @@ -2334,8 +2334,8 @@ async def mock_create(self, **kwargs): chunks.append(chunk) assert chunks == streaming_tool_call_chunks - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert call_args["event"] == "$ai_generation" assert props["$ai_provider"] == "openai" @@ -2369,7 +2369,7 @@ async def mock_create(self, **kwargs): received = [c async for c in stream] assert received == [chunk] - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 @pytest.mark.asyncio @@ -2400,25 +2400,7 @@ async def mock_create(self, **kwargs): break assert source.closed is True - assert mock_client.capture.call_count == 1 - - -def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_openai_response): - mock_client.enable_full_ai_capture = True - with patch( - "openai.resources.chat.completions.Completions.create", - return_value=mock_openai_response, - ): - client = OpenAI(api_key="test-key", posthog_client=mock_client) - client.chat.completions.create( - model="gpt-4", - messages=[{"role": "user", "content": "Hello"}], - posthog_distinct_id="test-id", - ) - - mock_client.capture.assert_not_called() assert mock_client.capture_ai.call_count == 1 - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_generation" def test_multimodal_client_skips_media_redaction(mock_client, mock_openai_response): @@ -2481,8 +2463,8 @@ def test_provider_override_chat_completions(mock_client, mock_openai_response): # The override must never reach the underlying OpenAI request. assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 - call_args = mock_client.capture.call_args[1] + assert mock_client.capture_ai.call_count == 1 + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "deepseek" @@ -2507,8 +2489,8 @@ def test_provider_override_omitted_defaults_to_openai( posthog_distinct_id="test-id", ) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "openai" @@ -2549,8 +2531,8 @@ def test_provider_override_streaming_chat_completions(mock_client): list(response_generator) assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "groq" assert props["$ai_model"] == "gpt-4" @@ -2572,9 +2554,9 @@ def test_provider_override_responses_api( assert response == mock_openai_response_with_responses_api assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "xai" assert props["$ai_model"] == "gpt-4o-mini" @@ -2595,9 +2577,9 @@ def test_provider_override_embeddings(mock_client, mock_embedding_response): assert response == mock_embedding_response assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "mistral" assert props["$ai_model"] == "text-embedding-3-small" @@ -2620,9 +2602,9 @@ def test_provider_override_chat_completions_parse(mock_client, mock_openai_respo assert response == mock_openai_response assert mock_parse.call_count == 1 assert "posthog_provider_override" not in mock_parse.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - call_args = mock_client.capture.call_args[1] + call_args = mock_client.capture_ai.call_args[1] props = call_args["properties"] assert props["$ai_provider"] == "cerebras" assert props["$ai_model"] == "gpt-4" @@ -2649,9 +2631,9 @@ async def test_async_provider_override_chat_completions( assert response == mock_openai_response mock_create.assert_awaited_once() assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "together" assert props["$ai_model"] == "gpt-4" @@ -2675,8 +2657,8 @@ async def test_async_provider_override_omitted_defaults_to_openai( posthog_distinct_id="test-id", ) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "openai" @@ -2715,8 +2697,8 @@ async def chunk_iterable(): assert "posthog_provider_override" not in captured_kwargs["kwargs"] assert len(chunks) == len(streaming_tool_call_chunks) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + assert mock_client.capture_ai.call_count == 1 + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "fireworks" assert props["$ai_model"] == "gpt-4" @@ -2738,9 +2720,9 @@ async def test_async_provider_override_embeddings(mock_client, mock_embedding_re assert response == mock_embedding_response assert mock_create.await_count == 1 assert "posthog_provider_override" not in mock_create.call_args.kwargs - assert mock_client.capture.call_count == 1 + assert mock_client.capture_ai.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_provider"] == "perplexity" assert props["$ai_model"] == "text-embedding-3-small" @@ -2760,7 +2742,7 @@ def test_served_service_tier_lands_in_model_parameters( posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert props["$ai_model_parameters"]["service_tier"] == "flex" assert props["$ai_service_tier"] == "flex" @@ -2777,7 +2759,7 @@ def test_response_without_service_tier_omits_it(mock_client, mock_openai_respons posthog_distinct_id="test-id", ) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] assert "service_tier" not in props["$ai_model_parameters"] assert "$ai_service_tier" not in props diff --git a/posthog/test/ai/openai_agents/test_processor.py b/posthog/test/ai/openai_agents/test_processor.py index 69745413b..203434abf 100644 --- a/posthog/test/ai/openai_agents/test_processor.py +++ b/posthog/test/ai/openai_agents/test_processor.py @@ -109,7 +109,7 @@ def test_on_trace_start_stores_metadata(self, processor, mock_client, mock_trace """Test that on_trace_start stores metadata but does not capture an event.""" processor.on_trace_start(mock_trace) - mock_client.capture.assert_not_called() + mock_client.capture_ai.assert_not_called() assert mock_trace.trace_id in processor._trace_metadata def test_on_trace_end_captures_ai_trace(self, processor, mock_client, mock_trace): @@ -117,8 +117,8 @@ def test_on_trace_end_captures_ai_trace(self, processor, mock_client, mock_trace processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - mock_client.capture.assert_called_once() - call_kwargs = mock_client.capture.call_args[1] + mock_client.capture_ai.assert_called_once() + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_trace" assert call_kwargs["distinct_id"] == "test-user" @@ -137,10 +137,10 @@ def test_group_id_becomes_session_id( mock_span.span_data = GenerationSpanData(model="gpt-4o") processor.on_span_start(mock_span) processor.on_span_end(mock_span) - span_kwargs = mock_client.capture.call_args[1] + span_kwargs = mock_client.capture_ai.call_args[1] processor.on_trace_end(mock_trace) - trace_kwargs = mock_client.capture.call_args[1] + trace_kwargs = mock_client.capture_ai.call_args[1] assert span_kwargs["properties"]["$ai_session_id"] == "group_123" assert span_kwargs["properties"]["$ai_group_id"] == "group_123" @@ -156,7 +156,7 @@ def test_personless_mode_when_no_distinct_id(self, mock_client, mock_trace): processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["options"] == {"process_person_profile": False} # Should fallback to trace_id as the distinct_id assert call_kwargs["distinct_id"] == mock_trace.trace_id @@ -170,7 +170,7 @@ def test_personless_mode_for_spans_when_no_distinct_id( ) processor.on_trace_start(mock_trace) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() span_data = GenerationSpanData(model="gpt-4o") mock_span.span_data = span_data @@ -178,7 +178,7 @@ def test_personless_mode_for_spans_when_no_distinct_id( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["options"] == {"process_person_profile": False} assert call_kwargs["distinct_id"] == mock_span.trace_id @@ -196,7 +196,7 @@ def resolver(trace): ) processor.on_trace_start(mock_trace) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() span_data = GenerationSpanData(model="gpt-4o") mock_span.span_data = span_data @@ -204,7 +204,7 @@ def resolver(trace): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["options"] == {"process_person_profile": False} assert call_kwargs["distinct_id"] == mock_span.trace_id @@ -218,7 +218,7 @@ def test_person_profile_when_distinct_id_provided(self, mock_client, mock_trace) processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert "options" not in call_kwargs def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace): @@ -229,7 +229,7 @@ def test_on_trace_end_clears_metadata(self, processor, mock_client, mock_trace): processor.on_trace_end(mock_trace) assert mock_trace.trace_id not in processor._trace_metadata # Also verify it captured the event - mock_client.capture.assert_called_once() + mock_client.capture_ai.assert_called_once() def test_on_span_start_tracks_time(self, processor, mock_span): """Test that on_span_start records start time.""" @@ -250,8 +250,8 @@ def test_generation_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - mock_client.capture.assert_called_once() - call_kwargs = mock_client.capture.call_args[1] + mock_client.capture_ai.assert_called_once() + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_generation" assert call_kwargs["properties"]["$ai_trace_id"] == "trace_123456789" @@ -285,7 +285,7 @@ def test_generation_span_with_reasoning_tokens( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_reasoning_tokens"] == 400 def test_function_span_mapping(self, processor, mock_client, mock_span): @@ -300,7 +300,7 @@ def test_function_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "get_weather" @@ -323,7 +323,7 @@ def test_agent_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "CustomerServiceAgent" @@ -345,7 +345,7 @@ def test_handoff_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_type"] == "handoff" @@ -367,7 +367,7 @@ def test_guardrail_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "ContentFilter" @@ -385,7 +385,7 @@ def test_custom_span_mapping(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_name"] == "database_query" @@ -414,7 +414,7 @@ def test_privacy_mode_redacts_content(self, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] # Content should be redacted assert call_kwargs["properties"]["$ai_input"] is None @@ -435,7 +435,7 @@ def test_client_privacy_mode_redacts_content(self, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - properties = mock_client.capture.call_args.kwargs["properties"] + properties = mock_client.capture_ai.call_args.kwargs["properties"] assert properties["$ai_input"] is None assert properties["$ai_output_choices"] is None @@ -480,7 +480,7 @@ def test_generation_span_image_input_is_redacted( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] captured_input = call_kwargs["properties"]["$ai_input"] assert captured_input[0]["content"][0]["image_url"] == "[base64 image redacted]" @@ -534,7 +534,7 @@ def test_response_span_image_output_is_redacted( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] captured_output = call_kwargs["properties"]["$ai_output_choices"] assert ( captured_output[0]["content"][0]["image_url"] == "[base64 image redacted]" @@ -549,7 +549,7 @@ def test_error_handling_in_span(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_is_error"] is True assert call_kwargs["properties"]["$ai_error"] == "Rate limit exceeded" @@ -567,7 +567,7 @@ def test_generation_span_includes_total_tokens( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_total_tokens"] == 150 def test_error_type_categorization_model_behavior( @@ -584,7 +584,7 @@ def test_error_type_categorization_model_behavior( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error" def test_error_type_categorization_user_error( @@ -598,7 +598,7 @@ def test_error_type_categorization_user_error( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "user_error" def test_error_type_categorization_input_guardrail( @@ -614,7 +614,7 @@ def test_error_type_categorization_input_guardrail( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert ( call_kwargs["properties"]["$ai_error_type"] == "input_guardrail_triggered" ) @@ -632,7 +632,7 @@ def test_error_type_categorization_output_guardrail( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert ( call_kwargs["properties"]["$ai_error_type"] == "output_guardrail_triggered" ) @@ -648,7 +648,7 @@ def test_error_type_categorization_max_turns( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "max_turns_exceeded" def test_error_type_categorization_unknown(self, processor, mock_client, mock_span): @@ -660,7 +660,7 @@ def test_error_type_categorization_unknown(self, processor, mock_client, mock_sp processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "unknown" def test_response_span_with_output_and_total_tokens( @@ -686,7 +686,7 @@ def test_response_span_with_output_and_total_tokens( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_generation" assert call_kwargs["properties"]["$ai_total_tokens"] == 35 @@ -715,7 +715,7 @@ def test_response_span_includes_total_cost( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_total_cost_usd"] == total_cost_usd def test_speech_span_with_pass_through_properties( @@ -735,7 +735,7 @@ def test_speech_span_with_pass_through_properties( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_type"] == "speech" @@ -768,7 +768,7 @@ def test_transcription_span_with_pass_through_properties( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["event"] == "$ai_span" assert call_kwargs["properties"]["$ai_span_type"] == "transcription" @@ -799,7 +799,7 @@ def test_transcription_span_audio_input_redacted( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_input"] == "[base64 audio redacted]" def test_transcription_span_audio_input_passthrough( @@ -834,7 +834,7 @@ def test_latency_calculation(self, processor, mock_client, mock_span): mock_time.return_value = 1001.5 # 1.5 seconds later processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_latency"] == pytest.approx(1.5, rel=0.01) def test_groups_included_in_events(self, mock_client, mock_trace, mock_span): @@ -848,7 +848,7 @@ def test_groups_included_in_events(self, mock_client, mock_trace, mock_span): processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["groups"] == {"company": "acme", "team": "engineering"} def test_additional_properties_included(self, mock_client, mock_trace): @@ -862,7 +862,7 @@ def test_additional_properties_included(self, mock_client, mock_trace): processor.on_trace_start(mock_trace) processor.on_trace_end(mock_trace) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["environment"] == "production" assert call_kwargs["properties"]["version"] == "1.0" @@ -889,7 +889,7 @@ def test_generation_span_with_no_usage(self, processor, mock_client, mock_span): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert "$ai_input_tokens" not in call_kwargs["properties"] assert "$ai_output_tokens" not in call_kwargs["properties"] assert "$ai_total_tokens" not in call_kwargs["properties"] @@ -907,7 +907,7 @@ def test_generation_span_with_partial_usage( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_input_tokens"] == 42 # The output side was never reported, so it is omitted rather than 0; # the total is the sum of the reported sides. @@ -928,7 +928,7 @@ def test_error_type_categorization_by_type_field_only( processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["properties"]["$ai_error_type"] == "model_behavior_error" def test_distinct_id_resolved_from_trace_for_spans( @@ -946,7 +946,7 @@ def resolver(trace): # Start trace - this resolves and stores distinct_id processor.on_trace_start(mock_trace) - mock_client.capture.reset_mock() + mock_client.capture_ai.reset_mock() # End a span - should use the stored distinct_id from trace span_data = GenerationSpanData(model="gpt-4o") @@ -955,7 +955,7 @@ def resolver(trace): processor.on_span_start(mock_span) processor.on_span_end(mock_span) - call_kwargs = mock_client.capture.call_args[1] + call_kwargs = mock_client.capture_ai.call_args[1] assert call_kwargs["distinct_id"] == "user-Test Workflow" def test_eviction_of_stale_entries(self, mock_client): @@ -1044,24 +1044,13 @@ def test_processor_groups_are_forwarded(self, mock_client): processor._capture_event("$ai_trace", {}) - assert mock_client.capture.call_args.kwargs["groups"] == {"company": "acme"} + assert mock_client.capture_ai.call_args.kwargs["groups"] == {"company": "acme"} def test_capture_errors_are_logged_and_suppressed(self, mock_client, caplog): - mock_client.capture.side_effect = RuntimeError("capture failed") + mock_client.capture_ai.side_effect = RuntimeError("capture failed") processor = PostHogTracingProcessor(client=mock_client) with caplog.at_level(logging.DEBUG, logger="posthog"): processor._capture_event("$ai_trace", {}) assert "Failed to capture PostHog event: capture failed" in caplog.text - - -def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_trace): - mock_client.enable_full_ai_capture = True - processor = PostHogTracingProcessor(client=mock_client, distinct_id="test-user") - processor.on_trace_start(mock_trace) - processor.on_trace_end(mock_trace) - - mock_client.capture.assert_not_called() - mock_client.capture_ai.assert_called_once() - assert mock_client.capture_ai.call_args[1]["event"] == "$ai_trace" diff --git a/posthog/test/ai/test_token_reporting.py b/posthog/test/ai/test_token_reporting.py index 790a3a6db..e06c2f075 100644 --- a/posthog/test/ai/test_token_reporting.py +++ b/posthog/test/ai/test_token_reporting.py @@ -52,7 +52,7 @@ def test_token_counts_trace_back_to_a_provider_report( ): capture_streaming_event(mock_client, _event_data(usage_stats)) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] for key in ("$ai_input_tokens", "$ai_output_tokens"): if key in expected: assert props[key] == expected[key] @@ -85,7 +85,7 @@ def test_aux_token_fields_trace_back_to_a_provider_report( ): capture_streaming_event(mock_client, _event_data(usage_stats, provider=provider)) - props = mock_client.capture.call_args[1]["properties"] + props = mock_client.capture_ai.call_args[1]["properties"] for key in ( "$ai_cache_read_input_tokens", "$ai_cache_creation_input_tokens", diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 2234de1b3..9034f8744 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -83,14 +83,6 @@ def test_sync_mode_capture_ai_posts_single_event_batch_to_ai_endpoint(self): batch = sent_batch(mock_post) self.assertEqual([e["event"] for e in batch], ["$ai_generation"]) - def test_multimodal_client_routes_wrapper_captures_to_ai_lane(self): - client = self._client(enable_full_ai_capture=True) - with patch_capture_send("consumer") as mock_post: - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.flush() - by_path = _events_by_path(mock_post) - self.assertEqual(set(by_path), {_CAPTURE_AI_V1_PATH}) - def test_disabled_client_never_starts_ai_lane(self): client = Client(TEST_API_KEY, disabled=True) client.capture_ai("$ai_generation", distinct_id="d") @@ -464,10 +456,13 @@ def test_fork_rebuild_replaces_sync_mode_queues(self): class TestCaptureAiEventHelper(unittest.TestCase): - """`_capture_ai_event` rides the AI lane only when the client opted in.""" + """`_capture_ai_event` always rides the AI lane; `enable_full_ai_capture` only controls content.""" - def test_opted_in_routes_through_ai_lane(self): - client = Client(TEST_API_KEY, flush_interval=0.05, enable_full_ai_capture=True) + @parameterized.expand([("default", False), ("full_capture", True)]) + def test_routes_only_through_ai_lane(self, _name, full_capture): + client = Client( + TEST_API_KEY, flush_interval=0.05, enable_full_ai_capture=full_capture + ) with patch_capture_send("consumer") as mock_post: _capture_ai_event( client, @@ -483,19 +478,6 @@ def test_opted_in_routes_through_ai_lane(self): ) client.join() - def test_default_keeps_capture_path(self): - client = Client(TEST_API_KEY, flush_interval=0.05) - with patch_capture_send("consumer") as mock_post: - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.flush() - - self.assertEqual( - [call.kwargs["path"] for call in mock_post.call_args_list], - [_CAPTURE_V1_PATH], - ) - self.assertEqual(client._ai_lane.consumers, []) - client.join() - def test_adds_ai_library_identity_and_preserves_provider_and_model(self): client = mock.Mock() _capture_ai_event( @@ -505,59 +487,16 @@ def test_adds_ai_library_identity_and_preserves_provider_and_model(self): properties={"$ai_provider": "openai", "$ai_model": "gpt-4o"}, ) - properties = client.capture.call_args.kwargs["properties"] + properties = client.capture_ai.call_args.kwargs["properties"] self.assertEqual(properties["$ai_lib"], "posthog-ai") self.assertEqual(properties["$ai_lib_version"], VERSION) self.assertEqual(properties["$ai_provider"], "openai") self.assertEqual(properties["$ai_model"], "gpt-4o") - def test_default_mock_clients_keep_seeing_capture(self): - # Downstream test suites pass Mock clients into the wrappers; without - # the opt-in they must keep seeing plain `capture()` calls. - client = mock.Mock() - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - client.capture_ai.assert_not_called() - - def test_opted_in_prefers_capture_ai(self): - client = mock.Mock(spec=["capture", "capture_ai", "enable_full_ai_capture"]) - client.enable_full_ai_capture = True - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture_ai.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - client.capture.assert_not_called() - - def test_opted_in_duck_typed_client_without_method_falls_back(self): + @parameterized.expand([("default", False), ("full_capture", True)]) + def test_client_without_capture_ai_falls_back_to_capture(self, _name, full_capture): client = mock.Mock(spec=["capture", "enable_full_ai_capture"]) - client.enable_full_ai_capture = True - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - - def test_client_multimodal_flag_prefers_capture_ai(self): - client = mock.Mock(spec=["capture", "capture_ai", "enable_full_ai_capture"]) - client.enable_full_ai_capture = True - _capture_ai_event(client, "$ai_generation", distinct_id="d") - client.capture_ai.assert_called_once_with( - event="$ai_generation", - distinct_id="d", - properties={"$ai_lib": "posthog-ai", "$ai_lib_version": VERSION}, - ) - client.capture.assert_not_called() - - def test_client_multimodal_flag_off_keeps_capture(self): - client = mock.Mock(spec=["capture", "capture_ai", "enable_full_ai_capture"]) - client.enable_full_ai_capture = False + client.enable_full_ai_capture = full_capture _capture_ai_event(client, "$ai_generation", distinct_id="d") client.capture.assert_called_once_with( event="$ai_generation", @@ -571,7 +510,7 @@ class TestLanesRefuseWorkAfterShutdown(unittest.TestCase): afterwards, even a lazy AI lane that never started before shutdown.""" def test_late_ai_capture_after_shutdown_starts_nothing_and_sends_nothing(self): - client = Client(TEST_API_KEY, enable_full_ai_capture=True, flush_interval=0.05) + client = Client(TEST_API_KEY, flush_interval=0.05) client.shutdown() with patch_capture_send("consumer") as mock_post: _capture_ai_event(client, "$ai_generation", distinct_id="d") From ed296dcb0ff00ddad061d54bb4a1543fa4d91a0a Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:47:22 -0700 Subject: [PATCH 13/14] feat: add an AI capture lane to AsyncPosthog (#1028) * feat: AsyncPosthog AI capture lane * fix: type capture v1 send timeouts as float * feat: AsyncPosthog privacy_mode and enable_full_ai_capture * fix: keep async capture_ai from raising on a non-string event name --- posthog/_async_consumer.py | 48 +++++- posthog/_async_request.py | 6 +- posthog/async_client.py | 217 ++++++++++++++++++++------- posthog/capture_send.py | 4 +- posthog/test/test_ai_capture_lane.py | 36 ++++- posthog/test/test_async_client.py | 121 ++++++++++++--- references/public_api_snapshot.txt | 7 +- sdk_compliance_adapter/adapter.py | 2 +- typings/requests/__init__.pyi | 2 +- 9 files changed, 350 insertions(+), 93 deletions(-) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index ce0b209c2..77e1294a4 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -101,9 +101,11 @@ def __init__( flush_at: int, flush_interval: float, retries: int, - timeout: int, + timeout: float, historical_migration: bool, capture_compression: CaptureCompression, + endpoint: str = _CAPTURE_V1_PATH, + max_msg_size: int = MAX_MSG_SIZE, ) -> None: self.queue = queue self.api_key = api_key @@ -116,6 +118,8 @@ def __init__( self.timeout = timeout self.historical_migration = historical_migration self.capture_compression = capture_compression + self.endpoint = endpoint + self.max_msg_size = max_msg_size self._carryover: Optional[tuple[dict[str, Any], int]] = None self._flush_event = asyncio.Event() @@ -176,7 +180,7 @@ async def upload(self, batch: list[dict[str, Any]]) -> None: await self.request(batch) except Exception as error: await _report_capture_failure( - self.on_error, self.log, error, batch, _CAPTURE_V1_PATH + self.on_error, self.log, error, batch, self.endpoint ) finally: for _ in batch: @@ -234,12 +238,15 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: self.queue.task_done() continue - if item_size > MAX_MSG_SIZE: + if item_size > self.max_msg_size: + # Log only name and size: AI events may carry unredacted + # multimodal payloads that must not leak into logs. self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", + "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", item.get("event"), item_size, - MAX_MSG_SIZE // 1024, + self.max_msg_size // 1024, + self.endpoint, ) self.queue.task_done() continue @@ -263,4 +270,35 @@ async def request(self, batch: list[dict[str, Any]]) -> None: timeout=self.timeout, max_retries=self.retries, historical_migration=self.historical_migration, + path=self.endpoint, ) + + +class _AsyncLane: + """One capture queue, the consumer tasks that drain it, and the endpoint they post to. + + The client owns one lane per traffic class (analytics, AI), so each gets + its own backpressure, timeout, size cap and compression. + """ + + def __init__( + self, + *, + name: str, + max_queue_size: int, + endpoint: str, + max_msg_size: int, + timeout: float, + capture_compression: CaptureCompression, + ) -> None: + self.name = name + self.queue: asyncio.Queue[Any] = asyncio.Queue(max_queue_size) + self.endpoint = endpoint + self.max_msg_size = max_msg_size + self.timeout = timeout + self.capture_compression = capture_compression + self.consumers: list[_AsyncConsumer] = [] + self.worker_tasks: list[asyncio.Task[None]] = [] + + def pending_items(self) -> int: + return int(getattr(self.queue, "_unfinished_tasks", self.queue.qsize())) diff --git a/posthog/_async_request.py b/posthog/_async_request.py index 9ddc7f7a2..1e10c3cd2 100644 --- a/posthog/_async_request.py +++ b/posthog/_async_request.py @@ -7,7 +7,7 @@ from urllib.parse import quote from .capture_compression import CaptureCompression -from .capture_send import _parse_retry_after, _send_v1_batch +from .capture_send import _CAPTURE_V1_PATH, _parse_retry_after, _send_v1_batch from .request import ( APIError, DatetimeSerializer, @@ -150,9 +150,10 @@ async def async_send_v1_batch( batch: list[dict[str, Any]], *, compression: CaptureCompression, - timeout: int, + timeout: float, max_retries: int, historical_migration: bool, + path: str = _CAPTURE_V1_PATH, ) -> None: """Run the existing capture-v1 submitter off-loop to preserve wire parity.""" await asyncio.to_thread( @@ -164,4 +165,5 @@ async def async_send_v1_batch( timeout=timeout, max_retries=max_retries, historical_migration=historical_migration, + path=path, ) diff --git a/posthog/async_client.py b/posthog/async_client.py index 00732d399..4956484de 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -19,6 +19,7 @@ from ._async_consumer import ( _STOP, _AsyncConsumer, + _AsyncLane, _invoke_callback, _report_capture_failure, _is_processing_event, @@ -34,6 +35,7 @@ from .args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs from .capture_compression import ( CaptureCompression, + _resolve_capture_ai_compression, _resolve_capture_compression, ) from .capture_event import ( @@ -44,7 +46,7 @@ _fill_event_defaults, _merge_groups, ) -from .capture_send import _CAPTURE_V1_PATH +from .capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from .client import ( MAX_DICT_SIZE as _MAX_DICT_SIZE, _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES, @@ -54,6 +56,7 @@ _add_context_session_id, _context_tag_defaults, _personless_options, + _positive_config_value, add_context_tags as _add_context_tags, get_identity_state as _get_identity_state, stringify_id as _stringify_id, @@ -82,7 +85,7 @@ mark_exception_as_captured, try_attach_code_variables_to_frames, ) -from .consumer import MAX_MSG_SIZE +from .consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE from .feature_flag_evaluations import ( FeatureFlagEvaluations, _EvaluatedFlagRecord, @@ -108,7 +111,12 @@ class AsyncClient: ``capture()`` is a synchronous, non-blocking queue write. Use ``await capture_immediate()`` when the caller must wait for delivery. + ``capture_ai()`` and ``await capture_ai_immediate()`` do the same for AI + events, on a separate queue that posts to the AI capture endpoint. ``flush()``, ``join()``, and ``shutdown()`` are awaitable lifecycle methods. + + The ``capture_ai_*``, ``privacy_mode`` and ``enable_full_ai_capture`` + options work the same way as on :class:`posthog.Client`. """ log = logging.getLogger("posthog") @@ -143,6 +151,12 @@ def __init__( code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, + capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, + capture_ai_max_queue_size: int = 1000, + capture_ai_timeout: float = 30, + capture_ai_max_event_bytes: int = AI_MAX_MSG_SIZE, + privacy_mode: bool = False, + enable_full_ai_capture: bool = False, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, @@ -153,6 +167,18 @@ def __init__( raise ValueError("flush_at must be greater than zero") if flush_interval <= 0: raise ValueError("flush_interval must be greater than zero") + capture_ai_max_queue_size = _positive_config_value( + "capture_ai_max_queue_size", capture_ai_max_queue_size, integer=True + ) + capture_ai_timeout = _positive_config_value( + "capture_ai_timeout", capture_ai_timeout + ) + capture_ai_max_event_bytes = _positive_config_value( + "capture_ai_max_event_bytes", + capture_ai_max_event_bytes, + integer=True, + maximum=AI_MAX_MSG_SIZE, + ) self.api_key = (project_api_key or "").strip() self.raw_host = normalize_host(host) @@ -170,7 +196,12 @@ def __init__( self.super_options = super_options self._release_id = _resolve_release_id() self.capture_compression = _resolve_capture_compression(capture_compression) + self.capture_ai_compression = _resolve_capture_ai_compression( + capture_ai_compression + ) self.capture_trace_context = capture_trace_context + self.privacy_mode = privacy_mode + self.enable_full_ai_capture = enable_full_ai_capture is True if personal_api_key is not None and secret_key is None: warnings.warn( "`personal_api_key` is deprecated; use `secret_key` instead.", @@ -245,12 +276,29 @@ def __init__( "api_key is empty after trimming whitespace; check your project API key" ) - self._queue: asyncio.Queue[Any] = asyncio.Queue(max_queue_size) + self._analytics_lane = _AsyncLane( + name="analytics", + max_queue_size=max_queue_size, + endpoint=_CAPTURE_V1_PATH, + max_msg_size=MAX_MSG_SIZE, + timeout=timeout, + capture_compression=self.capture_compression, + ) + # AI events post to their own endpoint, so multi-MB events stay off the + # analytics endpoint's smaller caps. Its workers start on the first AI + # event, so clients that never send one pay nothing. + self._ai_lane = _AsyncLane( + name="ai", + max_queue_size=capture_ai_max_queue_size, + endpoint=_CAPTURE_AI_V1_PATH, + max_msg_size=capture_ai_max_event_bytes, + timeout=capture_ai_timeout, + capture_compression=self.capture_ai_compression, + ) + self._lanes = (self._analytics_lane, self._ai_lane) self._worker_count = max(1, thread) self._flush_at = flush_at self._flush_interval = flush_interval - self._consumers: list[_AsyncConsumer] = [] - self._worker_tasks: list[asyncio.Task[None]] = [] self._immediate_callers: dict[asyncio.Task[Any], int] = {} self._inflight_operations: set[asyncio.Future[None]] = set() self._http_client: Optional[Any] = None @@ -264,7 +312,7 @@ def __init__( self._register_duplicate_client() async def __aenter__(self) -> AsyncClient: - self._ensure_workers_started() + self._ensure_workers_started(self._analytics_lane) return self async def __aexit__(self, exc_type, exc, tb) -> bool: @@ -325,9 +373,9 @@ def _get_http_client(self): self._http_client = _build_client(self.host) return self._http_client - def _new_consumer(self) -> _AsyncConsumer: + def _new_consumer(self, lane: _AsyncLane) -> _AsyncConsumer: return _AsyncConsumer( - self._queue, + lane.queue, self.api_key, host=self.host, on_error=self.on_error, @@ -335,22 +383,27 @@ def _new_consumer(self) -> _AsyncConsumer: flush_at=self._flush_at, flush_interval=self._flush_interval, retries=self.max_retries, - timeout=self.timeout, + timeout=lane.timeout, historical_migration=self.historical_migration, - capture_compression=self.capture_compression, + capture_compression=lane.capture_compression, + endpoint=lane.endpoint, + max_msg_size=lane.max_msg_size, ) - def _ensure_workers_started(self) -> None: - if self.disabled or not self.send or self._closed or self._worker_tasks: + def _ensure_workers_started(self, lane: _AsyncLane) -> None: + if self.disabled or not self.send or self._closed or lane.worker_tasks: return self._bind_loop() for _ in range(self._worker_count): - consumer = self._new_consumer() - self._consumers.append(consumer) - self._worker_tasks.append(asyncio.create_task(consumer.run())) + consumer = self._new_consumer(lane) + lane.consumers.append(consumer) + lane.worker_tasks.append(asyncio.create_task(consumer.run())) + + def _all_worker_tasks(self) -> list[asyncio.Task[None]]: + return [task for lane in self._lanes for task in lane.worker_tasks] def _enqueue_prepared_event( - self, prepared: dict[str, Any], defaults: _EventDefaults + self, prepared: dict[str, Any], lane: _AsyncLane, defaults: _EventDefaults ) -> bool: if not self._accepting or self._closed: return False @@ -368,13 +421,13 @@ def _enqueue_prepared_event( if self._loop is None: if running_loop is not None: - self._ensure_workers_started() - self._queue.put_nowait(queued_event) + self._ensure_workers_started(lane) + lane.queue.put_nowait(queued_event) return True if running_loop is self._loop: - self._ensure_workers_started() - self._queue.put_nowait(queued_event) + self._ensure_workers_started(lane) + lane.queue.put_nowait(queued_event) return True if running_loop is not None: raise RuntimeError("AsyncClient cannot be shared across event loops") @@ -388,8 +441,8 @@ def enqueue_on_bound_loop() -> None: if not self._accepting or self._closed: admitted.set_result(False) return - self._ensure_workers_started() - self._queue.put_nowait(queued_event) + self._ensure_workers_started(lane) + lane.queue.put_nowait(queued_event) except BaseException as error: admitted.set_exception(error) else: @@ -549,6 +602,30 @@ def capture( self, event: str, **kwargs: Unpack[OptionalCaptureArgs] ) -> Optional[str]: """Queue an event without blocking for network delivery.""" + return self._capture(event, self._analytics_lane, kwargs) + + def capture_ai( + self, event: str, **kwargs: Unpack[OptionalCaptureArgs] + ) -> Optional[str]: + """Queue an AI event for the AI capture endpoint without blocking. + + Takes the same arguments and returns the same value as ``capture()``. + The event uses a separate queue with a larger per-event size cap. The + payload is sent as given, with no redaction or truncation. + """ + self._log_non_ai_event(event) + return self._capture(event, self._ai_lane, kwargs) + + def _log_non_ai_event(self, event: str) -> None: + if isinstance(event, str) and not event.startswith("$ai_"): + self.log.debug( + "capture_ai called with non-AI event name %r; routing it to the AI endpoint anyway.", + event, + ) + + def _capture( + self, event: str, lane: _AsyncLane, kwargs: OptionalCaptureArgs + ) -> Optional[str]: try: msg, property_allowlist, defaults = self._build_capture_event(event, kwargs) prepared, sent_uuid = self._prepare_event(msg, property_allowlist) @@ -557,12 +634,12 @@ def capture( if not self.send: return sent_uuid - if not self._enqueue_prepared_event(prepared, defaults): + if not self._enqueue_prepared_event(prepared, lane, defaults): return None self.log.debug("queued async event %s", event) return sent_uuid except asyncio.QueueFull: - self.log.warning("PostHog async capture queue is full") + self.log.warning("PostHog async %s capture queue is full", lane.name) return None except Exception as error: if self.debug: @@ -584,6 +661,21 @@ async def capture_immediate( self, event: str, **kwargs: Unpack[OptionalCaptureArgs] ) -> Optional[str]: """Capture one event and wait until its delivery attempt completes.""" + return await self._capture_immediate(event, self._analytics_lane, kwargs) + + async def capture_ai_immediate( + self, event: str, **kwargs: Unpack[OptionalCaptureArgs] + ) -> Optional[str]: + """Capture one AI event and wait until its delivery attempt completes. + + Uses the AI capture endpoint and the ``capture_ai_*`` settings. + """ + self._log_non_ai_event(event) + return await self._capture_immediate(event, self._ai_lane, kwargs) + + async def _capture_immediate( + self, event: str, lane: _AsyncLane, kwargs: OptionalCaptureArgs + ) -> Optional[str]: current = asyncio.current_task() if current is None: # pragma: no cover - async functions always have a task return None @@ -610,21 +702,22 @@ async def capture_immediate( "unable to serialize immediate event for sizing, dropping" ) return None - if event_size > MAX_MSG_SIZE: + if event_size > lane.max_msg_size: self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", + "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", processed.get("event"), event_size, - MAX_MSG_SIZE // 1024, + lane.max_msg_size // 1024, + lane.endpoint, ) return None - consumer = self._new_consumer() + consumer = self._new_consumer(lane) await consumer.request(error_batch) return sent_uuid except Exception as error: await _report_capture_failure( - self.on_error, self.log, error, error_batch, _CAPTURE_V1_PATH + self.on_error, self.log, error, error_batch, lane.endpoint ) if self.debug: raise @@ -780,7 +873,7 @@ def _enqueue_built_event( defaults = self._event_defaults( context_options=context_options, disable_geoip=disable_geoip ) - if not self._enqueue_prepared_event(prepared, defaults): + if not self._enqueue_prepared_event(prepared, self._analytics_lane, defaults): return None return sent_uuid @@ -1107,40 +1200,43 @@ def _capture_feature_flag_called_if_needed( reported_flags.add(reported_key) def _pending_queue_items(self) -> int: - return int(getattr(self._queue, "_unfinished_tasks", self._queue.qsize())) + return sum(lane.pending_items() for lane in self._lanes) def _defer_lifecycle_call(self, awaitable) -> None: task = asyncio.create_task(_run_outside_processing_event(awaitable)) self._deferred_lifecycle_tasks.add(task) task.add_done_callback(self._deferred_lifecycle_tasks.discard) - def _discard_undrainable_queue(self) -> None: + def _discard_undrainable_queue(self, lane: _AsyncLane) -> None: discarded = 0 while True: try: - self._queue.get_nowait() + lane.queue.get_nowait() except asyncio.QueueEmpty: break - self._queue.task_done() + lane.queue.task_done() discarded += 1 - orphaned = self._pending_queue_items() + orphaned = lane.pending_items() for _ in range(orphaned): - self._queue.task_done() + lane.queue.task_done() discarded += orphaned if discarded: self.log.warning( - "discarded %d async capture items because all workers exited", + "discarded %d async %s capture items because all workers exited", discarded, + lane.name, ) - async def _wait_for_queue_drain(self, deadline: Optional[float]) -> None: - live_workers = [task for task in self._worker_tasks if not task.done()] + async def _wait_for_queue_drain( + self, lane: _AsyncLane, deadline: Optional[float] + ) -> None: + live_workers = [task for task in lane.worker_tasks if not task.done()] if not live_workers: - self._discard_undrainable_queue() + self._discard_undrainable_queue(lane) return - queue_join = asyncio.create_task(self._queue.join()) + queue_join = asyncio.create_task(lane.queue.join()) async def wait_for_workers() -> None: await asyncio.wait(live_workers, return_when=asyncio.ALL_COMPLETED) @@ -1163,7 +1259,7 @@ async def wait_for_workers() -> None: await queue_join return - self._discard_undrainable_queue() + self._discard_undrainable_queue(lane) await queue_join finally: for task in (queue_join, workers_finished): @@ -1172,22 +1268,28 @@ async def wait_for_workers() -> None: await asyncio.gather(queue_join, workers_finished, return_exceptions=True) async def flush(self, timeout_seconds: Optional[float] = 10) -> None: - if asyncio.current_task() in self._worker_tasks or _is_processing_event(): + if asyncio.current_task() in self._all_worker_tasks() or _is_processing_event(): self._defer_lifecycle_call(self.flush(timeout_seconds)) return if not self.send or self.disabled or self._pending_queue_items() == 0: return - self._ensure_workers_started() + pending_lanes = [lane for lane in self._lanes if lane.pending_items()] + for lane in pending_lanes: + self._ensure_workers_started(lane) deadline = ( None if timeout_seconds is None else asyncio.get_running_loop().time() + timeout_seconds ) try: - for consumer in self._consumers: - consumer.request_flush() - - await self._wait_for_queue_drain(deadline) + # Wake every lane before waiting on any, so one lane's partial + # batch does not wait out its flush_interval behind the other. + for lane in pending_lanes: + for consumer in lane.consumers: + consumer.request_flush() + + for lane in pending_lanes: + await self._wait_for_queue_drain(lane, deadline) except asyncio.TimeoutError: self.log.warning( "flush timed out after %s seconds with %s items pending", @@ -1204,7 +1306,7 @@ async def _close_transport(self) -> None: async def shutdown(self) -> None: current = asyncio.current_task() if ( - current in self._worker_tasks + current in self._all_worker_tasks() or current in self._immediate_callers or _is_processing_event() ): @@ -1234,17 +1336,20 @@ async def shutdown(self) -> None: errors.append(error) try: - live_workers = [task for task in self._worker_tasks if not task.done()] - for _ in live_workers: - await self._queue.put(_STOP) - if self._worker_tasks: - await asyncio.gather(*self._worker_tasks, return_exceptions=True) + for lane in self._lanes: + live_workers = [t for t in lane.worker_tasks if not t.done()] + for _ in live_workers: + await lane.queue.put(_STOP) + worker_tasks = self._all_worker_tasks() + if worker_tasks: + await asyncio.gather(*worker_tasks, return_exceptions=True) except Exception as error: self.log.exception("Failed to stop async capture workers") errors.append(error) finally: - self._worker_tasks.clear() - self._consumers.clear() + for lane in self._lanes: + lane.worker_tasks.clear() + lane.consumers.clear() try: await self._close_transport() diff --git a/posthog/capture_send.py b/posthog/capture_send.py index f254f165d..3c5f916d7 100644 --- a/posthog/capture_send.py +++ b/posthog/capture_send.py @@ -283,7 +283,7 @@ def _post_v1( attempt: int, request_id: str, compression: CaptureCompression = CaptureCompression.NONE, - timeout: int = 15, + timeout: float = 15, sdk_info: str = USER_AGENT, session: Optional["requests.Session"] = None, path: str = _CAPTURE_V1_PATH, @@ -405,7 +405,7 @@ def _send_v1_batch( batch: list[dict], *, compression: CaptureCompression = CaptureCompression.NONE, - timeout: int = 15, + timeout: float = 15, max_retries: int = 3, historical_migration: bool = False, sdk_info: str = USER_AGENT, diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index 9034f8744..d70bf703e 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -9,6 +9,7 @@ import posthog from posthog.ai.utils import _capture_ai_event, finalize_ai_content, with_privacy_mode +from posthog.async_client import AsyncPosthog from posthog.capture_compression import CAPTURE_COMPRESSION_ENV_VAR, CaptureCompression from posthog.client import Client from posthog.consumer import AI_MAX_MSG_SIZE, AI_MAX_PROPERTIES_SIZE, MAX_MSG_SIZE @@ -699,14 +700,16 @@ def drop_uuid(event): class TestCaptureAiPrivacyMode(unittest.TestCase): """Privacy mode always wins over `enable_full_ai_capture`.""" - def test_privacy_mode_strips_content_despite_full_ai_capture(self): - client = Client( + @parameterized.expand([("sync", Client), ("async", AsyncPosthog)]) + def test_privacy_mode_strips_content_despite_full_ai_capture( + self, _name, client_cls + ): + client = client_cls( TEST_API_KEY, send=False, enable_full_ai_capture=True, privacy_mode=True, ) - self.addCleanup(client.join) payload = {"role": "user", "content": "sensitive prompt"} sanitized = with_privacy_mode( @@ -715,6 +718,33 @@ def test_privacy_mode_strips_content_despite_full_ai_capture(self): self.assertIsNone(sanitized) + @parameterized.expand( + [ + ("sync_default", Client, False), + ("sync_full_capture", Client, True), + ("async_default", AsyncPosthog, False), + ("async_full_capture", AsyncPosthog, True), + ] + ) + def test_full_ai_capture_controls_media_redaction( + self, _name, client_cls, full_capture + ): + client = client_cls( + TEST_API_KEY, send=False, enable_full_ai_capture=full_capture + ) + image = "data:image/jpeg;base64," + "A" * 64 + payload = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": image}}], + } + ] + + sanitized = finalize_ai_content(payload, ph_client=client) + + sent_url = sanitized[0]["content"][0]["image_url"]["url"] + self.assertEqual(sent_url == image, full_capture) + if __name__ == "__main__": unittest.main() diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 02ce2ce4c..27c370a2f 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -12,7 +12,8 @@ import pytest from posthog import AsyncClient, AsyncPosthog, CaptureCompression -from posthog.consumer import MAX_MSG_SIZE +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH +from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE from posthog.contexts import ( new_context, set_capture_exception_code_variables_context, @@ -88,8 +89,8 @@ async def send_batch(api_key, host, batch, **kwargs): with patch_async_capture_send(side_effect=send_batch): client = AsyncPosthog("test-key", flush_interval=30) - client._ensure_workers_started() - while not client._queue._getters: + client._ensure_workers_started(client._analytics_lane) + while not client._analytics_lane.queue._getters: await asyncio.sleep(0) asyncio.get_running_loop().set_debug(True) @@ -124,7 +125,7 @@ async def send_batch(api_key, host, batch, **kwargs): @pytest.mark.asyncio async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): client = AsyncPosthog("test-key", flush_interval=30) - client._ensure_workers_started() + client._ensure_workers_started(client._analytics_lane) scheduled_callbacks = [] capture_result = [] @@ -148,9 +149,9 @@ async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): capture_thread.join(timeout=1) assert capture_result == [None] - for task in client._worker_tasks: + for task in client._analytics_lane.worker_tasks: task.cancel() - await asyncio.gather(*client._worker_tasks, return_exceptions=True) + await asyncio.gather(*client._analytics_lane.worker_tasks, return_exceptions=True) await client._close_transport() @@ -283,22 +284,41 @@ async def send_batch(api_key, host, batch, **kwargs): @pytest.mark.asyncio -async def test_capture_immediate_drops_oversized_event_after_before_send(): +@pytest.mark.parametrize( + ("method_name", "client_kwargs", "payload_size", "sent"), + [ + ("capture_immediate", {}, MAX_MSG_SIZE, False), + ("capture_ai_immediate", {}, MAX_MSG_SIZE, True), + ("capture_ai_immediate", {"capture_ai_max_event_bytes": 1024}, 2048, False), + ], +) +async def test_capture_immediate_applies_its_lane_size_cap_after_before_send( + method_name, client_kwargs, payload_size, sent +): def before_send(event): - event["properties"]["user_input"] = "x" * MAX_MSG_SIZE + event["properties"]["user_input"] = "x" * payload_size return event with patch_async_capture_send(new=mock.AsyncMock()) as send_batch: - client = AsyncPosthog("test-key", before_send=before_send) - result = await client.capture_immediate("event", distinct_id="user-1") + client = AsyncPosthog("test-key", before_send=before_send, **client_kwargs) + result = await getattr(client, method_name)("event", distinct_id="user-1") await client.shutdown() - assert result is None - send_batch.assert_not_awaited() + assert (result is not None) is sent + assert send_batch.await_count == (1 if sent else 0) @pytest.mark.asyncio -async def test_capture_immediate_uses_capture_v1_without_building_httpx_client(): +@pytest.mark.parametrize( + ("method_name", "path", "compression", "timeout"), + [ + ("capture_immediate", _CAPTURE_V1_PATH, CaptureCompression.GZIP, 15), + ("capture_ai_immediate", _CAPTURE_AI_V1_PATH, CaptureCompression.DEFLATE, 45), + ], +) +async def test_capture_immediate_uses_its_lane_without_building_httpx_client( + method_name, path, compression, timeout +): with ( mock.patch( "posthog._async_consumer.async_send_v1_batch", new=mock.AsyncMock() @@ -308,24 +328,72 @@ async def test_capture_immediate_uses_capture_v1_without_building_httpx_client() client = AsyncPosthog( "test-key", capture_compression=CaptureCompression.GZIP, + capture_ai_compression=CaptureCompression.DEFLATE, + capture_ai_timeout=45, + ) + event_uuid = await getattr(client, method_name)( + "$ai_generation", distinct_id="user-1" ) - event_uuid = await client.capture_immediate("event", distinct_id="user-1") await client.shutdown() assert event_uuid is not None build_client.assert_not_called() send_v1.assert_awaited_once() - assert send_v1.await_args.kwargs["compression"] == CaptureCompression.GZIP + assert send_v1.await_args.kwargs["path"] == path + assert send_v1.await_args.kwargs["compression"] == compression + assert send_v1.await_args.kwargs["timeout"] == timeout assert send_v1.await_args.args[2][0]["uuid"] == event_uuid +@pytest.mark.asyncio +async def test_ai_events_queue_on_their_own_lane_and_flush_drains_both(): + sent: dict[str, list[str]] = {} + + async def send_batch(api_key, host, batch, **kwargs): + sent.setdefault(kwargs["path"], []).extend(e["event"] for e in batch) + + with patch_async_capture_send(side_effect=send_batch): + client = AsyncPosthog("test-key", flush_at=100, flush_interval=30) + client.capture("pageview", distinct_id="user-1") + assert client._ai_lane.worker_tasks == [] + client.capture_ai( + "$ai_generation", + distinct_id="user-1", + properties={"$ai_input": "x" * MAX_MSG_SIZE}, + ) + await client.flush(timeout_seconds=1) + assert sent == { + _CAPTURE_V1_PATH: ["pageview"], + _CAPTURE_AI_V1_PATH: ["$ai_generation"], + } + + client.capture_ai("$ai_span", distinct_id="user-1") + await client.shutdown() + + assert sent[_CAPTURE_AI_V1_PATH] == ["$ai_generation", "$ai_span"] + assert client._all_worker_tasks() == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event", [None, 123]) +async def test_ai_capture_handles_a_non_string_event_name_like_capture(event): + client = AsyncPosthog("test-key", send=False) + try: + expected = client.capture(event, distinct_id="user-1") is None + assert (client.capture_ai(event, distinct_id="user-1") is None) is expected + immediate = await client.capture_ai_immediate(event, distinct_id="user-1") + assert (immediate is None) is expected + finally: + await client.shutdown() + + @pytest.mark.asyncio async def test_send_false_accepts_without_starting_workers_or_transport(): with mock.patch("posthog.async_client._build_client") as build_client: client = AsyncPosthog("test-key", send=False) assert client.capture("event", distinct_id="user-1") is not None assert await client.capture_immediate("event", distinct_id="user-1") is not None - assert client._worker_tasks == [] + assert client._all_worker_tasks() == [] await client.shutdown() build_client.assert_not_called() @@ -380,7 +448,7 @@ async def test_capture_after_shutdown_is_dropped_without_restarting_workers(): await client.shutdown() assert client.capture("event", distinct_id="user-1") is None assert await client.capture_immediate("event", distinct_id="user-1") is None - assert client._worker_tasks == [] + assert client._all_worker_tasks() == [] build_client.assert_not_called() @@ -521,10 +589,10 @@ async def send_batch(api_key, host, batch, **kwargs): @pytest.mark.asyncio async def test_shutdown_returns_when_all_workers_exited_with_queued_work(): client = AsyncPosthog("test-key", flush_interval=30) - client._ensure_workers_started() - for task in client._worker_tasks: + client._ensure_workers_started(client._analytics_lane) + for task in client._analytics_lane.worker_tasks: task.cancel() - await asyncio.gather(*client._worker_tasks, return_exceptions=True) + await asyncio.gather(*client._analytics_lane.worker_tasks, return_exceptions=True) assert client.capture("undrainable event", distinct_id="user-1") is not None try: @@ -625,8 +693,17 @@ async def flush_and_close(): assert batches[0][0]["event"] == "event" -@pytest.mark.parametrize(("option", "value"), [("flush_at", 0), ("flush_interval", 0)]) -def test_rejects_non_positive_batch_settings(option, value): +@pytest.mark.parametrize( + ("option", "value"), + [ + ("flush_at", 0), + ("flush_interval", 0), + ("capture_ai_max_queue_size", 0), + ("capture_ai_timeout", 0), + ("capture_ai_max_event_bytes", AI_MAX_MSG_SIZE + 1), + ], +) +def test_rejects_invalid_batch_settings(option, value): with pytest.raises(ValueError, match=option): AsyncPosthog("test-key", **{option: value}) diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index a0214f6f5..516b1e716 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -656,6 +656,7 @@ attribute posthog.args.OptionalSetArgs.timestamp: NotRequired[Optional[Union[dat attribute posthog.args.OptionalSetArgs.uuid: NotRequired[Optional[Union[str, UUID]]] attribute posthog.async_client.AsyncClient.api_key = (project_api_key or '').strip() attribute posthog.async_client.AsyncClient.before_send = before_send +attribute posthog.async_client.AsyncClient.capture_ai_compression = _resolve_capture_ai_compression(capture_ai_compression) attribute posthog.async_client.AsyncClient.capture_compression = _resolve_capture_compression(capture_compression) attribute posthog.async_client.AsyncClient.capture_exception_code_variables = capture_exception_code_variables attribute posthog.async_client.AsyncClient.capture_trace_context = capture_trace_context @@ -667,6 +668,7 @@ attribute posthog.async_client.AsyncClient.debug = debug attribute posthog.async_client.AsyncClient.disable_geoip = disable_geoip attribute posthog.async_client.AsyncClient.disabled = disabled or not self.api_key attribute posthog.async_client.AsyncClient.distinct_ids_feature_flags_reported = SizeLimitedDict(_MAX_DICT_SIZE, set) +attribute posthog.async_client.AsyncClient.enable_full_ai_capture = enable_full_ai_capture is True attribute posthog.async_client.AsyncClient.feature_flags_request_max_retries = max(0, feature_flags_request_max_retries) attribute posthog.async_client.AsyncClient.feature_flags_request_timeout_seconds = feature_flags_request_timeout_seconds attribute posthog.async_client.AsyncClient.historical_migration = historical_migration @@ -678,6 +680,7 @@ attribute posthog.async_client.AsyncClient.log_captured_exceptions = log_capture attribute posthog.async_client.AsyncClient.max_retries = max(0, max_retries) attribute posthog.async_client.AsyncClient.on_error = on_error attribute posthog.async_client.AsyncClient.personal_api_key = self.secret_key +attribute posthog.async_client.AsyncClient.privacy_mode = privacy_mode attribute posthog.async_client.AsyncClient.project_root = os.getcwd() attribute posthog.async_client.AsyncClient.raw_host = normalize_host(host) attribute posthog.async_client.AsyncClient.secret_key = (resolved_secret_key.strip() if isinstance(resolved_secret_key, str) else resolved_secret_key) or None @@ -1176,7 +1179,7 @@ class posthog.ai.types.TokenUsage class posthog.ai.types.ToolInProgress class posthog.args.OptionalCaptureArgs class posthog.args.OptionalSetArgs -class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] = None, *, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, max_retries: int = 3, timeout: int = 15, thread: int = 1, disabled: bool = False, disable_geoip: bool = True, is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, super_options: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, capture_exception_code_variables: bool = False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, feature_flags_request_timeout_seconds: int = 3, feature_flags_request_max_retries: int = 1) +class posthog.async_client.AsyncClient(project_api_key: str, host: Optional[str] = None, *, debug: bool = False, max_queue_size: int = 10000, send: bool = True, on_error=None, flush_at: int = 100, flush_interval: float = 5.0, max_retries: int = 3, timeout: int = 15, thread: int = 1, disabled: bool = False, disable_geoip: bool = True, is_server: bool = True, historical_migration: bool = False, super_properties: Optional[dict[str, Any]] = None, super_options: Optional[dict[str, Any]] = None, before_send=None, log_captured_exceptions: bool = False, project_root: Optional[str] = None, capture_exception_code_variables: bool = False, code_variables_mask_patterns=None, code_variables_ignore_patterns=None, code_variables_mask_url_credentials=None, code_variables_detect_secrets=None, in_app_modules: Optional[list[str]] = None, capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_compression: Optional[Union[CaptureCompression, str]] = None, capture_ai_max_queue_size: int = 1000, capture_ai_timeout: float = 30, capture_ai_max_event_bytes: int = AI_MAX_MSG_SIZE, privacy_mode: bool = False, enable_full_ai_capture: bool = False, capture_trace_context: bool = False, secret_key: Optional[str] = None, personal_api_key: Optional[str] = None, feature_flags_request_timeout_seconds: int = 3, feature_flags_request_max_retries: int = 1) class posthog.async_client.AsyncPosthog class posthog.bucketed_rate_limiter.BucketedRateLimiter(bucket_size: Number, refill_rate: Number, refill_interval_seconds: Number, on_bucket_rate_limited: Optional[Callable[[Hashable], None]] = None, clock: Callable[[], float] = time.monotonic) class posthog.capture_compression.CaptureCompression @@ -1568,6 +1571,8 @@ method posthog.ai.stream.AsyncStreamWrapper.aclose() -> None method posthog.ai.stream.AsyncStreamWrapper.close() -> None method posthog.async_client.AsyncClient.alias(previous_id: ID_TYPES, distinct_id: Optional[str], timestamp: Optional[Union[datetime, str]] = None, uuid: Optional[str] = None, disable_geoip: Optional[bool] = None, options: Optional[Dict[str, Any]] = None) -> Optional[str] method posthog.async_client.AsyncClient.capture(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] +method posthog.async_client.AsyncClient.capture_ai(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] +method posthog.async_client.AsyncClient.capture_ai_immediate(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.async_client.AsyncClient.capture_exception(exception: Optional[ExceptionArg] = None, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.async_client.AsyncClient.capture_immediate(event: str, **kwargs: Unpack[OptionalCaptureArgs]) -> Optional[str] method posthog.async_client.AsyncClient.evaluate_flags(distinct_id: Optional[ID_TYPES] = None, *, groups: Optional[Mapping[str, Union[str, int]]] = None, person_properties: Optional[Dict[str, Any]] = None, group_properties: Optional[Dict[str, Dict[str, Any]]] = None, disable_geoip: Optional[bool] = None, flag_keys: Optional[list[str]] = None, device_id: Optional[str] = None) -> FeatureFlagEvaluations diff --git a/sdk_compliance_adapter/adapter.py b/sdk_compliance_adapter/adapter.py index 19068152b..7a9ec437c 100644 --- a/sdk_compliance_adapter/adapter.py +++ b/sdk_compliance_adapter/adapter.py @@ -172,7 +172,7 @@ def patched_post_v1( attempt: int, request_id: str, compression: CaptureCompression = CaptureCompression.NONE, - timeout: int = 15, + timeout: float = 15, sdk_info: str = USER_AGENT, session: Any = None, path: str = _CAPTURE_V1_PATH, diff --git a/typings/requests/__init__.pyi b/typings/requests/__init__.pyi index 9e7892beb..4cc05e76f 100644 --- a/typings/requests/__init__.pyi +++ b/typings/requests/__init__.pyi @@ -34,7 +34,7 @@ class Session: *, data: str | bytes, headers: dict[str, str], - timeout: int, + timeout: float, stream: bool = ..., ) -> Response: ... def get( From 5ec93d0dc45586ab80b95ba5f8bd39c46a1f95e0 Mon Sep 17 00:00:00 2001 From: Eli Reisman Date: Tue, 13 Oct 2026 20:47:38 -0700 Subject: [PATCH 14/14] test: cover capture invariants on both lanes (#1029) * test: capture invariant hardening * test: assert exactly one capture loss line per failed batch --- posthog/test/test_ai_capture_lane.py | 12 +++ posthog/test/test_async_client.py | 54 ++++++++---- posthog/test/test_capture_invariants.py | 111 ++++++++++++++++++++++++ posthog/test/test_client.py | 59 ++++++++++--- 4 files changed, 209 insertions(+), 27 deletions(-) create mode 100644 posthog/test/test_capture_invariants.py diff --git a/posthog/test/test_ai_capture_lane.py b/posthog/test/test_ai_capture_lane.py index d70bf703e..a2c90f138 100644 --- a/posthog/test/test_ai_capture_lane.py +++ b/posthog/test/test_ai_capture_lane.py @@ -209,6 +209,18 @@ def test_sync_mode_ai_event_over_cap_is_not_sent(self): mock_send.assert_not_called() self.assertIn("exceeds the 1024KiB limit", "\n".join(logs.output)) + def test_sync_mode_ai_event_at_endpoint_ceiling_is_sent(self): + client = Client(TEST_API_KEY, sync_mode=True) + with patch_capture_send("client") as mock_send: + result = client.capture_ai( + "$ai_generation", + distinct_id="d", + properties={"p": "x" * AI_MAX_PROPERTIES_SIZE}, + ) + + self.assertIsNotNone(result) + mock_send.assert_called_once() + def test_analytics_lane_rejects_events_over_900kib(self): client = self._client() consumer = client.consumers[0] diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 27c370a2f..f111000b3 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -13,7 +13,7 @@ from posthog import AsyncClient, AsyncPosthog, CaptureCompression from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH -from posthog.consumer import AI_MAX_MSG_SIZE, MAX_MSG_SIZE +from posthog.consumer import AI_MAX_MSG_SIZE, AI_MAX_PROPERTIES_SIZE, MAX_MSG_SIZE from posthog.contexts import ( new_context, set_capture_exception_code_variables_context, @@ -123,7 +123,10 @@ async def send_batch(api_key, host, batch, **kwargs): @pytest.mark.asyncio -async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): +@pytest.mark.parametrize("method_name", ["capture", "capture_ai"]) +async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission( + method_name, +): client = AsyncPosthog("test-key", flush_interval=30) client._ensure_workers_started(client._analytics_lane) scheduled_callbacks = [] @@ -136,7 +139,7 @@ async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): ): capture_thread = threading.Thread( target=lambda: capture_result.append( - client.capture("threaded event", distinct_id="user-1") + getattr(client, method_name)("threaded event", distinct_id="user-1") ), daemon=True, ) @@ -149,9 +152,12 @@ async def test_cross_thread_capture_rechecks_shutdown_before_queue_admission(): capture_thread.join(timeout=1) assert capture_result == [None] - for task in client._analytics_lane.worker_tasks: + assert client._ai_lane.worker_tasks == [] + assert client._pending_queue_items() == 0 + worker_tasks = client._all_worker_tasks() + for task in worker_tasks: task.cancel() - await asyncio.gather(*client._analytics_lane.worker_tasks, return_exceptions=True) + await asyncio.gather(*worker_tasks, return_exceptions=True) await client._close_transport() @@ -288,7 +294,8 @@ async def send_batch(api_key, host, batch, **kwargs): ("method_name", "client_kwargs", "payload_size", "sent"), [ ("capture_immediate", {}, MAX_MSG_SIZE, False), - ("capture_ai_immediate", {}, MAX_MSG_SIZE, True), + ("capture_ai_immediate", {}, AI_MAX_PROPERTIES_SIZE, True), + ("capture_ai_immediate", {}, AI_MAX_MSG_SIZE, False), ("capture_ai_immediate", {"capture_ai_max_event_bytes": 1024}, 2048, False), ], ) @@ -359,7 +366,7 @@ async def send_batch(api_key, host, batch, **kwargs): client.capture_ai( "$ai_generation", distinct_id="user-1", - properties={"$ai_input": "x" * MAX_MSG_SIZE}, + properties={"$ai_input": "x" * AI_MAX_PROPERTIES_SIZE}, ) await client.flush(timeout_seconds=1) assert sent == { @@ -813,23 +820,38 @@ async def on_error(error, batch): @pytest.mark.asyncio -@pytest.mark.parametrize("immediate", [False, True]) -async def test_failed_capture_does_not_log_server_response_detail(caplog, immediate): +@pytest.mark.parametrize( + ("method_name", "path"), + [ + ("capture", _CAPTURE_V1_PATH), + ("capture_immediate", _CAPTURE_V1_PATH), + ("capture_ai", _CAPTURE_AI_V1_PATH), + ("capture_ai_immediate", _CAPTURE_AI_V1_PATH), + ], +) +async def test_failed_capture_logs_one_loss_line_without_response_detail( + caplog, method_name, path +): caplog.set_level(logging.DEBUG, logger="posthog") server_error = APIError(400, "password=server-secret") with patch_async_capture_send(side_effect=server_error): client = AsyncPosthog("test-key", flush_at=1, max_retries=0) - if immediate: - await client.capture_immediate("event", distinct_id="user-1") - else: - client.capture("event", distinct_id="user-1") - await client.flush(timeout_seconds=1) + result = getattr(client, method_name)("$ai_generation", distinct_id="user-1") + if method_name.endswith("_immediate"): + await result + await client.flush(timeout_seconds=1) await client.shutdown() assert "server-secret" not in caplog.text - assert "APIError" in caplog.text - assert "status=400" in caplog.text + loss_lines = [ + r.getMessage() + for r in caplog.records + if "event(s) not persisted by" in r.getMessage() + ] + assert loss_lines == [ + f"[PostHog] 1 event(s) not persisted by {path}: APIError (status=400)" + ] @pytest.mark.asyncio diff --git a/posthog/test/test_capture_invariants.py b/posthog/test/test_capture_invariants.py new file mode 100644 index 000000000..48f933bee --- /dev/null +++ b/posthog/test/test_capture_invariants.py @@ -0,0 +1,111 @@ +import json +from typing import Any, Optional +from unittest import mock + +import pytest +from requests import Response + +from posthog import AsyncPosthog, CaptureCompression, Client +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH + +_UUID = "01890000-0000-7000-8000-000000000001" +_TIMESTAMP = "2026-01-02T03:04:05+00:00" + + +def _event_kwargs() -> dict[str, Any]: + return { + "distinct_id": "user-1", + "uuid": _UUID, + "timestamp": _TIMESTAMP, + "properties": { + "$process_person_profile": False, + "$product_tour_id": "tour-1", + "empty_list": [], + "empty_map": {}, + "none": None, + "nested": {"a": {"b": [1, None, {}]}}, + }, + "options": { + "process_person_profile": None, + "cookieless_mode": False, + "custom_list": [], + "custom_map": {"inner": {}}, + }, + } + + +def _record_wire(sent: list[tuple[str, list[dict]]]): + def post(url, data=None, **kwargs): + events = json.loads(data)["batch"] + sent.append((url, events)) + response = Response() + response.status_code = 200 + response._content = json.dumps( + {"results": {event["uuid"]: {"result": "ok"} for event in events}} + ).encode() + return response + + session = mock.Mock() + session.post.side_effect = post + return mock.patch("posthog.capture_send._get_session", return_value=session) + + +def _send_sync(method_name: str, before_send) -> list[tuple[str, list[dict]]]: + sent: list[tuple[str, list[dict]]] = [] + with _record_wire(sent): + client = Client( + "test-key", + sync_mode=True, + before_send=before_send, + capture_compression=CaptureCompression.NONE, + ) + assert getattr(client, method_name)("$ai_generation", **_event_kwargs()) + client.shutdown() + return sent + + +async def _send_async(method_name: str, before_send) -> list[tuple[str, list[dict]]]: + sent: list[tuple[str, list[dict]]] = [] + with _record_wire(sent): + client = AsyncPosthog( + "test-key", + before_send=before_send, + capture_compression=CaptureCompression.NONE, + ) + assert await getattr(client, method_name)("$ai_generation", **_event_kwargs()) + await client.shutdown() + return sent + + +def _pass_through(event: dict) -> Optional[dict]: + return event + + +@pytest.mark.parametrize( + ("method_name", "path"), + [("capture", _CAPTURE_V1_PATH), ("capture_ai", _CAPTURE_AI_V1_PATH)], +) +def test_pass_through_before_send_leaves_sync_wire_unchanged(method_name, path): + without_hook = _send_sync(method_name, None) + with_hook = _send_sync(method_name, _pass_through) + + assert [url.endswith(path) for url, _ in without_hook] == [True] + assert with_hook == without_hook + assert without_hook[0][1][0]["options"]["process_person_profile"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method_name", "path"), + [ + ("capture_immediate", _CAPTURE_V1_PATH), + ("capture_ai_immediate", _CAPTURE_AI_V1_PATH), + ], +) +async def test_pass_through_before_send_leaves_async_wire_unchanged(method_name, path): + without_hook = await _send_async(method_name, None) + with_hook = await _send_async(method_name, _pass_through) + + assert [url.endswith(path) for url, _ in without_hook] == [True] + assert with_hook == without_hook + assert without_hook[0][1][0]["options"]["process_person_profile"] is False diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py index c2c3e9ad9..281ab3bfa 100644 --- a/posthog/test/test_client.py +++ b/posthog/test/test_client.py @@ -18,7 +18,7 @@ import pytest from posthog.capture_compression import CaptureCompression -from posthog.capture_send import _CAPTURE_V1_PATH +from posthog.capture_send import _CAPTURE_AI_V1_PATH, _CAPTURE_V1_PATH from posthog.client import Client from posthog.contexts import get_context_session_id, new_context, set_context_session from posthog.request import APIError, GetResponse @@ -2768,20 +2768,35 @@ def test_sync_send_failure_in_debug_reports_once_then_raises(self, method_name): self.assertIs(raised.exception, failure) self.assertEqual(errors, [failure]) - def test_sync_send_failure_without_on_error_logs_aggregate_line(self): + @parameterized.expand( + [ + ("analytics", "capture", _CAPTURE_V1_PATH), + ("ai", "capture_ai", _CAPTURE_AI_V1_PATH), + ] + ) + def test_sync_send_failure_without_on_error_logs_aggregate_line( + self, _name, method_name, path + ): client = Client(FAKE_TEST_API_KEY, sync_mode=True) with patch_capture_send( "client", side_effect=APIError(400, "password=server-secret") ): with self.assertLogs("posthog", level="ERROR") as logs: - result = client.capture("event", distinct_id="distinct_id") + result = getattr(client, method_name)( + "$ai_generation", distinct_id="distinct_id" + ) self.assertIsNone(result) output = "\n".join(logs.output) - self.assertIn( - "1 event(s) not persisted by /i/v1/analytics/events: APIError (status=400)", - output, + loss_lines = [ + r.getMessage() + for r in logs.records + if "event(s) not persisted by" in r.getMessage() + ] + self.assertEqual( + loss_lines, + [f"[PostHog] 1 event(s) not persisted by {path}: APIError (status=400)"], ) self.assertNotIn("server-secret", output) @@ -3405,12 +3420,21 @@ def test_shutdown_does_not_wait_for_idle_consumers_flush_interval(self): self.assertLess(time.monotonic() - start, 1) - def test_shutdown_waits_for_racing_enqueue_before_draining(self): + @parameterized.expand( + [ + ("analytics", "capture", "_analytics_lane", _CAPTURE_V1_PATH), + ("ai_first_event", "capture_ai", "_ai_lane", _CAPTURE_AI_V1_PATH), + ] + ) + def test_shutdown_waits_for_racing_enqueue_before_draining( + self, _name, method_name, lane_attr, path + ): client = Client(FAKE_TEST_API_KEY, flush_interval=0.01) + lane = getattr(client, lane_attr) put_started = threading.Event() release_put = threading.Event() shutdown_done = threading.Event() - original_put = client.queue.put + original_put = lane.queue.put capture_result = [] def blocking_put(*args, **kwargs): @@ -3420,14 +3444,17 @@ def blocking_put(*args, **kwargs): capture_thread = threading.Thread( target=lambda: capture_result.append( - client.capture("racing event", distinct_id="distinct_id") + getattr(client, method_name)("$ai_racing", distinct_id="distinct_id") ) ) shutdown_thread = threading.Thread( target=lambda: (client.shutdown(), shutdown_done.set()) ) - with mock.patch.object(client.queue, "put", side_effect=blocking_put): + with ( + patch_capture_send("consumer") as mock_send, + mock.patch.object(lane.queue, "put", side_effect=blocking_put), + ): capture_thread.start() self.assertTrue(put_started.wait(2)) shutdown_thread.start() @@ -3443,7 +3470,17 @@ def blocking_put(*args, **kwargs): self.assertFalse(shutdown_thread.is_alive()) self.assertTrue(shutdown_done.is_set()) self.assertIsNotNone(capture_result[0]) - self.assertTrue(client.queue.empty()) + self.assertTrue(lane.queue.empty()) + self.assertEqual( + [ + (call.kwargs["path"], len(sent_batch(mock_send, i))) + for i, call in enumerate(mock_send.call_args_list) + ], + [(path, 1)], + ) + self.assertFalse( + any(c.is_alive() for lane in client._lanes for c in lane.consumers) + ) def test_shutdown_waits_for_sync_send_and_rejects_later_sends(self): client = Client(FAKE_TEST_API_KEY, sync_mode=True)