Fix init logger state isolation

py/src/braintrust/logger.py

@@ -1429,6 +1429,22 @@ def _internal_get_global_state() -> BraintrustState:
     return _state
 
 
+def _resolve_state(
+    state: "BraintrustState | None", api_key: str | None, app_url: str | None
+) -> "BraintrustState":
+    """Resolve the state to use for a logger/experiment/dataset.
+
+    If an explicit state is provided, use it. Otherwise, if api_key or app_url
+    is specified, create a new isolated BraintrustState to avoid conflicts with
+    the global state. If neither is specified, use the global state.
+    """
+    if state is not None:
+        return state
+    if api_key is not None or app_url is not None:
+        return BraintrustState()
+    return _state
+
+
 _internal_reset_global_state()
 _logger = logging.getLogger("braintrust")
 
@@ -1576,7 +1592,7 @@ def init(
     :returns: The experiment object.
     """
 
-    state: BraintrustState = state or _state
+    state = _resolve_state(state, api_key, app_url)
 
     if project is None and project_id is None:
         raise ValueError("Must specify at least one of project or project_id")
@@ -1739,7 +1755,7 @@ def init_dataset(
     :returns: The dataset object.
     """
 
-    state = state or _state
+    state = _resolve_state(state, api_key, app_url)
 
     def compute_metadata():
         state.login(org_name=org_name, api_key=api_key, app_url=app_url)
@@ -1766,11 +1782,12 @@ def compute_metadata():
     )
 
 
-def _compute_logger_metadata(project_name: str | None = None, project_id: str | None = None):
-    login()
-    org_id = _state.org_id
+def _compute_logger_metadata(
+    state: "BraintrustState", project_name: str | None = None, project_id: str | None = None
+):
+    org_id = state.org_id
     if project_id is None:
