fix(exceptions): restore exception hooks safely

.sampo/changesets/resolute-seer-ukko.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Restore exception hooks safely

posthog/exception_capture.py

@@ -33,9 +33,13 @@ def __init__(
         refill_interval_seconds=DEFAULT_REFILL_INTERVAL_SECONDS,
     ):
         self.client = client
+        self._closed = False
         self.original_excepthook = sys.excepthook
-        sys.excepthook = self.exception_handler
-        threading.excepthook = self.thread_exception_handler
+        self._original_threading_excepthook = threading.excepthook
+        self._sys_excepthook = self.exception_handler
+        self._threading_excepthook = self.thread_exception_handler
+        sys.excepthook = self._sys_excepthook
+        threading.excepthook = self._threading_excepthook
         # opt-in client-side rate limiting: per exception type, allow a burst
         # of captures, then refill over time
         self._rate_limiter = None
@@ -47,17 +51,63 @@ def __init__(
             )
 
     def close(self):
-        sys.excepthook = self.original_excepthook
+        if self._closed:
+            return
+
+        self._closed = True
+        original_excepthook = self._resolve_hook(
+            self.original_excepthook,
+            "exception_handler",
+            "original_excepthook",
+        )
+        original_threading_excepthook = self._resolve_hook(
+            self._original_threading_excepthook,
+            "thread_exception_handler",
+            "_original_threading_excepthook",
+        )
+
+        # Keep each final ownership check and assignment together without a
+        # Python call between them. On supported GIL-enabled CPython builds,
+        # ordinary Python thread scheduling has no switch point inside either
+        # straight-line pair, minimizing the window for external hook writers.
+        if sys.excepthook is self._sys_excepthook:
+            sys.excepthook = original_excepthook
+        if threading.excepthook is self._threading_excepthook:
+            threading.excepthook = original_threading_excepthook
+
         if self._rate_limiter is not None:
             self._rate_limiter.stop()
 
     def exception_handler(self, exc_type, exc_value, exc_traceback):
-        # don't affect default behaviour.
-        self.capture_exception((exc_type, exc_value, exc_traceback))
-        self.original_excepthook(exc_type, exc_value, exc_traceback)
+        if not self._closed:
+            self.capture_exception((exc_type, exc_value, exc_traceback))
+        previous_hook = self._resolve_hook(
+            self.original_excepthook,
+            "exception_handler",
+            "original_excepthook",
+        )
+        previous_hook(exc_type, exc_value, exc_traceback)
 
     def thread_exception_handler(self, args):
-        self.capture_exception((args.exc_type, args.exc_value, args.exc_traceback))
+        if not self._closed:
+            self.capture_exception((args.exc_type, args.exc_value, args.exc_traceback))
+        previous_hook = self._resolve_hook(
+            self._original_threading_excepthook,
+            "thread_exception_handler",
+            "_original_threading_excepthook",
+        )
+        previous_hook(args)
+
+    @staticmethod
+    def _resolve_hook(hook, handler_name, previous_hook_name):
+        """Skip closed ExceptionCapture hooks while preserving the hook chain."""
+        while True:
+            owner = getattr(hook, "__self__", None)
+            if not isinstance(owner, ExceptionCapture) or not owner._closed:
+                return hook
+            if hook != getattr(owner, handler_name):
+                return hook
+            hook = getattr(owner, previous_hook_name)
 
     def exception_receiver(self, exc_info, extra_properties):
         if "distinct_id" in extra_properties:

posthog/test/test_exception_capture.py

@@ -1,6 +1,8 @@
 import subprocess
 import sys
+import threading
 from textwrap import dedent
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -133,6 +135,143 @@ def test_rate_limit_keys_on_outermost_of_chained_exceptions():
         capture.close()
 
 
