From ebb3a1fb658507f9995d044390d3a6e122c2d851 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 27 Oct 2026 19:56:34 +0000
Subject: [PATCH] Query MongoDB timestamps using normalized ISO string bounds

---
 great_ai/helper/normalise_datetime.py  |  8 +++++++
 great_ai/persistence/mongodb_driver.py |  9 ++++++--
 tests/test_mongodb_queries.py          | 32 ++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 great_ai/helper/normalise_datetime.py
 create mode 100644 tests/test_mongodb_queries.py

diff --git a/great_ai/helper/normalise_datetime.py b/great_ai/helper/normalise_datetime.py
new file mode 100644
index 0000000..7b12c4f
--- /dev/null
+++ b/great_ai/helper/normalise_datetime.py
@@ -0,0 +1,8 @@
+from datetime import datetime, timezone
+
+
+def normalise_datetime(value: datetime) -> datetime:
+    """Interpret naive timestamps as UTC and normalize aware timestamps to UTC."""
+    return (
+        value.astimezone(timezone.utc).replace(tzinfo=None) if value.tzinfo else value
+    )
diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py
index 3be8279..6ea4d94 100644
--- a/great_ai/persistence/mongodb_driver.py
+++ b/great_ai/persistence/mongodb_driver.py
@@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 from pymongo import ASCENDING, DESCENDING, MongoClient
 
+from ..helper.normalise_datetime import normalise_datetime
 from ..utilities import chunk
 from ..views import Filter, SortBy, Trace
 from .tracing_database_driver import TracingDatabaseDriver
@@ -125,10 +126,14 @@ class MongoDbDriver(TracingDatabaseDriver):
             and_query.append({})
 
         if since:
-            and_query.append({"created": {"$gte": since}})
+            and_query.append(
+                {"created": {"$gte": normalise_datetime(since).isoformat()}}
+            )
 
         if until:
-            and_query.append({"created": {"$lte": until}})
+            and_query.append(
+                {"created": {"$lte": normalise_datetime(until).isoformat()}}
+            )
 
         if has_feedback is not None:
             and_query.append(
diff --git a/tests/test_mongodb_queries.py b/tests/test_mongodb_queries.py
new file mode 100644
index 0000000..048fe0f
--- /dev/null
+++ b/tests/test_mongodb_queries.py
@@ -0,0 +1,32 @@
+from datetime import datetime, timedelta, timezone
+from typing import Any, Iterator, Tuple
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from great_ai import MongoDbDriver
+
+
+@pytest.fixture
+def mongo() -> Iterator[Tuple[MongoDbDriver, Any]]:
+    driver = MongoDbDriver.__new__(MongoDbDriver)
+    driver.mongo_connection_string = "mongodb://test"
+    driver.mongo_database = "test"
+    with patch("great_ai.persistence.mongodb_driver.MongoClient") as factory:
+        client = factory.__getitem__.return_value.return_value.__enter__.return_value
+        collection = client.__getitem__.return_value.traces
+        collection.count_documents.return_value = 3
+        collection.find.return_value.__enter__.return_value = []
+        yield driver, collection
+
+
+def test_mongo_date_bounds_match_stored_iso_strings(
+    mongo: Tuple[MongoDbDriver, Any],
+) -> None:
+    driver, collection = mongo
+    since = datetime(2026, 1, 1, 1, tzinfo=timezone(timedelta(hours=1)))
+    until = datetime(2026, 1, 2)
+    driver.query(since=since, until=until)
+    predicates = collection.find.call_args.kwargs["filter"]["$and"]
+    assert {"created": {"$gte": "2026-01-01T00:00:00"}} in predicates
+    assert {"created": {"$lte": "2026-01-02T00:00:00"}} in predicates