fix(observe): support async iterable streaming bodies

langfuse/_client/observe.py

@@ -5,7 +5,8 @@
 from functools import wraps
 from typing import (
     Any,
-    AsyncGenerator,
+    AsyncIterable,
+    AsyncIterator,
     Callable,
     Dict,
     Generator,
@@ -497,7 +498,7 @@ def _wrap_async_generator_result(
             LangfuseEmbedding,
             LangfuseGuardrail,
         ],
-        generator: AsyncGenerator,
+        generator: AsyncIterable[Any],
         capture_output: bool,
         transform_to_string: Optional[Callable[[Iterable], str]] = None,
     ) -> Any:
@@ -670,7 +671,7 @@ class _ContextPreservedAsyncGeneratorWrapper:
 
     def __init__(
         self,
-        generator: AsyncGenerator,
+        generator: AsyncIterable[Any],
         context: contextvars.Context,
         span: Union[
             LangfuseSpan,
@@ -687,6 +688,7 @@ def __init__(
         transform_fn: Optional[Callable[[Iterable], str]],
     ) -> None:
         self.generator = generator
+        self._iterator: Optional[AsyncIterator[Any]] = None
         self.context = context
         self.items: List[Any] = []
         self.span = span
@@ -699,14 +701,16 @@ def __aiter__(self) -> "_ContextPreservedAsyncGeneratorWrapper":
         return self
 
     def _generator_never_resumed(self) -> bool:
+        """Check if an async generator is still open after cancellation."""
+        generator = self.generator if self._iterator is None else self._iterator
+        if not inspect.isasyncgen(generator):
+            return False
         try:
-            state = inspect.getasyncgenstate(self.generator)
-        except (AttributeError, TypeError):
+            state = inspect.getasyncgenstate(generator)
+        except AttributeError:
             # getasyncgenstate is Python 3.12+; fall back to the attributes it reads.
-            frame = getattr(self.generator, "ag_frame", None)
-            return frame is not None and not getattr(
-                self.generator, "ag_running", False
-            )
+            frame = getattr(generator, "ag_frame", None)
+            return frame is not None and not getattr(generator, "ag_running", False)
 
         return state in ("AGEN_CREATED", "AGEN_SUSPENDED")
 
@@ -755,7 +759,11 @@ async def aclose(self) -> None:
                 self._finalize()
 
     async def _close_generator(self) -> None:
-        await asyncio.create_task(self.generator.aclose(), context=self.context)
+        """Close the active iterator if it offers async cleanup."""
+        iterator = self.generator if self._iterator is None else self._iterator
+        close = getattr(iterator, "aclose", None)
+        if close is not None:
+            await asyncio.create_task(close(), context=self.context)
 
     async def close(self) -> None:
         await self.aclose()
@@ -766,11 +774,17 @@ def __del__(self) -> None:
         else:
             self._finalize()
 
+    async def _next_item(self) -> Any:
+        """Start iteration lazily and advance it inside the observation context."""
+        if self._iterator is None:
+            self._iterator = aiter(self.generator)
+        return await anext(self._iterator)
+
     async def __anext__(self) -> Any:
         try:
             # Run the generator's __anext__ in the preserved context
             item = await asyncio.create_task(
-                self.generator.__anext__(),  # type: ignore
+                self._next_item(),
                 context=self.context,
             )
 

tests/unit/test_observe.py

@@ -3,7 +3,7 @@
 import gc
 import inspect
 import json
-from typing import Any, AsyncGenerator, Generator, cast
+from typing import Any, AsyncGenerator, AsyncIterable, Generator, cast
 
 import pytest
 
@@ -94,11 +94,16 @@ def body() -> Generator[str, None, None]:
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "use_async_iterable", [False, True], ids=["generator", "iterable"]
+)
 async def test_streaming_response_preserves_context_without_output_capture(
-    langfuse_memory_client: Any, memory_exporter: Any
+    langfuse_memory_client: Any, memory_exporter: Any, use_async_iterable: bool
 ) -> None:
+    """Stream both supported body shapes while retaining the endpoint span."""
+
     class StreamingResponse:
-        def __init__(self, body_iterator: AsyncGenerator[str, None]) -> None:
+        def __init__(self, body_iterator: AsyncIterable[str]) -> None:
             self.body_iterator = body_iterator
 
     @observe(name="stream_step")
@@ -109,9 +114,14 @@ async def body() -> AsyncGenerator[str, None]:
         for index in range(2):
             yield await stream_step(index)
 
+    class AsyncBody:
+        def __aiter__(self) -> AsyncGenerator[str, None]:
+            """Create the body iterator when streaming starts."""
+            return body()
+
     @observe(name="endpoint", capture_output=False)
     async def endpoint() -> StreamingResponse:
-        return StreamingResponse(body())
+        return StreamingResponse(AsyncBody() if use_async_iterable else body())
 
     response = await endpoint()
 
@@ -269,6 +279,71 @@ async def generator() -> AsyncGenerator[str, None]:
     assert span.ended == 1
 
 
+@pytest.mark.asyncio
+async def test_async_iterable_wrapper_closes_returned_iterator_in_context() -> None:
+    """Close the iterator created by the iterable within the saved context."""
+    marker = contextvars.ContextVar("marker", default="ambient")
+    seen: list[str] = []
+
+    class Body:
+        def __aiter__(self) -> AsyncGenerator[str, None]:
+            """Yield a chunk and record the context used for cleanup."""
+
+            async def generate() -> AsyncGenerator[str, None]:
+                try:
+                    yield "item_0"
+                    yield "item_1"
+                finally:
+                    seen.append(marker.get())
+
+            return generate()
+
+    span = SpanRecorder()
+    context = contextvars.copy_context()
+    context.run(marker.set, "preserved")
+    wrapper = _ContextPreservedAsyncGeneratorWrapper(
+        Body(), context, cast(Any, span), False, None
+    )
+
+    assert await anext(wrapper) == "item_0"
+    marker.set("ambient-now")
+    await wrapper.aclose()
+
+    assert seen == ["preserved"]
+    assert span.ended == 1
+
+
+@pytest.mark.asyncio
+async def test_async_iterable_wrapper_closes_without_iterator_aclose() -> None:
+    """Finish the span when the active iterator has no async close method."""
+
+    class Iterator:
+        async def __anext__(self) -> str:
+            """Return a chunk without exposing an aclose method."""
+            return "item_0"
+
+        def __aiter__(self) -> "Iterator":
+            """Return this iterator for async iteration."""
+            return self
+
+    class Body:
+        def __aiter__(self) -> Iterator:
+            """Create an iterator without an aclose method."""
+            return Iterator()
+
+    span = SpanRecorder()
+    wrapper = _ContextPreservedAsyncGeneratorWrapper(
+        Body(), contextvars.copy_context(), cast(Any, span), False, None
+    )
+
+    assert await anext(wrapper) == "item_0"
+    await wrapper.aclose()
+    await wrapper.aclose()
+
+    assert span.ended == 1
+    assert span.updates == []
+
+
 def test_sync_generator_wrapper_close_closes_generator_after_span_ended() -> None:
     marker = contextvars.ContextVar("marker", default="ambient")
     seen: list[str] = []