fix(celery): capture task exceptions with exception autocapture

.sampo/changesets/celery-exception-autocapture.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+Exception autocapture now captures Celery task failures and retry exceptions. `PosthogCeleryIntegration` also captures the exception of a retried task.

posthog/__init__.py

@@ -356,7 +356,9 @@ def get_tags() -> Dict[str, Any]:
         serverless handlers should call ``flush()`` before returning. Applied
         when ``setup()`` builds the global client, or on a later ``setup()``
         call if no span has been started yet.
-    enable_exception_autocapture: Automatically capture uncaught exceptions.
+    enable_exception_autocapture: Automatically capture uncaught exceptions,
+        and Celery task exceptions when Celery is imported before the client
+        is created.
     log_captured_exceptions: Also log exceptions captured by error tracking.
     project_root: Root path used to determine in-app exception stack frames.
     privacy_mode: Capture AI usage metadata without prompt inputs or outputs.

posthog/client.py

@@ -774,7 +774,8 @@ def __init__(
                 to 1. Set to 0 to disable retries.
             super_properties: Properties merged into every captured event.
             enable_exception_autocapture: Automatically capture uncaught
-                exceptions.
+                exceptions, and Celery task exceptions when Celery is
+                imported before the client is created.
             log_captured_exceptions: Also log exceptions captured by error
                 tracking.
             project_root: Root path used to determine in-app stack frames for

posthog/exception_capture.py

@@ -44,6 +44,13 @@ def __init__(
         self._threading_excepthook = self.thread_exception_handler
         sys.excepthook = self._sys_excepthook
         threading.excepthook = self._threading_excepthook
+        self._disconnect_celery = None
+        try:
+            from .integrations.celery import _instrument_exception_autocapture
+
+            self._disconnect_celery = _instrument_exception_autocapture(self)
+        except Exception:
+            self.log.exception("Failed to set up Celery exception autocapture")
         # opt-in client-side rate limiting: per exception type, allow a burst
         # of captures, then refill over time
         self._rate_limiter = None
@@ -79,6 +86,9 @@ def close(self):
         if threading.excepthook is self._threading_excepthook:
             threading.excepthook = original_threading_excepthook
 
+        if self._disconnect_celery is not None:
+            self._disconnect_celery()
+
         if self._rate_limiter is not None:
             self._rate_limiter.stop()
 

posthog/integrations/celery.py

@@ -60,14 +60,23 @@
     Additional properties on specific states:
 
     - **failure**: ``error_type``, ``error_message``
-    - **retry**: ``celery_reason``
+    - **retry**: ``celery_reason``, and ``error_type``, ``error_message`` when
+      the retry has an exception
+
+Exception autocapture:
+    When ``enable_exception_autocapture`` is on and Celery is imported before
+    the PostHog client is created, the client captures task failures and retry
+    exceptions without this integration. An instrumented integration with
+    ``capture_exceptions=True`` replaces that capture and applies its
+    ``task_filter``.
 """
 
 import atexit
 import json
 import logging
+import sys
 import time
-from typing import Any, Callable, Optional, cast
+from typing import TYPE_CHECKING, Any, Callable, Optional, cast
 
 from .. import contexts
 from ..client import Client
@@ -76,13 +85,20 @@
     _capture_exception_with_metadata,
 )
 
+if TYPE_CHECKING:
+    from ..exception_capture import ExceptionCapture
+
 
 CONTEXT_DISTINCT_ID_HEADER = "X-POSTHOG-DISTINCT-ID"
 CONTEXT_SESSION_ID_HEADER = "X-POSTHOG-SESSION-ID"
 CONTEXT_TAGS_HEADER = "X-POSTHOG-CONTEXT-TAGS"
 
 logger = logging.getLogger("posthog")
 
+# Instrumented integrations. Exception autocapture defers to them, so a task
+# exception is captured once and respects the integration's task filter.
+_instrumented_integrations: set["PosthogCeleryIntegration"] = set()
+
 
 class PosthogCeleryIntegration:
     """Celery integration that captures task lifecycle events and exceptions.
@@ -158,6 +174,7 @@ def instrument(self) -> None:
         atexit.register(self.shutdown)
 
         self._instrumented = True
+        _instrumented_integrations.add(self)
 
     def _disconnect_signals(self) -> None:
         if not self._instrumented or not self._signals:
@@ -176,6 +193,7 @@ def _disconnect_signals(self) -> None:
 
         self._signals = None
         self._instrumented = False
+        _instrumented_integrations.discard(self)
 
     def uninstrument(self) -> None:
         """Disconnect Celery signal handlers and unregister exit cleanup.
@@ -338,6 +356,7 @@ def _on_task_retry(self, *args, **kwargs):
             extra_properties={
                 "celery_reason": str(kwargs.get("reason")),
             },
+            exception=_retry_exception(kwargs.get("reason")),
             **kwargs,
         )
 
@@ -378,7 +397,7 @@ def _handle_task_end(
                 should_track = self._should_track(task_name, task_properties)
 
             if exception and self.capture_exceptions and should_track:
-                self._capture_exception(exception)
+                self._capture_exception(exception, _capture_metadata(state))
 
             if self.capture_task_lifecycle_events and should_track:
                 self._capture_event(f"celery task {state}", properties=task_properties)
@@ -478,12 +497,11 @@ def _capture_event(self, event: str, properties: dict[str, Any]) -> None:
 
             capture(event, properties=properties)
 
-    def _capture_exception(self, exception: Exception) -> None:
-        capture_metadata: _ExceptionCaptureMetadata = {
-            "level": "error",
-            "source": "celery.task_failure",
-            "mechanism": {"type": "task", "handled": False},
-        }
+    def _capture_exception(
+        self,
+        exception: BaseException,
+        capture_metadata: _ExceptionCaptureMetadata,
+    ) -> None:
         if self.client:
             _capture_exception_with_metadata(self.client, exception, capture_metadata)
         else:
@@ -492,6 +510,67 @@ def _capture_exception(self, exception: Exception) -> None:
             cast(Any, capture_exception)(exception, _capture_metadata=capture_metadata)
 
 
+def _retry_exception(reason: Any) -> Optional[BaseException]:
+    # Celery sends a ``Retry`` as the reason. Its ``exc`` is the exception
+    # that caused the retry, or ``None`` when the task retried without one.
+    exc = getattr(reason, "exc", None)
+    return exc if isinstance(exc, BaseException) else None
+
+
+def _capture_metadata(state: str) -> _ExceptionCaptureMetadata:
+    if state == "retry":
+        # Celery handles the exception and runs the task again.
+        return {
+            "level": "warning",
+            "source": "celery.task_retry",
+            "mechanism": {"type": "task", "handled": True},
+        }
+    return {
+        "level": "error",
+        "source": "celery.task_failure",
+        "mechanism": {"type": "task", "handled": False},
+    }
+
+
+def _instrument_exception_autocapture(
+    exception_capture: "ExceptionCapture",
+) -> Optional[Callable[[], None]]:
+    """Capture Celery task exceptions with exception autocapture.
+
+    Celery catches task exceptions, so ``sys.excepthook`` never sees them.
+    Returns a callable that disconnects the handlers, or ``None`` when the
+    process does not use Celery.
+    """
+    if "celery" not in sys.modules:
+        return None
+
+    from celery import signals
+
+    def capture(exception: Optional[BaseException], state: str) -> None:
+        if exception is None or exception_capture._closed:
+            return
+        if any(i.capture_exceptions for i in _instrumented_integrations):
+            return
+        exception_capture._capture_exception(
+            exception, capture_metadata=_capture_metadata(state)
+        )
+
+    def on_task_failure(*args, **kwargs):
+        capture(kwargs.get("exception"), "failure")
+
+    def on_task_retry(*args, **kwargs):
+        capture(_retry_exception(kwargs.get("reason")), "retry")
+
+    signals.task_failure.connect(on_task_failure, weak=False)
+    signals.task_retry.connect(on_task_retry, weak=False)
+
+    def disconnect() -> None:
+        signals.task_failure.disconnect(on_task_failure)
+        signals.task_retry.disconnect(on_task_retry)
+
+    return disconnect
+
+
 __all__ = [
     "PosthogCeleryIntegration",
 ]