trace is lost in spans coming from async generators

[chapati]

Issue

In the current version of SDK @observe loses context/trace id for spans that are coming from async generators and nested async generators (when observed async generator is called in another observed async generator). Such spans appear not on the current trace but added as a new trace and are displayed separately in the UI. The previous implementation have been prematurely removing observations from the stack/not properly restoring the stack state during async generator iteration.

Impact

My changes improve the reliability of observation tracking when using async generators with the @observe decorator. This is particularly important for applications that use async generators extensively, such as streaming responses from language models or other async data processing pipelines.

The fix ensures that observations are properly tracked and that the observation stack is maintained in a consistent state throughout the lifecycle of async generators.

Note

Tested for 2 weeks in-house

---

[!IMPORTANT] Fixes context/trace ID loss in async generators with @observe by managing observation stack correctly in langfuse_decorator.py.

• Behavior:
  • Fixes context/trace ID loss in async generators with @observe decorator in langfuse_decorator.py.
  • Ensures observations are properly tracked and stack is consistently maintained during async generator iteration.
• Functions:
  • Modifies _wrap_async_generator_result() to manage observation stack correctly by setting and restoring stack state around anext(generator) calls.
  • Removes commented-out code in _finalize_call() related to stack management.

^{This description was created by for 8fa20d60a9967a62fd4aedba0151fbd86ebbf320. You can customize this summary. It will automatically update as commits are pushed.}

Greptile Summary

Disclaimer: Experimental PR review

Fixes context/trace ID loss in async generators when using the @observe decorator by properly maintaining observation stack state during async generator iteration.

• Modified _wrap_async_generator_result in langfuse/decorators/langfuse_decorator.py to preserve observation stack between generator yields
• Added proper context restoration during anext calls to maintain parent-child relationships
• Fixed spans appearing as separate traces in UI by preventing premature observation stack removal

_{💡 (1/5) You can manually trigger the bot by mentioning @greptileai in a comment!}

[CLAassistant]

All committers have signed the CLA.

[chapati]

Hello, is it going to be merged? We really need this fix. Thanks.

[hassiebp]

Thanks again for your contribution @chapati - this should be fixed in our new Python SDK v3. If you still encounter this issue, feel free to raise a Github issue