diff --git a/strix/llm/provider_bans.py b/strix/llm/provider_bans.py index 3ac89fe1..be2bebb2 100644 --- a/strix/llm/provider_bans.py +++ b/strix/llm/provider_bans.py @@ -8,7 +8,6 @@ import threading import httpx from strix.config.loader import load_settings -from strix.telemetry.logging import current_scan_id logger = logging.getLogger(__name__) @@ -44,56 +43,30 @@ def _provider_slugs(model: str) -> dict[str, str]: return _SLUGS[model] -def ban_provider(model: str, provider: str, reason: str) -> bool: +def ban_provider(model: str, provider: str, reason: str) -> None: """Ban for the rest of the run; ``provider`` may be a display name or slug.""" if not load_settings().llm.openrouter_provider_bans: - return False + return model = _model_id(model) try: slugs = _provider_slugs(model) + slug = slugs[provider.strip().lower()] except Exception as exc: # noqa: BLE001 logger.warning( - "provider_ban_skipped upstream=%s model=%s reason=%s scan_id=%s: " - "could not list the model's providers: %s", - provider, - model, - reason, - current_scan_id(), - exc, + "provider_ban_skipped upstream=%s model=%s reason=%s: %r", provider, model, reason, exc ) - return False - slug = slugs.get(provider.strip().lower()) - if slug is None: - logger.warning( - "provider_ban_skipped upstream=%s model=%s reason=%s scan_id=%s: unknown provider", - provider, - model, - reason, - current_scan_id(), - ) - return False + return with _LOCK: banned = _BANNED.setdefault(model, set()) if slug in banned: - return True + return banned.add(slug) logger.warning( - "provider_banned upstream=%s slug=%s model=%s reason=%s scan_id=%s", - provider, - slug, - model, - reason, - current_scan_id(), + "provider_banned upstream=%s slug=%s model=%s reason=%s", provider, slug, model, reason ) if banned >= set(slugs.values()): banned.clear() - logger.warning( - "provider_bans_reset model=%s scan_id=%s: every provider was banned", - model, - current_scan_id(), - ) - return False - return True + logger.warning("provider_bans_reset model=%s: every provider was banned", model) def banned_providers(model: str) -> list[str]: diff --git a/strix/telemetry/logging.py b/strix/telemetry/logging.py index de3c5dab..670b2236 100644 --- a/strix/telemetry/logging.py +++ b/strix/telemetry/logging.py @@ -25,10 +25,6 @@ def set_scan_id(scan_id: str) -> None: _SCAN_ID.set(scan_id) -def current_scan_id() -> str: - return _SCAN_ID.get() or "-" - - def set_agent_id(agent_id: str | None) -> None: """Set or clear the agent_id seen on every log record from this point. diff --git a/tests/test_provider_bans.py b/tests/test_provider_bans.py index 0c3f3e4d..c082f7d7 100644 --- a/tests/test_provider_bans.py +++ b/tests/test_provider_bans.py @@ -38,40 +38,30 @@ def endpoints() -> Iterator[MagicMock]: def test_ban_maps_provider_names_to_slugs(endpoints: MagicMock) -> None: - with patch.object(provider_bans, "logger") as logger: - assert provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") - assert provider_bans.ban_provider("z-ai/glm-5.3", "google", "provider_unavailable") - assert provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") + provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") + provider_bans.ban_provider("z-ai/glm-5.3", "google", "provider_unavailable") + provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") assert provider_bans.banned_providers(MODEL) == ["google-vertex", "inference-net"] assert provider_bans.banned_providers("openrouter/other/model") == [] endpoints.assert_called_once_with( "https://openrouter.ai/api/v1/models/z-ai/glm-5.3/endpoints", timeout=5 ) - assert logger.warning.call_count == 2 - assert logger.warning.call_args_list[0].args[:5] == ( - "provider_banned upstream=%s slug=%s model=%s reason=%s scan_id=%s", - "InferenceNet", - "inference-net", - "z-ai/glm-5.3", - "context_length", - ) def test_banning_every_provider_unbans_them_all() -> None: - assert provider_bans.ban_provider(MODEL, "Relace", "incomplete_tool_call") - assert provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") - with patch.object(provider_bans, "logger") as logger: - assert not provider_bans.ban_provider(MODEL, "Google", "provider_unavailable") + provider_bans.ban_provider(MODEL, "Relace", "incomplete_tool_call") + provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") + assert provider_bans.banned_providers(MODEL) == ["inference-net", "relace"] + provider_bans.ban_provider(MODEL, "Google", "provider_unavailable") assert provider_bans.banned_providers(MODEL) == [] - assert logger.warning.call_args.args[0].startswith("provider_bans_reset ") def test_ban_is_skipped_for_an_unknown_provider_or_failed_lookup(endpoints: MagicMock) -> None: - assert not provider_bans.ban_provider(MODEL, "Nobody", "provider_unavailable") + provider_bans.ban_provider(MODEL, "Nobody", "provider_unavailable") endpoints.side_effect = httpx.ConnectError("offline") - assert not provider_bans.ban_provider("openrouter/other/model", "Relace", "x") + provider_bans.ban_provider("openrouter/other/model", "Relace", "x") assert provider_bans.banned_providers(MODEL) == [] @@ -79,7 +69,7 @@ def test_ban_is_skipped_for_an_unknown_provider_or_failed_lookup(endpoints: Magi def test_bans_can_be_turned_off(endpoints: MagicMock) -> None: with patch("strix.llm.provider_bans.load_settings") as settings: settings.return_value.llm.openrouter_provider_bans = False - assert not provider_bans.ban_provider(MODEL, "Relace", "incomplete_tool_call") + provider_bans.ban_provider(MODEL, "Relace", "incomplete_tool_call") assert provider_bans.banned_providers(MODEL) == [] endpoints.assert_not_called()