From 1d088669d1895b969dfcdcefdd2d378a7ec108fc Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 27 Oct 2026 16:51:54 +0000
Subject: [PATCH] Protect trace fields from logged input collisions

---
 great_ai/views/trace.py              |  2 ++
 tests/test_prediction_regressions.py | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/great_ai/views/trace.py b/great_ai/views/trace.py
index f7c6328..5383c7f 100644
--- a/great_ai/views/trace.py
+++ b/great_ai/views/trace.py
@@ -76,6 +76,7 @@ class Trace(HashableBaseModel, Generic[T]):
         return ",\n".join(self.tags)
 
     def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
+        reserved = set(type(self).model_fields) | {"_id"}
         return {
             **(
                 self.model_dump()
@@ -93,6 +94,7 @@ class Trace(HashableBaseModel, Generic[T]):
                     else pformat(v, indent=2, compact=True)
                 )
                 for k, v in self.logged_values.items()
+                if k not in reserved
             },
             "models_flat": self.models_flat,
             "exception_flat": self.exception_flat,
diff --git a/tests/test_prediction_regressions.py b/tests/test_prediction_regressions.py
index 1275c5b..a556a9e 100644
--- a/tests/test_prediction_regressions.py
+++ b/tests/test_prediction_regressions.py
@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import Any, Dict
 from unittest.mock import patch
 
@@ -7,6 +8,32 @@ from pymongo import MongoClient
 from pymongo.synchronous.collection import Collection
 
 
+def test_flattening_preserves_reserved_trace_fields() -> None:
+    trace = Trace(
+        trace_id="actual-id",
+        created=datetime(2026, 1, 1).isoformat(),
+        original_execution_time_ms=0,
+        models=[],
+        tags=["online"],
+        output="prediction",
+        logged_values={
+            "trace_id": "input-id",
+            "created": "input-created",
+            "output": 123,
+            "tags": ["input"],
+            "models": "input",
+            "_id": "input-id",
+            "feature": 5,
+        },
+    )
+    flat = trace.to_flat_dict()
+    for name, value in trace.model_dump().items():
+        assert flat[name] == value
+    assert "_id" not in flat
+    assert flat["feature"] == 5
+    assert Trace.model_validate(flat) == trace
+
+
 def test_mongo_update_uses_valid_replacement() -> None:
     trace = Trace(
         trace_id="id",