fix(async-streaming): harden context preservation

Author: hassiebp

langfuse/_client/observe.py

@@ -535,6 +535,20 @@ def _wrap_async_generator_result(
 observe = _decorator.observe
 
 
+def _get_generator_output(
+    items: List[Any],
+    transform_fn: Optional[Callable[[Iterable], str]],
+) -> Any:
+    output: Any = items
+
+    if transform_fn is not None:
+        output = transform_fn(items)
+    elif all(isinstance(item, str) for item in items):
+        output = "".join(items)
+
+    return output
+
+
 class _ContextPreservedSyncGeneratorWrapper:
     """Sync generator wrapper that ensures each iteration runs in preserved context."""
 
@@ -560,9 +574,17 @@ def __init__(
         self.items: List[Any] = []
         self.span = span
         self.transform_fn = transform_fn
+        self._finalized = False
 
-    def __iter__(self) -> "_ContextPreservedSyncGeneratorWrapper":
-        return self
+    def __iter__(self) -> Generator[Any, None, None]:
+        try:
+            while True:
+                yield self.__next__()
+        except StopIteration:
+            return
+        finally:
+            if not self._finalized:
+                self.close()
 
     def __next__(self) -> Any:
         try:
@@ -573,25 +595,65 @@ def __next__(self) -> Any:
             return item
 
         except StopIteration:
-            # Handle output and span cleanup when generator is exhausted
-            output: Any = self.items
+            self._finalize()
+            raise  # Re-raise StopIteration
 
-            if self.transform_fn is not None:
-                output = self.transform_fn(self.items)
+        except (Exception, asyncio.CancelledError) as e:
+            self._finalize(error=e)
+            raise
 
-            elif all(isinstance(item, str) for item in self.items):
-                output = "".join(self.items)
+    def close(self) -> None:
+        if self._finalized:
+            return
 
-            self.span.update(output=output).end()
+        try:
+            close_method = getattr(self.generator, "close", None)
+            if callable(close_method):
+                self.context.run(close_method)
+        except (Exception, asyncio.CancelledError) as e:
+            self._finalize(error=e)
+            raise
 
-            raise  # Re-raise StopIteration
+        self._finalize()
 
+    def throw(self, typ: Any, val: Any = None, tb: Any = None) -> Any:
+        throw_method = getattr(self.generator, "throw", None)
+        if not callable(throw_method):
+            raise AttributeError("Wrapped generator does not support throw()")
+
+        try:
+            if tb is not None:
+                item = self.context.run(throw_method, typ, val, tb)
+            elif val is not None:
+                item = self.context.run(throw_method, typ, val)
+            else:
+                item = self.context.run(throw_method, typ)
+
+            self.items.append(item)
+
+            return item
+        except StopIteration:
+            self._finalize()
+            raise
         except (Exception, asyncio.CancelledError) as e:
+            self._finalize(error=e)
+            raise
+
+    def _finalize(self, error: Optional[BaseException] = None) -> None:
+        if self._finalized:
+            return
+
+        self._finalized = True
+
+        if error is not None:
             self.span.update(
-                level="ERROR", status_message=str(e) or type(e).__name__
+                level="ERROR", status_message=str(error) or type(error).__name__
             ).end()
+            return
 
-            raise
+        self.span.update(
+            output=_get_generator_output(self.items, self.transform_fn)
+        ).end()
 
 
 class _ContextPreservedAsyncGeneratorWrapper:
@@ -619,43 +681,93 @@ def __init__(
         self.items: List[Any] = []
         self.span = span
         self.transform_fn = transform_fn
+        self._finalized = False
 
     def __aiter__(self) -> "_ContextPreservedAsyncGeneratorWrapper":
         return self
 
     async def __anext__(self) -> Any:
         try:
             # Run the generator's __anext__ in the preserved context
-            try:
-                # Python 3.10+ approach with context parameter
-                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__()
+            item = await self._run_in_preserved_context(self.generator.__anext__)
 
             self.items.append(item)
 
             return item
 
         except StopAsyncIteration:
-            # Handle output and span cleanup when generator is exhausted
-            output: Any = self.items
+            self._finalize()
+            raise  # Re-raise StopAsyncIteration
+        except (Exception, asyncio.CancelledError) as e:
+            self._finalize(error=e)
+            raise
 
-            if self.transform_fn is not None:
-                output = self.transform_fn(self.items)
+    async def close(self) -> None:
+        await self.aclose()
 
-            elif all(isinstance(item, str) for item in self.items):
-                output = "".join(self.items)
+    async def aclose(self) -> None:
+        if self._finalized:
+            return
 
-            self.span.update(output=output).end()
+        try:
+            close_method = getattr(self.generator, "aclose", None)
+            if callable(close_method):
+                await self._run_in_preserved_context(close_method)
+        except (Exception, asyncio.CancelledError) as e:
+            self._finalize(error=e)
+            raise
 
-            raise  # Re-raise StopAsyncIteration
+        self._finalize()
+
+    async def athrow(self, typ: Any, val: Any = None, tb: Any = None) -> Any:
+        throw_method = getattr(self.generator, "athrow", None)
+        if not callable(throw_method):
+            raise AttributeError("Wrapped async generator does not support athrow()")
+
+        try:
+            if tb is not None:
+                item = await self._run_in_preserved_context(
+                    lambda: throw_method(typ, val, tb)
+                )
+            elif val is not None:
+                item = await self._run_in_preserved_context(
+                    lambda: throw_method(typ, val)
+                )
+            else:
+                item = await self._run_in_preserved_context(lambda: throw_method(typ))
+
+            self.items.append(item)
+
+            return item
+        except StopAsyncIteration:
+            self._finalize()
+            raise
         except (Exception, asyncio.CancelledError) as e:
+            self._finalize(error=e)
+            raise
+
+    async def _run_in_preserved_context(self, factory: Callable[[], Any]) -> Any:
+        awaitable = self.context.run(factory)
+
+        try:
+            task = asyncio.create_task(awaitable, context=self.context)  # type: ignore[call-arg]
+        except TypeError:
+            task = self.context.run(asyncio.create_task, awaitable)
+
+        return await task
+
+    def _finalize(self, error: Optional[BaseException] = None) -> None:
+        if self._finalized:
+            return
+
+        self._finalized = True
+
+        if error is not None:
             self.span.update(
-                level="ERROR", status_message=str(e) or type(e).__name__
+                level="ERROR", status_message=str(error) or type(error).__name__
             ).end()
+            return
 
-            raise
+        self.span.update(
+            output=_get_generator_output(self.items, self.transform_fn)
+        ).end()

langfuse/openai.py

@@ -1010,6 +1010,7 @@ def __init__(
         self.response = response
         self.generation = generation
         self.completion_start_time: Optional[datetime] = None
+        self._finalized = False
 
     def __iter__(self) -> Any:
         try:
@@ -1039,12 +1040,31 @@ def __next__(self) -> Any:
             raise
 
     def __enter__(self) -> Any:
-        return self.__iter__()
+        return self
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        pass
+        self.close()
+
+    def close(self) -> None:
+        if self._finalized:
+            return
+
+        close_method = getattr(self.response, "close", None)
+        if callable(close_method):
+            try:
+                close_method()
+            finally:
+                self._finalize()
+            return
+
+        self._finalize()
 
     def _finalize(self) -> None:
+        if self._finalized:
+            return
+
+        self._finalized = True
+
         try:
             model, completion, usage, metadata = (
                 _extract_streamed_response_api_response(self.items)
@@ -1081,6 +1101,7 @@ def __init__(
         self.response = response
         self.generation = generation
         self.completion_start_time: Optional[datetime] = None
+        self._finalized = False
 
     async def __aiter__(self) -> Any:
         try:
@@ -1110,12 +1131,17 @@ async def __anext__(self) -> Any:
             raise
 
     async def __aenter__(self) -> Any:
-        return self.__aiter__()
+        return self
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        pass
+        await self.aclose()
 
     async def _finalize(self) -> None:
+        if self._finalized:
+            return
+
+        self._finalized = True
+
         try:
             model, completion, usage, metadata = (
                 _extract_streamed_response_api_response(self.items)
@@ -1142,11 +1168,40 @@ async def close(self) -> None:
 
         Automatically called if the response body is read to completion.
         """
-        await self.response.close()
+        if self._finalized:
+            return
+
+        close_method = getattr(self.response, "close", None)
+        if callable(close_method):
+            try:
+                await close_method()
+            finally:
+                await self._finalize()
+            return
+
+        await self._finalize()
 
     async def aclose(self) -> None:
         """Close the response and release the connection.
 
         Automatically called if the response body is read to completion.
         """
-        await self.response.aclose()
+        if self._finalized:
+            return
+
+        close_method = getattr(self.response, "aclose", None)
+        if callable(close_method):
+            try:
+                await close_method()
+            finally:
+                await self._finalize()
+        else:
+            close_method = getattr(self.response, "close", None)
+            if callable(close_method):
+                try:
+                    await close_method()
+                finally:
+                    await self._finalize()
+                return
+
+            await self._finalize()