"""模型 Key 兜底：402 识别、主 / 兜底选择、耗尽标记的持久化与自动恢复、Agent 用的包装模型。""" from __future__ import annotations from datetime import datetime, timedelta import pytest from backend.browser import llm_factory, llm_keys @pytest.fixture def clean_keys(monkeypatch): """把 settings 表换成内存字典，并复位模块内缓存，测试之间互不影响。""" store: dict[str, str] = {} monkeypatch.setattr(llm_keys.db, "get_setting", lambda key, default="": store.get(key, default)) monkeypatch.setattr(llm_keys.db, "set_setting", lambda key, value: store.__setitem__(key, value)) monkeypatch.setattr(llm_keys, "_exhausted_at", None) monkeypatch.setattr(llm_keys, "_loaded", True) from backend.config import settings monkeypatch.setattr(settings, "deepseek_api_key_fallback", "FALLBACK") monkeypatch.setattr(settings, "dashscope_api_key_fallback", "") monkeypatch.setattr(settings, "llm_primary_retry_hours", 1.0) return store class TestBalanceErrorDetection: @pytest.mark.parametrize("error", [ "Error code: 402 - {'error': {'message': 'Insufficient Balance'}}", "Insufficient Balance", "Arrearage", ]) def test_recognised_from_the_message(self, error): assert llm_keys.is_balance_error(error) is True def test_recognised_from_the_status_code(self): error = RuntimeError("something") error.status_code = 402 assert llm_keys.is_balance_error(error) is True @pytest.mark.parametrize("error", [ "Error code: 401 - Unauthorized", "Error code: 429 - Rate limit", "timeout waiting for selector", "", ]) def test_other_failures_are_not_balance_errors(self, error): # 限流、鉴权失败、页面超时都不该触发切 Key assert llm_keys.is_balance_error(error) is False class TestKeySelection: def test_primary_is_used_while_it_works(self, clean_keys): assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False) def test_fallback_is_used_after_the_primary_is_marked_exhausted(self, clean_keys): llm_keys.mark_primary_exhausted("402") assert llm_keys.choose("deepseek", "PRIMARY") == ("FALLBACK", True) def test_no_fallback_configured_keeps_using_the_primary(self, clean_keys, monkeypatch): from backend.config import settings monkeypatch.setattr(settings, "deepseek_api_key_fallback", "") llm_keys.mark_primary_exhausted("402") assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False) def test_fallback_identical_to_the_primary_is_not_a_fallback(self, clean_keys, monkeypatch): from backend.config import settings monkeypatch.setattr(settings, "deepseek_api_key_fallback", "PRIMARY") llm_keys.mark_primary_exhausted("402") assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False) def test_only_a_fallback_configured_is_still_usable(self, clean_keys): key, using = llm_keys.choose("deepseek", "") assert key == "FALLBACK" and using is True def test_qwen_reads_its_own_fallback(self, clean_keys, monkeypatch): from backend.config import settings monkeypatch.setattr(settings, "dashscope_api_key_fallback", "QWEN_FALLBACK") llm_keys.mark_primary_exhausted("402") assert llm_keys.choose("qwen", "PRIMARY") == ("QWEN_FALLBACK", True) class TestExhaustionMarker: def test_marker_is_persisted(self, clean_keys): llm_keys.mark_primary_exhausted("402") assert clean_keys[llm_keys.EXHAUSTED_KEY] assert llm_keys.primary_exhausted() is True def test_clearing_restores_the_primary(self, clean_keys): llm_keys.mark_primary_exhausted("402") llm_keys.clear_primary_exhausted() assert llm_keys.primary_exhausted() is False assert llm_keys.choose("deepseek", "PRIMARY") == ("PRIMARY", False) def test_primary_is_retried_after_the_retry_window(self, clean_keys, monkeypatch): # 充值之后不用重启：过了重试周期下一次调用先试主 Key monkeypatch.setattr(llm_keys, "_exhausted_at", datetime.now() - timedelta(hours=2)) assert llm_keys.primary_exhausted() is False def test_marker_still_holds_inside_the_retry_window(self, clean_keys, monkeypatch): monkeypatch.setattr(llm_keys, "_exhausted_at", datetime.now() - timedelta(minutes=10)) assert llm_keys.primary_exhausted() is True def test_status_reports_what_the_ui_shows(self, clean_keys): assert llm_keys.status("deepseek", "PRIMARY") == { "fallback_configured": True, "using_fallback": False, "primary_exhausted_at": None, "retry_hours": 1.0, } llm_keys.mark_primary_exhausted("402") status = llm_keys.status("deepseek", "PRIMARY") assert status["using_fallback"] is True and status["primary_exhausted_at"] class FakeModel: def __init__(self, name: str, fail: bool = False): self.name = name self.provider = "deepseek" self.calls = 0 self.fail = fail async def ainvoke(self, messages, output_format=None, **kwargs): self.calls += 1 if self.fail: from browser_use.llm.exceptions import ModelProviderError raise ModelProviderError("Error code: 402 - {'error': {'message': 'Insufficient Balance'}}", model=self.name) return f"answer-from-{self.name}" class TestFallbackChatModel: """Agent 拿到的是这个包装对象：跑到一半余额耗尽也要接着跑完，不能让任务失败。""" def _wrap(self, primary_fails: bool): primary, fallback = FakeModel("primary", fail=primary_fails), FakeModel("fallback") return llm_factory.FallbackChatModel(primary=primary, fallback=fallback, model="deepseek-chat"), primary, fallback async def test_healthy_primary_is_used_and_the_fallback_stays_idle(self, clean_keys): wrapper, primary, fallback = self._wrap(primary_fails=False) assert await wrapper.ainvoke([], None) == "answer-from-primary" assert primary.calls == 1 and fallback.calls == 0 async def test_402_switches_to_the_fallback_and_retries_the_same_call(self, clean_keys): wrapper, primary, fallback = self._wrap(primary_fails=True) assert await wrapper.ainvoke([], None) == "answer-from-fallback" assert primary.calls == 1 and fallback.calls == 1 assert wrapper.using_fallback is True assert llm_keys.primary_exhausted() is True async def test_later_calls_go_straight_to_the_fallback(self, clean_keys): wrapper, primary, _ = self._wrap(primary_fails=True) await wrapper.ainvoke([], None) await wrapper.ainvoke([], None) assert primary.calls == 1 # 不会每一步都再撞一次主 Key async def test_a_new_wrapper_starts_on_the_fallback_when_the_marker_is_set(self, clean_keys): llm_keys.mark_primary_exhausted("402") wrapper, primary, _ = self._wrap(primary_fails=True) assert wrapper.using_fallback is True assert await wrapper.ainvoke([], None) == "answer-from-fallback" assert primary.calls == 0 async def test_a_successful_primary_call_clears_a_stale_marker(self, clean_keys): llm_keys.mark_primary_exhausted("402") llm_keys.clear_primary_exhausted() # 重试周期到了 wrapper, _, fallback = self._wrap(primary_fails=False) await wrapper.ainvoke([], None) assert llm_keys.primary_exhausted() is False and fallback.calls == 0 async def test_non_balance_errors_are_not_swallowed(self, clean_keys): from browser_use.llm.exceptions import ModelProviderError class Boom(FakeModel): async def ainvoke(self, messages, output_format=None, **kwargs): self.calls += 1 raise ModelProviderError("Error code: 500 - server error", model=self.name) fallback = FakeModel("fallback") wrapper = llm_factory.FallbackChatModel(primary=Boom("primary"), fallback=fallback, model="m") with pytest.raises(ModelProviderError, match="500"): await wrapper.ainvoke([], None) assert fallback.calls == 0 # 服务端 500 不该消耗兜底 Key def test_it_looks_like_a_chat_model_to_browser_use(self, clean_keys): wrapper, _, _ = self._wrap(primary_fails=False) assert wrapper.provider == "deepseek" assert wrapper.name == "primary" assert wrapper.model_name == "deepseek-chat" assert wrapper._verified_api_keys is False class TestLlmConfig: def test_config_reports_both_keys_and_which_one_is_active(self, clean_keys, monkeypatch): monkeypatch.setattr(llm_factory, "_setting", lambda key, env_key, default="": { "llm_provider": "deepseek", "llm_model": "deepseek-chat", "llm_base_url": "https://api.deepseek.com/v1", "llm_api_key": "PRIMARY", }.get(key, default if isinstance(default, str) else "")) cfg = llm_factory.get_llm_config() assert cfg["primary_api_key"] == "PRIMARY" and cfg["fallback_api_key"] == "FALLBACK" assert cfg["api_key"] == "PRIMARY" and cfg["using_fallback"] is False llm_keys.mark_primary_exhausted("402") cfg = llm_factory.get_llm_config() assert cfg["api_key"] == "FALLBACK" and cfg["using_fallback"] is True