fix(observe): preserve context when closing wrappers

Author: hassiebp

langfuse/_client/observe.py

@@ -600,7 +600,7 @@ def close(self) -> None:
             return
 
         try:
-            self.generator.close()
+            self.context.run(self.generator.close)
         except (Exception, asyncio.CancelledError) as error:
             self._finalize_with_error(error)
             raise
@@ -695,7 +695,13 @@ async def aclose(self) -> None:
             return
 
         try:
-            await self.generator.aclose()
+            try:
+                await asyncio.create_task(
+                    self.generator.aclose(),
+                    context=self.context,
+                )  # type: ignore
+            except TypeError:
+                await self.generator.aclose()
         except (Exception, asyncio.CancelledError) as error:
             self._finalize_with_error(error)
             raise

tests/unit/test_observe.py

@@ -137,6 +137,37 @@ def generator() -> Generator[str, None, None]:
     assert span.updates == []
 
 
+def test_sync_generator_wrapper_close_preserves_context() -> None:
+    marker = contextvars.ContextVar("marker", default="ambient")
+    seen: list[str] = []
+
+    def generator() -> Generator[str, None, None]:
+        try:
+            yield "item_0"
+            yield "item_1"
+        finally:
+            seen.append(marker.get())
+
+    span = SpanRecorder()
+    context = contextvars.copy_context()
+    context.run(marker.set, "preserved")
+    wrapper = _ContextPreservedSyncGeneratorWrapper(
+        generator(),
+        context,
+        cast(Any, span),
+        False,
+        None,
+    )
+
+    assert next(wrapper) == "item_0"
+    marker.set("ambient-now")
+
+    wrapper.close()
+
+    assert seen == ["preserved"]
+    assert span.ended == 1
+
+
 def test_sync_generator_wrapper_del_ends_span_when_abandoned() -> None:
     def generator() -> Generator[str, None, None]:
         yield "item_0"
@@ -184,6 +215,38 @@ async def generator() -> AsyncGenerator[str, None]:
     assert span.updates == []
 
 
+@pytest.mark.asyncio
+async def test_async_generator_wrapper_aclose_preserves_context() -> None:
+    marker = contextvars.ContextVar("marker", default="ambient")
+    seen: list[str] = []
+
+    async def generator() -> AsyncGenerator[str, None]:
+        try:
+            yield "item_0"
+            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__() == "item_0"
+    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]: