fix(observe): preserve async wrapper context on py310

Author: hassiebp

langfuse/_client/observe.py

@@ -701,7 +701,7 @@ async def aclose(self) -> None:
                     context=self.context,
                 )  # type: ignore
             except TypeError:
-                await self.generator.aclose()
+                await self.context.run(asyncio.create_task, self.generator.aclose())
         except (Exception, asyncio.CancelledError) as error:
             self._finalize_with_error(error)
             raise
@@ -718,14 +718,17 @@ async def __anext__(self) -> Any:
         try:
             # Run the generator's __anext__ in the preserved context
             try:
-                # Python 3.10+ approach with context parameter
+                # Python 3.11+ approach with explicit task context
                 item = await asyncio.create_task(
                     self.generator.__anext__(),  # type: ignore
                     context=self.context,
                 )  # type: ignore
             except TypeError:
-                # Python < 3.10 fallback - context parameter not supported
-                item = await self.generator.__anext__()
+                # Python 3.10 fallback - create the task inside the preserved context.
+                item = await self.context.run(
+                    asyncio.create_task,
+                    self.generator.__anext__(),  # type: ignore
+                )
 
             if self.capture_output:
                 self.items.append(item)

tests/unit/test_observe.py

@@ -247,6 +247,49 @@ async def generator() -> AsyncGenerator[str, None]:
     assert span.ended == 1
 
 
+@pytest.mark.asyncio
+async def test_async_generator_wrapper_fallback_preserves_context(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    marker = contextvars.ContextVar("marker", default="ambient")
+    seen: list[str] = []
+    original_create_task = asyncio.create_task
+
+    def create_task_with_type_error(*args: Any, **kwargs: Any) -> asyncio.Task[Any]:
+        if "context" in kwargs:
+            raise TypeError("context argument unsupported")
+
+        return original_create_task(*args, **kwargs)
+
+    monkeypatch.setattr(asyncio, "create_task", create_task_with_type_error)
+
+    async def generator() -> AsyncGenerator[str, None]:
+        try:
+            yield marker.get()
+            yield "item_1"
+        finally:
+            seen.append(marker.get())
+
+    span = SpanRecorder()
+    context = contextvars.copy_context()
+    context.run(marker.set, "preserved")
+    wrapper = _ContextPreservedAsyncGeneratorWrapper(
+        generator(),
+        context,
+        cast(Any, span),
+        False,
+        None,
+    )
+
+    assert await wrapper.__anext__() == "preserved"
+    marker.set("ambient-now")
+
+    await wrapper.aclose()
+
+    assert seen == ["preserved"]
+    assert span.ended == 1
+
+
 @pytest.mark.asyncio
 async def test_async_generator_wrapper_del_ends_span_when_abandoned() -> None:
     async def generator() -> AsyncGenerator[str, None]: