From aeb4f452ffb89ebfd18ca7c7610787365eea5c5f Mon Sep 17 00:00:00 2001
From: Nickita Khylkouski
 <90287684+nickita-khylkouski@users.noreply.github.com>
Date: Thu, 10 Sep 2026 09:04:24 -0700
Subject: [PATCH] fix: clean up async consumer waiters on cancellation

---
 .../changesets/async-consumer-cancellation.md |  5 ++
 posthog/_async_consumer.py                    | 20 ++++---
 posthog/test/test_async_consumer.py           | 55 +++++++++++++++++++
 3 files changed, 71 insertions(+), 9 deletions(-)
 create mode 100644 .sampo/changesets/async-consumer-cancellation.md

diff --git a/.sampo/changesets/async-consumer-cancellation.md b/.sampo/changesets/async-consumer-cancellation.md
new file mode 100644
index 000000000..8ffcf9051
--- /dev/null
+++ b/.sampo/changesets/async-consumer-cancellation.md
@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Clean up queue and flush waiters when the async capture consumer is cancelled.
diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py
index 387caab44..04eeb082b 100644
--- a/posthog/_async_consumer.py
+++ b/posthog/_async_consumer.py
@@ -117,15 +117,17 @@ def request_flush(self) -> None:
     async def _get_or_flush(self, timeout: float) -> tuple[Any, bool]:
         get_task = asyncio.create_task(self.queue.get())
         flush_task = asyncio.create_task(self._flush_event.wait())
-        done, pending = await asyncio.wait(
-            {get_task, flush_task},
-            timeout=timeout,
-            return_when=asyncio.FIRST_COMPLETED,
-        )
-        for task in pending:
-            task.cancel()
-        if pending:
-            await asyncio.gather(*pending, return_exceptions=True)
+        try:
+            done, _ = await asyncio.wait(
+                {get_task, flush_task},
+                timeout=timeout,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        finally:
+            for task in (get_task, flush_task):
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(get_task, flush_task, return_exceptions=True)
 
         if get_task in done:
             return get_task.result(), False
diff --git a/posthog/test/test_async_consumer.py b/posthog/test/test_async_consumer.py
index 290ce687d..9bb3f682f 100644
--- a/posthog/test/test_async_consumer.py
+++ b/posthog/test/test_async_consumer.py
@@ -152,3 +152,58 @@ async def test_request_stops_after_configured_retry_limit():
 
     assert batch_post.await_count == 3
     assert [call.args[0] for call in sleep.await_args_list] == [1, 2]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("run_worker", [False, True], ids=["wait", "worker"])
+async def test_get_or_flush_cancels_waiters_on_cancellation(run_worker):
+    consumer = make_consumer(retries=0)
+    consumer.flush_interval = 60
+    wait_started = asyncio.Event()
+    waiters = []
+    real_wait = asyncio.wait
+
+    async def observe_wait(tasks, **kwargs):
+        waiters.extend(tasks)
+        wait_started.set()
+        return await real_wait(tasks, **kwargs)
+
+    with mock.patch("posthog._async_consumer.asyncio.wait", side_effect=observe_wait):
+        task = asyncio.create_task(
+            consumer.run() if run_worker else consumer._get_or_flush(60)
+        )
+        try:
+            await wait_started.wait()
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+            assert len(waiters) == 2
+            assert all(waiter.cancelled() for waiter in waiters)
+        finally:
+            task.cancel()
+            for waiter in waiters:
+                waiter.cancel()
+            await asyncio.gather(task, *waiters, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("queued", "flush"), [(True, False), (False, True), (False, False), (True, True)]
+)
+async def test_get_or_flush_preserves_results_and_cleans_up_waiters(queued, flush):
+    consumer = make_consumer(retries=0)
+    event = {"event": "test"}
+    if queued:
+        consumer.queue.put_nowait(event)
+    if flush:
+        consumer.request_flush()
+    tasks_before = asyncio.all_tasks()
+
+    result = await consumer._get_or_flush(60 if queued or flush else 0)
+
+    assert result == (event if queued else None, flush and not queued)
+    assert consumer._flush_event.is_set() == (queued and flush)
+    assert not (asyncio.all_tasks() - tasks_before)
+    if queued:
+        consumer.queue.task_done()