docs(tracing): fix custom span output example (#4196)

src/agents/tracing/spans.py

@@ -46,7 +46,7 @@ class Span(abc.ABC, Generic[TSpanData]):
             "table": "users"
         }) as span:
             results = await db.query("SELECT * FROM users")
-            span.set_output({"count": len(results)})
+            span.span_data.data["output"] = {"count": len(results)}
 
         # Handling errors in spans
         with custom_span("risky_operation") as span:

tests/test_tracing.py

@@ -273,6 +273,24 @@ def test_spans_with_setters() -> None:
     )
 
 
+def test_custom_span_records_output_in_data() -> None:
+    with trace(workflow_name="test"):
+        with custom_span("database_query", {"operation": "SELECT", "table": "users"}) as span:
+            span.span_data.data["output"] = {"count": 2}
+
+    exported = span.export()
+    assert exported is not None
+    assert exported["span_data"] == {
+        "type": "custom",
+        "name": "database_query",
+        "data": {
+            "operation": "SELECT",
+            "table": "users",
+            "output": {"count": 2},
+        },
+    }
+
+
 def disabled_tracing():
     with trace(workflow_name="test", trace_id="123", group_id="456", disabled=True):
         with agent_span(name="agent_1"):