From 6f42542755f31feeeaa98e4178ca43bc4790f025 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 13 May 2026 08:30:18 -0400 Subject: [PATCH 1/3] fix: support scoped async functions --- .sampo/changesets/chivalrous-baron-tapio.md | 5 +++ posthog/contexts.py | 17 +++++++++ posthog/test/test_contexts.py | 42 +++++++++++++++++++++ 3 files changed, 64 insertions(+) create mode 100644 .sampo/changesets/chivalrous-baron-tapio.md diff --git a/.sampo/changesets/chivalrous-baron-tapio.md b/.sampo/changesets/chivalrous-baron-tapio.md new file mode 100644 index 00000000..9524b28a --- /dev/null +++ b/.sampo/changesets/chivalrous-baron-tapio.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: patch +--- + +Fix scoped context support for async functions diff --git a/posthog/contexts.py b/posthog/contexts.py index 326916da..56bed4a9 100644 --- a/posthog/contexts.py +++ b/posthog/contexts.py @@ -391,12 +391,29 @@ def process_payment(payment_id): # and then re-raised some_risky_function() + # When stacking decorators, the posthog.scoped decorator must be + # closest to the function. For example, with FastAPI middleware: + @app.middleware("http") + @posthog.scoped() + async def middleware(request, call_next): + return await call_next(request) + Category: Contexts """ def decorator(func: F) -> F: from functools import wraps + from inspect import iscoroutinefunction + + if iscoroutinefunction(func): + + @wraps(func) + async def async_wrapper(*args, **kwargs): + with new_context(fresh=fresh, capture_exceptions=capture_exceptions): + return await func(*args, **kwargs) + + return cast(F, async_wrapper) @wraps(func) def wrapper(*args, **kwargs): diff --git a/posthog/test/test_contexts.py b/posthog/test/test_contexts.py index 3ff6fddd..7c0a4a73 100644 --- a/posthog/test/test_contexts.py +++ b/posthog/test/test_contexts.py @@ -1,3 +1,4 @@ +import asyncio import unittest from unittest.mock import patch @@ -231,3 +232,44 @@ def function_with_context(): # Context should be cleared after function execution assert get_context_distinct_id() is None assert get_context_session_id() is None + + def test_scoped_decorator_async_function(self): + @scoped() + async def async_function_with_context(): + identify_context("user_async") + set_context_session("session_async") + return get_context_distinct_id(), get_context_session_id() + + distinct_id, session_id = asyncio.run(async_function_with_context()) + assert distinct_id == "user_async" + assert session_id == "session_async" + + # Context should be cleared after function execution + assert get_context_distinct_id() is None + assert get_context_session_id() is None + + @patch("posthog.capture_exception") + def test_scoped_decorator_async_exception(self, mock_capture): + test_exception = ValueError("Test async exception") + + def check_context_on_capture(exception, **kwargs): + # Assert context IDs are available when capture_exception is called + assert get_context_distinct_id() == "user_async" + assert get_context_session_id() == "session_async" + + mock_capture.side_effect = check_context_on_capture + + @scoped() + async def failing_async_function(): + identify_context("user_async") + set_context_session("session_async") + raise test_exception + + with self.assertRaises(ValueError): + asyncio.run(failing_async_function()) + + mock_capture.assert_called_once_with(test_exception) + + # Context should be cleared after function execution + assert get_context_distinct_id() is None + assert get_context_session_id() is None From 8c1d72c3f7dfe8a3af3c4d8ca46723e575ae5b69 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 13 May 2026 08:51:18 -0400 Subject: [PATCH 2/3] test: cover concurrent scoped async isolation --- posthog/test/test_contexts.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/posthog/test/test_contexts.py b/posthog/test/test_contexts.py index 7c0a4a73..1720ef78 100644 --- a/posthog/test/test_contexts.py +++ b/posthog/test/test_contexts.py @@ -248,6 +248,33 @@ async def async_function_with_context(): assert get_context_distinct_id() is None assert get_context_session_id() is None + def test_scoped_decorator_async_concurrent_context_isolation(self): + first_ready = asyncio.Event() + second_ready = asyncio.Event() + first_checked = asyncio.Event() + + @scoped() + async def first(): + identify_context("user_1") + first_ready.set() + await second_ready.wait() + distinct_id = get_context_distinct_id() + first_checked.set() + return distinct_id + + @scoped() + async def second(): + await first_ready.wait() + identify_context("user_2") + second_ready.set() + await first_checked.wait() + return get_context_distinct_id() + + async def run(): + return await asyncio.wait_for(asyncio.gather(first(), second()), timeout=1) + + assert asyncio.run(run()) == ["user_1", "user_2"] + @patch("posthog.capture_exception") def test_scoped_decorator_async_exception(self, mock_capture): test_exception = ValueError("Test async exception") From 8b612c98c5872535a4cc8785d76d547ab14ee463 Mon Sep 17 00:00:00 2001 From: Manoel Aranda Neto Date: Wed, 13 May 2026 08:56:17 -0400 Subject: [PATCH 3/3] test: parameterize scoped async cases --- posthog/test/test_contexts.py | 104 +++++++++++++++------------------- 1 file changed, 47 insertions(+), 57 deletions(-) diff --git a/posthog/test/test_contexts.py b/posthog/test/test_contexts.py index 1720ef78..a2f81fba 100644 --- a/posthog/test/test_contexts.py +++ b/posthog/test/test_contexts.py @@ -88,8 +88,6 @@ def successful_function(x, y): @patch("posthog.capture_exception") def test_scoped_decorator_exception(self, mock_capture): - test_exception = ValueError("Test exception") - def check_context_on_capture(exception, **kwargs): # Assert tags are available when capture_exception is called current_tags = get_tags() @@ -97,20 +95,40 @@ def check_context_on_capture(exception, **kwargs): mock_capture.side_effect = check_context_on_capture - @scoped() - def failing_function(): - tag("important_context", "value") - raise test_exception + for name, is_async in [("sync", False), ("async", True)]: + with self.subTest(name=name): + test_exception = ValueError(f"Test {name} exception") - # Function should raise the exception - with self.assertRaises(ValueError): - failing_function() + if is_async: - # Verify capture_exception was called - mock_capture.assert_called_once_with(test_exception) + @scoped() + async def failing_function(): + tag("important_context", "value") + raise test_exception - # Context should be cleared after function execution - assert get_tags() == {} + def run(): + return asyncio.run(failing_function()) + + else: + + @scoped() + def failing_function(): + tag("important_context", "value") + raise test_exception + + run = failing_function + + # Function should raise the exception + with self.assertRaises(ValueError): + run() + + # Verify capture_exception was called + mock_capture.assert_called_once_with(test_exception) + + # Context should be cleared after function execution + assert get_tags() == {} + + mock_capture.reset_mock() @patch("posthog.capture_exception") def test_new_context_exception_handling(self, mock_capture): @@ -220,33 +238,31 @@ def test_child_tags_override_parent_tags_in_non_fresh_context(self): def test_scoped_decorator_with_context_ids(self): @scoped() - def function_with_context(): + def sync_function_with_context(): identify_context("user456") set_context_session("session789") return get_context_distinct_id(), get_context_session_id() - distinct_id, session_id = function_with_context() - assert distinct_id == "user456" - assert session_id == "session789" - - # Context should be cleared after function execution - assert get_context_distinct_id() is None - assert get_context_session_id() is None - - def test_scoped_decorator_async_function(self): @scoped() async def async_function_with_context(): - identify_context("user_async") - set_context_session("session_async") + identify_context("user456") + set_context_session("session789") return get_context_distinct_id(), get_context_session_id() - distinct_id, session_id = asyncio.run(async_function_with_context()) - assert distinct_id == "user_async" - assert session_id == "session_async" + cases = [ + ("sync", sync_function_with_context, lambda func: func()), + ("async", async_function_with_context, lambda func: asyncio.run(func())), + ] - # Context should be cleared after function execution - assert get_context_distinct_id() is None - assert get_context_session_id() is None + for name, func, run in cases: + with self.subTest(name=name): + distinct_id, session_id = run(func) + assert distinct_id == "user456" + assert session_id == "session789" + + # Context should be cleared after function execution + assert get_context_distinct_id() is None + assert get_context_session_id() is None def test_scoped_decorator_async_concurrent_context_isolation(self): first_ready = asyncio.Event() @@ -274,29 +290,3 @@ async def run(): return await asyncio.wait_for(asyncio.gather(first(), second()), timeout=1) assert asyncio.run(run()) == ["user_1", "user_2"] - - @patch("posthog.capture_exception") - def test_scoped_decorator_async_exception(self, mock_capture): - test_exception = ValueError("Test async exception") - - def check_context_on_capture(exception, **kwargs): - # Assert context IDs are available when capture_exception is called - assert get_context_distinct_id() == "user_async" - assert get_context_session_id() == "session_async" - - mock_capture.side_effect = check_context_on_capture - - @scoped() - async def failing_async_function(): - identify_context("user_async") - set_context_session("session_async") - raise test_exception - - with self.assertRaises(ValueError): - asyncio.run(failing_async_function()) - - mock_capture.assert_called_once_with(test_exception) - - # Context should be cleared after function execution - assert get_context_distinct_id() is None - assert get_context_session_id() is None