diff --git a/langfuse/_client/observe.py b/langfuse/_client/observe.py index e8786a0e0..50d8f59ef 100644 --- a/langfuse/_client/observe.py +++ b/langfuse/_client/observe.py @@ -625,18 +625,19 @@ def __init__( def __aiter__(self) -> "_ContextPreservedAsyncGeneratorWrapper": return self + async def __anext__(self) -> Any: + tokens = [] + + # Load the context. + # We were using asyncio.create_task before, but this would confuse anyio context managers + # RuntimeError "Attempted to exit cancel scope in a different task than it was entered in" + if self.context: + for var, value in self.context.items(): + tokens.append((var, var.set(value))) 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.generator.__anext__() self.items.append(item) @@ -661,3 +662,7 @@ async def __anext__(self) -> Any: ).end() raise + finally: + # Put back the original context + for var, token in tokens: + var.reset(token)