fix(sdk): backpropagate association property to nearest workflow/task (…

…#1300)

packages/traceloop-sdk/tests/test_association_properties.py

@@ -0,0 +1,71 @@
+from opentelemetry.semconv.ai import SpanAttributes
+from traceloop.sdk import Traceloop
+from traceloop.sdk.decorators import task, workflow
+
+
+def test_association_properties(exporter):
+    @workflow(name="test_workflow")
+    def test_workflow():
+        return test_task()
+
+    @task(name="test_task")
+    def test_task():
+        return
+
+    Traceloop.set_association_properties({"user_id": 1, "user_name": "John Doe"})
+    test_workflow()
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == [
+        "test_task.task",
+        "test_workflow.workflow",
+    ]
+
+    some_task_span = spans[0]
+    some_workflow_span = spans[1]
+    assert (
+        some_workflow_span.attributes[
+            f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
+        ]
+        == 1
+    )
+    assert (
+        some_workflow_span.attributes[
+            f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_name"
+        ]
+        == "John Doe"
+    )
+    assert (
+        some_task_span.attributes[
+            f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_id"
+        ]
+        == 1
+    )
+    assert (
+        some_task_span.attributes[
+            f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.user_name"
+        ]
+        == "John Doe"
+    )
+
+
+def test_association_properties_within_workflow(exporter):
+    @workflow(name="test_workflow_within")
+    def test_workflow():
+        Traceloop.set_association_properties({"session_id": 15})
+        return
+
+    test_workflow()
+
+    spans = exporter.get_finished_spans()
+    assert [span.name for span in spans] == [
+        "test_workflow_within.workflow",
+    ]
+
+    some_workflow_span = spans[0]
+    assert (
+        some_workflow_span.attributes[
+            f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.session_id"
+        ]
+        == 15
+    )

packages/traceloop-sdk/tests/test_nested_tasks.py

@@ -1,9 +1,7 @@
-import pytest
 from opentelemetry.semconv.ai import SpanAttributes
 from traceloop.sdk.decorators import task, workflow
 
 
-@pytest.mark.vcr
 def test_nested_tasks(exporter):
     @workflow(name="some_workflow")
     def some_workflow():

packages/traceloop-sdk/traceloop/sdk/tracing/tracing.py

@@ -327,16 +327,9 @@ def _span_processor_on_start(self, span, parent_context):
         if entity_name is not None:
             span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, entity_name)
 
-        correlation_id = get_value("correlation_id")
-        if correlation_id is not None:
-            span.set_attribute(SpanAttributes.TRACELOOP_CORRELATION_ID, correlation_id)
-
         association_properties = get_value("association_properties")
         if association_properties is not None:
-            for key, value in association_properties.items():
-                span.set_attribute(
-                    f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.{key}", value
-                )
+            _set_association_properties_attributes(span, association_properties)
 
             if not self.enable_content_tracing:
                 if self.__content_allow_list.is_allowed(association_properties):
@@ -409,6 +402,18 @@ def get_tracer(self):
 def set_association_properties(properties: dict) -> None:
     attach(set_value("association_properties", properties))
 
+    # Attach association properties to the current span, if it's a workflow or a task
+    span = trace.get_current_span()
+    if get_value("workflow_name") is not None or get_value("entity_name") is not None:
+        _set_association_properties_attributes(span, properties)
+
+
+def _set_association_properties_attributes(span, properties: dict) -> None:
+    for key, value in properties.items():
+        span.set_attribute(
+            f"{SpanAttributes.TRACELOOP_ASSOCIATION_PROPERTIES}.{key}", value
+        )
+
 
 def set_workflow_name(workflow_name: str) -> None:
     attach(set_value("workflow_name", workflow_name))