From b45de74be2c6a8dbe323894554b849a49d963f4d Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 14:11:29 +0200 Subject: [PATCH 01/13] feat: add asyncio capture client --- .sampo/changesets/wily-windweaver-kullervo.md | 5 + posthog/__init__.py | 2 + posthog/_async_consumer.py | 202 +++++ posthog/_async_request.py | 145 ++++ posthog/async_client.py | 773 ++++++++++++++++++ posthog/test/test_async_client.py | 369 +++++++++ posthog/test/test_async_request.py | 127 +++ pyproject.toml | 1 + references/public_api_snapshot.txt | 42 + uv.lock | 6 +- 10 files changed, 1671 insertions(+), 1 deletion(-) create mode 100644 .sampo/changesets/wily-windweaver-kullervo.md create mode 100644 posthog/_async_consumer.py create mode 100644 posthog/_async_request.py create mode 100644 posthog/async_client.py create mode 100644 posthog/test/test_async_client.py create mode 100644 posthog/test/test_async_request.py diff --git a/.sampo/changesets/wily-windweaver-kullervo.md b/.sampo/changesets/wily-windweaver-kullervo.md new file mode 100644 index 000000000..a952cfdfa --- /dev/null +++ b/.sampo/changesets/wily-windweaver-kullervo.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +Add an asyncio-native client for buffered and immediate event capture diff --git a/posthog/__init__.py b/posthog/__init__.py index 9729d2a64..01b10f53f 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -12,6 +12,8 @@ from posthog.capture_compression import CaptureCompression as CaptureCompression from posthog.capture_mode import CaptureMode as CaptureMode from posthog.client import Client +from posthog.async_client import AsyncClient as AsyncClient +from posthog.async_client import AsyncPosthog as AsyncPosthog from posthog.exception_capture import ExceptionCapture from posthog.contexts import ( identify_context as inner_identify_context, diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py new file mode 100644 index 000000000..570460858 --- /dev/null +++ b/posthog/_async_consumer.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import asyncio +import inspect +import json +import logging +from collections.abc import Awaitable, Callable +from typing import Any, Optional + +from ._async_request import async_batch_post, 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 + +_FLUSH = object() +_STOP = object() + + +class _AsyncConsumer: + """Consume an asyncio queue and upload capture batches.""" + + log = logging.getLogger("posthog") + + def __init__( + self, + queue: asyncio.Queue[Any], + api_key: str, + *, + host: Optional[str], + on_error: Optional[Callable[[Exception, list[dict[str, Any]]], Any]], + 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 + self.host = host + self.on_error = on_error + 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 + + async def run(self) -> None: + self.log.debug("async consumer is running") + try: + while True: + batch, stop = await self.next() + if batch: + await self.upload(batch) + if stop: + return + except asyncio.CancelledError: + raise + except Exception: + self.log.exception("async consumer stopped after an unexpected error") + finally: + self.log.debug("async consumer exited") + + async def upload(self, batch: list[dict[str, Any]]) -> None: + try: + await self.request(batch) + except Exception as error: + self.log.error("error uploading async capture batch: %s", error) + if self.on_error: + try: + result = self.on_error(error, batch) + if inspect.isawaitable(result): + await result + except Exception as callback_error: + self.log.error("on_error handler failed: %s", callback_error) + finally: + for _ in batch: + self.queue.task_done() + + async def next(self) -> tuple[list[dict[str, Any]], bool]: + items: list[dict[str, Any]] = [] + total_size = 0 + stop = False + started = asyncio.get_running_loop().time() + + while len(items) < self.flush_at: + remaining = self.flush_interval - ( + asyncio.get_running_loop().time() - started + ) + if remaining <= 0: + break + + try: + queued = await asyncio.wait_for(self.queue.get(), timeout=remaining) + except asyncio.TimeoutError: + break + + if queued is _FLUSH: + self.queue.task_done() + break + if queued is _STOP: + self.queue.task_done() + stop = True + break + + try: + item = await self.process_event(queued) + except Exception: + self.log.exception("unable to process queued event, dropping") + self.queue.task_done() + continue + + if item is None: + self.queue.task_done() + continue + + try: + serialized = await asyncio.to_thread( + json.dumps, item, cls=DatetimeSerializer + ) + item_size = len(serialized.encode()) + except Exception: + self.log.error("unable to serialize queued event for sizing, dropping") + self.queue.task_done() + continue + + if item_size > MAX_MSG_SIZE: + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", + item.get("event"), + item_size, + MAX_MSG_SIZE // 1024, + ) + self.queue.task_done() + continue + + items.append(item) + total_size += item_size + if total_size >= BATCH_SIZE_LIMIT: + self.log.debug("hit async batch size limit (size: %d)", total_size) + break + + 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)) diff --git a/posthog/_async_request.py b/posthog/_async_request.py new file mode 100644 index 000000000..79ce3a386 --- /dev/null +++ b/posthog/_async_request.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +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 .capture_compression import CaptureCompression +from .capture_v1 import _send_v1_batch +from .request import APIError, DatetimeSerializer, USER_AGENT, normalize_host +from .utils import remove_trailing_slash + +try: # pragma: no cover - exercised when the optional dependency is absent + import httpx +except ImportError: # pragma: no cover + httpx = None + + +def _require_httpx(): + if httpx is None: # pragma: no cover + raise RuntimeError( + "Async PostHog support requires httpx. Install it with `posthog[async]`." + ) + return httpx + + +def _build_client(host: Optional[str] = None): + httpx_module = _require_httpx() + base_url = remove_trailing_slash(normalize_host(host)) + return httpx_module.AsyncClient(base_url=base_url, follow_redirects=True) + + +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 _parse_retry_after(response: Any) -> Optional[float]: + value = response.headers.get("Retry-After") + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _process_response(response: Any) -> None: + if response.status_code == 200: + return + + retry_after = _parse_retry_after(response) + try: + payload = response.json() + detail = payload["detail"] + except (KeyError, TypeError, ValueError): + detail = response.text + raise APIError(response.status_code, detail, retry_after=retry_after) + + +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") + response = await http_client.post( + path, content=data, headers=headers, timeout=timeout + ) + _process_response(response) + finally: + if owns_client: + await http_client.aclose() + + +async def async_send_v1_batch( + api_key: str, + host: Optional[str], + batch: list[dict[str, Any]], + *, + compression: CaptureCompression, + timeout: int, + max_retries: int, + historical_migration: bool, +) -> None: + """Run the existing capture-v1 submitter off-loop to preserve wire parity.""" + await asyncio.to_thread( + _send_v1_batch, + api_key, + host, + batch, + compression=compression, + timeout=timeout, + max_retries=max_retries, + historical_migration=historical_migration, + ) diff --git a/posthog/async_client.py b/posthog/async_client.py new file mode 100644 index 000000000..8e6267f42 --- /dev/null +++ b/posthog/async_client.py @@ -0,0 +1,773 @@ +from __future__ import annotations + +import asyncio +import inspect +import logging +import os +import sys +import warnings +import weakref +from datetime import datetime, timezone +from typing import Any, Dict, Optional, Union +from uuid import UUID, uuid4 + +from typing_extensions import Unpack + +from ._async_consumer import _FLUSH, _STOP, _AsyncConsumer +from ._async_request import _build_client +from .args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs +from .capture_compression import ( + CaptureCompression, + _resolve_capture_compression, +) +from .capture_mode import CaptureMode, _resolve_capture_mode +from .client import ( + Client as _SyncClient, + add_context_tags as _add_context_tags, + get_identity_state as _get_identity_state, + stringify_id as _stringify_id, +) +from .contexts import get_context_session_id as _get_context_session_id +from .exception_utils import ( + DEFAULT_CODE_VARIABLES_DETECT_SECRETS, + DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS, + DEFAULT_CODE_VARIABLES_MASK_PATTERNS, + DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS, + _get_current_otel_span_properties, + exc_info_from_error, + exception_is_already_captured, + exceptions_from_error_tuple, + handle_in_app, + mark_exception_as_captured, + try_attach_code_variables_to_frames, +) +from .request import determine_server_host, normalize_host +from .utils import _normalize_timestamp, clean, system_context +from .version import VERSION + +__all__ = ["AsyncClient", "AsyncPosthog"] + + +class AsyncClient: + """Asyncio-native PostHog capture client. + + ``capture()`` is a synchronous, non-blocking queue write. Use + ``await capture_immediate()`` when the caller must wait for delivery. + ``flush()``, ``join()``, and ``shutdown()`` are awaitable lifecycle methods. + """ + + log = logging.getLogger("posthog") + + 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, + 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, + ) -> None: + self.api_key = (project_api_key or "").strip() + self.raw_host = normalize_host(host) + self.host = determine_server_host(host) + 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 + self.disable_geoip = disable_geoip + self.is_server = is_server + self.historical_migration = historical_migration + self.super_properties = super_properties + self.capture_mode = _resolve_capture_mode(capture_mode) + self.capture_compression = _resolve_capture_compression( + capture_compression, gzip_fallback=gzip + ) + self.capture_trace_context = capture_trace_context + self.log_captured_exceptions = log_captured_exceptions + self.capture_exception_code_variables = capture_exception_code_variables + self.code_variables_mask_patterns = ( + code_variables_mask_patterns + if code_variables_mask_patterns is not None + else DEFAULT_CODE_VARIABLES_MASK_PATTERNS + ) + self.code_variables_ignore_patterns = ( + code_variables_ignore_patterns + if code_variables_ignore_patterns is not None + else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS + ) + self.code_variables_mask_url_credentials = ( + code_variables_mask_url_credentials + if code_variables_mask_url_credentials is not None + else DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS + ) + self.code_variables_detect_secrets = ( + code_variables_detect_secrets + if code_variables_detect_secrets is not None + else DEFAULT_CODE_VARIABLES_DETECT_SECRETS + ) + self.in_app_modules = in_app_modules + self.project_root = project_root + if self.project_root is None: + try: + self.project_root = os.getcwd() + except Exception: + self.project_root = None + + if before_send is not None and not callable(before_send): + self.log.warning("before_send is not callable, it will be ignored") + before_send = None + self.before_send = before_send + + if debug: + logging.basicConfig() + self.log.setLevel(logging.DEBUG) + + if not self.api_key: + self.log.error( + "api_key is empty after trimming whitespace; check your project API key" + ) + + self._queue: asyncio.Queue[Any] = asyncio.Queue(max_queue_size) + 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_tasks: set[asyncio.Task[Any]] = set() + self._http_client: Optional[Any] = None + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._accepting = True + self._closed = False + self._shutdown_lock = asyncio.Lock() + self._deferred_lifecycle_tasks: set[asyncio.Task[Any]] = set() + self._duplicate_client_registry_key: Optional[tuple[str, str]] = None + self._register_duplicate_client() + + async def __aenter__(self) -> AsyncClient: + self._ensure_workers_started() + return self + + async def __aexit__(self, exc_type, exc, tb) -> bool: + await self.shutdown() + return False + + def _register_duplicate_client(self) -> None: + if self.disabled or not self.send or not self.api_key: + return + + registry_key = (self.api_key, self.host) + should_warn = False + with _SyncClient._client_registry_lock: + clients = _SyncClient._client_registry.setdefault( + registry_key, weakref.WeakSet() + ) + has_existing_client = len(clients) > 0 + clients.add(self) + self._duplicate_client_registry_key = registry_key + if ( + has_existing_client + and registry_key not in _SyncClient._duplicate_client_warnings + ): + _SyncClient._duplicate_client_warnings.add(registry_key) + should_warn = True + + if should_warn: + self.log.warning( + "Multiple active PostHog clients detected for the same project API key " + "and host. Reuse one client per application when possible." + ) + + def _unregister_duplicate_client(self) -> None: + registry_key = self._duplicate_client_registry_key + if registry_key is None: + return + with _SyncClient._client_registry_lock: + clients = _SyncClient._client_registry.get(registry_key) + if clients is not None: + clients.discard(self) + if not clients: + del _SyncClient._client_registry[registry_key] + _SyncClient._duplicate_client_warnings.discard(registry_key) + self._duplicate_client_registry_key = None + + def _bind_loop(self) -> asyncio.AbstractEventLoop: + loop = asyncio.get_running_loop() + if self._loop is None: + self._loop = loop + elif self._loop is not loop: + raise RuntimeError("AsyncClient cannot be shared across event loops") + return loop + + def _get_http_client(self): + self._bind_loop() + if self._http_client is None: + self._http_client = _build_client(self.host) + 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, + host=self.host, + on_error=self.on_error, + 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 _ensure_workers_started(self) -> None: + if self.disabled or not self.send or self._closed or self._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())) + + 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): + self.log.error( + "Invalid UUID %r. Falling back to a generated UUID.", raw_uuid + ) + else: + msg["uuid"] = normalized + return normalized + + normalized = str(uuid4()) + msg["uuid"] = normalized + return normalized + + 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: + return None, None + + timestamp = msg.get("timestamp") + if timestamp is None: + timestamp = datetime.now(tz=timezone.utc) + try: + msg["timestamp"] = _normalize_timestamp(timestamp) + except ValueError: + self.log.warning( + "Invalid timestamp %r. Falling back to the current UTC time.", + timestamp, + ) + msg["timestamp"] = datetime.now(tz=timezone.utc).isoformat() + + sent_uuid = self._normalize_uuid(msg) + properties = msg.setdefault("properties", {}) + 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.is_server: + msg["properties"]["$is_server"] = True + if property_allowlist is not None: + msg["properties"] = { + key: value + for key, value in msg["properties"].items() + if key in property_allowlist + } + + msg["distinct_id"] = _stringify_id(msg.get("distinct_id")) + cleaned = clean(msg) + return cleaned, sent_uuid + + async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: + if self.before_send is None: + return msg + + original_uuid = msg["uuid"] + try: + result = self.before_send(msg) + if inspect.isawaitable(result): + result = await result + if result is None: + self.log.debug("event dropped by before_send callback") + return None + if not isinstance(result, dict): + raise TypeError("before_send must return a dict or None") + processed = clean(result) + processed["uuid"] = original_uuid + if not self._accepting: + return None + return processed + except Exception as error: + self.log.exception("Error in before_send callback: %s", error) + return None + + def _build_capture_event( + self, event: str, kwargs: OptionalCaptureArgs + ) -> tuple[dict[str, Any], Optional[bool], Any]: + properties = {**(kwargs.get("properties") or {}), **system_context()} + if self.capture_trace_context: + properties = {**_get_current_otel_span_properties(), **properties} + properties = _add_context_tags(properties) + assert properties is not None + + 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 + + flags_snapshot = kwargs.get("flags") + send_feature_flags = kwargs.get("send_feature_flags") + if flags_snapshot is not None: + properties = {**flags_snapshot._get_event_properties(), **properties} + elif send_feature_flags: + warnings.warn( + "AsyncClient does not support deprecated send_feature_flags. Pass a " + "flags snapshot from evaluate_flags() instead.", + DeprecationWarning, + stacklevel=3, + ) + + return ( + { + "properties": properties, + "timestamp": kwargs.get("timestamp"), + "distinct_id": distinct_id, + "event": event, + "uuid": kwargs.get("uuid"), + }, + kwargs.get("disable_geoip"), + kwargs.get("_property_allowlist"), + ) + + def capture( + self, event: str, **kwargs: Unpack[OptionalCaptureArgs] + ) -> 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 + ) + if prepared is None or sent_uuid is None: + return None + if not self.send: + return sent_uuid + + try: + self._ensure_workers_started() + except RuntimeError: + # Construction and capture before the loop starts are supported. + # flush()/shutdown() will bind the client and start the workers. + pass + + self._queue.put_nowait(prepared) + self.log.debug("queued async event %s", event) + return sent_uuid + except asyncio.QueueFull: + self.log.warning("PostHog async capture queue is full") + return None + except Exception as error: + if self.debug: + raise + self.log.exception("Error in async capture: %s", error) + return None + + async def capture_immediate( + self, event: str, **kwargs: Unpack[OptionalCaptureArgs] + ) -> Optional[str]: + """Capture one event and wait until its delivery attempt completes.""" + current = asyncio.current_task() + if current is None: # pragma: no cover - async functions always have a task + return None + if not self._accepting: + return None + self._bind_loop() + self._immediate_tasks.add(current) + 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 + ) + if prepared is None or sent_uuid is None: + return None + processed = await self._process_event(prepared) + if processed is None: + return None + error_batch = [processed] + if not self.send: + return sent_uuid + + consumer = self._new_consumer() + await consumer.request(error_batch) + return sent_uuid + except Exception as error: + if self.on_error: + try: + callback_result = self.on_error(error, error_batch) + if inspect.isawaitable(callback_result): + await callback_result + except Exception as callback_error: + self.log.error("on_error handler failed: %s", callback_error) + if self.debug: + raise + self.log.exception("Error in immediate async capture: %s", error) + return None + finally: + self._immediate_tasks.discard(current) + + def _build_person_properties_event( + self, event: str, property_key: str, kwargs: OptionalSetArgs + ) -> Optional[dict[str, Any]]: + properties = _add_context_tags(kwargs.get("properties") or {}) + distinct_id, personless = _get_identity_state(kwargs.get("distinct_id")) + if personless or not properties: + return None + return { + "timestamp": kwargs.get("timestamp"), + "distinct_id": distinct_id, + property_key: properties, + "event": event, + "uuid": kwargs.get("uuid"), + } + + def set(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: + try: + 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")) + except Exception as error: + if self.debug: + raise + self.log.exception("Error in async set: %s", error) + return None + + def set_once(self, **kwargs: Unpack[OptionalSetArgs]) -> Optional[str]: + try: + 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")) + except Exception as error: + if self.debug: + raise + self.log.exception("Error in async set_once: %s", error) + return None + + def group_identify( + self, + 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]: + try: + if not _stringify_id(group_type): + self.log.warning( + "group_identify() called without a group_type, dropping event" + ) + return None + if not _stringify_id(group_key): + self.log.warning( + "group_identify() called without a group_key, dropping event" + ) + return None + + resolved_distinct_id = _get_identity_state(distinct_id)[0] + msg: dict[str, Any] = { + "event": "$groupidentify", + "properties": { + "$group_type": group_type, + "$group_key": group_key, + "$group_set": properties or {}, + }, + "distinct_id": resolved_distinct_id, + "timestamp": timestamp, + "uuid": uuid, + } + session_id = _get_context_session_id() + if session_id: + msg["properties"]["$session_id"] = str(session_id) + return self._enqueue_built_event(msg, disable_geoip) + except Exception as error: + if self.debug: + raise + self.log.exception("Error in async group_identify: %s", error) + return None + + def alias( + self, + 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]: + try: + resolved_previous_id = _stringify_id(previous_id) + if not resolved_previous_id: + self.log.warning("alias() called without a previous_id, dropping event") + return None + resolved_distinct_id, personless = _get_identity_state(distinct_id) + if personless: + self.log.warning("alias() called without a distinct_id, dropping event") + return None + msg: dict[str, Any] = { + "properties": { + "distinct_id": resolved_previous_id, + "alias": resolved_distinct_id, + }, + "timestamp": timestamp, + "event": "$create_alias", + "distinct_id": resolved_previous_id, + "uuid": uuid, + } + session_id = _get_context_session_id() + if session_id: + msg["properties"]["$session_id"] = str(session_id) + return self._enqueue_built_event(msg, disable_geoip) + except Exception as error: + if self.debug: + raise + self.log.exception("Error in async alias: %s", error) + return None + + def _enqueue_built_event( + self, msg: dict[str, Any], disable_geoip: Optional[bool] + ) -> Optional[str]: + prepared, sent_uuid = self._prepare_event(msg, disable_geoip) + if prepared is None or sent_uuid is None: + return None + if not self.send: + return sent_uuid + try: + self._ensure_workers_started() + except RuntimeError: + pass + self._queue.put_nowait(prepared) + return sent_uuid + + def capture_exception( + self, + exception: Optional[ExceptionArg] = None, + **kwargs: Unpack[OptionalCaptureArgs], + ) -> Optional[str]: + """Capture an exception. This method never raises, including in debug mode.""" + try: + if exception is not None and exception_is_already_captured(exception): + self.log.debug("Exception already captured, skipping") + return None + exc_info = ( + exc_info_from_error(exception) + if exception is not None + else sys.exc_info() + ) + if exc_info is None or exc_info == (None, None, None): + self.log.warning("No exception information available") + return None + + exceptions = exceptions_from_error_tuple(exc_info) + event = handle_in_app( + {"exception": {"values": exceptions}}, + in_app_include=self.in_app_modules, + project_root=self.project_root, + ) + exceptions = event["exception"]["values"] + properties = { + "$exception_list": exceptions, + **(kwargs.get("properties") or {}), + } + if self.capture_exception_code_variables: + try_attach_code_variables_to_frames( + exceptions, + exc_info, + mask_patterns=self.code_variables_mask_patterns, + ignore_patterns=self.code_variables_ignore_patterns, + mask_url_credentials=self.code_variables_mask_url_credentials, + detect_secrets=self.code_variables_detect_secrets, + ) + if self.log_captured_exceptions: + self.log.exception(exception, extra=kwargs) + + result = self.capture( + "$exception", + distinct_id=kwargs.get("distinct_id"), + properties=properties, + timestamp=kwargs.get("timestamp"), + uuid=kwargs.get("uuid"), + groups=kwargs.get("groups"), + flags=kwargs.get("flags"), + disable_geoip=kwargs.get("disable_geoip"), + ) + if exception is not None and result is not None: + mark_exception_as_captured(exception, result) + return result + except Exception as error: + self.log.exception("Failed to capture exception: %s", error) + return None + + def _pending_queue_items(self) -> int: + return int(getattr(self._queue, "_unfinished_tasks", self._queue.qsize())) + + def _defer_lifecycle_call(self, awaitable) -> None: + task = asyncio.create_task(awaitable) + self._deferred_lifecycle_tasks.add(task) + task.add_done_callback(self._deferred_lifecycle_tasks.discard) + + async def flush(self, timeout_seconds: Optional[float] = 10) -> None: + if asyncio.current_task() in self._worker_tasks: + 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() + deadline = ( + None + if timeout_seconds is None + else asyncio.get_running_loop().time() + timeout_seconds + ) + try: + for _ in self._worker_tasks: + if deadline is None: + await self._queue.put(_FLUSH) + else: + remaining = max(0.0, deadline - asyncio.get_running_loop().time()) + await asyncio.wait_for(self._queue.put(_FLUSH), remaining) + + if deadline is None: + await self._queue.join() + else: + remaining = max(0.0, deadline - asyncio.get_running_loop().time()) + await asyncio.wait_for(self._queue.join(), remaining) + except asyncio.TimeoutError: + self.log.warning( + "flush timed out after %s seconds with %s items pending", + timeout_seconds, + self._pending_queue_items(), + ) + + async def _close_transport(self) -> None: + http_client = self._http_client + self._http_client = None + if http_client is not None: + await http_client.aclose() + + async def shutdown(self) -> None: + current = asyncio.current_task() + if current in self._worker_tasks or current in self._immediate_tasks: + self._accepting = False + self._defer_lifecycle_call(self.shutdown()) + return + + async with self._shutdown_lock: + if self._closed: + return + self._bind_loop() + self._accepting = False + errors: list[Exception] = [] + + current = asyncio.current_task() + pending_immediate = [ + task + for task in self._immediate_tasks + if task is not current and not task.done() + ] + if pending_immediate: + await asyncio.gather(*pending_immediate, return_exceptions=True) + + try: + await self.flush(timeout_seconds=None) + except Exception as error: + self.log.exception("Failed to flush async capture queue") + errors.append(error) + + try: + for _ in self._worker_tasks: + await self._queue.put(_STOP) + if self._worker_tasks: + await asyncio.gather(*self._worker_tasks, return_exceptions=False) + except Exception as error: + self.log.exception("Failed to stop async capture workers") + errors.append(error) + finally: + self._worker_tasks.clear() + self._consumers.clear() + + try: + await self._close_transport() + except Exception as error: + self.log.exception("Failed to close async capture transport") + errors.append(error) + + try: + self._unregister_duplicate_client() + except Exception as error: + self.log.exception("Failed to unregister async client") + errors.append(error) + + self._closed = True + if errors and self.debug: + raise errors[0] + + async def join(self) -> None: + await self.shutdown() + + +class AsyncPosthog(AsyncClient): + """Customer-facing name for :class:`AsyncClient`.""" + + pass diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py new file mode 100644 index 000000000..af16caf75 --- /dev/null +++ b/posthog/test/test_async_client.py @@ -0,0 +1,369 @@ +from __future__ import annotations + +import asyncio +import logging +from unittest import mock + +import pytest + +from posthog import AsyncClient, AsyncPosthog, CaptureCompression, CaptureMode + + +@pytest.mark.asyncio +async def test_async_posthog_is_the_customer_facing_async_client(): + client = AsyncPosthog("test-key", send=False) + assert isinstance(client, AsyncClient) + await client.shutdown() + + +@pytest.mark.asyncio +async def test_capture_is_a_synchronous_queue_write_and_flushes(): + batches = [] + + async def batch_post(*args, **kwargs): + batches.append(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + async with AsyncPosthog("test-key", flush_at=100, flush_interval=30) as client: + event_uuid = client.capture( + "async event", + distinct_id="user-1", + properties={"plan": "pro"}, + ) + assert isinstance(event_uuid, str) + await client.flush(timeout_seconds=1) + + assert len(batches) == 1 + event = batches[0][0] + assert event["event"] == "async event" + assert event["distinct_id"] == "user-1" + assert event["properties"]["plan"] == "pro" + assert event["properties"]["$lib"] == "posthog-python" + assert event["properties"]["$geoip_disable"] is True + assert event["properties"]["$is_server"] is True + assert event["uuid"] == event_uuid + + +@pytest.mark.asyncio +async def test_capture_runs_async_before_send_in_consumer(): + batches = [] + + async def before_send(event): + await asyncio.sleep(0) + event["properties"]["from_before_send"] = True + return event + + async def batch_post(*args, **kwargs): + batches.append(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + async with AsyncPosthog( + "test-key", before_send=before_send, flush_interval=30 + ) as client: + event_uuid = client.capture("event", distinct_id="user-1") + await client.flush(timeout_seconds=1) + + assert batches[0][0]["properties"]["from_before_send"] is True + assert batches[0][0]["uuid"] == event_uuid + + +@pytest.mark.asyncio +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: + async with AsyncPosthog( + "test-key", before_send=before_send, flush_interval=0.01 + ) as client: + accepted_uuid = client.capture("event", distinct_id="user-1") + await client.flush(timeout_seconds=1) + + assert accepted_uuid is not None + batch_post.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" + delivered.set() + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key") + event_uuid = await client.capture_immediate( + "immediate event", distinct_id="user-1" + ) + assert delivered.is_set() + assert event_uuid is not None + await client.shutdown() + + +@pytest.mark.asyncio +async def test_capture_immediate_supports_async_before_send(): + batches = [] + + async def before_send(event): + await asyncio.sleep(0) + event["properties"]["processed"] = True + return event + + async def batch_post(*args, **kwargs): + batches.append(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key", before_send=before_send) + result = await client.capture_immediate("event", distinct_id="user-1") + await client.shutdown() + + assert result is not None + assert batches[0][0]["properties"]["processed"] is True + + +@pytest.mark.asyncio +async def test_capture_immediate_uses_capture_v1_without_building_httpx_client(): + with ( + mock.patch( + "posthog._async_consumer.async_send_v1_batch", new=mock.AsyncMock() + ) as send_v1, + mock.patch("posthog.async_client._build_client") as build_client, + ): + client = AsyncPosthog( + "test-key", + capture_mode=CaptureMode.V1, + capture_compression=CaptureCompression.GZIP, + ) + 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.args[2][0]["uuid"] == event_uuid + + +@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 == [] + await client.shutdown() + build_client.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method_name", "method_kwargs", "expected_event"), + [ + ( + "set", + {"distinct_id": "user-1", "properties": {"email": "a@example.com"}}, + "$set", + ), + ( + "set_once", + {"distinct_id": "user-1", "properties": {"first_seen": True}}, + "$set_once", + ), + ( + "alias", + {"previous_id": "anon-1", "distinct_id": "user-1"}, + "$create_alias", + ), + ( + "group_identify", + {"group_type": "company", "group_key": "company-1"}, + "$groupidentify", + ), + ], +) +async def test_identify_methods_enqueue_events( + method_name, method_kwargs, expected_event +): + batches = [] + + async def batch_post(*args, **kwargs): + batches.append(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + async with AsyncPosthog("test-key", flush_interval=30) as client: + method = getattr(client, method_name) + assert method(**method_kwargs) is not None + await client.flush(timeout_seconds=1) + + assert batches[0][0]["event"] == expected_event + + +@pytest.mark.asyncio +async def test_capture_after_shutdown_is_dropped_without_restarting_workers(): + with mock.patch("posthog.async_client._build_client") as build_client: + client = AsyncPosthog("test-key") + 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 == [] + build_client.assert_not_called() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_an_in_flight_batch_instead_of_cancelling_it(): + upload_started = asyncio.Event() + allow_upload = asyncio.Event() + delivered = [] + + async def batch_post(*args, **kwargs): + upload_started.set() + await allow_upload.wait() + delivered.extend(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key", flush_at=1) + client.capture("event", distinct_id="user-1") + await upload_started.wait() + shutdown = asyncio.create_task(client.shutdown()) + await asyncio.sleep(0) + assert not shutdown.done() + allow_upload.set() + await shutdown + + assert [event["event"] for event in delivered] == ["event"] + + +@pytest.mark.asyncio +async def test_shutdown_called_from_before_send_is_deferred_without_deadlock(): + callback_finished = asyncio.Event() + client: AsyncPosthog | None = None + + async def before_send(event): + assert client is not None + await client.shutdown() + callback_finished.set() + return event + + with mock.patch( + "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() + ) as batch_post: + 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) + await asyncio.wait_for(client.shutdown(), timeout=1) + + batch_post.assert_not_awaited() + assert client.capture("after shutdown", distinct_id="user-1") is None + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_in_flight_immediate_capture(): + upload_started = asyncio.Event() + allow_upload = asyncio.Event() + + async def batch_post(*args, **kwargs): + upload_started.set() + await allow_upload.wait() + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key") + capture = asyncio.create_task( + client.capture_immediate("event", distinct_id="user-1") + ) + await upload_started.wait() + shutdown = asyncio.create_task(client.shutdown()) + await asyncio.sleep(0) + assert not shutdown.done() + allow_upload.set() + assert await capture is not None + await shutdown + + +@pytest.mark.asyncio +async def test_reuses_and_closes_instance_owned_http_client(): + http_client = mock.Mock() + http_client.aclose = mock.AsyncMock() + + with ( + mock.patch( + "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, + ): + client = AsyncPosthog("test-key") + await client.capture_immediate("first", distinct_id="user-1") + await client.capture_immediate("second", distinct_id="user-1") + await client.shutdown() + + build.assert_called_once_with(client.host) + assert [call.kwargs["client"] for call in batch_post.await_args_list] == [ + http_client, + http_client, + ] + http_client.aclose.assert_awaited_once_with() + + +def test_capture_before_loop_starts_is_flushed_when_loop_runs(): + batches = [] + 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 flush_and_close(): + with mock.patch( + "posthog._async_consumer.async_batch_post", side_effect=batch_post + ): + await client.shutdown() + + asyncio.run(flush_and_close()) + assert batches[0][0]["event"] == "event" + + +@pytest.mark.asyncio +async def test_capture_exception_never_raises_in_debug_mode(): + client = AsyncPosthog("test-key", send=False, debug=True) + with mock.patch.object(client, "capture", side_effect=RuntimeError("broken")): + assert client.capture_exception(ValueError("boom")) is None + await client.shutdown() + + +@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()): + async with AsyncPosthog("test-key", flush_interval=30) as client: + client.capture( + "event", + distinct_id="user-1", + properties={"password": "super-secret"}, + ) + await client.flush(timeout_seconds=1) + + assert "super-secret" not in caplog.text + assert "test-key" not in caplog.text + + +@pytest.mark.asyncio +async def test_flush_timeout_reports_unfinished_items(caplog): + upload_started = asyncio.Event() + allow_upload = asyncio.Event() + + async def batch_post(*args, **kwargs): + upload_started.set() + await allow_upload.wait() + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key", flush_at=1) + client.capture("event", distinct_id="user-1") + await upload_started.wait() + await client.flush(timeout_seconds=0.01) + assert "items pending" in caplog.text + allow_upload.set() + await client.shutdown() diff --git a/posthog/test/test_async_request.py b/posthog/test/test_async_request.py new file mode 100644 index 000000000..4a07f0d14 --- /dev/null +++ b/posthog/test/test_async_request.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import asyncio +import logging +import subprocess +import sys +from unittest import mock + +import pytest + +from posthog._async_request import ( + _build_client, + _process_response, + async_batch_post, +) +from posthog.request import APIError + + +class FakeResponse: + def __init__(self, status_code=200, payload=None): + self.status_code = status_code + self._payload = payload if payload is not None else {"ok": True} + self.headers = {} + self.text = str(self._payload) + + def json(self): + return self._payload + + +class FakeAsyncClient: + def __init__(self, response=None): + self.response = response or FakeResponse() + self.calls = [] + self.closed = False + + async def post(self, *args, **kwargs): + self.calls.append((args, kwargs)) + return self.response + + async def aclose(self): + self.closed = True + + +def test_import_posthog_does_not_require_async_extra(): + script = """ +import builtins +original_import = builtins.__import__ +def blocked_import(name, *args, **kwargs): + if name == 'httpx' or name.startswith('httpx.'): + raise ImportError('httpx intentionally unavailable') + return original_import(name, *args, **kwargs) +builtins.__import__ = blocked_import +import posthog +assert posthog.Client +assert posthog.Posthog +assert posthog.AsyncClient +""" + subprocess.run([sys.executable, "-c", script], check=True) + + +def test_build_client_scopes_requests_to_host_and_follows_redirects(): + with mock.patch("posthog._async_request.httpx.AsyncClient") as async_client: + _build_client("https://example.com/") + async_client.assert_called_once_with( + base_url="https://example.com", follow_redirects=True + ) + + +@pytest.mark.asyncio +async def test_async_batch_post_uses_relative_path_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][0] == ("/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_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(), + ) + + +def test_process_response_raises_api_error_without_logging_payload(caplog): + caplog.set_level(logging.DEBUG, logger="posthog") + response = FakeResponse(400, {"detail": "password=secret"}) + + with pytest.raises(APIError): + _process_response(response) + + assert "password=secret" not in caplog.text diff --git a/pyproject.toml b/pyproject.toml index fccfe3663..5543e46f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ Repository = "https://github.com/posthog/posthog-python" [project.optional-dependencies] langchain = ["langchain>=1.3.9"] +async = ["httpx>=0.27.0,<1.0"] # Opt-in zstd support for capture-v1 request compression. Kept out of the core # dependencies so existing installs are unaffected; Python gains stdlib zstd # only in 3.14 (compression.zstd), so the third-party package is needed until diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 1f92d60e7..c302dbe66 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -3,6 +3,8 @@ # Public API scope: public posthog modules (excluding tests) and their exported # members. Modules with __all__ use it; other modules include non-underscore # names. External imports are excluded. +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 @@ -515,6 +517,33 @@ attribute posthog.args.OptionalSetArgs.distinct_id: NotRequired[Optional[ID_TYPE 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]]] +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_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 +attribute posthog.async_client.AsyncClient.code_variables_mask_patterns = code_variables_mask_patterns if code_variables_mask_patterns is not None else DEFAULT_CODE_VARIABLES_MASK_PATTERNS +attribute posthog.async_client.AsyncClient.code_variables_mask_url_credentials = code_variables_mask_url_credentials if code_variables_mask_url_credentials is not None else DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS +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.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 +attribute posthog.async_client.AsyncClient.is_server = is_server +attribute posthog.async_client.AsyncClient.log = logging.getLogger('posthog') +attribute posthog.async_client.AsyncClient.log_captured_exceptions = log_captured_exceptions +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.project_root = os.getcwd() +attribute posthog.async_client.AsyncClient.raw_host = normalize_host(host) +attribute posthog.async_client.AsyncClient.send = send +attribute posthog.async_client.AsyncClient.super_properties = super_properties +attribute posthog.async_client.AsyncClient.timeout = timeout attribute posthog.before_send = None attribute posthog.bucketed_rate_limiter.Number = Union[int, float] attribute posthog.bucketed_rate_limiter.ONE_DAY_IN_SECONDS = 86400.0 @@ -918,6 +947,8 @@ 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) +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 @@ -1250,6 +1281,16 @@ method posthog.ai.prompts.Prompts.compile(prompt: str, variables: PromptVariable method posthog.ai.prompts.Prompts.get(name: str, *, with_metadata: Optional[bool] = None, cache_ttl_seconds: Optional[int] = None, fallback: Optional[str] = None, version: Optional[int] = None, label: Optional[str] = None) -> Union[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.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.flush(timeout_seconds: Optional[float] = 10) -> None +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.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] +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.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] @@ -1412,6 +1453,7 @@ module posthog.ai.stream module posthog.ai.types module posthog.ai.utils module posthog.args +module posthog.async_client module posthog.bucketed_rate_limiter module posthog.capture_compression module posthog.capture_mode diff --git a/uv.lock b/uv.lock index da40814a3..212c1a852 100644 --- a/uv.lock +++ b/uv.lock @@ -2743,6 +2743,9 @@ dependencies = [ ] [package.optional-dependencies] +async = [ + { name = "httpx" }, +] dev = [ { name = "django-stubs" }, { name = "griffe" }, @@ -2821,6 +2824,7 @@ requires-dist = [ { name = "gevent", marker = "implementation_name == 'cpython' and extra == 'test'", specifier = ">=25.4.1" }, { name = "google-genai", marker = "extra == 'test'" }, { name = "griffe", marker = "extra == 'dev'" }, + { name = "httpx", marker = "extra == 'async'", specifier = ">=0.27.0,<1.0" }, { name = "langchain", marker = "extra == 'langchain'", specifier = ">=1.3.9" }, { name = "langchain-anthropic", marker = "extra == 'test'", specifier = ">=1.0" }, { name = "langchain-community", marker = "extra == 'test'", specifier = ">=0.4" }, @@ -2862,7 +2866,7 @@ requires-dist = [ { name = "zstandard", marker = "extra == 'test'", specifier = ">=0.23.0" }, { name = "zstandard", marker = "extra == 'zstd'", specifier = ">=0.23.0" }, ] -provides-extras = ["langchain", "zstd", "otel", "dev", "test"] +provides-extras = ["langchain", "async", "zstd", "otel", "dev", "test"] [package.metadata.requires-dev] dev = [{ name = "claude-agent-sdk", specifier = ">=0.1.50" }] From dcb1331c74fde768a274cd235874794fc1bbf8b4 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 14:18:56 +0200 Subject: [PATCH 02/13] fix: preserve in-flight async capture during shutdown --- posthog/async_client.py | 7 +++++-- posthog/test/test_async_client.py | 34 ++++++++++++++++++++++++++++++- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/posthog/async_client.py b/posthog/async_client.py index 8e6267f42..f922229c5 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -90,6 +90,11 @@ def __init__( capture_compression: Optional[Union[CaptureCompression, str]] = None, capture_trace_context: bool = False, ) -> None: + if flush_at <= 0: + raise ValueError("flush_at must be greater than zero") + if flush_interval <= 0: + raise ValueError("flush_interval must be greater than zero") + self.api_key = (project_api_key or "").strip() self.raw_host = normalize_host(host) self.host = determine_server_host(host) @@ -338,8 +343,6 @@ async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: raise TypeError("before_send must return a dict or None") processed = clean(result) processed["uuid"] = original_uuid - if not self._accepting: - return None return processed except Exception as error: self.log.exception("Error in before_send callback: %s", error) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index af16caf75..236512558 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -255,10 +255,36 @@ async def before_send(event): await asyncio.wait_for(callback_finished.wait(), timeout=1) await asyncio.wait_for(client.shutdown(), timeout=1) - batch_post.assert_not_awaited() + batch_post.assert_awaited_once() assert client.capture("after shutdown", distinct_id="user-1") is None +@pytest.mark.asyncio +async def test_external_shutdown_delivers_event_already_in_before_send(): + callback_started = asyncio.Event() + allow_callback = asyncio.Event() + delivered = [] + + async def before_send(event): + callback_started.set() + await allow_callback.wait() + return event + + async def batch_post(*args, **kwargs): + delivered.extend(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key", before_send=before_send, flush_at=1) + client.capture("event", distinct_id="user-1") + await callback_started.wait() + shutdown = asyncio.create_task(client.shutdown()) + await asyncio.sleep(0) + allow_callback.set() + await shutdown + + assert [event["event"] for event in delivered] == ["event"] + + @pytest.mark.asyncio async def test_shutdown_waits_for_in_flight_immediate_capture(): upload_started = asyncio.Event() @@ -326,6 +352,12 @@ 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): + with pytest.raises(ValueError, match=option): + AsyncPosthog("test-key", **{option: value}) + + @pytest.mark.asyncio async def test_capture_exception_never_raises_in_debug_mode(): client = AsyncPosthog("test-key", send=False, debug=True) From 3116cdecd29f2e7c3ef02d5726963543e07fff88 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 14:25:17 +0200 Subject: [PATCH 03/13] fix: sanitize async capture failure logs --- posthog/_async_consumer.py | 23 +++++++++++++++++------ posthog/async_client.py | 12 +++++++++--- posthog/test/test_async_client.py | 23 +++++++++++++++++++++++ 3 files changed, 49 insertions(+), 9 deletions(-) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 570460858..ddae8c7c9 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -66,8 +66,10 @@ async def run(self) -> None: return except asyncio.CancelledError: raise - except Exception: - self.log.exception("async consumer stopped after an unexpected error") + except Exception as error: + self.log.error( + "async consumer stopped after an unexpected %s", type(error).__name__ + ) finally: self.log.debug("async consumer exited") @@ -75,14 +77,20 @@ async def upload(self, batch: list[dict[str, Any]]) -> None: try: await self.request(batch) except Exception as error: - self.log.error("error uploading async capture batch: %s", error) + self.log.error( + "async capture upload failed (%s, status=%s)", + type(error).__name__, + getattr(error, "status", None), + ) if self.on_error: try: result = self.on_error(error, batch) if inspect.isawaitable(result): await result except Exception as callback_error: - self.log.error("on_error handler failed: %s", callback_error) + self.log.error( + "on_error handler failed (%s)", type(callback_error).__name__ + ) finally: for _ in batch: self.queue.task_done() @@ -115,8 +123,11 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: try: item = await self.process_event(queued) - except Exception: - self.log.exception("unable to process queued event, dropping") + except Exception as error: + self.log.error( + "unable to process queued event, dropping (%s)", + type(error).__name__, + ) self.queue.task_done() continue diff --git a/posthog/async_client.py b/posthog/async_client.py index f922229c5..b5a4f819b 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -345,7 +345,7 @@ async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: processed["uuid"] = original_uuid return processed except Exception as error: - self.log.exception("Error in before_send callback: %s", error) + self.log.error("Error in before_send callback (%s)", type(error).__name__) return None def _build_capture_event( @@ -461,10 +461,16 @@ async def capture_immediate( if inspect.isawaitable(callback_result): await callback_result except Exception as callback_error: - self.log.error("on_error handler failed: %s", callback_error) + self.log.error( + "on_error handler failed (%s)", type(callback_error).__name__ + ) if self.debug: raise - self.log.exception("Error in immediate async capture: %s", error) + self.log.error( + "Immediate async capture failed (%s, status=%s)", + type(error).__name__, + getattr(error, "status", None), + ) return None finally: self._immediate_tasks.discard(current) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 236512558..0498a3cbc 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -7,6 +7,7 @@ import pytest from posthog import AsyncClient, AsyncPosthog, CaptureCompression, CaptureMode +from posthog.request import APIError @pytest.mark.asyncio @@ -382,6 +383,28 @@ async def test_queued_payload_is_not_written_to_debug_logs(caplog): assert "test-key" not 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): + 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 + ): + 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) + await client.shutdown() + + assert "server-secret" not in caplog.text + assert "APIError" in caplog.text + assert "status=400" in caplog.text + + @pytest.mark.asyncio async def test_flush_timeout_reports_unfinished_items(caplog): upload_started = asyncio.Event() From a96cc23618731ef2beebd90b638fb9e1f14a6408 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 14:39:49 +0200 Subject: [PATCH 04/13] fix: reset async client logging level --- posthog/async_client.py | 2 ++ posthog/test/test_async_client.py | 11 +++++++++++ 2 files changed, 13 insertions(+) diff --git a/posthog/async_client.py b/posthog/async_client.py index b5a4f819b..d91f51b23 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -152,6 +152,8 @@ def __init__( if debug: logging.basicConfig() self.log.setLevel(logging.DEBUG) + else: + self.log.setLevel(logging.WARNING) if not self.api_key: self.log.error( diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 0498a3cbc..012791a9c 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -10,6 +10,17 @@ from posthog.request import APIError +@pytest.mark.asyncio +async def test_debug_logging_does_not_leak_to_later_clients(): + debug_client = AsyncPosthog("test-key", send=False, debug=True) + assert debug_client.log.level == logging.DEBUG + await debug_client.shutdown() + + normal_client = AsyncPosthog("test-key", send=False) + assert normal_client.log.level == logging.WARNING + await normal_client.shutdown() + + @pytest.mark.asyncio async def test_async_posthog_is_the_customer_facing_async_client(): client = AsyncPosthog("test-key", send=False) From ed5ae7c609a7c8915af5878c28b6acdda07f02f4 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 14:44:08 +0200 Subject: [PATCH 05/13] fix: keep async capture batches within size limit --- posthog/_async_consumer.py | 15 ++++++++++++--- posthog/test/test_async_client.py | 23 +++++++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index ddae8c7c9..3eeb90fa3 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -54,6 +54,7 @@ def __init__( self.capture_mode = capture_mode self.capture_compression = capture_compression self.http_client = http_client + self._carryover: Optional[tuple[dict[str, Any], int]] = None async def run(self) -> None: self.log.debug("async consumer is running") @@ -101,6 +102,12 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: stop = False started = asyncio.get_running_loop().time() + if self._carryover is not None: + carried_item, carried_size = self._carryover + self._carryover = None + items.append(carried_item) + total_size = carried_size + while len(items) < self.flush_at: remaining = self.flush_interval - ( asyncio.get_running_loop().time() - started @@ -155,12 +162,14 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: self.queue.task_done() continue - items.append(item) - total_size += item_size - if total_size >= BATCH_SIZE_LIMIT: + if items and total_size + item_size > BATCH_SIZE_LIMIT: + self._carryover = (item, item_size) self.log.debug("hit async batch size limit (size: %d)", total_size) break + items.append(item) + total_size += item_size + return items, stop async def request(self, batch: list[dict[str, Any]]) -> None: diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 012791a9c..cde20ce40 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -224,6 +224,29 @@ async def test_capture_after_shutdown_is_dropped_without_restarting_workers(): build_client.assert_not_called() +@pytest.mark.asyncio +async def test_batch_size_overflow_event_is_sent_in_the_next_batch(): + batches = [] + + async def batch_post(*args, **kwargs): + batches.append(kwargs["batch"]) + + with ( + mock.patch("posthog._async_consumer.BATCH_SIZE_LIMIT", 800), + mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post), + ): + client = AsyncPosthog("test-key", flush_at=10, flush_interval=30) + client.capture("first", distinct_id="user-1", properties={"value": "a" * 400}) + client.capture("second", distinct_id="user-1", properties={"value": "b" * 400}) + await client.flush(timeout_seconds=1) + await client.shutdown() + + assert [[event["event"] for event in batch] for batch in batches] == [ + ["first"], + ["second"], + ] + + @pytest.mark.asyncio async def test_shutdown_waits_for_an_in_flight_batch_instead_of_cancelling_it(): upload_started = asyncio.Event() From 04965487e38e05706ddfc7d4fb64d237acb24f1e Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 14:59:48 +0200 Subject: [PATCH 06/13] fix: track immediate capture operation completion --- posthog/async_client.py | 27 ++++++++++++++++++--------- posthog/test/test_async_client.py | 25 +++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/posthog/async_client.py b/posthog/async_client.py index d91f51b23..25ed6b2dd 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -166,7 +166,8 @@ def __init__( self._flush_interval = flush_interval self._consumers: list[_AsyncConsumer] = [] self._worker_tasks: list[asyncio.Task[None]] = [] - self._immediate_tasks: set[asyncio.Task[Any]] = set() + self._immediate_callers: dict[asyncio.Task[Any], int] = {} + self._immediate_completions: set[asyncio.Future[None]] = set() self._http_client: Optional[Any] = None self._loop: Optional[asyncio.AbstractEventLoop] = None self._accepting = True @@ -434,8 +435,10 @@ async def capture_immediate( return None if not self._accepting: return None - self._bind_loop() - self._immediate_tasks.add(current) + loop = self._bind_loop() + completion: asyncio.Future[None] = loop.create_future() + self._immediate_completions.add(completion) + 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( @@ -475,7 +478,14 @@ async def capture_immediate( ) return None finally: - self._immediate_tasks.discard(current) + remaining_calls = self._immediate_callers[current] - 1 + if remaining_calls: + self._immediate_callers[current] = remaining_calls + else: + del self._immediate_callers[current] + if not completion.done(): + completion.set_result(None) + self._immediate_completions.discard(completion) def _build_person_properties_event( self, event: str, property_key: str, kwargs: OptionalSetArgs @@ -719,7 +729,7 @@ async def _close_transport(self) -> None: async def shutdown(self) -> None: current = asyncio.current_task() - if current in self._worker_tasks or current in self._immediate_tasks: + if current in self._worker_tasks or current in self._immediate_callers: self._accepting = False self._defer_lifecycle_call(self.shutdown()) return @@ -731,11 +741,10 @@ async def shutdown(self) -> None: self._accepting = False errors: list[Exception] = [] - current = asyncio.current_task() pending_immediate = [ - task - for task in self._immediate_tasks - if task is not current and not task.done() + completion + for completion in self._immediate_completions + if not completion.done() ] if pending_immediate: await asyncio.gather(*pending_immediate, return_exceptions=True) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index cde20ce40..ac9d633a8 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -320,6 +320,31 @@ async def batch_post(*args, **kwargs): assert [event["event"] for event in delivered] == ["event"] +@pytest.mark.asyncio +async def test_shutdown_waits_for_immediate_operation_not_its_long_lived_caller(): + upload_started = asyncio.Event() + allow_upload = asyncio.Event() + shutdown_task = None + + async def batch_post(*args, **kwargs): + upload_started.set() + await allow_upload.wait() + + async def capture_then_await_shutdown(client): + result = await client.capture_immediate("event", distinct_id="user-1") + assert result is not None + assert shutdown_task is not None + await shutdown_task + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key") + caller = asyncio.create_task(capture_then_await_shutdown(client)) + await upload_started.wait() + shutdown_task = asyncio.create_task(client.shutdown()) + allow_upload.set() + await asyncio.wait_for(caller, timeout=1) + + @pytest.mark.asyncio async def test_shutdown_waits_for_in_flight_immediate_capture(): upload_started = asyncio.Event() From 25fafd189c5af406100b32d9355482f986b25b2e Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 15:07:22 +0200 Subject: [PATCH 07/13] fix: coordinate async worker flushes --- posthog/_async_consumer.py | 33 +++++++++++++++----- posthog/async_client.py | 30 +++++++++++------- posthog/test/test_async_client.py | 52 +++++++++++++++++++++++++++++++ 3 files changed, 96 insertions(+), 19 deletions(-) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 3eeb90fa3..0fc26065d 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -13,7 +13,6 @@ from .consumer import BATCH_SIZE_LIMIT, MAX_MSG_SIZE from .request import APIError, DatetimeSerializer, EVENTS_ENDPOINT -_FLUSH = object() _STOP = object() @@ -55,6 +54,7 @@ def __init__( self.capture_compression = capture_compression self.http_client = http_client self._carryover: Optional[tuple[dict[str, Any], int]] = None + self._flush_event = asyncio.Event() async def run(self) -> None: self.log.debug("async consumer is running") @@ -74,6 +74,29 @@ async def run(self) -> None: finally: self.log.debug("async consumer exited") + def request_flush(self) -> None: + self._flush_event.set() + + async def _get_or_flush(self, timeout: float) -> tuple[Any, bool]: + get_task = asyncio.create_task(self.queue.get()) + flush_task = asyncio.create_task(self._flush_event.wait()) + done, pending = await asyncio.wait( + {get_task, flush_task}, + timeout=timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + if get_task in done: + return get_task.result(), False + if flush_task in done: + self._flush_event.clear() + return None, True + return None, False + async def upload(self, batch: list[dict[str, Any]]) -> None: try: await self.request(batch) @@ -115,14 +138,10 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: if remaining <= 0: break - try: - queued = await asyncio.wait_for(self.queue.get(), timeout=remaining) - except asyncio.TimeoutError: + queued, flush_requested = await self._get_or_flush(remaining) + if flush_requested or queued is None: break - if queued is _FLUSH: - self.queue.task_done() - break if queued is _STOP: self.queue.task_done() stop = True diff --git a/posthog/async_client.py b/posthog/async_client.py index 25ed6b2dd..d243de49d 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -13,8 +13,8 @@ from typing_extensions import Unpack -from ._async_consumer import _FLUSH, _STOP, _AsyncConsumer -from ._async_request import _build_client +from ._async_consumer import _STOP, _AsyncConsumer +from ._async_request import _build_client, _require_httpx from .args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs from .capture_compression import ( CaptureCompression, @@ -259,6 +259,10 @@ def _new_consumer(self) -> _AsyncConsumer: 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 @@ -407,12 +411,15 @@ def capture( if not self.send: return sent_uuid + self._validate_transport_available() try: - self._ensure_workers_started() + asyncio.get_running_loop() except RuntimeError: - # Construction and capture before the loop starts are supported. - # flush()/shutdown() will bind the client and start the workers. + # Capture before the loop starts is supported. flush()/shutdown() + # will bind the client and start the workers. pass + else: + self._ensure_workers_started() self._queue.put_nowait(prepared) self.log.debug("queued async event %s", event) @@ -615,10 +622,13 @@ def _enqueue_built_event( return None if not self.send: return sent_uuid + self._validate_transport_available() try: - self._ensure_workers_started() + asyncio.get_running_loop() except RuntimeError: pass + else: + self._ensure_workers_started() self._queue.put_nowait(prepared) return sent_uuid @@ -702,12 +712,8 @@ async def flush(self, timeout_seconds: Optional[float] = 10) -> None: else asyncio.get_running_loop().time() + timeout_seconds ) try: - for _ in self._worker_tasks: - if deadline is None: - await self._queue.put(_FLUSH) - else: - remaining = max(0.0, deadline - asyncio.get_running_loop().time()) - await asyncio.wait_for(self._queue.put(_FLUSH), remaining) + for consumer in self._consumers: + consumer.request_flush() if deadline is None: await self._queue.join() diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index ac9d633a8..a903fbc76 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -159,6 +159,23 @@ 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") + 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: @@ -247,6 +264,41 @@ async def batch_post(*args, **kwargs): ] +@pytest.mark.asyncio +async def test_flush_wakes_each_worker_with_a_partial_batch(): + slow_callback_started = asyncio.Event() + allow_slow_callback = asyncio.Event() + delivered = [] + + async def before_send(event): + if event["event"] == "slow": + slow_callback_started.set() + await allow_slow_callback.wait() + return event + + async def batch_post(*args, **kwargs): + delivered.extend(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog( + "test-key", + thread=2, + flush_at=100, + flush_interval=30, + before_send=before_send, + ) + client.capture("slow", distinct_id="user-1") + await slow_callback_started.wait() + client.capture("fast", distinct_id="user-1") + flush = asyncio.create_task(client.flush(timeout_seconds=1)) + await asyncio.sleep(0) + allow_slow_callback.set() + await flush + await client.shutdown() + + assert sorted(event["event"] for event in delivered) == ["fast", "slow"] + + @pytest.mark.asyncio async def test_shutdown_waits_for_an_in_flight_batch_instead_of_cancelling_it(): upload_started = asyncio.Event() From 6163ec85dde2bf6bcc6996167d7a006faf906586 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 15:58:36 +0200 Subject: [PATCH 08/13] docs: clarify async alternative to sync mode --- posthog/__init__.py | 3 ++- posthog/client.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index 01b10f53f..0431bc883 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -318,7 +318,8 @@ def get_tags() -> Dict[str, Any]: debug: Enable verbose SDK logging and re-raise errors from public APIs. send: If False, queueing succeeds but events are not sent to PostHog. sync_mode: If True, send events synchronously instead of using background - worker threads. + worker threads. This blocks the calling thread; in asyncio applications + such as FastAPI, use ``AsyncPosthog`` instead. disabled: If True, disable captures and API requests. Useful in tests. secret_key: A Personal API Key or Project Secret API Key used for local feature flag evaluation and remote config payloads. diff --git a/posthog/client.py b/posthog/client.py index dcdae7cd1..437eca2da 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -732,7 +732,9 @@ def __init__( 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. + background worker threads. This blocks the calling thread; in + asyncio applications such as FastAPI, use ``AsyncPosthog`` + instead. timeout: HTTP request timeout in seconds for event uploads. thread: Number of background consumer threads. poll_interval: Seconds between local feature flag definition refreshes. From 0fbf6bea86bb4ce2f6661b14fca3fe1e9a85ef5d Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 26 Aug 2026 16:35:50 +0200 Subject: [PATCH 09/13] address async capture review feedback --- posthog/_async_consumer.py | 10 +-- posthog/_async_request.py | 2 +- posthog/async_client.py | 74 +++++++++++++++++--- posthog/test/test_async_client.py | 78 +++++++++++++++++++++ posthog/test/test_async_consumer.py | 101 ++++++++++++++++++++++++++++ posthog/test/test_async_request.py | 4 +- 6 files changed, 252 insertions(+), 17 deletions(-) create mode 100644 posthog/test/test_async_consumer.py diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 0fc26065d..7ccfdd32a 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -16,6 +16,11 @@ _STOP = object() +async def _serialized_event_size(event: dict[str, Any]) -> int: + serialized = await asyncio.to_thread(json.dumps, event, cls=DatetimeSerializer) + return len(serialized.encode()) + + class _AsyncConsumer: """Consume an asyncio queue and upload capture batches.""" @@ -162,10 +167,7 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: continue try: - serialized = await asyncio.to_thread( - json.dumps, item, cls=DatetimeSerializer - ) - item_size = len(serialized.encode()) + item_size = await _serialized_event_size(item) except Exception: self.log.error("unable to serialize queued event for sizing, dropping") self.queue.task_done() diff --git a/posthog/_async_request.py b/posthog/_async_request.py index 79ce3a386..b71cf6c3d 100644 --- a/posthog/_async_request.py +++ b/posthog/_async_request.py @@ -31,7 +31,7 @@ def _require_httpx(): def _build_client(host: Optional[str] = None): httpx_module = _require_httpx() base_url = remove_trailing_slash(normalize_host(host)) - return httpx_module.AsyncClient(base_url=base_url, follow_redirects=True) + return httpx_module.AsyncClient(base_url=base_url, follow_redirects=False) def _serialize_v0_body( diff --git a/posthog/async_client.py b/posthog/async_client.py index d243de49d..d91bd6baa 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -13,7 +13,7 @@ from typing_extensions import Unpack -from ._async_consumer import _STOP, _AsyncConsumer +from ._async_consumer import _STOP, _AsyncConsumer, _serialized_event_size from ._async_request import _build_client, _require_httpx from .args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs from .capture_compression import ( @@ -27,7 +27,14 @@ get_identity_state as _get_identity_state, stringify_id as _stringify_id, ) -from .contexts import get_context_session_id as _get_context_session_id +from .contexts import ( + get_capture_exception_code_variables_context, + get_code_variables_detect_secrets_context, + get_code_variables_ignore_patterns_context, + get_code_variables_mask_patterns_context, + get_code_variables_mask_url_credentials_context, + get_context_session_id as _get_context_session_id, +) from .exception_utils import ( DEFAULT_CODE_VARIABLES_DETECT_SECRETS, DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS, @@ -41,6 +48,7 @@ mark_exception_as_captured, try_attach_code_variables_to_frames, ) +from .consumer import MAX_MSG_SIZE from .request import determine_server_host, normalize_host from .utils import _normalize_timestamp, clean, system_context from .version import VERSION @@ -340,9 +348,12 @@ async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: original_uuid = msg["uuid"] try: - result = self.before_send(msg) - if inspect.isawaitable(result): - result = await result + if inspect.iscoroutinefunction(self.before_send): + result = await self.before_send(msg) + else: + result = await asyncio.to_thread(self.before_send, msg) + if inspect.isawaitable(result): + result = await result if result is None: self.log.debug("event dropped by before_send callback") return None @@ -462,6 +473,21 @@ async def capture_immediate( error_batch = [processed] if not self.send: return sent_uuid + try: + event_size = await _serialized_event_size(processed) + except Exception: + self.log.error( + "unable to serialize immediate event for sizing, dropping" + ) + return None + if event_size > MAX_MSG_SIZE: + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", + processed.get("event"), + event_size, + MAX_MSG_SIZE // 1024, + ) + return None consumer = self._new_consumer() await consumer.request(error_batch) @@ -662,14 +688,42 @@ def capture_exception( "$exception_list": exceptions, **(kwargs.get("properties") or {}), } - if self.capture_exception_code_variables: + context_enabled = get_capture_exception_code_variables_context() + context_mask = get_code_variables_mask_patterns_context() + context_ignore = get_code_variables_ignore_patterns_context() + context_mask_url_credentials = ( + get_code_variables_mask_url_credentials_context() + ) + context_detect_secrets = get_code_variables_detect_secrets_context() + enabled = ( + context_enabled + if context_enabled is not None + else self.capture_exception_code_variables + ) + if enabled: try_attach_code_variables_to_frames( exceptions, exc_info, - mask_patterns=self.code_variables_mask_patterns, - ignore_patterns=self.code_variables_ignore_patterns, - mask_url_credentials=self.code_variables_mask_url_credentials, - detect_secrets=self.code_variables_detect_secrets, + mask_patterns=( + context_mask + if context_mask is not None + else self.code_variables_mask_patterns + ), + ignore_patterns=( + context_ignore + if context_ignore is not None + else self.code_variables_ignore_patterns + ), + mask_url_credentials=( + context_mask_url_credentials + if context_mask_url_credentials is not None + else self.code_variables_mask_url_credentials + ), + detect_secrets=( + context_detect_secrets + if context_detect_secrets is not None + else self.code_variables_detect_secrets + ), ) if self.log_captured_exceptions: self.log.exception(exception, extra=kwargs) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index a903fbc76..0198919df 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -2,11 +2,21 @@ import asyncio import logging +import threading from unittest import mock import pytest from posthog import AsyncClient, AsyncPosthog, CaptureCompression, CaptureMode +from posthog.consumer import MAX_MSG_SIZE +from posthog.contexts import ( + new_context, + set_capture_exception_code_variables_context, + set_code_variables_detect_secrets_context, + set_code_variables_ignore_patterns_context, + set_code_variables_mask_patterns_context, + set_code_variables_mask_url_credentials_context, +) from posthog.request import APIError @@ -79,6 +89,24 @@ async def batch_post(*args, **kwargs): assert batches[0][0]["uuid"] == event_uuid +@pytest.mark.asyncio +async def test_capture_offloads_synchronous_before_send(): + callback_thread = None + + def before_send(event): + nonlocal callback_thread + callback_thread = threading.get_ident() + return event + + with mock.patch("posthog._async_consumer.async_batch_post", new=mock.AsyncMock()): + client = AsyncPosthog("test-key", before_send=before_send) + result = await client.capture_immediate("event", distinct_id="user-1") + await client.shutdown() + + assert result is not None + assert callback_thread != threading.get_ident() + + @pytest.mark.asyncio async def test_capture_drops_event_when_before_send_raises(): async def before_send(_event): @@ -136,6 +164,23 @@ async def batch_post(*args, **kwargs): assert batches[0][0]["properties"]["processed"] is True +@pytest.mark.asyncio +async def test_capture_immediate_drops_oversized_event_after_before_send(): + 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: + 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() + + @pytest.mark.asyncio async def test_capture_immediate_uses_capture_v1_without_building_httpx_client(): with ( @@ -478,6 +523,39 @@ async def test_capture_exception_never_raises_in_debug_mode(): await client.shutdown() +@pytest.mark.asyncio +async def test_capture_exception_uses_context_code_variable_settings(): + client = AsyncPosthog( + "test-key", + send=False, + capture_exception_code_variables=False, + ) + with ( + new_context(), + mock.patch( + "posthog.async_client.try_attach_code_variables_to_frames" + ) as attach, + ): + set_capture_exception_code_variables_context(True) + set_code_variables_mask_patterns_context(["mask-me"]) + set_code_variables_ignore_patterns_context(["ignore-me"]) + set_code_variables_mask_url_credentials_context(False) + set_code_variables_detect_secrets_context(False) + try: + raise ValueError("boom") + except ValueError as error: + assert client.capture_exception(error, distinct_id="user-1") is not None + + attach.assert_called_once() + assert attach.call_args.kwargs == { + "mask_patterns": ["mask-me"], + "ignore_patterns": ["ignore-me"], + "mask_url_credentials": False, + "detect_secrets": False, + } + await client.shutdown() + + @pytest.mark.asyncio async def test_queued_payload_is_not_written_to_debug_logs(caplog): caplog.set_level(logging.DEBUG, logger="posthog") diff --git a/posthog/test/test_async_consumer.py b/posthog/test/test_async_consumer.py new file mode 100644 index 000000000..4fb5c613b --- /dev/null +++ b/posthog/test/test_async_consumer.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import asyncio +from unittest import mock + +import pytest + +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: + return _AsyncConsumer( + asyncio.Queue(), + "test-key", + host="https://example.com", + on_error=None, + 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("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] diff --git a/posthog/test/test_async_request.py b/posthog/test/test_async_request.py index 4a07f0d14..ed34e1157 100644 --- a/posthog/test/test_async_request.py +++ b/posthog/test/test_async_request.py @@ -58,11 +58,11 @@ def blocked_import(name, *args, **kwargs): subprocess.run([sys.executable, "-c", script], check=True) -def test_build_client_scopes_requests_to_host_and_follows_redirects(): +def test_build_client_scopes_requests_to_host_without_following_redirects(): with mock.patch("posthog._async_request.httpx.AsyncClient") as async_client: _build_client("https://example.com/") async_client.assert_called_once_with( - base_url="https://example.com", follow_redirects=True + base_url="https://example.com", follow_redirects=False ) From 61a3d12b0cdb2c001e50315a6679f2f15a48400b Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 27 Aug 2026 06:55:17 +0200 Subject: [PATCH 10/13] fix: preserve context for buffered async events --- posthog/_async_consumer.py | 42 +++++++++++++++++++++--- posthog/async_client.py | 34 +++++++++++--------- posthog/test/test_async_client.py | 53 +++++++++++++++++++++++++++++++ 3 files changed, 110 insertions(+), 19 deletions(-) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 7ccfdd32a..9c3a892f2 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -1,10 +1,12 @@ from __future__ import annotations import asyncio +import contextvars import inspect import json import logging from collections.abc import Awaitable, Callable +from dataclasses import dataclass from typing import Any, Optional from ._async_request import async_batch_post, async_send_v1_batch @@ -14,6 +16,28 @@ from .request import APIError, DatetimeSerializer, EVENTS_ENDPOINT _STOP = object() +_PROCESSING_EVENT = contextvars.ContextVar( + "posthog_async_processing_event", default=False +) + + +def _is_processing_event() -> bool: + return _PROCESSING_EVENT.get() + + +@dataclass(frozen=True) +class _QueuedEvent: + event: dict[str, Any] + context: contextvars.Context + + +async def _invoke_callback(callback, *args): + if inspect.iscoroutinefunction(callback): + return await callback(*args) + result = await asyncio.to_thread(callback, *args) + if inspect.isawaitable(result): + return await result + return result async def _serialized_event_size(event: dict[str, Any]) -> int: @@ -102,6 +126,15 @@ async def _get_or_flush(self, timeout: float) -> tuple[Any, bool]: return None, True return None, False + async def _process_queued_event( + self, event: dict[str, Any] + ) -> Optional[dict[str, Any]]: + token = _PROCESSING_EVENT.set(True) + try: + return await self.process_event(event) + finally: + _PROCESSING_EVENT.reset(token) + async def upload(self, batch: list[dict[str, Any]]) -> None: try: await self.request(batch) @@ -113,9 +146,7 @@ async def upload(self, batch: list[dict[str, Any]]) -> None: ) if self.on_error: try: - result = self.on_error(error, batch) - if inspect.isawaitable(result): - await result + 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__ @@ -153,7 +184,10 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: break try: - item = await self.process_event(queued) + process_task = queued.context.run( + asyncio.create_task, self._process_queued_event(queued.event) + ) + item = await process_task except Exception as error: self.log.error( "unable to process queued event, dropping (%s)", diff --git a/posthog/async_client.py b/posthog/async_client.py index d91bd6baa..b3d881f48 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -import inspect +import contextvars import logging import os import sys @@ -13,7 +13,14 @@ from typing_extensions import Unpack -from ._async_consumer import _STOP, _AsyncConsumer, _serialized_event_size +from ._async_consumer import ( + _STOP, + _AsyncConsumer, + _invoke_callback, + _is_processing_event, + _QueuedEvent, + _serialized_event_size, +) from ._async_request import _build_client, _require_httpx from .args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs from .capture_compression import ( @@ -348,12 +355,7 @@ async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: original_uuid = msg["uuid"] try: - if inspect.iscoroutinefunction(self.before_send): - result = await self.before_send(msg) - else: - result = await asyncio.to_thread(self.before_send, msg) - if inspect.isawaitable(result): - result = await result + result = await _invoke_callback(self.before_send, msg) if result is None: self.log.debug("event dropped by before_send callback") return None @@ -432,7 +434,7 @@ def capture( else: self._ensure_workers_started() - self._queue.put_nowait(prepared) + self._queue.put_nowait(_QueuedEvent(prepared, contextvars.copy_context())) self.log.debug("queued async event %s", event) return sent_uuid except asyncio.QueueFull: @@ -495,9 +497,7 @@ async def capture_immediate( except Exception as error: if self.on_error: try: - callback_result = self.on_error(error, error_batch) - if inspect.isawaitable(callback_result): - await callback_result + 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__ @@ -655,7 +655,7 @@ def _enqueue_built_event( pass else: self._ensure_workers_started() - self._queue.put_nowait(prepared) + self._queue.put_nowait(_QueuedEvent(prepared, contextvars.copy_context())) return sent_uuid def capture_exception( @@ -754,7 +754,7 @@ def _defer_lifecycle_call(self, awaitable) -> None: task.add_done_callback(self._deferred_lifecycle_tasks.discard) async def flush(self, timeout_seconds: Optional[float] = 10) -> None: - if asyncio.current_task() in self._worker_tasks: + if asyncio.current_task() in self._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: @@ -789,7 +789,11 @@ async def _close_transport(self) -> None: async def shutdown(self) -> None: current = asyncio.current_task() - if current in self._worker_tasks or current in self._immediate_callers: + if ( + current in self._worker_tasks + or current in self._immediate_callers + or _is_processing_event() + ): self._accepting = False self._defer_lifecycle_call(self.shutdown()) return diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 0198919df..f13c8d6b7 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextvars import logging import threading from unittest import mock @@ -89,6 +90,38 @@ async def batch_post(*args, **kwargs): assert batches[0][0]["uuid"] == event_uuid +@pytest.mark.asyncio +async def test_capture_preserves_context_for_each_buffered_event(): + request_id = contextvars.ContextVar("request_id") + batches = [] + + def before_send(event): + event["properties"]["request_id"] = request_id.get() + return event + + async def batch_post(*args, **kwargs): + batches.append(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + 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") + request_id.reset(token) + + token = request_id.set("request-B") + client.capture("event-B", distinct_id="user-1") + request_id.reset(token) + + await client.flush(timeout_seconds=1) + await client.shutdown() + + events = [event for batch in batches for event in batch] + assert [event["properties"]["request_id"] for event in events] == [ + "request-A", + "request-B", + ] + + @pytest.mark.asyncio async def test_capture_offloads_synchronous_before_send(): callback_thread = None @@ -572,6 +605,26 @@ async def test_queued_payload_is_not_written_to_debug_logs(caplog): assert "test-key" not in caplog.text +@pytest.mark.asyncio +async def test_capture_immediate_offloads_synchronous_on_error(): + callback_thread = None + + def on_error(error, batch): + nonlocal callback_thread + callback_thread = threading.get_ident() + + with mock.patch( + "posthog._async_consumer.async_batch_post", + side_effect=APIError(400, "failed"), + ): + client = AsyncPosthog("test-key", on_error=on_error, max_retries=0) + result = await client.capture_immediate("event", distinct_id="user-1") + await client.shutdown() + + assert result is None + assert callback_thread != threading.get_ident() + + @pytest.mark.asyncio @pytest.mark.parametrize("immediate", [False, True]) async def test_failed_capture_does_not_log_server_response_detail(caplog, immediate): From d59c99a390e3c1898fa7f22b6cce59fbbf8b1d74 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 27 Aug 2026 07:00:02 +0200 Subject: [PATCH 11/13] fix: clear event context for deferred lifecycle --- posthog/_async_consumer.py | 8 ++++++++ posthog/async_client.py | 3 ++- posthog/test/test_async_client.py | 7 ++++++- 3 files changed, 16 insertions(+), 2 deletions(-) diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 9c3a892f2..387caab44 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -25,6 +25,14 @@ def _is_processing_event() -> bool: return _PROCESSING_EVENT.get() +async def _run_outside_processing_event(awaitable): + token = _PROCESSING_EVENT.set(False) + try: + return await awaitable + finally: + _PROCESSING_EVENT.reset(token) + + @dataclass(frozen=True) class _QueuedEvent: event: dict[str, Any] diff --git a/posthog/async_client.py b/posthog/async_client.py index b3d881f48..c151ae5a8 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -19,6 +19,7 @@ _invoke_callback, _is_processing_event, _QueuedEvent, + _run_outside_processing_event, _serialized_event_size, ) from ._async_request import _build_client, _require_httpx @@ -749,7 +750,7 @@ def _pending_queue_items(self) -> int: return int(getattr(self._queue, "_unfinished_tasks", self._queue.qsize())) def _defer_lifecycle_call(self, awaitable) -> None: - task = asyncio.create_task(awaitable) + task = asyncio.create_task(_run_outside_processing_event(awaitable)) self._deferred_lifecycle_tasks.add(task) task.add_done_callback(self._deferred_lifecycle_tasks.discard) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index f13c8d6b7..04af127c8 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -418,7 +418,12 @@ async def before_send(event): 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) - await asyncio.wait_for(client.shutdown(), timeout=1) + + async def wait_until_closed(): + while not client._closed: + await asyncio.sleep(0) + + await asyncio.wait_for(wait_until_closed(), timeout=1) batch_post.assert_awaited_once() assert client.capture("after shutdown", distinct_id="user-1") is None From b7f70b5f803760a032d2298b9b10fa8eeb03aeed Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 27 Aug 2026 21:52:26 +0200 Subject: [PATCH 12/13] fix: harden async capture delivery lifecycle --- posthog/_async_request.py | 46 ++++++++++- posthog/async_client.py | 127 +++++++++++++++++++++++------ posthog/test/test_async_client.py | 69 ++++++++++++++++ posthog/test/test_async_request.py | 56 ++++++++++++- 4 files changed, 266 insertions(+), 32 deletions(-) diff --git a/posthog/_async_request.py b/posthog/_async_request.py index b71cf6c3d..e5a4c8cc6 100644 --- a/posthog/_async_request.py +++ b/posthog/_async_request.py @@ -8,6 +8,7 @@ from gzip import GzipFile from io import BytesIO from typing import Any, Optional +from urllib.parse import urljoin, urlsplit from .capture_compression import CaptureCompression from .capture_v1 import _send_v1_batch @@ -61,6 +62,24 @@ def _serialize_v0_body( 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_path( + base_url: str, current_path: str, location: str +) -> Optional[str]: + target = urlsplit(urljoin(urljoin(f"{base_url}/", current_path), location)) + if _origin(target.geturl()) != _origin(base_url): + return None + path = target.path or "/" + return f"{path}?{target.query}" if target.query else path + + def _parse_retry_after(response: Any) -> Optional[float]: value = response.headers.get("Retry-After") if value is None: @@ -113,10 +132,29 @@ async def async_batch_post( http_client = client or _build_client(host) try: logging.getLogger("posthog").debug("making async capture request") - response = await http_client.post( - path, content=data, headers=headers, timeout=timeout - ) - _process_response(response) + base_url = remove_trailing_slash(normalize_host(host)) + request_path = path + for redirect_count in range(6): + response = await http_client.post( + request_path, 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_path = ( + _same_origin_redirect_path(base_url, request_path, location) + if location + else None + ) + if redirect_path is None: + raise APIError(400, "Cross-origin or invalid redirect blocked") + if redirect_count >= 5: + raise APIError(400, "Too many capture redirects") + request_path = redirect_path finally: if owns_client: await http_client.aclose() diff --git a/posthog/async_client.py b/posthog/async_client.py index c151ae5a8..65ca61c9e 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -1,10 +1,12 @@ from __future__ import annotations import asyncio +import concurrent.futures import contextvars import logging import os import sys +import threading import warnings import weakref from datetime import datetime, timezone @@ -186,6 +188,7 @@ def __init__( self._immediate_completions: set[asyncio.Future[None]] = set() self._http_client: Optional[Any] = None self._loop: Optional[asyncio.AbstractEventLoop] = None + self._loop_thread_id: Optional[int] = None self._accepting = True self._closed = False self._shutdown_lock = asyncio.Lock() @@ -244,6 +247,7 @@ def _bind_loop(self) -> asyncio.AbstractEventLoop: loop = asyncio.get_running_loop() if self._loop is None: self._loop = loop + self._loop_thread_id = threading.get_ident() elif self._loop is not loop: raise RuntimeError("AsyncClient cannot be shared across event loops") return loop @@ -288,6 +292,42 @@ 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]) -> None: + queued_event = _QueuedEvent(prepared, contextvars.copy_context()) + try: + running_loop = asyncio.get_running_loop() + except RuntimeError: + running_loop = None + + if self._loop is None: + if running_loop is not None: + self._ensure_workers_started() + self._queue.put_nowait(queued_event) + return + + if running_loop is self._loop: + self._ensure_workers_started() + self._queue.put_nowait(queued_event) + return + if running_loop is not None: + raise RuntimeError("AsyncClient cannot be shared across event loops") + if self._loop_thread_id == threading.get_ident() or not self._loop.is_running(): + raise RuntimeError("AsyncClient event loop is not running") + + admitted: concurrent.futures.Future[None] = concurrent.futures.Future() + + def enqueue_on_bound_loop() -> None: + try: + self._ensure_workers_started() + self._queue.put_nowait(queued_event) + except BaseException as error: + admitted.set_exception(error) + else: + admitted.set_result(None) + + self._loop.call_soon_threadsafe(enqueue_on_bound_loop) + admitted.result() + def _normalize_uuid(self, msg: dict[str, Any]) -> str: raw_uuid = msg.pop("uuid", None) if raw_uuid is not None: @@ -426,16 +466,7 @@ def capture( return sent_uuid self._validate_transport_available() - try: - asyncio.get_running_loop() - except RuntimeError: - # Capture before the loop starts is supported. flush()/shutdown() - # will bind the client and start the workers. - pass - else: - self._ensure_workers_started() - - self._queue.put_nowait(_QueuedEvent(prepared, contextvars.copy_context())) + self._enqueue_prepared_event(prepared) self.log.debug("queued async event %s", event) return sent_uuid except asyncio.QueueFull: @@ -650,13 +681,7 @@ def _enqueue_built_event( if not self.send: return sent_uuid self._validate_transport_available() - try: - asyncio.get_running_loop() - except RuntimeError: - pass - else: - self._ensure_workers_started() - self._queue.put_nowait(_QueuedEvent(prepared, contextvars.copy_context())) + self._enqueue_prepared_event(prepared) return sent_uuid def capture_exception( @@ -754,6 +779,63 @@ def _defer_lifecycle_call(self, awaitable) -> None: self._deferred_lifecycle_tasks.add(task) task.add_done_callback(self._deferred_lifecycle_tasks.discard) + def _discard_undrainable_queue(self) -> None: + discarded = 0 + while True: + try: + self._queue.get_nowait() + except asyncio.QueueEmpty: + break + self._queue.task_done() + discarded += 1 + + orphaned = self._pending_queue_items() + for _ in range(orphaned): + self._queue.task_done() + discarded += orphaned + if discarded: + self.log.warning( + "discarded %d async capture items because all workers exited", + discarded, + ) + + async def _wait_for_queue_drain(self, deadline: Optional[float]) -> None: + live_workers = [task for task in self._worker_tasks if not task.done()] + if not live_workers: + self._discard_undrainable_queue() + return + + queue_join = asyncio.create_task(self._queue.join()) + + async def wait_for_workers() -> None: + await asyncio.wait(live_workers, return_when=asyncio.ALL_COMPLETED) + + workers_finished = asyncio.create_task(wait_for_workers()) + try: + timeout = ( + None + if deadline is None + else max(0.0, deadline - asyncio.get_running_loop().time()) + ) + done, _ = await asyncio.wait( + {queue_join, workers_finished}, + timeout=timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + raise asyncio.TimeoutError + if queue_join in done: + await queue_join + return + + self._discard_undrainable_queue() + await queue_join + finally: + for task in (queue_join, workers_finished): + if not task.done(): + task.cancel() + 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(): self._defer_lifecycle_call(self.flush(timeout_seconds)) @@ -770,11 +852,7 @@ async def flush(self, timeout_seconds: Optional[float] = 10) -> None: for consumer in self._consumers: consumer.request_flush() - if deadline is None: - await self._queue.join() - else: - remaining = max(0.0, deadline - asyncio.get_running_loop().time()) - await asyncio.wait_for(self._queue.join(), remaining) + await self._wait_for_queue_drain(deadline) except asyncio.TimeoutError: self.log.warning( "flush timed out after %s seconds with %s items pending", @@ -821,10 +899,11 @@ async def shutdown(self) -> None: errors.append(error) try: - for _ in self._worker_tasks: + 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=False) + await asyncio.gather(*self._worker_tasks, return_exceptions=True) except Exception as error: self.log.exception("Failed to stop async capture workers") errors.append(error) diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 04af127c8..05b5c9d11 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -3,6 +3,8 @@ import asyncio import contextvars import logging +import subprocess +import sys import threading from unittest import mock @@ -67,6 +69,57 @@ async def batch_post(*args, **kwargs): assert event["uuid"] == event_uuid +def test_capture_from_worker_thread_wakes_loop_bound_queue(): + script = r""" +import asyncio +import os +import threading +from unittest import mock +from posthog import AsyncPosthog + +async def main(): + delivered = [] + + async def batch_post(*args, **kwargs): + delivered.extend(kwargs["batch"]) + + with mock.patch( + "posthog._async_consumer.async_batch_post", side_effect=batch_post + ): + client = AsyncPosthog("test-key", flush_interval=30) + client._ensure_workers_started() + while not client._queue._getters: + await asyncio.sleep(0) + + asyncio.get_running_loop().set_debug(True) + capture_result = [] + capture_thread = threading.Thread( + target=lambda: capture_result.append( + client.capture("threaded event", distinct_id="user-1") + ), + daemon=True, + ) + capture_thread.start() + for _ in range(100): + if not capture_thread.is_alive(): + break + await asyncio.sleep(0.01) + + if capture_thread.is_alive(): + os._exit(2) + if len(capture_result) != 1 or capture_result[0] is None: + os._exit(3) + + await asyncio.wait_for(client.flush(), timeout=1) + await client.shutdown() + if [event["event"] for event in delivered] != ["threaded event"]: + os._exit(4) + +asyncio.run(main()) +""" + subprocess.run([sys.executable, "-c", script], check=True, timeout=5) + + @pytest.mark.asyncio async def test_capture_runs_async_before_send_in_consumer(): batches = [] @@ -455,6 +508,22 @@ async def batch_post(*args, **kwargs): assert [event["event"] for event in delivered] == ["event"] +@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: + task.cancel() + await asyncio.gather(*client._worker_tasks, return_exceptions=True) + + assert client.capture("undrainable event", distinct_id="user-1") is not None + try: + await asyncio.wait_for(client.shutdown(), timeout=0.1) + finally: + if client._http_client is not None: + await client._http_client.aclose() + + @pytest.mark.asyncio async def test_shutdown_waits_for_immediate_operation_not_its_long_lived_caller(): upload_started = asyncio.Event() diff --git a/posthog/test/test_async_request.py b/posthog/test/test_async_request.py index ed34e1157..8b4da4b7e 100644 --- a/posthog/test/test_async_request.py +++ b/posthog/test/test_async_request.py @@ -17,10 +17,10 @@ class FakeResponse: - def __init__(self, status_code=200, payload=None): + def __init__(self, status_code=200, payload=None, headers=None): self.status_code = status_code self._payload = payload if payload is not None else {"ok": True} - self.headers = {} + self.headers = headers or {} self.text = str(self._payload) def json(self): @@ -29,13 +29,17 @@ def json(self): class FakeAsyncClient: def __init__(self, response=None): - self.response = response or FakeResponse() + self.responses = ( + list(response) + if isinstance(response, list) + else [response or FakeResponse()] + ) self.calls = [] self.closed = False async def post(self, *args, **kwargs): self.calls.append((args, kwargs)) - return self.response + return self.responses.pop(0) async def aclose(self): self.closed = True @@ -85,6 +89,50 @@ async def test_async_batch_post_uses_relative_path_and_sanitized_logs(caplog): 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[0] for call in client.calls] == [ + ("/batch/",), + ("/redirected-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() From 54cf24c1782b8bcae0d9d223cd62a0f13ecf7437 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 27 Aug 2026 21:57:35 +0200 Subject: [PATCH 13/13] fix: reject cross-thread capture during shutdown --- posthog/async_client.py | 23 +++++++++++++-------- posthog/test/test_async_client.py | 33 +++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/posthog/async_client.py b/posthog/async_client.py index 65ca61c9e..52bec2fc3 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -292,7 +292,9 @@ 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]) -> None: + def _enqueue_prepared_event(self, prepared: dict[str, Any]) -> bool: + if not self._accepting or self._closed: + return False queued_event = _QueuedEvent(prepared, contextvars.copy_context()) try: running_loop = asyncio.get_running_loop() @@ -303,30 +305,33 @@ def _enqueue_prepared_event(self, prepared: dict[str, Any]) -> None: if running_loop is not None: self._ensure_workers_started() self._queue.put_nowait(queued_event) - return + return True if running_loop is self._loop: self._ensure_workers_started() self._queue.put_nowait(queued_event) - return + return True if running_loop is not None: raise RuntimeError("AsyncClient cannot be shared across event loops") if self._loop_thread_id == threading.get_ident() or not self._loop.is_running(): raise RuntimeError("AsyncClient event loop is not running") - admitted: concurrent.futures.Future[None] = concurrent.futures.Future() + admitted: concurrent.futures.Future[bool] = concurrent.futures.Future() def enqueue_on_bound_loop() -> None: try: + if not self._accepting or self._closed: + admitted.set_result(False) + return self._ensure_workers_started() self._queue.put_nowait(queued_event) except BaseException as error: admitted.set_exception(error) else: - admitted.set_result(None) + admitted.set_result(True) self._loop.call_soon_threadsafe(enqueue_on_bound_loop) - admitted.result() + return admitted.result() def _normalize_uuid(self, msg: dict[str, Any]) -> str: raw_uuid = msg.pop("uuid", None) @@ -466,7 +471,8 @@ def capture( return sent_uuid self._validate_transport_available() - self._enqueue_prepared_event(prepared) + if not self._enqueue_prepared_event(prepared): + return None self.log.debug("queued async event %s", event) return sent_uuid except asyncio.QueueFull: @@ -681,7 +687,8 @@ def _enqueue_built_event( if not self.send: return sent_uuid self._validate_transport_available() - self._enqueue_prepared_event(prepared) + if not self._enqueue_prepared_event(prepared): + return None return sent_uuid def capture_exception( diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index 05b5c9d11..fcb6e2298 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -120,6 +120,39 @@ async def batch_post(*args, **kwargs): subprocess.run([sys.executable, "-c", script], check=True, timeout=5) +@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() + scheduled_callbacks = [] + capture_result = [] + + with mock.patch.object( + client._loop, + "call_soon_threadsafe", + side_effect=lambda callback: scheduled_callbacks.append(callback), + ): + capture_thread = threading.Thread( + target=lambda: capture_result.append( + client.capture("threaded event", distinct_id="user-1") + ), + daemon=True, + ) + capture_thread.start() + while not scheduled_callbacks: + await asyncio.sleep(0) + + client._accepting = False + scheduled_callbacks.pop()() + capture_thread.join(timeout=1) + + assert capture_result == [None] + for task in client._worker_tasks: + task.cancel() + await asyncio.gather(*client._worker_tasks, return_exceptions=True) + await client._close_transport() + + @pytest.mark.asyncio async def test_capture_runs_async_before_send_in_consumer(): batches = []