From 7b7ac6684b395b55bceb29420fc026c2fd6ee0e6 Mon Sep 17 00:00:00 2001 From: Aafimalek <91587101+Aafimalek@users.noreply.github.com> Date: Fri, 6 Nov 2026 13:16:44 +0530 Subject: [PATCH 1/2] fix(observe): support async iterable streaming bodies --- langfuse/_client/observe.py | 36 +++++++++++++++++++++++++----------- tests/unit/test_observe.py | 18 ++++++++++++++---- 2 files changed, 39 insertions(+), 15 deletions(-) diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index 5aa88ce2e..ffcf5254a 100644 --- a/langfuse/_client/observe.py +++ b/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, ) diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index 21ff1418d..33b6f8670 100644 --- a/tests/unit/test_observe.py +++ b/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() From 05f14d211dea70fedb525ebd3f0fb458d0b3683c Mon Sep 17 00:00:00 2001 From: Aafimalek <91587101+Aafimalek@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:09:04 +0530 Subject: [PATCH 2/2] test(observe): cover early close for async iterables --- tests/unit/test_observe.py | 65 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index 33b6f8670..6e2779fe6 100644 --- a/tests/unit/test_observe.py +++ b/tests/unit/test_observe.py @@ -279,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] = []