From a13659ca5d0abd85979277eb2bb4ea17c545891d Mon Sep 17 00:00:00 2001 From: kzahiri1 Date: Thu, 1 Oct 2026 14:58:34 -0700 Subject: [PATCH] fix: prevent selected loads from replacing fallback cache --- py/src/braintrust/logger.py | 8 +++++- py/src/braintrust/test_logger.py | 46 ++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 21155a2dd..d3b21111c 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -2066,6 +2066,7 @@ def load_prompt( :returns: The prompt object. """ effective_environment = None if version is not None else environment + should_fall_back_to_cache = version is None and effective_environment is None if id: # When loading by ID, we don't need project or slug @@ -2102,7 +2103,7 @@ def compute_metadata(): if not _is_loader_cache_fallback_error(server_error): raise # If environment or version was specified, don't fall back to cache - if effective_environment is not None or version is not None: + if not should_fall_back_to_cache: raise ValueError(f"Prompt not found with specified parameters") from server_error eprint(f"Failed to load prompt, attempting to fall back to cache: {server_error}") @@ -2138,6 +2139,9 @@ def compute_metadata(): ) resp_prompt = response["objects"][0] prompt = PromptSchema.from_dict_deep(resp_prompt) + # Selected versions/environments must not replace the default fallback entry. + if not should_fall_back_to_cache: + return prompt try: if id: _state._prompt_cache.set( @@ -2322,6 +2326,8 @@ def load_parameters( ) parameters = RemoteEvalParameters.from_function_row(response["objects"][0]) + if not should_fall_back_to_cache: + return parameters try: if id: _state._parameters_cache.set(parameters, id=id, cache_namespace=cache_namespace) diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 597ab4f77..e0c57ccc9 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -705,6 +705,52 @@ def test_load_prefers_version_over_environment(load, endpoint, response, lookup, assert endpoint_mock.call_args_list == [expected_call] +@pytest.mark.parametrize( + ("load", "response_factory", "resource", "cache_name", "cache_type"), + [ + (braintrust.load_prompt, _prompt_response, "prompts.get_prompt", "_prompt_cache", PromptCache), + ( + braintrust.load_parameters, + _parameters_response, + "functions.get_function", + "_parameters_cache", + ParametersCache, + ), + ], + ids=["prompt", "parameters"], +) +@pytest.mark.parametrize("by_id", [False, True], ids=["by-slug", "by-id"]) +@pytest.mark.parametrize("selector", [{"environment": "staging"}, {"version": "v2"}], ids=["environment", "version"]) +def test_load_selected_values_do_not_replace_latest_cache( + load, response_factory, resource, cache_name, cache_type, by_id, selector +): + simulate_login() + latest_response = response_factory("saved") + selected_response = copy.deepcopy(latest_response) + selected_response["objects"][0]["_xact_id"] = "v2" + lookup = {"project": "test-project", "slug": "saved"} + if by_id: + resource += "_id" + latest_response = latest_response["objects"][0] + selected_response = selected_response["objects"][0] + lookup = {"id": latest_response["id"]} + + mock_api_client = MagicMock() + endpoint_mock = operator.attrgetter(resource)(mock_api_client) + endpoint_mock.side_effect = [latest_response, selected_response, _http_error(503)] + cache = cache_type(memory_cache=LRUCache(max_size=10)) + + with ( + patch.object(logger._state, "api_client", return_value=mock_api_client), + patch.object(logger._state, cache_name, cache), + ): + assert load(**lookup).version == "v1" + assert load(**lookup, **selector).version == "v2" + assert load(**lookup).version == "v1" + + assert endpoint_mock.call_count == 3 + + def test_load_parameters_returns_remote_object(): simulate_login() mock_api_client = MagicMock()