From 87e4673ac60327bc67cc61a1f2fd8afdea8353b0 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 27 Oct 2026 19:51:37 +0000
Subject: [PATCH] Honor explicitly selected custom storage backends

---
 great_ai/context.py                     |  4 ++++
 tests/test_configuration_regressions.py | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+)
 create mode 100644 tests/test_configuration_regressions.py

diff --git a/great_ai/context.py b/great_ai/context.py
index 24412a5..d6d611f 100644
--- a/great_ai/context.py
+++ b/great_ai/context.py
@@ -202,6 +202,8 @@ def _initialize_tracing_database(
                     )
                     tracing_driver.configure_credentials_from_file(p)
                     return tracing_driver
+    if selected is not None:
+        return selected
     logger.warning(
         "Cannot find credentials files, defaulting to using ParallelTinyDbDriver"
     )
@@ -227,6 +229,8 @@ def _initialize_large_file(
                     )
                     large_file.configure_credentials_from_file(p)
                     return large_file
+    if selected is not None:
+        return selected
     logger.warning("Cannot find credentials files, defaulting to using LargeFileLocal")
     return LargeFileLocal
 
diff --git a/tests/test_configuration_regressions.py b/tests/test_configuration_regressions.py
new file mode 100644
index 0000000..88c76cb
--- /dev/null
+++ b/tests/test_configuration_regressions.py
@@ -0,0 +1,19 @@
+from great_ai import ParallelTinyDbDriver, configure
+from great_ai.context import get_context
+from great_ai.large_file import LargeFileLocal
+
+
+def test_explicit_custom_backends_are_preserved(isolated_context: None) -> None:
+    class CustomDatabase(ParallelTinyDbDriver):
+        pass
+
+    class CustomStorage(LargeFileLocal):
+        pass
+
+    configure(
+        tracing_database_factory=CustomDatabase,
+        large_file_implementation=CustomStorage,
+        disable_se4ml_banner=True,
+    )
+    assert isinstance(get_context().tracing_database, CustomDatabase)
+    assert get_context().large_file_implementation is CustomStorage