From 4c0dc0b88b8428e1ffc3f5646dcef03493829943 Mon Sep 17 00:00:00 2001 From: David Traina <44659830+DavidTraina@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:41:29 -0500 Subject: [PATCH 1/4] fix(observe): defer span end when a cancellation never resumed the generator --- langfuse/_client/observe.py | 59 +++++++++++++--- tests/unit/test_observe.py | 133 ++++++++++++++++++++++++++++++++++++ 2 files changed, 182 insertions(+), 10 deletions(-) diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index bd0a3edee..7472cf8ca 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -616,6 +616,12 @@ def _finalize_with_error(self, error: BaseException) -> None: def close(self) -> None: if self._span_ended: + # The span can end while the generator is still suspended (e.g. + # an error surfaced from __next__ without resuming the generator). + # Still close the generator so its cleanup runs deterministically + # in the preserved context instead of at GC time under an + # arbitrary ambient context. + self.context.run(self.generator.close) return try: @@ -678,10 +684,19 @@ def __init__( self.capture_output = capture_output self.transform_fn = transform_fn self._span_ended = False + self._pending_error: Optional[BaseException] = None def __aiter__(self) -> "_ContextPreservedAsyncGeneratorWrapper": return self + def _generator_never_resumed(self) -> bool: + try: + state = inspect.getasyncgenstate(self.generator) + except TypeError: + return False + + return state in ("AGEN_CREATED", "AGEN_SUSPENDED") + def _finalize(self) -> None: if self._span_ended: return @@ -711,27 +726,41 @@ def _finalize_with_error(self, error: BaseException) -> None: async def aclose(self) -> None: if self._span_ended: + # The span can end while the generator is still suspended. Still + # close the generator so its cleanup runs deterministically in + # the preserved context instead of at GC time under an arbitrary + # ambient context. + await self._close_generator() return try: - try: - await asyncio.create_task( - self.generator.aclose(), - context=self.context, - ) # type: ignore - except TypeError: - await self.context.run(asyncio.create_task, self.generator.aclose()) + await self._close_generator() except (Exception, asyncio.CancelledError) as error: self._finalize_with_error(error) raise else: - self._finalize() + if self._pending_error is not None: + self._finalize_with_error(self._pending_error) + else: + self._finalize() + + async def _close_generator(self) -> None: + try: + await asyncio.create_task( + self.generator.aclose(), + context=self.context, + ) # type: ignore + except TypeError: + await self.context.run(asyncio.create_task, self.generator.aclose()) async def close(self) -> None: await self.aclose() def __del__(self) -> None: - self._finalize() + if self._pending_error is not None: + self._finalize_with_error(self._pending_error) + else: + self._finalize() async def __anext__(self) -> Any: try: @@ -757,6 +786,16 @@ async def __anext__(self) -> Any: except StopAsyncIteration: self._finalize() raise # Re-raise StopAsyncIteration - except (Exception, asyncio.CancelledError) as e: + except asyncio.CancelledError as e: + if self._generator_never_resumed(): + # The cancellation was delivered before the inner task resumed + # the generator. Defer the span end so aclose() can run the + # generator's cleanup (including any span updates it makes) + # before the span is finalized with this error. + self._pending_error = e + raise + self._finalize_with_error(e) + raise + except Exception as e: self._finalize_with_error(e) raise diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index 5527be9b9..ea5c49bac 100644 --- a/tests/unit/test_observe.py +++ b/tests/unit/test_observe.py @@ -1,6 +1,7 @@ import asyncio import contextvars import gc +import inspect import json import sys from typing import Any, AsyncGenerator, Generator, cast @@ -270,6 +271,138 @@ async def generator() -> AsyncGenerator[str, None]: assert span.ended == 1 +def test_sync_generator_wrapper_close_closes_generator_after_span_ended() -> 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" + + # __next__ raising without resuming the generator (here: re-entering the + # preserved context) ends the span while the generator is still suspended. + with pytest.raises(RuntimeError): + context.run(lambda: next(wrapper)) + + assert span.ended == 1 + assert seen == [] + + marker.set("ambient-now") + wrapper.close() + + assert seen == ["preserved"] + assert span.ended == 1 + + +@pytest.mark.asyncio +async def test_async_generator_wrapper_aclose_closes_generator_after_span_ended() -> ( + 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" + + # Span ends while the generator is still suspended. + wrapper._finalize_with_error(RuntimeError("ended early")) + assert span.ended == 1 + assert seen == [] + + marker.set("ambient-now") + await wrapper.aclose() + + assert seen == ["preserved"] + assert span.ended == 1 + + +@pytest.mark.asyncio +async def test_async_generator_wrapper_defers_span_end_for_unresumed_cancel() -> 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") + raw = generator() + wrapper = _ContextPreservedAsyncGeneratorWrapper( + raw, + context, + cast(Any, span), + False, + None, + ) + + async def consume() -> None: + async for _ in wrapper: + await asyncio.sleep(0) + + consumer = asyncio.create_task(consume()) + for _ in range(4): + await asyncio.sleep(0) + # The cancel lands between chunks, before the inner __anext__ task's + # first step: the generator is never resumed. + asyncio.get_running_loop().call_soon(consumer.cancel) + with pytest.raises(asyncio.CancelledError): + await consumer + + assert inspect.getasyncgenstate(raw) == "AGEN_SUSPENDED" + # Span end is deferred so the generator's cleanup can still update it. + assert span.ended == 0 + + marker.set("ambient-now") + await wrapper.aclose() + + assert inspect.getasyncgenstate(raw) == "AGEN_CLOSED" + assert seen == ["preserved"] + assert span.ended == 1 + # The retained cancellation still finalizes the span as an error. + assert span.updates[-1] == { + "level": "ERROR", + "status_message": "CancelledError", + } + + @pytest.mark.asyncio async def test_async_generator_wrapper_fallback_preserves_context( monkeypatch: pytest.MonkeyPatch, From 545ddc745cd5e87fbd176cc0fcf9a7ece29fe22e Mon Sep 17 00:00:00 2001 From: David Traina <44659830+DavidTraina@users.noreply.github.com> Date: Thu, 30 Jul 2026 01:00:35 -0500 Subject: [PATCH 2/4] fix(observe): only treat create_task signature errors as fallback trigger --- langfuse/_client/observe.py | 9 +++++++-- tests/unit/test_observe.py | 26 ++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index 7472cf8ca..58f7944c1 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -746,12 +746,17 @@ async def aclose(self) -> None: async def _close_generator(self) -> None: try: - await asyncio.create_task( + close_task = asyncio.create_task( self.generator.aclose(), context=self.context, ) # type: ignore except TypeError: - await self.context.run(asyncio.create_task, self.generator.aclose()) + # Python 3.10: create_task has no context parameter. Only the + # create_task call is guarded so a TypeError raised from the + # generator's own cleanup is not mistaken for it. + close_task = self.context.run(asyncio.create_task, self.generator.aclose()) + + await close_task async def close(self) -> None: await self.aclose() diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index ea5c49bac..2d37b82ed 100644 --- a/tests/unit/test_observe.py +++ b/tests/unit/test_observe.py @@ -403,6 +403,32 @@ async def consume() -> None: } +@pytest.mark.asyncio +async def test_async_generator_wrapper_aclose_propagates_cleanup_type_error() -> None: + async def generator() -> AsyncGenerator[str, None]: + try: + yield "item_0" + finally: + raise TypeError("cleanup failed") + + span = SpanRecorder() + wrapper = _ContextPreservedAsyncGeneratorWrapper( + generator(), + contextvars.copy_context(), + cast(Any, span), + False, + None, + ) + + assert await wrapper.__anext__() == "item_0" + + with pytest.raises(TypeError, match="cleanup failed"): + await wrapper.aclose() + + assert span.ended == 1 + assert span.updates[-1] == {"level": "ERROR", "status_message": "cleanup failed"} + + @pytest.mark.asyncio async def test_async_generator_wrapper_fallback_preserves_context( monkeypatch: pytest.MonkeyPatch, From f955baa58941b60db917aced4794bbf179b59dd7 Mon Sep 17 00:00:00 2001 From: David Traina <44659830+DavidTraina@users.noreply.github.com> Date: Thu, 30 Jul 2026 01:07:08 -0500 Subject: [PATCH 3/4] fix(observe): support generator state check on python < 3.12 --- langfuse/_client/observe.py | 28 +++++++----------- tests/unit/test_observe.py | 58 +++++++++++++++++++++++++++++++++---- 2 files changed, 62 insertions(+), 24 deletions(-) diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index 58f7944c1..5158023eb 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -616,11 +616,7 @@ def _finalize_with_error(self, error: BaseException) -> None: def close(self) -> None: if self._span_ended: - # The span can end while the generator is still suspended (e.g. - # an error surfaced from __next__ without resuming the generator). - # Still close the generator so its cleanup runs deterministically - # in the preserved context instead of at GC time under an - # arbitrary ambient context. + # Still close the generator so cleanup runs in the preserved context, not at GC time. self.context.run(self.generator.close) return @@ -692,8 +688,12 @@ def __aiter__(self) -> "_ContextPreservedAsyncGeneratorWrapper": def _generator_never_resumed(self) -> bool: try: state = inspect.getasyncgenstate(self.generator) - except TypeError: - return False + except (AttributeError, TypeError): + # 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 + ) return state in ("AGEN_CREATED", "AGEN_SUSPENDED") @@ -726,10 +726,7 @@ def _finalize_with_error(self, error: BaseException) -> None: async def aclose(self) -> None: if self._span_ended: - # The span can end while the generator is still suspended. Still - # close the generator so its cleanup runs deterministically in - # the preserved context instead of at GC time under an arbitrary - # ambient context. + # Still close the generator so cleanup runs in the preserved context, not at GC time. await self._close_generator() return @@ -751,9 +748,7 @@ async def _close_generator(self) -> None: context=self.context, ) # type: ignore except TypeError: - # Python 3.10: create_task has no context parameter. Only the - # create_task call is guarded so a TypeError raised from the - # generator's own cleanup is not mistaken for it. + # Python 3.10 create_task has no context param; guard only the call itself. close_task = self.context.run(asyncio.create_task, self.generator.aclose()) await close_task @@ -793,10 +788,7 @@ async def __anext__(self) -> Any: raise # Re-raise StopAsyncIteration except asyncio.CancelledError as e: if self._generator_never_resumed(): - # The cancellation was delivered before the inner task resumed - # the generator. Defer the span end so aclose() can run the - # generator's cleanup (including any span updates it makes) - # before the span is finalized with this error. + # Defer span end so aclose() can run the generator's cleanup first. self._pending_error = e raise self._finalize_with_error(e) diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index 2d37b82ed..e080abd0c 100644 --- a/tests/unit/test_observe.py +++ b/tests/unit/test_observe.py @@ -295,8 +295,7 @@ def generator() -> Generator[str, None, None]: assert next(wrapper) == "item_0" - # __next__ raising without resuming the generator (here: re-entering the - # preserved context) ends the span while the generator is still suspended. + # An error from __next__ that never resumed the generator ends the span. with pytest.raises(RuntimeError): context.run(lambda: next(wrapper)) @@ -380,20 +379,19 @@ async def consume() -> None: consumer = asyncio.create_task(consume()) for _ in range(4): await asyncio.sleep(0) - # The cancel lands between chunks, before the inner __anext__ task's - # first step: the generator is never resumed. + # Cancel lands before the inner __anext__ task's first step. asyncio.get_running_loop().call_soon(consumer.cancel) with pytest.raises(asyncio.CancelledError): await consumer - assert inspect.getasyncgenstate(raw) == "AGEN_SUSPENDED" + assert raw.ag_frame is not None # still suspended, never resumed # Span end is deferred so the generator's cleanup can still update it. assert span.ended == 0 marker.set("ambient-now") await wrapper.aclose() - assert inspect.getasyncgenstate(raw) == "AGEN_CLOSED" + assert raw.ag_frame is None # closed assert seen == ["preserved"] assert span.ended == 1 # The retained cancellation still finalizes the span as an error. @@ -403,6 +401,54 @@ async def consume() -> None: } +@pytest.mark.asyncio +async def test_async_generator_wrapper_defers_unresumed_cancel_without_inspect_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Python < 3.12 has no inspect.getasyncgenstate; the fallback must still defer. + monkeypatch.delattr(inspect, "getasyncgenstate", raising=False) + + seen: list[str] = [] + + async def generator() -> AsyncGenerator[str, None]: + try: + yield "item_0" + yield "item_1" + finally: + seen.append("closed") + + span = SpanRecorder() + wrapper = _ContextPreservedAsyncGeneratorWrapper( + generator(), + contextvars.copy_context(), + cast(Any, span), + False, + None, + ) + + async def consume() -> None: + async for _ in wrapper: + await asyncio.sleep(0) + + consumer = asyncio.create_task(consume()) + for _ in range(4): + await asyncio.sleep(0) + asyncio.get_running_loop().call_soon(consumer.cancel) + with pytest.raises(asyncio.CancelledError): + await consumer + + assert span.ended == 0 + + await wrapper.aclose() + + assert seen == ["closed"] + assert span.ended == 1 + assert span.updates[-1] == { + "level": "ERROR", + "status_message": "CancelledError", + } + + @pytest.mark.asyncio async def test_async_generator_wrapper_aclose_propagates_cleanup_type_error() -> None: async def generator() -> AsyncGenerator[str, None]: From bb4ea207ba1044cf25ab45ea99a776f1bf38fff0 Mon Sep 17 00:00:00 2001 From: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:31:22 +0200 Subject: [PATCH 4/4] fix(observe): avoid unawaited generator tasks --- langfuse/_client/observe.py | 14 +++++++------- tests/unit/test_observe.py | 24 ++++++++++-------------- 2 files changed, 17 insertions(+), 21 deletions(-) diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index 5158023eb..53aa79a58 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -2,6 +2,7 @@ import contextvars import inspect import os +import sys from functools import wraps from typing import ( Any, @@ -48,6 +49,8 @@ P = ParamSpec("P") R = TypeVar("R") +_ASYNCIO_CREATE_TASK_SUPPORTS_CONTEXT = sys.version_info >= (3, 11) + class LangfuseDecorator: """Implementation of the @observe decorator for seamless Langfuse tracing integration. @@ -742,13 +745,12 @@ async def aclose(self) -> None: self._finalize() async def _close_generator(self) -> None: - try: + if _ASYNCIO_CREATE_TASK_SUPPORTS_CONTEXT: close_task = asyncio.create_task( self.generator.aclose(), context=self.context, ) # type: ignore - except TypeError: - # Python 3.10 create_task has no context param; guard only the call itself. + else: close_task = self.context.run(asyncio.create_task, self.generator.aclose()) await close_task @@ -765,14 +767,12 @@ def __del__(self) -> None: async def __anext__(self) -> Any: try: # Run the generator's __anext__ in the preserved context - try: - # Python 3.11+ approach with explicit task context + if _ASYNCIO_CREATE_TASK_SUPPORTS_CONTEXT: item = await asyncio.create_task( self.generator.__anext__(), # type: ignore context=self.context, ) # type: ignore - except TypeError: - # Python 3.10 fallback - create the task inside the preserved context. + else: item = await self.context.run( asyncio.create_task, self.generator.__anext__(), # type: ignore diff --git a/tests/unit/test_observe.py b/tests/unit/test_observe.py index e080abd0c..f2ff11789 100644 --- a/tests/unit/test_observe.py +++ b/tests/unit/test_observe.py @@ -9,6 +9,7 @@ import pytest from langfuse import observe +from langfuse._client import observe as observe_module from langfuse._client.attributes import LangfuseOtelSpanAttributes from langfuse._client.observe import ( _ContextPreservedAsyncGeneratorWrapper, @@ -352,6 +353,7 @@ async def generator() -> AsyncGenerator[str, None]: async def test_async_generator_wrapper_defers_span_end_for_unresumed_cancel() -> None: marker = contextvars.ContextVar("marker", default="ambient") seen: list[str] = [] + cleanup_span_states: list[int] = [] async def generator() -> AsyncGenerator[str, None]: try: @@ -359,6 +361,8 @@ async def generator() -> AsyncGenerator[str, None]: yield "item_1" finally: seen.append(marker.get()) + cleanup_span_states.append(span.ended) + span.update(cleanup=True) span = SpanRecorder() context = contextvars.copy_context() @@ -393,12 +397,12 @@ async def consume() -> None: assert raw.ag_frame is None # closed assert seen == ["preserved"] + assert cleanup_span_states == [0] assert span.ended == 1 - # The retained cancellation still finalizes the span as an error. - assert span.updates[-1] == { - "level": "ERROR", - "status_message": "CancelledError", - } + assert span.updates == [ + {"cleanup": True}, + {"level": "ERROR", "status_message": "CancelledError"}, + ] @pytest.mark.asyncio @@ -481,15 +485,7 @@ async def test_async_generator_wrapper_fallback_preserves_context( ) -> 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) + monkeypatch.setattr(observe_module, "_ASYNCIO_CREATE_TASK_SUPPORTS_CONTEXT", False) async def generator() -> AsyncGenerator[str, None]: try: