diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py
index 506d74d..3be8279 100644
--- a/great_ai/persistence/mongodb_driver.py
+++ b/great_ai/persistence/mongodb_driver.py
@@ -158,7 +158,7 @@ class MongoDbDriver(TracingDatabaseDriver):
         serialized["_id"] = new_version.trace_id
 
         with MongoClient[Any](self.mongo_connection_string) as client:
-            client[self.mongo_database].traces.update_one({"_id": id}, serialized)
+            client[self.mongo_database].traces.replace_one({"_id": id}, serialized)
 
     def delete(self, id: str) -> None:
         with MongoClient[Any](self.mongo_connection_string) as client:
diff --git a/tests/test_prediction_regressions.py b/tests/test_prediction_regressions.py
new file mode 100644
index 0000000..1275c5b
--- /dev/null
+++ b/tests/test_prediction_regressions.py
@@ -0,0 +1,37 @@
+from typing import Any, Dict
+from unittest.mock import patch
+
+from great_ai import Trace
+from great_ai.persistence.mongodb_driver import MongoDbDriver
+from pymongo import MongoClient
+from pymongo.synchronous.collection import Collection
+
+
+def test_mongo_update_uses_valid_replacement() -> None:
+    trace = Trace(
+        trace_id="id",
+        created="2026-01-01",
+        original_execution_time_ms=0,
+        logged_values={},
+        models=[],
+        tags=[],
+        feedback="correct",
+    )
+    driver = MongoDbDriver.__new__(MongoDbDriver)
+    driver.mongo_connection_string = "mongodb://127.0.0.1:1"
+    driver.mongo_database = "test"
+    client: MongoClient[Dict[str, Any]] = MongoClient(
+        driver.mongo_connection_string, connect=False
+    )
+    with (
+        patch("great_ai.persistence.mongodb_driver.MongoClient") as factory,
+        patch.object(
+            Collection, "_update_retryable", return_value={"n": 1, "nModified": 1}
+        ) as update,
+    ):
+        factory.__getitem__.return_value.return_value = client
+        driver.update("id", trace)
+    assert update.call_args.args[0] == {"_id": "id"}
+    document = update.call_args.args[1]
+    assert document["_id"] == "id"
+    assert Trace.model_validate(document) == trace