fix(mcp): reset background capture after fork (#816)

Author: marandaneto

.sampo/changesets/ardent-queen-vainamoinen.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Reset MCP background capture state after fork

posthog/mcp/_instrumentation.py

@@ -11,6 +11,7 @@
 
 import asyncio
 import concurrent.futures
+import os
 import threading
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Set
@@ -36,6 +37,23 @@
 _bg_loop_lock = threading.Lock()
 
 
+def _reinit_background_loop_after_fork() -> None:
+    """Drop background-loop state inherited by a forked child.
+
+    The loop's daemon thread does not survive ``fork()``, and its lock may have
+    been held by a vanished thread. Replace the state without acquiring the old
+    lock or trying to close the inherited loop, which can no longer be driven.
+    """
+    global _BACKGROUND_TASKS, _bg_loop, _bg_loop_lock
+    _BACKGROUND_TASKS = set()
+    _bg_loop = None
+    _bg_loop_lock = threading.Lock()
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reinit_background_loop_after_fork)
+
+
 def _get_background_loop() -> asyncio.AbstractEventLoop:
     global _bg_loop
     if _bg_loop is None:

posthog/test/mcp/test_instrumentation_fork.py

@@ -0,0 +1,81 @@
+import asyncio
+import os
+import signal
+import threading
+import warnings
+
+import pytest
+
+import posthog.mcp._instrumentation as instrumentation
+
+
+@pytest.mark.skipif(
+    not hasattr(os, "fork") or not hasattr(os, "register_at_fork"),
+    reason="requires os.fork and os.register_at_fork",
+)
+def test_sync_capture_completes_after_fork():
+    parent_capture_started = threading.Event()
+    finish_parent_capture = threading.Event()
+
+    async def pending_parent_capture():
+        parent_capture_started.set()
+        while not finish_parent_capture.is_set():
+            await asyncio.sleep(0.01)
+
+    instrumentation.fire_and_forget(pending_parent_capture())
+    assert parent_capture_started.wait(timeout=2)
+    parent_loop = instrumentation._bg_loop
+    assert parent_loop is not None
+    assert instrumentation._BACKGROUND_TASKS
+
+    read_fd, write_fd = os.pipe()
+    instrumentation._bg_loop_lock.acquire()
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", DeprecationWarning)
+            pid = os.fork()
+        if pid == 0:
+            os.close(read_fd)
+            signal.alarm(5)
+            try:
+                inherited_work_cleared = not instrumentation._BACKGROUND_TASKS
+                child_capture_completed = []
+
+                async def child_capture():
+                    child_capture_completed.append(True)
+
+                instrumentation.fire_and_forget(child_capture())
+                instrumentation.drain_pending_sync(timeout=2)
+                new_loop_created = instrumentation._bg_loop is not parent_loop
+
+                if (
+                    inherited_work_cleared
+                    and child_capture_completed == [True]
+                    and new_loop_created
+                ):
+                    result = "ok"
+                else:
+                    result = (
+                        f"inherited_work_cleared={inherited_work_cleared}, "
+                        f"child_capture_completed={child_capture_completed}, "
+                        f"new_loop_created={new_loop_created}"
+                    )
+            except BaseException as error:
+                result = f"exception: {error!r}"
+            finally:
+                signal.alarm(0)
+                os.write(write_fd, result.encode())
+                os.close(write_fd)
+                os._exit(0)
+
+        os.close(write_fd)
+        result = os.read(read_fd, 4096).decode()
+        os.close(read_fd)
+        _, status = os.waitpid(pid, 0)
+    finally:
+        instrumentation._bg_loop_lock.release()
+        finish_parent_capture.set()
+        instrumentation.drain_pending_sync(timeout=2)
+
+    assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, result
+    assert result == "ok"