truncate large spans

Author: nikhilc-microsoft

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/exporters/agent365_exporter.py

@@ -23,6 +23,7 @@
     kind_name,
     partition_by_identity,
     status_name,
+    truncate_span_if_needed,
 )
 
 # ---- Exporter ---------------------------------------------------------------
@@ -295,7 +296,7 @@ def _map_span(self, sp: ReadableSpan) -> dict[str, Any]:
         start_ns = sp.start_time
         end_ns = sp.end_time
 
-        return {
+        span_dict = {
             "traceId": hex_trace_id(ctx.trace_id),
             "spanId": hex_span_id(ctx.span_id),
             "parentSpanId": parent_span_id,
@@ -308,3 +309,6 @@ def _map_span(self, sp: ReadableSpan) -> dict[str, Any]:
             "links": links,
             "status": status,
         }
+
+        # Apply truncation if needed
+        return truncate_span_if_needed(span_dict)

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/exporters/utils.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import json
+import logging
 import os
 from collections.abc import Sequence
 from typing import Any
@@ -13,6 +15,11 @@
     TENANT_ID_KEY,
 )
 
+logger = logging.getLogger(__name__)
+
+# Maximum allowed span size in bytes (250KB)
+MAX_SPAN_SIZE_BYTES = 250 * 1024
+
 
 def hex_trace_id(value: int) -> str:
     # 128-bit -> 32 hex chars
@@ -46,6 +53,77 @@ def status_name(code: StatusCode) -> str:
         return str(code)
 
 
+def truncate_span_if_needed(span_dict: dict[str, Any]) -> dict[str, Any]:
+    """
+    Truncate span attributes if the serialized span exceeds MAX_SPAN_SIZE_BYTES.
+    This mirrors the .NET ExportFormatter behavior.
+
+    Args:
+        span_dict: The span dictionary to potentially truncate
+
+    Returns:
+        The potentially truncated span dictionary
+    """
+    try:
+        # Serialize the span to check its size
+        serialized = json.dumps(span_dict, separators=(",", ":"))
+        current_size = len(serialized.encode("utf-8"))
+
+        if current_size <= MAX_SPAN_SIZE_BYTES:
+            return span_dict
+
+        logger.warning(
+            f"Span size ({current_size} bytes) exceeds limit ({MAX_SPAN_SIZE_BYTES} bytes). "
+            "Truncating large payload attributes."
+        )
+
+        # Create a deep copy to modify (shallow copy would still reference original attributes)
+        truncated_span = span_dict.copy()
+        if "attributes" in truncated_span:
+            truncated_span["attributes"] = truncated_span["attributes"].copy()
+        attributes = truncated_span.get("attributes", {})
+
+        # Track what was truncated for logging
+        truncated_keys = []
+
+        # Sort attributes by size (largest first) and truncate until size is acceptable
+        if attributes:
+            # Calculate size of each attribute value when serialized
+            attr_sizes = []
+            for key, value in attributes.items():
+                try:
+                    value_size = len(json.dumps(value, separators=(",", ":")).encode("utf-8"))
+                    attr_sizes.append((key, value_size))
+                except Exception:
+                    # If we can't serialize the value, assume it's small
+                    attr_sizes.append((key, 0))
+
+            # Sort by size (descending - largest first)
+            attr_sizes.sort(key=lambda x: x[1], reverse=True)
+
+            # Truncate largest attributes first until size is acceptable
+            for key, _ in attr_sizes:
+                if key in attributes:
+                    attributes[key] = "TRUNCATED"
+                    truncated_keys.append(key)
+
+                    # Check size after truncation
+                    serialized = json.dumps(truncated_span, separators=(",", ":"))
+                    current_size = len(serialized.encode("utf-8"))
+
+                    if current_size <= MAX_SPAN_SIZE_BYTES:
+                        break
+
+        if truncated_keys:
+            logger.info(f"Truncated attributes: {', '.join(truncated_keys)}")
+
+        return truncated_span
+
+    except Exception as e:
+        logger.error(f"Error during span truncation: {e}")
+        return span_dict
+
+
 def partition_by_identity(
     spans: Sequence[ReadableSpan],
 ) -> dict[tuple[str, str], list[ReadableSpan]]:

tests/observability/core/exporters/test_utils.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import unittest
+
+from microsoft_agents_a365.observability.core.exporters.utils import (
+    truncate_span_if_needed,
+)
+
+
+class TestUtils(unittest.TestCase):
+    """Unit tests for utility functions."""
+
+    def test_truncate_span_if_needed(self):
+        """Test truncate_span_if_needed with various span sizes."""
+        # Small span - should return unchanged
+        small_span = {
+            "traceId": "abc123",
+            "spanId": "def456",
+            "name": "small_span",
+            "attributes": {"key1": "value1", "key2": "value2"},
+        }
+        result = truncate_span_if_needed(small_span)
+        self.assertIsNotNone(result)
+        self.assertEqual(result["name"], "small_span")
+        self.assertEqual(result["attributes"]["key1"], "value1")
+
+        # Large span with large payload attributes - should truncate attributes
+        large_span = {
+            "traceId": "abc123",
+            "spanId": "def456",
+            "name": "large_span",
+            "attributes": {
+                "gen_ai.system": "openai",
+                "gen_ai.request.model": "gpt-4",
+                "gen_ai.response.model": "gpt-4",
+                "gen_ai.input.messages": "x" * 150000,  # Large payload
+                "gen_ai.output.messages": "y" * 150000,  # Large payload
+                "gen_ai.sample.attribute": "x" * 250000,  # Large payload
+                "small_attr": "small_value",
+            },
+        }
+        result = truncate_span_if_needed(large_span)
+        self.assertIsNotNone(result)
+        # The largest attributes should be truncated first
+        self.assertEqual(result["attributes"]["gen_ai.input.messages"], "TRUNCATED")
+        self.assertEqual(result["attributes"]["small_attr"], "small_value")  # Unchanged
+        self.assertEqual(result["attributes"]["gen_ai.sample.attribute"], "TRUNCATED")
+
+        # Extremely large span - should return truncated span even if still large
+        extreme_span = {
+            "traceId": "abc123",
+            "spanId": "def456",
+            "name": "extreme_span",
+            "attributes": {f"attr_{i}": "x" * 10000 for i in range(100)},  # Many large attributes
+            "events": [
+                {"name": f"event_{i}", "attributes": {"data": "y" * 10000}} for i in range(50)
+            ],
+        }
+        result = truncate_span_if_needed(extreme_span)
+        self.assertIsNotNone(result)  # Should always return a span, even if still large
+        # All attributes should be truncated due to size
+        for key in result["attributes"]:
+            self.assertEqual(result["attributes"][key], "TRUNCATED")
+
+
+if __name__ == "__main__":
+    unittest.main()