fix(mcp): isolate pending capture tasks (#815)

* fix(mcp): isolate pending capture tasks

* test(mcp): await background capture futures

* address mcp pending task review feedback

* fix(mcp): reset pending task state after fork

Author: marandaneto

.sampo/changesets/regal-witch-kullervo.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Isolate MCP pending capture tasks by owner and loop

posthog/mcp/__init__.py

@@ -125,10 +125,12 @@ async def capture(self, event: str, properties: Optional[dict] = None) -> None:
             await coro
 
     async def flush(self) -> None:
-        """Await in-flight auto-captured events scheduled on the current event loop.
+        """Await this server's in-flight auto-captures on the current event loop.
         Call this before ``posthog.shutdown()`` on exit so trailing tool-call events
         aren't dropped. (Then call ``posthog.flush()``/``shutdown()`` to send them.)"""
-        await drain_pending()
+        data = get_server_tracking_data(self._key)
+        if data is not None:
+            await drain_pending(data)
 
 
 class _NoopAnalytics(McpAnalytics):
@@ -138,6 +140,10 @@ def __init__(self) -> None:  # noqa: D401 - graceful degradation handle
     async def capture(self, event: str, properties: Optional[dict] = None) -> None:
         return None
 
+    async def flush(self) -> None:
+        # There is no tracking key to look up or pending work to drain.
+        return None
+
 
 def _resolve_client(posthog_client: Optional[Client]) -> Optional[Client]:
     if posthog_client is not None:

posthog/mcp/_instrumentation.py

@@ -14,7 +14,7 @@
 import os
 import threading
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Set
+from typing import Any, Dict, List, Optional
 
 from ._capture import capture_event
 from ._event_types import MCPAnalyticsEventType
@@ -26,10 +26,10 @@
 from .session import resolve_session_id
 from .session_token import SessionTokenPayload, decode_session_id
 
-# Keep strong refs to in-flight capture tasks/futures so they aren't GC'd mid-flight,
-# and so the asyncio ones can be awaited via drain_pending() before shutdown. Holds
-# asyncio.Task (running-loop path) or concurrent.futures.Future (sync background-loop path).
-_BACKGROUND_TASKS: Set[Any] = set()
+# Keep strong refs to in-flight capture tasks/futures and their lifecycle owners so
+# they aren't GC'd mid-flight and lifecycle drains can select only their own work.
+_BACKGROUND_TASKS: Dict[Any, Any] = {}
+_tasks_lock = threading.Lock()
 
 # A single daemon event loop for hosts with no running loop (sync dispatchers
 # like PostHogMCP). Created lazily and reused, so we never leak a loop per call.
@@ -44,8 +44,9 @@ def _reinit_background_loop_after_fork() -> None:
     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()
+    global _BACKGROUND_TASKS, _tasks_lock, _bg_loop, _bg_loop_lock
+    _BACKGROUND_TASKS = {}
+    _tasks_lock = threading.Lock()
     _bg_loop = None
     _bg_loop_lock = threading.Lock()
 
@@ -67,60 +68,74 @@ def _get_background_loop() -> asyncio.AbstractEventLoop:
     return _bg_loop
 
 
+def _track_task(task: Any, owner: Any) -> None:
+    with _tasks_lock:
+        _BACKGROUND_TASKS[task] = owner
+    task.add_done_callback(_on_task_done)
+
+
 def _on_task_done(task: Any) -> None:
-    _BACKGROUND_TASKS.discard(task)
+    with _tasks_lock:
+        _BACKGROUND_TASKS.pop(task, None)
     try:
         if not task.cancelled() and task.exception() is not None:
             log(f"background capture task failed: {task.exception()}")
     except Exception:  # noqa: BLE001 - never let bookkeeping raise
         pass
 
 
-def fire_and_forget(coro: Optional[Any]) -> None:
-    """Schedule a capture coroutine without blocking the tool path. No-ops if the
-    coroutine is ``None`` (no sink). Runs on the current loop when there is one,
-    otherwise on a shared daemon loop (sync hosts) — never creates a throwaway loop."""
+def fire_and_forget(
+    coro: Optional[Any], owner: Any, *, background: bool = False
+) -> None:
+    """Schedule capture work and associate it with its lifecycle owner.
+
+    Async instrumentation uses its current loop. Sync-only owners can request the
+    shared background loop so their synchronous lifecycle methods can safely drain
+    captures even when invoked by a host that also has a running event loop.
+    """
     if coro is None:
         return
     try:
-        asyncio.get_running_loop()
+        running_loop = asyncio.get_running_loop()
     except RuntimeError:
-        # No running loop (sync host) — schedule on the shared background loop.
-        future = asyncio.run_coroutine_threadsafe(coro, _get_background_loop())
-        _BACKGROUND_TASKS.add(future)
-        future.add_done_callback(_on_task_done)
-        return
-    task = asyncio.ensure_future(coro)
-    _BACKGROUND_TASKS.add(task)
-    task.add_done_callback(_on_task_done)
+        running_loop = None
 
+    if background or running_loop is None:
+        loop = _get_background_loop()
+        future = asyncio.run_coroutine_threadsafe(coro, loop)
+        _track_task(future, owner)
+        return
 
-async def drain_pending() -> None:
-    """Await in-flight capture work before ``posthog.shutdown()`` instead of racing a
-    sleep. Covers both paths: ``asyncio.Task`` (running-loop hosts) and the
-    ``concurrent.futures.Future`` scheduled on the background loop (sync hosts like
-    PostHogMCP) — the latter wrapped so it can be awaited on the current loop."""
-    awaitables: List[Any] = []
-    for t in list(_BACKGROUND_TASKS):
-        if isinstance(t, asyncio.Task):
-            if not t.done():
-                awaitables.append(t)
-        elif isinstance(t, concurrent.futures.Future):
-            if not t.done():
-                awaitables.append(asyncio.wrap_future(t))
-    if awaitables:
-        await asyncio.gather(*awaitables, return_exceptions=True)
-
-
-def drain_pending_sync(timeout: Optional[float] = None) -> None:
-    """Block until background-loop captures finish. For sync hosts (PostHogMCP) that
-    can't await :func:`drain_pending` — call it before ``flush()``/``shutdown()`` so
-    trailing events aren't still in flight when the client tears down."""
-    futures = [
-        t
-        for t in list(_BACKGROUND_TASKS)
-        if isinstance(t, concurrent.futures.Future) and not t.done()
-    ]
+    task = running_loop.create_task(coro)
+    _track_task(task, owner)
+
+
+async def drain_pending(owner: Any) -> None:
+    """Await this owner's in-flight captures bound to the current event loop."""
+    loop = asyncio.get_running_loop()
+    with _tasks_lock:
+        tasks = [
+            task
+            for task, task_owner in _BACKGROUND_TASKS.items()
+            if task_owner is owner
+            and isinstance(task, asyncio.Task)
+            and task.get_loop() is loop
+            and not task.done()
+        ]
+    if tasks:
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+
+def drain_pending_sync(owner: Any, timeout: Optional[float] = None) -> None:
+    """Block until this owner's shared-background-loop captures finish."""
+    with _tasks_lock:
+        futures = [
+            task
+            for task, task_owner in _BACKGROUND_TASKS.items()
+            if task_owner is owner
+            and isinstance(task, concurrent.futures.Future)
+            and not task.done()
+        ]
     if futures:
         concurrent.futures.wait(futures, timeout=timeout)
 
@@ -188,7 +203,7 @@ async def _maybe_emit_initialize(
     await _apply_event_properties(
         data, event, {"method": "initialize", "params": {}}, extra
     )
-    fire_and_forget(capture_event(data, event))
+    fire_and_forget(capture_event(data, event), data)
 
 
 async def _apply_event_properties(
@@ -251,7 +266,7 @@ async def prepare_request(
     session_id = await resolve_session_id(data, mcp_session_id, token=token)
     identify_event = await handle_identify(data, session_id, request, extra)
     if identify_event:
-        fire_and_forget(capture_event(data, identify_event))
+        fire_and_forget(capture_event(data, identify_event), data)
     await _maybe_emit_initialize(
         data, session_id, client_name, client_version, extra, protocol_version
     )
@@ -306,7 +321,7 @@ async def record_tool_call(
         if props is not None:
             event["properties"] = props
 
-        fire_and_forget(capture_event(data, event))
+        fire_and_forget(capture_event(data, event), data)
     except Exception as err:  # noqa: BLE001 - isolate analytics from the tool path
         log(f"record_tool_call failed (event dropped, tool unaffected): {err}")
 
@@ -389,7 +404,7 @@ async def record_missing_capability(
             event["user_intent"] = context.strip()
             event["user_intent_source"] = "context_parameter"
         await _apply_event_properties(data, event, request, extra)
-        fire_and_forget(capture_event(data, event))
+        fire_and_forget(capture_event(data, event), data)
     except Exception as err:  # noqa: BLE001 - isolate analytics from the tool path
         log(f"record_missing_capability failed (event dropped): {err}")
 
@@ -426,6 +441,6 @@ async def record_tools_list(
         if error is not None:
             event["error"] = capture_exception(error)
         await _apply_event_properties(data, event, request, extra)
-        fire_and_forget(capture_event(data, event))
+        fire_and_forget(capture_event(data, event), data)
     except Exception as err:  # noqa: BLE001 - isolate analytics from the tool path
         log(f"record_tools_list failed (event dropped): {err}")

posthog/mcp/posthog_mcp.py

@@ -64,14 +64,14 @@ def __init__(
 
     def flush(self, timeout_seconds: Optional[float] = 10) -> None:
         """Drain in-flight MCP captures scheduled on the background loop, then flush
-        the underlying client. The capture methods are fire-and-forget on a sync host,
-        so without this drain a trailing event could still be in flight at flush time."""
-        drain_pending_sync(timeout=timeout_seconds)
+        the underlying client. The capture methods are fire-and-forget, so without
+        this drain a trailing event could still be in flight at flush time."""
+        drain_pending_sync(self, timeout=timeout_seconds)
         return super().flush(timeout_seconds=timeout_seconds)
 
     def shutdown(self) -> None:
         """Drain in-flight MCP captures, then shut the underlying client down."""
-        drain_pending_sync()
+        drain_pending_sync(self)
         return super().shutdown()
 
     # --- capture methods -----------------------------------------------------
@@ -300,7 +300,10 @@ def _emit(self, event: Dict[str, Any]) -> None:
         options = McpCaptureOptions(
             enable_exception_autocapture=self._mcp_exception_autocapture
         )
-        fire_and_forget(self._mcp_sink.capture(event, options))
+        # PostHogMCP exposes synchronous lifecycle methods, so always use the shared
+        # background loop even when capture is called by an async host. This keeps
+        # flush()/shutdown() able to drain without blocking their own event loop's tasks.
+        fire_and_forget(self._mcp_sink.capture(event, options), self, background=True)
 
     def _inject_context(self, tool: Any, description: Optional[str]) -> Any:
         if isinstance(tool, dict):

posthog/test/mcp/_helpers.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import concurrent.futures
 
 
 class FakeClient:
@@ -32,9 +33,15 @@ async def flush_background():
     """Let fire-and-forget capture tasks run to completion."""
     import posthog.mcp._instrumentation as instr
 
+    loop = asyncio.get_running_loop()
     for _ in range(10):
         await asyncio.sleep(0)
-        pending = [t for t in list(instr._BACKGROUND_TASKS) if not t.done()]
+        pending = []
+        for task in list(instr._BACKGROUND_TASKS):
+            if isinstance(task, asyncio.Task) and task.get_loop() is loop:
+                pending.append(task)
+            elif isinstance(task, concurrent.futures.Future):
+                pending.append(asyncio.wrap_future(task))
         if pending:
             await asyncio.gather(*pending, return_exceptions=True)
     await asyncio.sleep(0)

posthog/test/mcp/test_fastmcp.py

@@ -1,5 +1,7 @@
 """End-to-end tests for the FastMCP adapter (Milestone 2)."""
 
+import asyncio
+
 import pytest
 
 import mcp.types as mcp_types
@@ -101,6 +103,27 @@ def spy_add(a: int, b: int) -> int:
     assert "context" not in props["$mcp_parameters"]["request"]["params"]["arguments"]
 
 
+async def test_analytics_flush_drains_its_own_captures():
+    async def slow_before_send(event):
+        await asyncio.sleep(0.05)
+        return event
+
+    server = make_server()
+    client = FakeClient()
+    analytics = instrument(
+        server, client, MCPAnalyticsOptions(before_send=slow_before_send)
+    )
+
+    await server._tool_manager.call_tool(
+        "add", {"a": 2, "b": 3, "context": "summing two numbers"}
+    )
+    assert _events(client, "$mcp_tool_call") == []
+
+    await analytics.flush()
+
+    assert len(_events(client, "$mcp_tool_call")) == 1
+
+
 async def test_initialize_emitted_once_per_session():
     server = make_server()
     client = FakeClient()
@@ -179,5 +202,6 @@ async def test_instrument_is_idempotent():
 
 async def test_unsupported_server_returns_noop_handle():
     handle = instrument(object(), FakeClient())
-    # graceful no-op: capture does nothing and does not raise
+    # graceful no-op: capture and flush do nothing and do not raise
     await handle.capture("anything")
+    await handle.flush()

posthog/test/mcp/test_instrumentation_fork.py

@@ -22,14 +22,16 @@ async def pending_parent_capture():
         while not finish_parent_capture.is_set():
             await asyncio.sleep(0.01)
 
-    instrumentation.fire_and_forget(pending_parent_capture())
+    owner = object()
+    instrumentation.fire_and_forget(pending_parent_capture(), owner)
     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()
+    instrumentation._tasks_lock.acquire()
     try:
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", DeprecationWarning)
@@ -44,8 +46,8 @@ async def pending_parent_capture():
                 async def child_capture():
                     child_capture_completed.append(True)
 
-                instrumentation.fire_and_forget(child_capture())
-                instrumentation.drain_pending_sync(timeout=2)
+                instrumentation.fire_and_forget(child_capture(), owner)
+                instrumentation.drain_pending_sync(owner, timeout=2)
                 new_loop_created = instrumentation._bg_loop is not parent_loop
 
                 if (
@@ -73,9 +75,10 @@ async def child_capture():
         os.close(read_fd)
         _, status = os.waitpid(pid, 0)
     finally:
+        instrumentation._tasks_lock.release()
         instrumentation._bg_loop_lock.release()
         finish_parent_capture.set()
-        instrumentation.drain_pending_sync(timeout=2)
+        instrumentation.drain_pending_sync(owner, timeout=2)
 
     assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, result
     assert result == "ok"