fix(sdk): aworkflow decorator for async generators (#2292)

packages/traceloop-sdk/tests/test_workflows.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 
 import pytest
@@ -259,3 +260,23 @@ async def unserializable_workflow(obj: object):
 
     spans = exporter.get_finished_spans()
     assert [span.name for span in spans] == ["unserializable_task.task"]
+
+
+@pytest.mark.asyncio
+async def test_async_generator_workflow(exporter):
+
+    @aworkflow(name="async generator")
+    async def stream_numbers():
+        for i in range(3):
+            yield i
+            await asyncio.sleep(0.1)
+
+    results = []
+
+    async for num in await stream_numbers():
+        results.append(num)
+
+    spans = exporter.get_finished_spans()
+
+    assert results == [0, 1, 2]
+    assert [span.name for span in spans] == ['async generator.workflow']

packages/traceloop-sdk/traceloop/sdk/decorators/base.py

@@ -175,11 +175,14 @@ async def wrap(*args, **kwargs):
                 except TypeError as e:
                     Telemetry().log_exception(e)
 
-                res = await fn(*args, **kwargs)
+                res = fn(*args, **kwargs)
 
-                # span will be ended in the generator
+                # If it's an async generator, return a new async generator that handles the span
                 if isinstance(res, types.AsyncGeneratorType):
-                    return await _ahandle_generator(span, ctx_token, res)
+                    return _ahandle_generator(span, ctx_token, res)
+
+                # Await here for non-generator async functions
+                res = await res
 
                 try:
                     if _should_send_prompts():