From b0b1eef7a1944412569b600f6bf62afb84218ba0 Mon Sep 17 00:00:00 2001
From: Manoel Aranda Neto <marandaneto@gmail.com>
Date: Thu, 30 Jul 2026 16:40:09 +0200
Subject: [PATCH 1/4] fix(mcp): isolate pending capture tasks

---
 .sampo/changesets/regal-witch-kullervo.md |   5 +
 posthog/mcp/__init__.py                   |   8 +-
 posthog/mcp/_instrumentation.py           | 118 +++++++++++--------
 posthog/mcp/posthog_mcp.py                |  13 ++-
 posthog/test/mcp/_helpers.py              |   9 +-
 posthog/test/mcp/test_pending_tasks.py    | 136 ++++++++++++++++++++++
 posthog/test/mcp/test_review_fixes.py     |   7 +-
 7 files changed, 237 insertions(+), 59 deletions(-)
 create mode 100644 .sampo/changesets/regal-witch-kullervo.md
 create mode 100644 posthog/test/mcp/test_pending_tasks.py

diff --git a/.sampo/changesets/regal-witch-kullervo.md b/.sampo/changesets/regal-witch-kullervo.md
new file mode 100644
index 000000000..e697b8b4e
--- /dev/null
+++ b/.sampo/changesets/regal-witch-kullervo.md
@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Isolate MCP pending capture tasks by owner and loop
diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py
index 3da91f7d8..d4778ccb1 100644
--- a/posthog/mcp/__init__.py
+++ b/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):
diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py
index 865e54733..17c24d545 100644
--- a/posthog/mcp/_instrumentation.py
+++ b/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,8 +52,19 @@ 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()}")
@@ -58,51 +72,59 @@ def _on_task_done(task: Any) -> None:
         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}")
diff --git a/posthog/mcp/posthog_mcp.py b/posthog/mcp/posthog_mcp.py
index 67dcb2101..487545e80 100644
--- a/posthog/mcp/posthog_mcp.py
+++ b/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):
diff --git a/posthog/test/mcp/_helpers.py b/posthog/test/mcp/_helpers.py
index 17a6b496a..bb5cdf503 100644
--- a/posthog/test/mcp/_helpers.py
+++ b/posthog/test/mcp/_helpers.py
@@ -32,9 +32,16 @@ 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 = [
+            task
+            for task in list(instr._BACKGROUND_TASKS)
+            if isinstance(task, asyncio.Task)
+            and task.get_loop() is loop
+            and not task.done()
+        ]
         if pending:
             await asyncio.gather(*pending, return_exceptions=True)
     await asyncio.sleep(0)
diff --git a/posthog/test/mcp/test_pending_tasks.py b/posthog/test/mcp/test_pending_tasks.py
new file mode 100644
index 000000000..a9abf09e3
--- /dev/null
+++ b/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]
diff --git a/posthog/test/mcp/test_review_fixes.py b/posthog/test/mcp/test_review_fixes.py
index 584f72496..f5f975728 100644
--- a/posthog/test/mcp/test_review_fixes.py
+++ b/posthog/test/mcp/test_review_fixes.py
@@ -200,8 +200,9 @@ async def slow_capture():
         await asyncio.sleep(0.05)
         done.append(1)
 
-    instr.fire_and_forget(slow_capture())
-    instr.drain_pending_sync(timeout=2)
+    owner = object()
+    instr.fire_and_forget(slow_capture(), owner)
+    instr.drain_pending_sync(owner, timeout=2)
 
     assert done == [1]
 
@@ -311,7 +312,7 @@ def test_posthogmcp_can_disable_exception_fanout():
     client.capture = lambda event, **kw: captured.append({"event": event, **kw})
 
     client.capture_tool_call("boom", is_error=True, error="kaboom")
-    instr.drain_pending_sync(timeout=2)
+    instr.drain_pending_sync(client, timeout=2)
 
     names = [c["event"] for c in captured]
     assert "$mcp_tool_call" in names

From 8094c344f2818e432573f6ff1898fc2ec1ebff04 Mon Sep 17 00:00:00 2001
From: Manoel Aranda Neto <marandaneto@gmail.com>
Date: Thu, 30 Jul 2026 18:49:39 +0200
Subject: [PATCH 2/4] test(mcp): await background capture futures

---
 posthog/test/mcp/_helpers.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/posthog/test/mcp/_helpers.py b/posthog/test/mcp/_helpers.py
index bb5cdf503..21afc62a9 100644
--- a/posthog/test/mcp/_helpers.py
+++ b/posthog/test/mcp/_helpers.py
@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import concurrent.futures
 
 
 class FakeClient:
@@ -35,13 +36,12 @@ async def flush_background():
     loop = asyncio.get_running_loop()
     for _ in range(10):
         await asyncio.sleep(0)