-        response = _state.app_conn().post_json(
+        response = state.app_conn().post_json(
             "api/project/register",
             {
                 "project_name": project_name or GLOBAL_PROJECT,
@@ -1783,7 +1800,7 @@ def _compute_logger_metadata(project_name: str | None = None, project_id: str |
             project=ObjectMetadata(id=resp_project["id"], name=resp_project["name"], full_info=resp_project),
         )
     elif project_name is None:
-        response = _state.app_conn().get_json("api/project", {"id": project_id})
+        response = state.app_conn().get_json("api/project", {"id": project_id})
         return OrgProjectMetadata(
             org_id=org_id, project=ObjectMetadata(id=project_id, name=response["name"], full_info=response)
         )
@@ -1819,7 +1836,7 @@ def init_logger(
     :returns: The newly created Logger.
     """
 
-    state = state or _state
+    state = _resolve_state(state, api_key, app_url)
     compute_metadata_args = dict(project_name=project, project_id=project_id)
 
     link_args = {
@@ -1831,7 +1848,7 @@ def init_logger(
 
     def compute_metadata():
         state.login(org_name=org_name, api_key=api_key, app_url=app_url, force_login=force_login)
-        return _compute_logger_metadata(**compute_metadata_args)
+        return _compute_logger_metadata(state, **compute_metadata_args)
 
     # For loggers, enable queue size limit enforcement (bounded queue)
     state.enforce_queue_size_limit(True)
@@ -3400,7 +3417,12 @@ def _span_components_to_object_id_lambda(components: SpanComponentsV4) -> Callab
         raise Exception("Impossible: compute_object_metadata_args not supported for experiments")
     elif components.object_type == SpanObjectTypeV3.PROJECT_LOGS:
         captured_compute_object_metadata_args = components.compute_object_metadata_args
-        return lambda: _compute_logger_metadata(**captured_compute_object_metadata_args).project.id
+
+        def compute_project_id():
+            login()
+            return _compute_logger_metadata(_state, **captured_compute_object_metadata_args).project.id
+
+        return compute_project_id
     else:
         raise Exception(f"Unknown object type: {components.object_type}")
 

py/src/braintrust/test_helpers.py

@@ -103,7 +103,7 @@ def init_test_logger(project_name: str):
     l._lazy_metadata = lazy_metadata  # Skip actual login by setting fake metadata directly
 
     # Replace the global _compute_logger_metadata function with a resolved LazyValue
-    def fake_compute_logger_metadata(project_name=None, project_id=None):
+    def fake_compute_logger_metadata(state, project_name=None, project_id=None):
         if project_id:
             project_metadata = ObjectMetadata(id=project_id, name=project_name, full_info=dict())
         else:

py/src/braintrust/test_logger.py

@@ -3174,3 +3174,104 @@ def test_multiple_attachment_types_tracked(with_memory_logger, with_simulate_log
     assert attachment in with_memory_logger.upload_attempts
     assert json_attachment in with_memory_logger.upload_attempts
     assert ext_attachment in with_memory_logger.upload_attempts
+
+
+class TestImplicitStateIsolation(TestCase):
+    """Test that passing api_key or app_url creates an isolated BraintrustState."""
+
+    def test_init_logger_with_api_key_creates_isolated_state(self):
+        """Test that init_logger with api_key creates a new isolated state."""
+        from braintrust.logger import BraintrustState, _state
+
+        # Create a logger with an explicit api_key
+        logger_a = init_logger(
+            project="test-project",
+            project_id="test-project-id",
+            api_key="test-api-key-a",
+            set_current=False,
+        )
+
+        # The logger should have its own state, not the global _state
+        assert logger_a.state is not _state
+        assert isinstance(logger_a.state, BraintrustState)
+
+    def test_init_logger_with_app_url_creates_isolated_state(self):
+        """Test that init_logger with app_url creates a new isolated state."""
+        from braintrust.logger import BraintrustState, _state
+
+        # Create a logger with an explicit app_url
+        logger_a = init_logger(
+            project="test-project",
+            project_id="test-project-id",
+            app_url="https://custom.braintrust.dev",
+            set_current=False,
+        )
+
+        # The logger should have its own state, not the global _state
+        assert logger_a.state is not _state
+        assert isinstance(logger_a.state, BraintrustState)
+
+    def test_init_logger_without_api_key_uses_global_state(self):
+        """Test that init_logger without api_key uses the global state."""
+        from braintrust.logger import _state
+
+        # Create a logger without api_key or app_url
+        logger_a = init_logger(
+            project="test-project",
+            project_id="test-project-id",
+            set_current=False,
+        )
+
+        # The logger should use the global _state
+        assert logger_a.state is _state
+
+    def test_multiple_loggers_with_different_api_keys_have_separate_states(self):
+        """Test that multiple loggers with different api_keys have separate states."""
+        # Create two loggers with different api_keys
+        logger_a = init_logger(
+            project="test-project-a",
+            project_id="test-project-id-a",
+            api_key="test-api-key-a",
+            set_current=False,
+        )
+
+        logger_b = init_logger(
+            project="test-project-b",
+            project_id="test-project-id-b",
+            api_key="test-api-key-b",
+            set_current=False,
+        )
+
+        # Each logger should have its own separate state
+        assert logger_a.state is not logger_b.state
+
+    def test_init_with_api_key_creates_isolated_state(self):
+        """Test that init (experiment) with api_key creates a new isolated state."""
+        from braintrust.logger import BraintrustState, _state
+
+        # Create an experiment with an explicit api_key
+        experiment = braintrust.init(
+            project="test-project",
+            experiment="test-experiment",
+            api_key="test-api-key",
+            set_current=False,
+        )
+
+        # The experiment should have its own state, not the global _state
+        assert experiment.state is not _state
+        assert isinstance(experiment.state, BraintrustState)
+
+    def test_init_dataset_with_api_key_creates_isolated_state(self):
+        """Test that init_dataset with api_key creates a new isolated state."""
+        from braintrust.logger import BraintrustState, _state
+
+        # Create a dataset with an explicit api_key
+        dataset = braintrust.init_dataset(
+            project="test-project",
+            name="test-dataset",
+            api_key="test-api-key",
+        )
+
+        # The dataset should have its own state, not the global _state
+        assert dataset.state is not _state
+        assert isinstance(dataset.state, BraintrustState)