diff --git a/py/src/braintrust/api/test_projects.py b/py/src/braintrust/api/test_projects.py index ca36baba..3821c534 100644 --- a/py/src/braintrust/api/test_projects.py +++ b/py/src/braintrust/api/test_projects.py @@ -1,3 +1,5 @@ +from unittest.mock import patch + import pytest from braintrust import init_logger from braintrust.api import BraintrustClient @@ -32,6 +34,17 @@ def test_projects_end_to_end_with_real_backend(api_key): assert registered_logger.id == created["id"] assert looked_up_logger.project.name == project_name + # Separate loggers share resolved project metadata without another HTTP + # request. Keep this on the real API path using the existing cassette. + with patch.object( + state._client.transport, "request_json", wraps=state._client.transport.request_json + ) as request: + another_registered = init_logger(project=project_name, state=state, set_current=False) + another_looked_up = init_logger(project_id=created["id"], state=state, set_current=False) + assert another_registered.id == registered_logger.id + assert another_looked_up.project.name == looked_up_logger.project.name + request.assert_not_called() + deleted = client.openapi.projects.delete_project_id(created["id"]) assert create_project == { diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 45e8355f..4c6dae50 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -8,6 +8,7 @@ import concurrent.futures import contextlib import contextvars +import copy import dataclasses import datetime import hashlib @@ -25,7 +26,7 @@ import types import uuid from abc import ABC, abstractmethod -from collections import Counter +from collections import Counter, OrderedDict from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence from functools import partial, wraps from types import TracebackType @@ -578,6 +579,23 @@ def _close(self) -> None: conn.close() +_LOGIN_TIMEOUT_SECONDS = 30.0 +_PROJECT_LOOKUP_TIMEOUT_SECONDS = 30.0 +_PROJECT_CACHE_TTL_SECONDS = 15 * 60 +_PROJECT_CACHE_MAX_SIZE = 1000 +_RequestResult = TypeVar("_RequestResult") +_LoginKey = tuple[str, str, str | None, int] +_ProjectCacheKey = tuple[str | None, str | None, str | None, str, str, str | None] + + +@dataclasses.dataclass +class _ProjectMetadataEntry: + future: "concurrent.futures.Future[OrgProjectMetadata]" = dataclasses.field( + default_factory=concurrent.futures.Future + ) + expires_at: float = float("inf") + + class BraintrustState: def __init__(self): self.id = str(uuid.uuid4()) @@ -608,6 +626,11 @@ def __init__(self): self._context_manager = None self._context_manager_lock = threading.Lock() self._client_lock = threading.RLock() + self._project_metadata_cache: OrderedDict[_ProjectCacheKey, _ProjectMetadataEntry] = OrderedDict() + self._pending_logins: dict[_LoginKey, concurrent.futures.Future[BraintrustState]] = {} + self._login_requests: set[concurrent.futures.Future[BraintrustState]] = set() + self._login_queue_tail: concurrent.futures.Future[BraintrustState] | None = None + self._login_generation = 0 def default_get_api_conn(): self.login() @@ -662,6 +685,17 @@ def default_get_api_conn(): self._otel_flush_callback: Any | None = None def reset_login_info(self): + with self._client_lock: + self._reset_login_info() + + def _reset_login_info(self): + self._login_generation += 1 + for request in list(self._login_requests): + if not request.done(): + request.set_exception(RuntimeError("Login cancelled by a reset.")) + self._pending_logins.clear() + self._login_queue_tail = None + self._project_metadata_cache.clear() self._loader_api_client_cache.clear() self.app_url: str | None = None @@ -743,6 +777,11 @@ async def flush_otel(self) -> None: def copy_state(self, other: "BraintrustState"): """Copy login information from another BraintrustState instance.""" + with self._client_lock: + self._copy_state(other) + + def _copy_state(self, other: "BraintrustState"): + self._project_metadata_cache.clear() self._loader_api_client_cache.clear() self.__dict__.update( { @@ -761,6 +800,11 @@ def copy_state(self, other: "BraintrustState"): "_last_otel_setting", "_context_manager_lock", "_client_lock", + "_project_metadata_cache", + "_pending_logins", + "_login_requests", + "_login_queue_tail", + "_login_generation", "_loader_api_client_cache", ) } @@ -774,6 +818,7 @@ def login( force_login: bool = False, ) -> None: with self._client_lock: + generation = self._login_generation if not force_login and self.logged_in: # We have already logged in. If any provided login inputs disagree # with our existing settings, raise an Exception warning the user to @@ -790,23 +835,121 @@ def check_updated_param(varname, arg, orig): check_updated_param("org_name", org_name, self.org_name) return - state = login_to_state( - app_url=app_url, - api_key=api_key, - org_name=org_name, - ) - self.copy_state(state) + # Resolve credentials outside the lock; dotenv discovery can perform I/O. + app_url = _get_app_url(app_url) + api_key = _require_api_key(api_key) + org_name = _get_org_name(org_name) + key = (app_url, api_key, org_name, id(_http_adapter)) + while True: + recheck_login = False + with self._client_lock: + if generation != self._login_generation: + raise RuntimeError("Login cancelled by a reset.") + if not force_login and self.logged_in: + self.login(app_url=app_url, api_key=api_key, org_name=org_name) + return + pending = self._pending_logins.get(key) if not force_login else None + previous = self._login_queue_tail + if not force_login and pending is None and previous is not None: + pending = previous + recheck_login = True + elif pending is None: + pending = concurrent.futures.Future() + self._login_queue_tail = pending + self._pending_logins[key] = pending + self._login_requests.add(pending) + + def cleanup(request): + with self._client_lock: + self._login_requests.discard(request) + if self._pending_logins.get(key) is request: + self._pending_logins.pop(key) + if self._login_queue_tail is request: + self._login_queue_tail = None + + pending.add_done_callback(cleanup) + self._start_request( + pending, + lambda: login_to_state(app_url=app_url, api_key=api_key, org_name=org_name), + timeout=_LOGIN_TIMEOUT_SECONDS, + label="login", + previous=previous, + on_success=self.copy_state, + on_discard=_close_login_state, + ) + if recheck_login: + # Validate against the preceding login's credentials on success; + # a failed predecessor leaves this caller free to try its own. + pending.exception() + else: + pending.result() + return + + def _start_request( + self, + future: "concurrent.futures.Future[_RequestResult]", + compute: "Callable[[], _RequestResult]", + *, + timeout: float, + label: str, + previous: "concurrent.futures.Future[_RequestResult] | None" = None, + on_success: "Callable[[_RequestResult], None] | None" = None, + on_discard: "Callable[[_RequestResult], None] | None" = None, + ) -> None: + """Bound shared requests, including adapters that ignore HTTP timeouts. + + A daemon worker lets callers recover from a stalled custom adapter. Late + results cannot publish login state or replace newer cache entries. + Future completion and state changes always hold the same lock. + """ + + def expire(): + with self._client_lock: + if not future.done(): + future.set_exception(TimeoutError(f"Braintrust {label} timed out after {timeout:g} seconds.")) + + def run(): + timer = threading.Timer(timeout, expire) + timer.daemon = True + timer.start() + if previous is not None: + try: + previous.result() + except BaseException: + pass + with self._client_lock: + if future.done(): + timer.cancel() + return + try: + value = compute() + with self._client_lock: + if future.done(): + if on_discard is not None: + on_discard(value) + return + if on_success is not None: + on_success(value) + future.set_result(value) + except BaseException as exc: + with self._client_lock: + if not future.done(): + future.set_exception(exc) + finally: + timer.cancel() + + ctx = contextvars.copy_context() + threading.Thread(target=lambda: ctx.run(run), name=f"braintrust-{label}", daemon=True).start() def api_client(self) -> BraintrustOpenApiClient: """Return the lazily bootstrapped OpenAPI client.""" if self._client is None: - with self._client_lock: - if self._client is None: - self.login() - if self._client is None: - raise RuntimeError("Braintrust API client was not initialized during login") - return self._client.openapi + self.login() + with self._client_lock: + if self._client is None: + raise RuntimeError("Braintrust API client was not initialized during login") + return self._client.openapi def app_conn(self): if not self._app_conn: @@ -919,6 +1062,9 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: global _http_adapter + if _http_adapter is not adapter: + with _state._client_lock: + _state._project_metadata_cache.clear() _http_adapter = adapter if _state._app_conn: @@ -1908,26 +2054,64 @@ def _compute_logger_metadata( state = state or _state state.login() org_id = state.org_id - if project_id is None: - body: _CreateProjectRequest = {"name": project_name or GLOBAL_PROJECT, "org_name": state.org_name} - if project_group_name is not None: - body["project_group_name"] = project_group_name - response = state.api_client().projects.post_project(body=body) - return OrgProjectMetadata( - org_id=org_id, - project=ObjectMetadata(id=response["id"], name=response["name"], full_info=dict(response)), - ) - elif project_name is None: - response = state.api_client().projects.get_project_id(project_id) - return OrgProjectMetadata( - org_id=org_id, - project=ObjectMetadata(id=response["id"], name=response["name"], full_info=dict(response)), - ) - else: + if project_id is not None and project_name is not None: return OrgProjectMetadata( org_id=org_id, project=ObjectMetadata(id=project_id, name=project_name, full_info=dict()) ) + with state._client_lock: + if not state.logged_in: + raise RuntimeError("Project lookup cancelled by a login reset.") + org_id = state.org_id + key = ( + state.app_url, + org_id, + state.login_token, + "name" if project_id is None else "id", + project_name or GLOBAL_PROJECT if project_id is None else project_id, + project_group_name if project_id is None else None, + ) + cache = state._project_metadata_cache + entry = cache.get(key) + if entry is None or time.monotonic() >= entry.expires_at: + # Capture the client and organization together so a forced login cannot + # redirect an already-started lookup to another credential or organization. + api_client = state.api_client() + org_name = state.org_name + entry = _ProjectMetadataEntry() + cache[key] = entry + cache.move_to_end(key) + if len(cache) > _PROJECT_CACHE_MAX_SIZE: + cache.popitem(last=False) + + def compute(): + if project_id is None: + body: _CreateProjectRequest = {"name": project_name or GLOBAL_PROJECT, "org_name": org_name} + if project_group_name is not None: + body["project_group_name"] = project_group_name + response = api_client.projects.post_project(body=body) + else: + response = api_client.projects.get_project_id(project_id) + return OrgProjectMetadata( + org_id=org_id, + project=ObjectMetadata(id=response["id"], name=response["name"], full_info=dict(response)), + ) + + def complete(request): + with state._client_lock: + if request.exception() is None: + entry.expires_at = time.monotonic() + _PROJECT_CACHE_TTL_SECONDS + elif cache.get(key) is entry: + cache.pop(key) + + entry.future.add_done_callback(complete) + state._start_request( + entry.future, compute, timeout=_PROJECT_LOOKUP_TIMEOUT_SECONDS, label="project lookup" + ) + else: + cache.move_to_end(key) + return copy.deepcopy(entry.future.result()) + def init_logger( project: str | None = None, @@ -2354,9 +2538,6 @@ def load_parameters( return parameters -login_lock = threading.RLock() - - def login( app_url: str | None = None, api_key: str | None = None, @@ -2378,9 +2559,7 @@ def login( global _state - # Only permit one thread to login at a time - with login_lock: - _state.login(app_url=app_url, api_key=api_key, org_name=org_name, force_login=force_login) + _state.login(app_url=app_url, api_key=api_key, org_name=org_name, force_login=force_login) def register_otel_flush(callback: Any) -> None: @@ -2459,6 +2638,13 @@ def _login_loader_client(options: _LoaderLoginOptions) -> BraintrustClient: return client +def _close_login_state(state: BraintrustState) -> None: + """Close resources created for a login result that can no longer be published.""" + for resource in (state._client, state._api_conn, state._app_conn, state._proxy_conn): + if resource is not None: + resource.close() + + def login_to_state( app_url: str | None = None, api_key: str | None = None, diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 48e4d416..bd064c16 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -2907,11 +2907,13 @@ def setUp(self): metadata_patch = patch.object(logger, "_compute_logger_metadata", _compute_logger_metadata) metadata_patch.start() self.addCleanup(metadata_patch.stop) - self.state = MagicMock() + self.state = logger.login_to_state(api_key=logger.TEST_API_KEY) self.state.org_name = "test-org" - self.state.context_manager = BraintrustState().context_manager - self.state.id_generator = logger.id_gen.get_id_generator() - self.projects = self.state.api_client.return_value.projects + api_client = MagicMock() + client_patch = patch.object(self.state, "api_client", return_value=api_client) + client_patch.start() + self.addCleanup(client_patch.stop) + self.projects = api_client.projects self.projects.post_project.return_value = {"id": "project-id", "name": "project"} self.projects.get_project_id.return_value = {"id": "project-id", "name": "project"} @@ -2954,7 +2956,7 @@ def test_project_id_takes_precedence(self): else: self.projects.get_project_id.assert_not_called() - def test_group_survives_unresolved_exports(self): + def test_grouped_loggers_and_unresolved_exports_share_project_lookup(self): instance = init_logger( project="project", _create_in_project_group="my-group", @@ -2968,7 +2970,6 @@ def test_group_survives_unresolved_exports(self): self.projects.post_project.assert_not_called() for exported in exports: with self.subTest(exported=exported): - self.projects.reset_mock() components = logger.SpanComponentsV4.from_str(exported) self.assertEqual( components.compute_object_metadata_args, @@ -2980,9 +2981,14 @@ def test_group_survives_unresolved_exports(self): ) with patch.object(logger, "_state", self.state): self.assertEqual(logger.span_components_to_object_id(components), "project-id") - self.projects.post_project.assert_called_once_with( - body={"name": "project", "org_name": "test-org", "project_group_name": "my-group"} - ) + other = init_logger( + project="project", _create_in_project_group="my-group", state=self.state, set_current=False + ) + self.assertEqual(instance.id, other.id) + self.assertEqual(instance.id, "project-id") + self.projects.post_project.assert_called_once_with( + body={"name": "project", "org_name": "test-org", "project_group_name": "my-group"} + ) def test_otel_parent_resolves_project_id_when_grouped(self): for group, expected_parent in [("my-group", "project_id:project-id"), (None, "project_name:project")]: @@ -2999,7 +3005,7 @@ def test_otel_parent_resolves_project_id_when_grouped(self): patch.object(logger, "current_logger", return_value=current), ): self.assertEqual(logger._current_braintrust_parent(), expected_parent) - expected_calls = 0 if group is None else 2 + expected_calls = 0 if group is None else 1 self.assertEqual(self.projects.post_project.call_count, expected_calls) def test_otel_parent_tolerates_failed_group_registration(self): @@ -3047,67 +3053,43 @@ def test_lazy_project_uses_generated_registration_and_lookup(self): class TestDatasetGeneratedAPI(TestCase): def test_init_dataset_uses_generated_project_and_dataset_resources(self): - mock_state = MagicMock() - mock_state.org_name = "test-org" - api_client = mock_state.api_client.return_value - api_client.projects.post_project.return_value = { - "id": "test-project-id", - "name": "test-project", - } - api_client.datasets.post_dataset.return_value = { - "id": "test-dataset-id", - "project_id": "test-project-id", - "name": "test-dataset", - } - - dataset = braintrust.init_dataset( - project="test-project", - name="test-dataset", - description="description", - metadata={"purpose": "test"}, - use_output=False, - state=mock_state, - ) - - self.assertEqual(dataset.id, "test-dataset-id") - api_client.projects.post_project.assert_called_once_with(body={"name": "test-project", "org_name": "test-org"}) - api_client.datasets.post_dataset.assert_called_once_with( - body={ - "project_id": "test-project-id", - "name": "test-dataset", - "description": "description", - "metadata": {"purpose": "test"}, - } - ) - mock_state.app_conn.assert_not_called() - - def test_init_dataset_without_name_uses_logs_name_with_generated_resources(self): - mock_state = MagicMock() - mock_state.org_name = "test-org" - api_client = mock_state.api_client.return_value - api_client.projects.post_project.return_value = { - "id": "test-project-id", - "name": "test-project", - } - api_client.datasets.post_dataset.return_value = { - "id": "test-dataset-id", - "project_id": "test-project-id", - "name": "logs", - } + for options, expected_name in [ + ({"name": "test-dataset", "metadata": {"purpose": "test"}}, "test-dataset"), + ({}, "logs"), + ]: + with self.subTest(options=options): + mock_state = MagicMock() + mock_state.org_name = "test-org" + api_client = mock_state.api_client.return_value + api_client.projects.post_project.return_value = {"id": "test-project-id", "name": "test-project"} + api_client.datasets.post_dataset.return_value = { + "id": "test-dataset-id", + "project_id": "test-project-id", + "name": expected_name, + } - dataset = braintrust.init_dataset( - project="test-project", - description="description", - use_output=False, - state=mock_state, - ) + dataset = braintrust.init_dataset( + project="test-project", + description="description", + use_output=False, + state=mock_state, + **options, + ) - self.assertEqual(dataset.name, "logs") - api_client.projects.post_project.assert_called_once_with(body={"name": "test-project", "org_name": "test-org"}) - api_client.datasets.post_dataset.assert_called_once_with( - body={"project_id": "test-project-id", "name": "logs", "description": "description"} - ) - mock_state.app_conn.assert_not_called() + self.assertEqual(dataset.id, "test-dataset-id") + self.assertEqual(dataset.name, expected_name) + api_client.projects.post_project.assert_called_once_with( + body={"name": "test-project", "org_name": "test-org"} + ) + api_client.datasets.post_dataset.assert_called_once_with( + body={ + "project_id": "test-project-id", + "name": expected_name, + "description": "description", + **options, + } + ) + mock_state.app_conn.assert_not_called() def test_init_dataset_looks_up_explicit_dataset_id(self): mock_state = MagicMock() @@ -3650,23 +3632,18 @@ def failing_function(): _assert_test_exception_group_contents(logs[0].get("error", "")) -def test_check_org_info_no_git_metadata_leaves_settings_none(): - """When org has no git_metadata, state.git_metadata_settings should be None (no org restriction).""" - state = BraintrustState() - org_info = [ - { - "id": "org-1", - "name": "org1", - "api_url": "https://api.example.com", - "proxy_url": "https://proxy.example.com", - } - ] - _check_org_info(state, org_info, None) - assert state.git_metadata_settings is None - - -def test_check_org_info_with_git_metadata_uses_server_settings(): - """When org provides git_metadata, it is used as-is.""" +@pytest.mark.parametrize( + ("org_settings", "expected"), + [ + pytest.param({}, None, id="no-org-restriction"), + pytest.param( + {"git_metadata": {"collect": "some", "fields": ["commit", "branch"]}}, + GitMetadataSettings(collect="some", fields=["commit", "branch"]), + id="server-settings", + ), + ], +) +def test_check_org_info_git_metadata(org_settings, expected): state = BraintrustState() org_info = [ { @@ -3674,12 +3651,11 @@ def test_check_org_info_with_git_metadata_uses_server_settings(): "name": "org1", "api_url": "https://api.example.com", "proxy_url": "https://proxy.example.com", - "git_metadata": {"collect": "some", "fields": ["commit", "branch"]}, + **org_settings, } ] _check_org_info(state, org_info, None) - assert state.git_metadata_settings.collect == "some" - assert set(state.git_metadata_settings.fields) == {"commit", "branch"} + assert state.git_metadata_settings == expected @pytest.mark.parametrize( diff --git a/py/src/braintrust/test_logger_caching.py b/py/src/braintrust/test_logger_caching.py new file mode 100644 index 00000000..67581a59 --- /dev/null +++ b/py/src/braintrust/test_logger_caching.py @@ -0,0 +1,459 @@ +"""Regression tests for shared login and project resolution.""" + +import concurrent.futures +import threading +from unittest.mock import MagicMock, patch + +import pytest +from braintrust import logger +from braintrust.logger import BraintrustState, _compute_logger_metadata +from braintrust.util import GLOBAL_PROJECT + + +@pytest.fixture(autouse=True) +def real_metadata_resolution(): + # Some older logger helpers replace the module function with a fake. + with patch.object(logger, "_compute_logger_metadata", _compute_logger_metadata): + yield + + +@pytest.fixture +def project_state(): + state = logger.login_to_state(api_key=logger.TEST_API_KEY) + client = MagicMock() + client.projects.post_project.side_effect = lambda *, body: { + "id": "project-id", + "name": body["name"], + "metadata": {"nested": [1]}, + } + client.projects.get_project_id.side_effect = lambda project_id: {"id": project_id, "name": "project"} + with patch.object(state, "api_client", return_value=client): + yield state, client + + +def _observe_request_waiters(state, count): + joined = threading.Event() + lock = threading.Lock() + waiting = 0 + start_request = state._start_request + + def start(future, *args, **kwargs): + result = future.result + exception = future.exception + + def wait(method, *args, **kwargs): + nonlocal waiting + with lock: + waiting += 1 + if waiting == count: + joined.set() + return method(*args, **kwargs) + + future.result = lambda *args, **kwargs: wait(result, *args, **kwargs) + future.exception = lambda *args, **kwargs: wait(exception, *args, **kwargs) + return start_request(future, *args, **kwargs) + + return patch.object(state, "_start_request", side_effect=start), joined + + +def _observe_discard(state): + discarded = threading.Event() + start_request = state._start_request + + def start(*args, on_discard=None, **kwargs): + def discard(value): + if on_discard is not None: + on_discard(value) + discarded.set() + + return start_request(*args, on_discard=discard, **kwargs) + + return patch.object(state, "_start_request", side_effect=start), discarded + + +def test_project_cache_default_id_name_and_group_keys(project_state): + state, client = project_state + for name in (None, "", GLOBAL_PROJECT): + _compute_logger_metadata(project_name=name, state=state) + assert client.projects.post_project.call_count == 1 + for group in (None, "one", "two", "one"): + _compute_logger_metadata(project_name="project", project_group_name=group, state=state) + assert client.projects.post_project.call_count == 4 + for _ in range(2): + _compute_logger_metadata(project_id="project", state=state) + client.projects.get_project_id.assert_called_once_with("project") + + +def test_project_cache_shares_pending_lookup_and_expires_after_success(project_state, monkeypatch): + state, client = project_state + now = [100.0] + monkeypatch.setattr(logger.time, "monotonic", lambda: now[0]) + entered = threading.Event() + release = threading.Event() + first = logger.init_logger(project="project", state=state, set_current=False) + response = {"id": "project-id", "name": "project", "metadata": {"nested": [1]}} + + def lookup(*, body): + entered.set() + assert release.wait(5) + return response + + client.projects.post_project.side_effect = lookup + observe, joined = _observe_request_waiters(state, 2) + with observe, concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + original = pool.submit(lambda: first.id) + assert entered.wait(5) + now[0] += 901 + shared = pool.submit(_compute_logger_metadata, project_name="project", state=state) + assert joined.wait(5) + release.set() + assert original.result(5) == shared.result(5).project.id == "project-id" + client.projects.post_project.assert_called_once() + + # Each logger owns a copy, including nested metadata. + first.project.name = "changed" + first.project.full_info["metadata"]["nested"].append(2) + assert shared.result().project.name == "project" + assert shared.result().project.full_info["metadata"]["nested"] == [1] + now[0] += 899 + cached = _compute_logger_metadata(project_name="project", state=state) + assert cached.project.name == "project" + assert cached.project.full_info["metadata"]["nested"] == [1] + client.projects.post_project.assert_called_once() + + # Cache hits do not slide expiry; resolved loggers keep their original ID. + now[0] += 1 + client.projects.post_project.side_effect = None + client.projects.post_project.return_value = {"id": "new-id", "name": "project"} + assert _compute_logger_metadata(project_name="project", state=state).project.id == "new-id" + assert first.id == "project-id" + assert client.projects.post_project.call_count == 2 + + +def test_project_cache_shares_errors_and_retries_later(project_state): + state, client = project_state + entered = threading.Event() + release = threading.Event() + error = ValueError("lookup failed") + + def lookup(*, body): + entered.set() + assert release.wait(5) + raise error + + client.projects.post_project.side_effect = lookup + observe, joined = _observe_request_waiters(state, 4) + with observe, concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + results = [pool.submit(_compute_logger_metadata, project_name="project", state=state) for _ in range(4)] + assert entered.wait(5) + assert joined.wait(5) + release.set() + assert all(result.exception(5) is error for result in results) + client.projects.post_project.assert_called_once() + client.projects.post_project.side_effect = None + client.projects.post_project.return_value = {"id": "project-id", "name": "project"} + for _ in range(2): + assert _compute_logger_metadata(project_name="project", state=state).project.id == "project-id" + assert client.projects.post_project.call_count == 2 + + +def test_old_project_failure_does_not_remove_replacement_after_reset(project_state): + state, client = project_state + entered = threading.Event() + release = threading.Event() + + def lookup(*, body): + entered.set() + assert release.wait(5) + raise ValueError("old failure") + + client.projects.post_project.side_effect = lookup + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + first = pool.submit(_compute_logger_metadata, project_name="project", state=state) + assert entered.wait(5) + state.copy_state(logger.login_to_state(api_key=logger.TEST_API_KEY)) + client.projects.post_project.side_effect = None + client.projects.post_project.return_value = {"id": "new-id", "name": "project"} + assert _compute_logger_metadata(project_name="project", state=state).project.id == "new-id" + release.set() + with pytest.raises(ValueError, match="old failure"): + first.result(5) + assert _compute_logger_metadata(project_name="project", state=state).project.id == "new-id" + assert client.projects.post_project.call_count == 2 + + +@pytest.mark.parametrize("invalidate", ["reset", "copy", "adapter"]) +def test_project_cache_invalidated_by_login_context(project_state, invalidate): + state, client = project_state + _compute_logger_metadata(project_name="project", state=state) + if invalidate == "reset": + state.reset_login_info() + state.login(api_key=logger.TEST_API_KEY) + elif invalidate == "copy": + state.copy_state(logger.login_to_state(api_key=logger.TEST_API_KEY)) + else: + with patch.object(logger, "_state", state), patch.object(logger, "_http_adapter", None): + logger.set_http_adapter(MagicMock()) + _compute_logger_metadata(project_name="project", state=state) + assert client.projects.post_project.call_count == 2 + + +def test_project_cache_lru_bound_and_state_isolation(project_state): + state, client = project_state + for index in range(1000): + _compute_logger_metadata(project_name=str(index), state=state) + _compute_logger_metadata(project_name="0", state=state) + _compute_logger_metadata(project_name="1000", state=state) + _compute_logger_metadata(project_name="0", state=state) + assert client.projects.post_project.call_count == 1001 + _compute_logger_metadata(project_name="1", state=state) + assert client.projects.post_project.call_count == 1002 + other = logger.login_to_state(api_key=logger.TEST_API_KEY) + with patch.object(other, "api_client", return_value=client): + _compute_logger_metadata(project_name="0", state=other) + assert client.projects.post_project.call_count == 1003 + + +def test_login_shares_cold_requests_and_preserves_forced_logins(): + state = BraintrustState() + logged_in = logger.login_to_state(api_key=logger.TEST_API_KEY) + entered = threading.Event() + release = threading.Event() + + def login(**kwargs): + entered.set() + assert release.wait(5) + return logged_in + + with patch.object(logger, "login_to_state", side_effect=login) as request: + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + results = [pool.submit(state.login, api_key=logger.TEST_API_KEY) for _ in range(8)] + assert entered.wait(5) + release.set() + for result in results: + result.result(5) + request.assert_called_once() + state.login(api_key=logger.TEST_API_KEY) + state.login(api_key=logger.TEST_API_KEY, force_login=True) + state.login(api_key=logger.TEST_API_KEY, force_login=True) + assert request.call_count == 3 + + +@pytest.mark.parametrize( + ("changed_param", "changed_value", "conflicts"), + [ + pytest.param("api_key", "second", True, id="api-key"), + pytest.param("org_name", "second", True, id="organization"), + pytest.param("app_url", "https://second.example.com", True, id="app-url"), + pytest.param("org_name", None, False, id="reuse-active-org"), + ], +) +@pytest.mark.parametrize("force_login", [False, True], ids=["non-forced", "forced"]) +def test_queued_login_requires_force_to_switch_credentials( + changed_param, changed_value, conflicts, force_login, monkeypatch +): + monkeypatch.delenv("BRAINTRUST_ORG_NAME", raising=False) + state = BraintrustState() + first_options = {"api_key": logger.TEST_API_KEY, "org_name": "first-org", "app_url": "https://first.example.com"} + second_options = {**first_options, changed_param: changed_value} + first_state = logger.login_to_state(**first_options) + second_state = logger.login_to_state( + api_key=logger.TEST_API_KEY, org_name=second_options["org_name"], app_url=second_options["app_url"] + ) + second_state.login_token = second_options["api_key"] + entered = threading.Event() + release = threading.Event() + + def login(**options): + if options == first_options: + entered.set() + assert release.wait(5) + return first_state + return second_state + + observe, joined = _observe_request_waiters(state, 2) + with observe, patch.object(logger, "login_to_state", side_effect=login) as request: + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(state.login, **first_options) + assert entered.wait(5) + second = pool.submit(state.login, **second_options, force_login=force_login) + assert joined.wait(5) + release.set() + first.result(5) + if force_login or not conflicts: + second.result(5) + else: + with pytest.raises(Exception, match="force_login=True"): + second.result(5) + assert request.call_count == (2 if force_login else 1) + expected = second_state if force_login else first_state + assert (state.app_url, state.login_token, state.org_name) == ( + expected.app_url, + expected.login_token, + expected.org_name, + ) + + +def test_login_shares_failure_and_retries_later(): + state = BraintrustState() + entered = threading.Event() + release = threading.Event() + error = ValueError("login failed") + + def fail(**kwargs): + entered.set() + assert release.wait(5) + raise error + + with patch.object(logger, "login_to_state", side_effect=fail) as request: + observe, joined = _observe_request_waiters(state, 4) + with observe, concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + results = [pool.submit(state.login, api_key=logger.TEST_API_KEY) for _ in range(4)] + assert entered.wait(5) + assert joined.wait(5) + release.set() + for result in results: + assert result.exception(5) is error + request.assert_called_once() + state.login(api_key=logger.TEST_API_KEY) + assert state.logged_in + + +@pytest.mark.parametrize("force_login", [False, True], ids=["non-forced", "forced"]) +def test_reset_cancels_pending_and_queued_logins(force_login): + state = BraintrustState() + old_state = logger.login_to_state(api_key=logger.TEST_API_KEY, org_name="old") + new_state = logger.login_to_state(api_key=logger.TEST_API_KEY, org_name="new") + entered = threading.Event() + release = threading.Event() + + def blocked(**kwargs): + entered.set() + assert release.wait(5) + return old_state + + observe, discarded = _observe_discard(state) + with observe, patch.object(logger, "login_to_state", side_effect=blocked) as request: + observe_waiters, joined = _observe_request_waiters(state, 2) + with observe_waiters, concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(state.login, api_key=logger.TEST_API_KEY) + assert entered.wait(5) + queued = pool.submit( + state.login, api_key=logger.TEST_API_KEY, org_name="queued-org", force_login=force_login + ) + assert joined.wait(5) + state.reset_login_info() + for result in (first, queued): + with pytest.raises(RuntimeError, match="cancelled by a reset"): + result.result(5) + request.side_effect = None + request.return_value = new_state + state.login(api_key=logger.TEST_API_KEY) + release.set() + assert discarded.wait(5) + assert state.org_name == "new" + assert request.call_count == 2 + + +@pytest.mark.parametrize("force_login", [False, True], ids=["non-forced", "forced"]) +def test_failed_login_does_not_block_queued_login(force_login): + state = BraintrustState() + logged_in = logger.login_to_state(api_key=logger.TEST_API_KEY, org_name="retry-org") + entered = threading.Event() + release = threading.Event() + + def fail(**kwargs): + entered.set() + assert release.wait(5) + raise ValueError("login failed") + + with patch.object(logger, "login_to_state") as request: + + def dispatch(**kwargs): + return fail(**kwargs) if request.call_count == 1 else logged_in + + request.side_effect = dispatch + observe, joined = _observe_request_waiters(state, 2) + with observe, concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(state.login, api_key=logger.TEST_API_KEY) + assert entered.wait(5) + second = pool.submit( + state.login, api_key=logger.TEST_API_KEY, org_name="retry-org", force_login=force_login + ) + assert joined.wait(5) + release.set() + with pytest.raises(ValueError, match="login failed"): + first.result(5) + second.result(5) + assert state.logged_in + assert request.call_count == 2 + + +def test_login_reset_during_credential_discovery(): + state = BraintrustState() + + def credentials(api_key): + state.reset_login_info() + return logger.TEST_API_KEY + + with patch.object(logger, "_require_api_key", side_effect=credentials): + with pytest.raises(RuntimeError, match="cancelled by a reset"): + state.login() + assert not state.logged_in + + +def test_project_lookup_reset_before_client_capture(project_state): + state, client = project_state + with patch.object(state, "login", side_effect=state.reset_login_info): + with pytest.raises(RuntimeError, match="cancelled by a login reset"): + _compute_logger_metadata(project_name="project", state=state) + client.projects.post_project.assert_not_called() + + +@pytest.mark.parametrize("kind", ["login", "project"]) +def test_stalled_requests_timeout_and_allow_retry(project_state, monkeypatch, kind): + state, client = project_state + entered = threading.Event() + release = threading.Event() + + def blocked(*args, **kwargs): + entered.set() + assert release.wait(5) + if kind == "login": + return old_state + return {"id": "stale-id", "name": "project"} + + if kind == "login": + old_state = logger.login_to_state(api_key=logger.TEST_API_KEY, org_name="old") + old_state._client = MagicMock() + old_state._api_conn = MagicMock() + old_state._app_conn = MagicMock() + old_state._proxy_conn = MagicMock() + state.reset_login_info() + monkeypatch.setattr(logger, "_LOGIN_TIMEOUT_SECONDS", 0.5) + target = patch.object(logger, "login_to_state", side_effect=blocked) + invoke = lambda: state.login(api_key=logger.TEST_API_KEY) + else: + monkeypatch.setattr(logger, "_PROJECT_LOOKUP_TIMEOUT_SECONDS", 0.05) + target = patch.object(client.projects, "post_project", side_effect=blocked) + invoke = lambda: _compute_logger_metadata(project_name="project", state=state) + observe, discarded = _observe_discard(state) + with observe: + with target: + with pytest.raises(TimeoutError, match="timed out"): + invoke() + assert entered.is_set() + if kind == "login": + state.login(api_key=logger.TEST_API_KEY, org_name="new") + else: + assert invoke().project.id == "project-id" + release.set() + assert discarded.wait(5) + if kind == "login": + assert state.org_name == "new" + old_state._client.close.assert_called_once() + old_state._api_conn.close.assert_called_once() + old_state._app_conn.close.assert_called_once() + old_state._proxy_conn.close.assert_called_once() + else: + assert invoke().project.id == "project-id"