fix(traceloop-sdk): avoid calling async json methods in JSONEncoder (#3968)

Author: trinhchien

packages/traceloop-sdk/tests/test_workflows.py

@@ -496,3 +496,55 @@ def dataclass_workflow(data: TestDataClass):
     assert task_span.parent.span_id == workflow_span.context.span_id
     assert workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "dataclass_workflow"
     assert task_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "dataclass_task"
+
+
+class FakeAsyncRequest:
+    async def json(self):
+        return {"ok": True}
+
+
+class FakeSyncRequest:
+    def json(self):
+        return {"ok": True}
+
+
+@pytest.mark.asyncio
+async def test_async_workflow_with_async_json_method_argument(exporter, recwarn):
+    @workflow(name="request_workflow")
+    async def request_workflow(request: FakeAsyncRequest):
+        return {"ok": True}
+
+    await request_workflow(FakeAsyncRequest())
+
+    assert not any(
+        w.category is RuntimeWarning and "was never awaited" in str(w.message)
+        for w in recwarn
+    )
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == ["request_workflow.workflow"]
+
+    workflow_span = spans[0]
+    assert json.loads(workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT]) == {
+        "args": ["FakeAsyncRequest"],
+        "kwargs": {},
+    }
+    assert json.loads(workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]) == {"ok": True}
+
+
+def test_workflow_with_sync_json_method_argument(exporter):
+    @workflow(name="request_workflow")
+    def request_workflow(request: FakeSyncRequest):
+        return {"ok": True}
+
+    request_workflow(FakeSyncRequest())
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == ["request_workflow.workflow"]
+
+    workflow_span = spans[0]
+    assert json.loads(workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_INPUT]) == {
+        "args": [{"ok": True}],
+        "kwargs": {},
+    }
+    assert json.loads(workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]) == {"ok": True}

packages/traceloop-sdk/traceloop/sdk/utils/json_encoder.py

@@ -1,4 +1,5 @@
 import dataclasses
+import inspect
 import json
 
 
@@ -15,7 +16,12 @@ def default(self, o):
             return o.to_json()
 
         if hasattr(o, "json"):
-            return o.json()
+            json_method = o.json
+            if callable(json_method) and not inspect.iscoroutinefunction(json_method):
+                result = json_method()
+                if not inspect.iscoroutine(result):
+                    return result
+                result.close()
 
         if hasattr(o, "__class__"):
             return o.__class__.__name__