Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .sampo/changesets/ardent-queen-vainamoinen.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
pypi/posthog: patch
---

Reset MCP background capture state after fork
18 changes: 18 additions & 0 deletions posthog/mcp/_instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down
81 changes: 81 additions & 0 deletions posthog/test/mcp/test_instrumentation_fork.py
Original file line number Diff line number Diff line change
@@ -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"