From 3d6679b854eaef5446be259fae3d90fe0e4091a5 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Thu, 30 Jul 2026 16:40:41 +0200 Subject: [PATCH] fix(mcp): reset background capture after fork --- .sampo/changesets/ardent-queen-vainamoinen.md | 5 ++ posthog/mcp/_instrumentation.py | 18 +++++ posthog/test/mcp/test_instrumentation_fork.py | 81 +++++++++++++++++++ 3 files changed, 104 insertions(+) create mode 100644 .sampo/changesets/ardent-queen-vainamoinen.md create mode 100644 posthog/test/mcp/test_instrumentation_fork.py diff --git a/.sampo/changesets/ardent-queen-vainamoinen.md b/.sampo/changesets/ardent-queen-vainamoinen.md new file mode 100644 index 000000000..ecdf61641 --- /dev/null +++ b/.sampo/changesets/ardent-queen-vainamoinen.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Reset MCP background capture state after fork diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 865e54733..940628524 100644 --- a/posthog/mcp/_instrumentation.py +++ b/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: diff --git a/posthog/test/mcp/test_instrumentation_fork.py b/posthog/test/mcp/test_instrumentation_fork.py new file mode 100644 index 000000000..bde5f079a --- /dev/null +++ b/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"