feat(openai-agents): Truncate long messages (#5141)

Author: alexander-alderman-webb

sentry_sdk/integrations/openai_agents/spans/invoke_agent.py

@@ -3,6 +3,7 @@
     get_start_span_function,
     set_data_normalized,
     normalize_message_roles,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.scope import should_send_default_pii
@@ -61,12 +62,17 @@ def invoke_agent_span(context, agent, kwargs):
 
         if len(messages) > 0:
             normalized_messages = normalize_message_roles(messages)
-            set_data_normalized(
-                span,
-                SPANDATA.GEN_AI_REQUEST_MESSAGES,
-                normalized_messages,
-                unpack=False,
+            scope = sentry_sdk.get_current_scope()
+            messages_data = truncate_and_annotate_messages(
+                normalized_messages, span, scope
             )
+            if messages_data is not None:
+                set_data_normalized(
+                    span,
+                    SPANDATA.GEN_AI_REQUEST_MESSAGES,
+                    messages_data,
+                    unpack=False,
+                )
 
     _set_agent_data(span, agent)
 

sentry_sdk/integrations/openai_agents/utils.py

@@ -4,6 +4,7 @@
     normalize_message_roles,
     set_data_normalized,
     normalize_message_role,
+    truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import SPANDATA, SPANSTATUS, OP
 from sentry_sdk.integrations import DidNotEnable
@@ -135,12 +136,16 @@ def _set_input_data(span, get_response_kwargs):
                     }
                 )
 
-    set_data_normalized(
-        span,
-        SPANDATA.GEN_AI_REQUEST_MESSAGES,
-        normalize_message_roles(request_messages),
-        unpack=False,
-    )
+    normalized_messages = normalize_message_roles(request_messages)
+    scope = sentry_sdk.get_current_scope()
+    messages_data = truncate_and_annotate_messages(normalized_messages, span, scope)
+    if messages_data is not None:
+        set_data_normalized(
+            span,
+            SPANDATA.GEN_AI_REQUEST_MESSAGES,
+            messages_data,
+            unpack=False,
+        )
 
 
 def _set_output_data(span, result):

tests/integrations/openai_agents/test_openai_agents.py

@@ -3,9 +3,13 @@
 import pytest
 from unittest.mock import MagicMock, patch
 import os
+import json
 
+import sentry_sdk
+from sentry_sdk import start_span
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.openai_agents import OpenAIAgentsIntegration
-from sentry_sdk.integrations.openai_agents.utils import safe_serialize
+from sentry_sdk.integrations.openai_agents.utils import _set_input_data, safe_serialize
 from sentry_sdk.utils import parse_version
 
 import agents
@@ -1225,3 +1229,46 @@ def failing_tool(message: str) -> str:
     # The span should be marked as error because the tool execution failed
     assert execute_tool_span["status"] == "internal_error"
     assert execute_tool_span["tags"]["status"] == "internal_error"
+
+
+def test_openai_agents_message_truncation(sentry_init, capture_events):
+    """Test that large messages are truncated properly in OpenAI Agents integration."""
+
+    large_content = (
+        "This is a very long message that will exceed our size limits. " * 1000
+    )
+
+    sentry_init(
+        integrations=[OpenAIAgentsIntegration()],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+    )
+
+    test_messages = [
+        {"role": "system", "content": "small message 1"},
+        {"role": "user", "content": large_content},
+        {"role": "assistant", "content": large_content},
+        {"role": "user", "content": "small message 4"},
+        {"role": "assistant", "content": "small message 5"},
+    ]
+
+    get_response_kwargs = {"input": test_messages}
+
+    with start_span(op="gen_ai.chat") as span:
+        scope = sentry_sdk.get_current_scope()
+        _set_input_data(span, get_response_kwargs)
+        if hasattr(scope, "_gen_ai_original_message_count"):
+            truncated_count = scope._gen_ai_original_message_count.get(span.span_id)
+            assert truncated_count == 5, (
+                f"Expected 5 original messages, got {truncated_count}"
+            )
+
+        assert SPANDATA.GEN_AI_REQUEST_MESSAGES in span._data
+        messages_data = span._data[SPANDATA.GEN_AI_REQUEST_MESSAGES]
+        assert isinstance(messages_data, str)
+
+        parsed_messages = json.loads(messages_data)
+        assert isinstance(parsed_messages, list)
+        assert len(parsed_messages) == 2
+        assert "small message 4" in str(parsed_messages[0])
+        assert "small message 5" in str(parsed_messages[1])