From c55326dc82f06b6bd0067bee240fc22ddfedefd3 Mon Sep 17 00:00:00 2001
From: Uma Annamalai <uannamalai@newrelic.com>
Date: Tue, 17 Oct 2023 18:20:02 -0700
Subject: [PATCH] Add 4096 char truncation for ML events.

---
 newrelic/api/transaction.py            |  2 +-
 newrelic/core/attribute.py             |  1 +
 newrelic/core/custom_event.py          | 12 +++++++++---
 tests/agent_features/test_ml_events.py | 10 ++++++++++
 4 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/newrelic/api/transaction.py b/newrelic/api/transaction.py
index 988b56be6e..ae2fa0dd50 100644
--- a/newrelic/api/transaction.py
+++ b/newrelic/api/transaction.py
@@ -1648,7 +1648,7 @@ def record_ml_event(self, event_type, params):
         if not settings.ml_insights_events.enabled:
             return
 
-        event = create_custom_event(event_type, params)
+        event = create_custom_event(event_type, params, is_ml_event=True)
         if event:
             self._ml_events.add(event, priority=self.priority)
 
diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py
index 10ae8e4597..80533df57e 100644
--- a/newrelic/core/attribute.py
+++ b/newrelic/core/attribute.py
@@ -89,6 +89,7 @@
 
 MAX_NUM_USER_ATTRIBUTES = 128
 MAX_ATTRIBUTE_LENGTH = 255
+MAX_ML_ATTRIBUTE_LENGTH = 4096
 MAX_64_BIT_INT = 2**63 - 1
 MAX_LOG_MESSAGE_LENGTH = 32768
 
diff --git a/newrelic/core/custom_event.py b/newrelic/core/custom_event.py
index 206fb84e68..a7d53b1684 100644
--- a/newrelic/core/custom_event.py
+++ b/newrelic/core/custom_event.py
@@ -18,7 +18,7 @@
 
 from newrelic.core.attribute import (check_name_is_string, check_name_length,
         process_user_attribute, NameIsNotStringException, NameTooLongException,
-        MAX_NUM_USER_ATTRIBUTES)
+        MAX_NUM_USER_ATTRIBUTES, MAX_ML_ATTRIBUTE_LENGTH)
 
 
 _logger = logging.getLogger(__name__)
@@ -72,7 +72,8 @@ def process_event_type(name):
     else:
         return name
 
-def create_custom_event(event_type, params):
+
+def create_custom_event(event_type, params, is_ml_event=False):
     """Creates a valid custom event.
 
     Ensures that the custom event has a valid name, and also checks
@@ -83,6 +84,8 @@ def create_custom_event(event_type, params):
     Args:
         event_type (str): The type (name) of the custom event.
         params (dict): Attributes to add to the event.
+        is_ml_event (bool): Boolean indicating whether create_custom_event was called from
+        record_ml_event for truncation purposes
 
     Returns:
         Custom event (list of 2 dicts), if successful.
@@ -99,7 +102,10 @@ def create_custom_event(event_type, params):
 
     try:
         for k, v in params.items():
-            key, value = process_user_attribute(k, v)
+            if is_ml_event:
+                key, value = process_user_attribute(k, v, max_length=MAX_ML_ATTRIBUTE_LENGTH)
+            else:
+                key, value = process_user_attribute(k, v)
             if key:
                 if len(attributes) >= MAX_NUM_USER_ATTRIBUTES:
                     _logger.debug('Maximum number of attributes already '
diff --git a/tests/agent_features/test_ml_events.py b/tests/agent_features/test_ml_events.py
index 5720224bbe..9a35f0cd00 100644
--- a/tests/agent_features/test_ml_events.py
+++ b/tests/agent_features/test_ml_events.py
@@ -102,6 +102,16 @@ def _test():
     _test()
 
 
+@reset_core_stats_engine()
+def test_record_ml_event_truncation():
+    @validate_ml_events([(_intrinsics, {"a": "a" * 4096})])
+    @background_task()
+    def _test():
+        record_ml_event("LabelEvent", {"a": "a" * 4100})
+
+    _test()
+
+
 @pytest.mark.parametrize(
     "params,expected",
     [