From 5389ba358009c106a068a93c85385011c592e413 Mon Sep 17 00:00:00 2001
From: Mahhheshh <100200105+Mahhheshh@users.noreply.github.com>
Date: Fri, 13 Mar 2026 17:56:56 +0000
Subject: [PATCH] Feat: Make atexit flush registration optional

---
 py/src/braintrust/logger.py      |  9 ++++++++-
 py/src/braintrust/test_logger.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py
index f2d20863..7ba250c0 100644
--- a/py/src/braintrust/logger.py
+++ b/py/src/braintrust/logger.py
@@ -1035,7 +1035,14 @@ def __init__(self, api_conn: LazyValue[HTTPConnection]):
         # Counter for tracking overflow uploads (useful for testing)
         self._overflow_upload_count = 0
 
-        atexit.register(self._finalize)
+        disable_atexit_flush = False
+        try:
+            disable_atexit_flush = os.environ["BRAINTRUST_DISABLE_ATEXIT_FLUSH"] in ("True", "1")
+        except:
+            pass
+
+        if not disable_atexit_flush:
+            atexit.register(self._finalize)
 
     def enforce_queue_size_limit(self, enforce: bool) -> None:
         """
diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py
index 9792ed2a..cecd5f88 100644
--- a/py/src/braintrust/test_logger.py
+++ b/py/src/braintrust/test_logger.py
@@ -125,6 +125,36 @@ def test_init_with_repo_info_does_not_raise(self):
         assert metadata.project.id == "test-project-id"
         assert metadata.experiment.name == "test-exp"
 
+    def test_init_enable_atexit_flush(self):
+        from braintrust.logger import _HTTPBackgroundLogger
+
+        api_con_response = lambda : {
+            "project": {"id": "test-project-id", "name": "test-project"},
+            "experiment": {"id": "test-exp-id", "name": "test-exp"},
+        }
+
+        with patch("atexit.register") as mock_register:
+            _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False))  # type: ignore
+            mock_register.assert_called()
+
+    def test_init_disable_atexit_flush(self):
+        from braintrust.logger import _HTTPBackgroundLogger
+
+        api_con_response = lambda : {
+            "project": {"id": "test-project-id", "name": "test-project"},
+            "experiment": {"id": "test-exp-id", "name": "test-exp"},
+        }
+
+        with patch.dict(os.environ, {"BRAINTRUST_DISABLE_ATEXIT_FLUSH": "True"}):
+            with patch("atexit.register") as mock_register:
+                _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False))  # type: ignore
+                mock_register.assert_not_called()
+
+        with patch.dict(os.environ, {"BRAINTRUST_DISABLE_ATEXIT_FLUSH": "1"}):
+            with patch("atexit.register") as mock_register:
+                _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False))  # type: ignore
+                mock_register.assert_not_called()
+
 class TestLogger(TestCase):
     def test_extract_attachments_no_op(self):
         attachments: List[BaseAttachment] = []