-        pending = [
-            task
-            for task in list(instr._BACKGROUND_TASKS)
-            if isinstance(task, asyncio.Task)
-            and task.get_loop() is loop
-            and not task.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)

From c719044d7edb8eb5c7d1de3723ccd5b575cd181c Mon Sep 17 00:00:00 2001
From: Manoel Aranda Neto <marandaneto@gmail.com>
Date: Tue, 4 Aug 2026 10:47:27 +0200
Subject: [PATCH 3/4] address mcp pending task review feedback

---
 posthog/mcp/__init__.py          |  6 ++++--
 posthog/mcp/_instrumentation.py  | 36 +++++++++++++-------------------
 posthog/test/mcp/test_fastmcp.py | 26 ++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 25 deletions(-)

diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py
index d4778ccb1..dbab60cdb 100644
--- a/posthog/mcp/__init__.py
+++ b/posthog/mcp/__init__.py
@@ -128,8 +128,6 @@ async def flush(self) -> None:
         """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.)"""
-        if self._key is None:
-            return
         data = get_server_tracking_data(self._key)
         if data is not None:
             await drain_pending(data)
@@ -142,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:
diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py
index 17c24d545..b99d5defc 100644
--- a/posthog/mcp/_instrumentation.py
+++ b/posthog/mcp/_instrumentation.py
@@ -13,7 +13,7 @@
 import concurrent.futures
 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
@@ -25,12 +25,9 @@
 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.
-# 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] = {}
+# 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
@@ -52,19 +49,15 @@ def _get_background_loop() -> asyncio.AbstractEventLoop:
     return _bg_loop
 
 
-def _track_task(task: Any, owner: Any, loop: asyncio.AbstractEventLoop) -> None:
+def _track_task(task: Any, owner: Any) -> None:
     with _tasks_lock:
-        _BACKGROUND_TASKS.add(task)
-        _TASK_OWNERS[task] = owner
-        _TASK_LOOPS[task] = loop
+        _BACKGROUND_TASKS[task] = owner
     task.add_done_callback(_on_task_done)
 
 
 def _on_task_done(task: Any) -> None:
     with _tasks_lock:
-        _BACKGROUND_TASKS.discard(task)
-        _TASK_OWNERS.pop(task, None)
-        _TASK_LOOPS.pop(task, None)
+        _BACKGROUND_TASKS.pop(task, None)
     try:
         if not task.cancelled() and task.exception() is not None:
             log(f"background capture task failed: {task.exception()}")
@@ -91,11 +84,11 @@ def fire_and_forget(
     if background or running_loop is None:
         loop = _get_background_loop()
         future = asyncio.run_coroutine_threadsafe(coro, loop)
-        _track_task(future, owner, loop)
+        _track_task(future, owner)
         return
 
     task = running_loop.create_task(coro)
-    _track_task(task, owner, running_loop)
+    _track_task(task, owner)
 
 
 async def drain_pending(owner: Any) -> None:
@@ -104,10 +97,10 @@ async def drain_pending(owner: Any) -> None:
     with _tasks_lock:
         tasks = [
             task
-            for task in _BACKGROUND_TASKS
-            if _TASK_OWNERS.get(task) is owner
-            and _TASK_LOOPS.get(task) is loop
+            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:
@@ -119,9 +112,8 @@ def drain_pending_sync(owner: Any, timeout: Optional[float] = None) -> None:
     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
+            for task, task_owner in _BACKGROUND_TASKS.items()
+            if task_owner is owner
             and isinstance(task, concurrent.futures.Future)
             and not task.done()
         ]
diff --git a/posthog/test/mcp/test_fastmcp.py b/posthog/test/mcp/test_fastmcp.py
index 87b46ea46..1784835ff 100644
--- a/posthog/test/mcp/test_fastmcp.py
+++ b/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()

From e7d5c41bc45ad0dd40dd78daec603694e724755a Mon Sep 17 00:00:00 2001
From: Manoel Aranda Neto <marandaneto@gmail.com>
Date: Tue, 4 Aug 2026 11:06:47 +0200
Subject: [PATCH 4/4] fix(mcp): reset pending task state after fork

---
 posthog/mcp/_instrumentation.py               |  5 +++--
 posthog/test/mcp/test_instrumentation_fork.py | 11 +++++++----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py
index 79e4719f7..ff5380abd 100644
--- a/posthog/mcp/_instrumentation.py
+++ b/posthog/mcp/_instrumentation.py
@@ -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()
 
diff --git a/posthog/test/mcp/test_instrumentation_fork.py b/posthog/test/mcp/test_instrumentation_fork.py
index bde5f079a..5e3d02f34 100644
--- a/posthog/test/mcp/test_instrumentation_fork.py
+++ b/posthog/test/mcp/test_instrumentation_fork.py
@@ -22,7 +22,8 @@ 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
@@ -30,6 +31,7 @@ async def pending_parent_capture():
 
     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"