fix(mcp): isolate pending capture tasks

.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,14 @@ 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()
+        if self._key is None:
+            return
+        data = get_server_tracking_data(self._key)
+        if data is not None:
+            await drain_pending(data)
 
 
 class _NoopAnalytics(McpAnalytics):

posthog/mcp/_instrumentation.py

@@ -25,10 +25,13 @@
 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).
+# Keep strong refs to in-flight capture tasks/futures so they aren't GC'd mid-flight.
+# The metadata lets lifecycle drains select only work belonging to their analytics
+# handle/client and, for asyncio tasks, only work bound to the current event loop.
 _BACKGROUND_TASKS: Set[Any] = set()
+_TASK_OWNERS: Dict[Any, Any] = {}
+_TASK_LOOPS: Dict[Any, asyncio.AbstractEventLoop] = {}
+_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.
@@ -49,60 +52,79 @@ def _get_background_loop() -> asyncio.AbstractEventLoop:
     return _bg_loop
 
 
+def _track_task(task: Any, owner: Any, loop: asyncio.AbstractEventLoop) -> None:
+    with _tasks_lock:
+        _BACKGROUND_TASKS.add(task)
+        _TASK_OWNERS[task] = owner
+        _TASK_LOOPS[task] = loop
+    task.add_done_callback(_on_task_done)
+
+
 def _on_task_done(task: Any) -> None:
-    _BACKGROUND_TASKS.discard(task)
+    with _tasks_lock:
+        _BACKGROUND_TASKS.discard(task)
+        _TASK_OWNERS.pop(task, None)
+        _TASK_LOOPS.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, loop)
+        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, running_loop)
+
+
+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 in _BACKGROUND_TASKS
+            if _TASK_OWNERS.get(task) is owner
+            and _TASK_LOOPS.get(task) is loop
+            and isinstance(task, asyncio.Task)
+            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 in _BACKGROUND_TASKS
+            if _TASK_OWNERS.get(task) is owner
+            and _TASK_LOOPS.get(task) is _bg_loop
+            and isinstance(task, concurrent.futures.Future)
+            and not task.done()
+        ]
     if futures:
         concurrent.futures.wait(futures, timeout=timeout)
 
@@ -170,7 +192,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(
@@ -233,7 +255,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
     )
@@ -288,7 +310,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}")
 
@@ -371,7 +393,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}")
 
@@ -408,6 +430,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_pending_tasks.py

@@ -0,0 +1,136 @@
+import asyncio
+import threading
+
+from posthog.mcp import PostHogMCP
+from posthog.mcp import _instrumentation as instrumentation
+
+
+async def test_async_drain_is_scoped_to_owner():
+    first_owner = object()
+    second_owner = object()
+    first_done = []
+    second_started = asyncio.Event()
+    release_second = asyncio.Event()
+
+    async def first_capture():
+        await asyncio.sleep(0)
+        first_done.append(True)
+
+    async def second_capture():
+        second_started.set()
+        await release_second.wait()
+
+    instrumentation.fire_and_forget(first_capture(), first_owner)
+    instrumentation.fire_and_forget(second_capture(), second_owner)
+    await second_started.wait()
+
+    await asyncio.wait_for(instrumentation.drain_pending(first_owner), timeout=1)
+
+    assert first_done == [True]
+    assert not release_second.is_set()
+
+    release_second.set()
+    await instrumentation.drain_pending(second_owner)
+
+
+async def test_async_drain_ignores_same_owner_tasks_on_another_loop():
+    owner = object()
+    foreign_started = threading.Event()
+    release_foreign = threading.Event()
+    foreign_done = threading.Event()
+    thread_errors = []
+
+    async def foreign_capture():
+        foreign_started.set()
+        while not release_foreign.is_set():
+            await asyncio.sleep(0.01)
+        foreign_done.set()
+
+    def run_foreign_loop():
+        async def run():
+            instrumentation.fire_and_forget(foreign_capture(), owner)
+            while not release_foreign.is_set():
+                await asyncio.sleep(0.01)
+            await instrumentation.drain_pending(owner)
+
+        try:
+            asyncio.run(run())
+        except BaseException as error:  # noqa: BLE001 - surfaced in the test thread
+            thread_errors.append(error)
+
+    thread = threading.Thread(target=run_foreign_loop)
+    thread.start()
+    assert foreign_started.wait(timeout=1)
+
+    local_done = []
+
+    async def local_capture():
+        await asyncio.sleep(0)
+        local_done.append(True)
+
+    try:
+        instrumentation.fire_and_forget(local_capture(), owner)
+        await asyncio.wait_for(instrumentation.drain_pending(owner), timeout=1)
+
+        assert local_done == [True]
+        assert not foreign_done.is_set()
+        assert thread_errors == []
+    finally:
+        release_foreign.set()
+        thread.join(timeout=2)
+
+    assert not thread.is_alive()
+    assert foreign_done.is_set()
+    assert thread_errors == []
+
+
+def test_sync_drain_is_scoped_to_owner():
+    first_owner = object()
+    second_owner = object()
+    first_done = []
+    second_started = threading.Event()
+    release_second = threading.Event()
+    second_done = threading.Event()
+
+    async def first_capture():
+        await asyncio.sleep(0)
+        first_done.append(True)
+
+    async def second_capture():
+        second_started.set()
+        while not release_second.is_set():
+            await asyncio.sleep(0.01)
+        second_done.set()
+
+    instrumentation.fire_and_forget(first_capture(), first_owner, background=True)
+    instrumentation.fire_and_forget(second_capture(), second_owner, background=True)
+    assert second_started.wait(timeout=1)
+
+    try:
+        instrumentation.drain_pending_sync(first_owner, timeout=1)
+
+        assert first_done == [True]
+        assert not second_done.is_set()
+    finally:
+        release_second.set()
+        instrumentation.drain_pending_sync(second_owner, timeout=2)
+
+    assert second_done.is_set()
+
+
+async def test_posthog_mcp_sync_flush_drains_capture_from_async_host(monkeypatch):
+    client = PostHogMCP("phc_test", disabled=True)
+    captured = []
+    flushed_after = []
+    client.capture = lambda event, **kwargs: captured.append({"event": event, **kwargs})
+
+    def record_flush(self, timeout_seconds=10):
+        flushed_after.append(list(captured))
+
+    monkeypatch.setattr("posthog.client.Client.flush", record_flush)
+
+    client.capture_tool_call("search")
+    client.flush(timeout_seconds=1)
+
+    assert captured[0]["event"] == "$mcp_tool_call"
+    assert flushed_after == [captured]