+def test_exception_hooks_delegate_and_restore_previous_hooks(monkeypatch):
+    from posthog.exception_capture import ExceptionCapture
+
+    sys_hook = MagicMock()
+    thread_hook = MagicMock()
+    monkeypatch.setattr(sys, "excepthook", sys_hook)
+    monkeypatch.setattr(threading, "excepthook", thread_hook)
+    client = MagicMock()
+    capture = ExceptionCapture(client)
+    exc_info = _exc_info(ValueError("boom"))
+    thread_args = SimpleNamespace(
+        exc_type=exc_info[0],
+        exc_value=exc_info[1],
+        exc_traceback=exc_info[2],
+        thread=threading.current_thread(),
+    )
+
+    capture.exception_handler(*exc_info)
+    capture.thread_exception_handler(thread_args)
+    capture.close()
+
+    assert client.capture_exception.call_count == 2
+    sys_hook.assert_called_once_with(*exc_info)
+    thread_hook.assert_called_once_with(thread_args)
+    assert sys.excepthook is sys_hook
+    assert threading.excepthook is thread_hook
+
+
+def test_close_does_not_overwrite_hooks_installed_later(monkeypatch):
+    from posthog.exception_capture import ExceptionCapture
+
+    monkeypatch.setattr(sys, "excepthook", MagicMock())
+    monkeypatch.setattr(threading, "excepthook", MagicMock())
+    capture = ExceptionCapture(MagicMock())
+    replacement_sys_hook = MagicMock()
+    replacement_thread_hook = MagicMock()
+    sys.excepthook = replacement_sys_hook
+    threading.excepthook = replacement_thread_hook
+
+    capture.close()
+
+    assert sys.excepthook is replacement_sys_hook
+    assert threading.excepthook is replacement_thread_hook
+
+
+def test_close_does_not_overwrite_hooks_replaced_while_resolving(monkeypatch):
+    from posthog.exception_capture import ExceptionCapture
+
+    monkeypatch.setattr(sys, "excepthook", MagicMock())
+    monkeypatch.setattr(threading, "excepthook", MagicMock())
+    capture = ExceptionCapture(MagicMock())
+    replacement_sys_hook = MagicMock()
+    replacement_thread_hook = MagicMock()
+    original_resolve_hook = capture._resolve_hook
+
+    def replace_hook_while_resolving(hook, handler_name, previous_hook_name):
+        if handler_name == "exception_handler":
+            sys.excepthook = replacement_sys_hook
+        else:
+            threading.excepthook = replacement_thread_hook
+        return original_resolve_hook(hook, handler_name, previous_hook_name)
+
+    monkeypatch.setattr(capture, "_resolve_hook", replace_hook_while_resolving)
+
+    capture.close()
+
+    assert sys.excepthook is replacement_sys_hook
+    assert threading.excepthook is replacement_thread_hook
+
+
+def test_multiple_captures_support_out_of_order_close(monkeypatch):
+    from posthog.exception_capture import ExceptionCapture
+
+    sys_hook = MagicMock()
+    thread_hook = MagicMock()
+    monkeypatch.setattr(sys, "excepthook", sys_hook)
+    monkeypatch.setattr(threading, "excepthook", thread_hook)
+    first_client = MagicMock()
+    second_client = MagicMock()
+    first = ExceptionCapture(first_client)
+    second = ExceptionCapture(second_client)
+
+    first.close()
+    assert sys.excepthook == second.exception_handler
+    assert threading.excepthook == second.thread_exception_handler
+
+    exc_info = _exc_info(RuntimeError("boom"))
+    thread_args = SimpleNamespace(
+        exc_type=exc_info[0],
+        exc_value=exc_info[1],
+        exc_traceback=exc_info[2],
+        thread=threading.current_thread(),
+    )
+    second.exception_handler(*exc_info)
+    second.thread_exception_handler(thread_args)
+
+    assert first_client.capture_exception.call_count == 0
+    assert second_client.capture_exception.call_count == 2
+    sys_hook.assert_called_once_with(*exc_info)
+    thread_hook.assert_called_once_with(thread_args)
+
+    second.close()
+    assert sys.excepthook is sys_hook
+    assert threading.excepthook is thread_hook
+
+
+def test_uncaught_thread_exception_preserves_default_diagnostic():
+    script = dedent(
+        """
+        import threading
+        from posthog.exception_capture import ExceptionCapture
+
+        class Client:
+            def capture_exception(self, exception, distinct_id=None):
+                print(f"captured:{exception[0].__name__}")
+
+        capture = ExceptionCapture(Client())
+
+        def fail():
+            raise RuntimeError("thread failure")
+
+        thread = threading.Thread(target=fail, name="failing-thread")
+        thread.start()
+        thread.join()
+        capture.close()
+        """
+    )
+
+    result = subprocess.run(
+        [sys.executable, "-c", script], capture_output=True, text=True, check=True
+    )
+
+    assert "captured:RuntimeError" in result.stdout
+    assert "Exception in thread failing-thread" in result.stderr
+    assert "RuntimeError: thread failure" in result.stderr
+
+
 def test_excepthook(tmpdir):
     app = tmpdir.join("app.py")
     app.write(