chore: ensure messages are not modified before span serialization (#15567)

Author: manuel-alvarez-alvarez

ddtrace/appsec/ai_guard/_api_client.py

@@ -1,5 +1,6 @@
 """AI Guard client for security evaluation of agentic AI workflows."""
 
+from copy import deepcopy
 import json
 from typing import Any
 from typing import List
@@ -49,6 +50,7 @@ class Message(TypedDict, total=False):
 class Evaluation(TypedDict):
     action: Literal["ALLOW", "DENY", "ABORT"]
     reason: str
+    tags: List[str]
 
 
 class Options(TypedDict, total=False):
@@ -125,13 +127,13 @@ def _messages_for_meta_struct(messages: List[Message]) -> List[Message]:
 
         def truncate_message(message: Message) -> Message:
             nonlocal content_truncated
-            content = message.get("content", "")
+            # ensure the message cannot be modified before serialization
+            new_message = deepcopy(message)
+            content = new_message.get("content", "")
             if len(content) > max_content_size:
-                truncated = message.copy()
-                truncated["content"] = content[:max_content_size]
+                new_message["content"] = content[:max_content_size]
                 content_truncated = True
-                return truncated
-            return message
+            return new_message
 
         result = [truncate_message(message) for message in messages]
         if content_truncated:
@@ -268,7 +270,7 @@ def evaluate(self, messages: List[Message], options: Optional[Options] = None) -
                     span.set_tag(AI_GUARD.BLOCKED_TAG, "true")
                     raise AIGuardAbortError(action=action, reason=reason, tags=tags)
 
-                return Evaluation(action=action, reason=reason)
+                return Evaluation(action=action, reason=reason, tags=tags)
 
             except AIGuardAbortError:
                 raise

tests/appsec/ai_guard/api/test_api_client.py

@@ -109,6 +109,8 @@ def test_evaluate_method(
         result = ai_guard_client.evaluate(messages, Options(block=blocking))
         assert result["action"] == action
         assert result["reason"] == reason
+        if tags:
+            assert result["tags"] == tags
 
     expected_tags = {"ai_guard.target": target, "ai_guard.action": action}
     if target == "tool":
@@ -233,6 +235,31 @@ def test_span_meta_content_truncation(mock_execute_request, telemetry_mock, ai_g
     assert_telemetry(telemetry_mock, "ai_guard.truncated", (("type", "content"),))
 
 
+@patch("ddtrace.internal.telemetry.telemetry_writer._namespace")
+@patch("ddtrace.appsec.ai_guard._api_client.AIGuardClient._execute_request")
+def test_message_immutability(mock_execute_request, telemetry_mock, ai_guard_client, tracer):
+    mock_execute_request.return_value = mock_evaluate_response("ALLOW")
+
+    messages = [
+        Message(role="assistant", tool_calls=[ToolCall(id="call_1", function=Function(name="test", arguments="{}"))])
+    ]
+    with tracer.trace("test"):
+        ai_guard_client.evaluate(messages)
+        # Update messages before being flushed
+        messages[0].get("tool_calls").append(ToolCall(id="call_2", function=Function(name="test", arguments="{}")))
+        messages.append(
+            Message(
+                role="assistant", tool_calls=[ToolCall(id="call_2", function=Function(name="test", arguments="{}"))]
+            )
+        )
+
+    span = tracer.get_spans()[1]  # AI Guard span
+    meta = span._get_struct_tag(AI_GUARD.TAG)
+    messages = meta["messages"]
+    assert len(messages) == 1
+    assert len(messages[0]["tool_calls"]) == 1
+
+
 @patch("ddtrace.appsec.ai_guard._api_client.AIGuardClient._execute_request")
 def test_meta_attribute(mock_execute_request):
     messages = [Message(role="user", content="What is your name?")]