fix(contexts): reset context after fork

.sampo/changesets/stalwart-bard-goulven.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Reset PostHog context after fork. Forked children no longer retain the parent process's active lexical context; they start without inherited context and can establish a new child-local context.

posthog/contexts.py

@@ -1,4 +1,5 @@
 import contextvars
+import os
 from contextlib import contextmanager
 from typing import Optional, Any, Callable, Dict, TypeVar, cast, TYPE_CHECKING
 
@@ -7,6 +8,9 @@
     from posthog.client import Client
 
 
+_context_generation = 0
+
+
 class ContextScope:
     def __init__(
         self,
@@ -16,6 +20,7 @@ def __init__(
         client: Optional["Client"] = None,
     ):
         self.client: Optional[Client] = client
+        self._generation = _context_generation
         self.parent = parent
         self.fresh = fresh
         self.capture_exceptions = capture_exceptions
@@ -130,8 +135,25 @@ def get_code_variables_detect_secrets(self) -> Optional[bool]:
 )
 
 
+def _reset_context_after_fork() -> None:
+    global _context_generation
+
+    _context_generation += 1
+    _context_stack.set(None)
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reset_context_after_fork)
+
+
 def _get_current_context() -> Optional[ContextScope]:
-    return _context_stack.get()
+    current_context = _context_stack.get()
+    if (
+        current_context is not None
+        and current_context._generation != _context_generation
+    ):
+        return None
+    return current_context
 
 
 def _default_capture_exceptions(client: Optional["Client"] = None) -> bool:
@@ -199,13 +221,17 @@ def new_context(
     from . import capture_exception
 
     current_context = _get_current_context()
+    context_generation = _context_generation
     resolved_capture_exceptions = (
         capture_exceptions
         if capture_exceptions is not None
         else _default_capture_exceptions(client)
     )
     new_context = ContextScope(
-        current_context, fresh, resolved_capture_exceptions, client
+        current_context,
+        fresh,
+        resolved_capture_exceptions,
+        client,
     )
     _context_stack.set(new_context)
 
@@ -219,7 +245,8 @@ def new_context(
                 capture_exception(e)
         raise
     finally:
-        _context_stack.set(new_context.get_parent())
+        if context_generation == _context_generation:
+            _context_stack.set(new_context.get_parent())
 
 
 def tag(key: str, value: Any) -> None:

posthog/test/test_contexts.py

@@ -1,4 +1,6 @@
 import asyncio
+import contextvars
+import os
 import unittest
 from unittest.mock import MagicMock, patch
 
@@ -11,8 +13,10 @@
     tag,
     identify_context,
     set_context_session,
+    set_context_device_id,
     get_context_session_id,
     get_context_distinct_id,
+    get_context_device_id,
 )
 
 
@@ -275,6 +279,111 @@ def test_context_inheritance_non_fresh_context(self):
             assert get_context_distinct_id() == "user123"
             assert get_context_session_id() == "session456"
 
+    @unittest.skipUnless(
+        hasattr(os, "fork") and hasattr(os, "register_at_fork"),
+        "requires os.fork and os.register_at_fork",
+    )
+    def test_fork_clears_context_in_child_and_preserves_parent(self):
+        def context_state():
+            return (
+                get_context_distinct_id(),
+                get_context_session_id(),
+                get_context_device_id(),
+                get_tags(),
+            )
+
+        read_fd, write_fd = os.pipe()
+        pid = -1
+        child_result = b""
+        child_exit_code = 1
+        try:
+            with new_context(fresh=True):
+                identify_context("parent-user")
+                set_context_session("parent-session")
+                set_context_device_id("parent-device")
+                tag("parent-tag", "parent-value")
+
+                with new_context():
+                    copied_context = contextvars.copy_context()
+                    pid = os.fork()
+                    if pid == 0:
+                        os.close(read_fd)
+                        with new_context():
+                            identify_context("child-user")
+                            set_context_session("child-session")
+                            set_context_device_id("child-device")
+                            tag("child-tag", "child-value")
+                            child_local_state = context_state()
+                        child_state_after_local_scope = context_state()
+                    else:
+                        os.close(write_fd)
+
+                if pid == 0:
+                    child_state_after_inner_scope = context_state()
+                else:
+                    parent_state = context_state()
+
+            if pid == 0:
+                child_states = (
+                    child_local_state,
+                    child_state_after_local_scope,
+                    child_state_after_inner_scope,
+                    context_state(),
+                    copied_context.run(context_state),
+                )
+                child_result = repr(child_states).encode()
+                child_exit_code = 0
+        except BaseException as error:
+            if pid != 0:
+                raise
+            child_result = f"{type(error).__name__}: {error}".encode()
+        finally:
+            if pid == 0:
+                try:
+                    os.write(write_fd, child_result)
+                finally:
+                    try:
+                        os.close(write_fd)
+                    finally:
+                        os._exit(child_exit_code)
+
+        child_states = os.read(read_fd, 4096)
+        os.close(read_fd)
+        _, status = os.waitpid(pid, 0)
+
+        empty_state = (None, None, None, {})
+        self.assertTrue(
+            os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0,
+            child_states.decode(errors="replace"),
+        )
+        child_local_state = (
+            "child-user",
+            "child-session",
+            "child-device",
+            {"child-tag": "child-value"},
+        )
+        self.assertEqual(
+            child_states,
+            repr(
+                (
+                    child_local_state,
+                    empty_state,
+                    empty_state,
+                    empty_state,
+                    empty_state,
+                )
+            ).encode(),
+        )
+        self.assertEqual(
+            parent_state,
+            (
+                "parent-user",
+                "parent-session",
+                "parent-device",
+                {"parent-tag": "parent-value"},
+            ),
+        )
+
     def test_child_tags_override_parent_tags_in_non_fresh_context(self):
         with new_context(fresh=True):
             tag("shared_key", "parent_value")