diff --git a/strix/config/models.py b/strix/config/models.py index 5e567d4b..6a0b8976 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -793,19 +793,16 @@ def _install_openrouter_stream_cost_capture() -> None: # Pin each agent's calls to one upstream provider so its prompt cache # survives between turns. body = super().transform_request(*args, **kwargs) + agent_id = request_log.current_call_context().agent_id + if agent_id and load_settings().llm.openrouter_sticky_sessions: + session_id = _OPENROUTER_SESSION_IDS.setdefault(agent_id, str(uuid.uuid4())) + body.setdefault("session_id", session_id) model = kwargs.get("model", args[0] if args else "") if banned := provider_bans.banned_providers(model): # Copy: strix-pro's hook sets ``provider`` to a dict it caches. provider = dict(body.get("provider") or {}) provider["ignore"] = sorted({*provider.get("ignore", ()), *banned}) body["provider"] = provider - agent_id = request_log.current_call_context().agent_id - if agent_id and load_settings().llm.openrouter_sticky_sessions: - session_id = _OPENROUTER_SESSION_IDS.setdefault(agent_id, str(uuid.uuid4())) - # A new session after each ban lets OpenRouter pick a new provider. - if bans := provider_bans.ban_count(model): - session_id = f"{session_id}-{bans}" - body.setdefault("session_id", session_id) return body # LiteLLM's provider-config factory reads litellm.OpenrouterConfig at call diff --git a/strix/llm/provider_bans.py b/strix/llm/provider_bans.py index 65bd769f..3ac89fe1 100644 --- a/strix/llm/provider_bans.py +++ b/strix/llm/provider_bans.py @@ -19,7 +19,6 @@ _LOCK = threading.Lock() # OpenRouter reports display names ("InferenceNet"); provider.ignore takes slugs ("inference-net"). _SLUGS: dict[str, dict[str, str]] = {} _BANNED: dict[str, set[str]] = {} -_BAN_COUNTS: dict[str, int] = {} def _model_id(model: str) -> str: @@ -78,7 +77,6 @@ def ban_provider(model: str, provider: str, reason: str) -> bool: if slug in banned: return True banned.add(slug) - _BAN_COUNTS[model] = _BAN_COUNTS.get(model, 0) + 1 logger.warning( "provider_banned upstream=%s slug=%s model=%s reason=%s scan_id=%s", provider, @@ -101,8 +99,3 @@ def ban_provider(model: str, provider: str, reason: str) -> bool: def banned_providers(model: str) -> list[str]: with _LOCK: return sorted(_BANNED.get(_model_id(model), ())) - - -def ban_count(model: str) -> int: - with _LOCK: - return _BAN_COUNTS.get(_model_id(model), 0) diff --git a/tests/test_cost_tracking.py b/tests/test_cost_tracking.py index 961dd415..cae4a0a4 100644 --- a/tests/test_cost_tracking.py +++ b/tests/test_cost_tracking.py @@ -400,27 +400,20 @@ def test_openrouter_request_ignores_banned_providers() -> None: assert config is not None messages: list[AllMessageValues] = [{"role": "user", "content": "hi"}] pro_preferences = {"order": ["novita", "relace"], "allow_fallbacks": True, "ignore": ["a"]} - token = request_log.bind_call_context("e5f6a7b8", "root") - try: - with ( - patch( - "litellm.llms.openrouter.chat.transformation.OpenrouterConfig.transform_request", - return_value={"provider": pro_preferences}, - ), - patch("strix.config.models.load_settings") as settings, - patch("strix.llm.provider_bans.banned_providers", return_value=["relace"]), - patch("strix.llm.provider_bans.ban_count", return_value=1), - ): - settings.return_value.llm.openrouter_sticky_sessions = True - body = config.transform_request( - model="moonshotai/kimi-k3", - messages=messages, - optional_params={}, - litellm_params={}, - headers={}, - ) - finally: - request_log.reset_call_context(token) + with ( + patch( + "litellm.llms.openrouter.chat.transformation.OpenrouterConfig.transform_request", + return_value={"provider": pro_preferences}, + ), + patch("strix.llm.provider_bans.banned_providers", return_value=["relace"]), + ): + body = config.transform_request( + model="moonshotai/kimi-k3", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) assert body["provider"] == { "order": ["novita", "relace"], @@ -428,9 +421,6 @@ def test_openrouter_request_ignores_banned_providers() -> None: "ignore": ["a", "relace"], } assert pro_preferences["ignore"] == ["a"] - session_id, bans = body["session_id"].rsplit("-", 1) - assert bans == "1" - assert str(uuid.UUID(session_id)) == session_id def _openrouter_config() -> Any: diff --git a/tests/test_provider_bans.py b/tests/test_provider_bans.py index 333373fc..0c3f3e4d 100644 --- a/tests/test_provider_bans.py +++ b/tests/test_provider_bans.py @@ -26,7 +26,7 @@ ENDPOINTS = [ @pytest.fixture(autouse=True) def endpoints() -> Iterator[MagicMock]: - for state in (provider_bans._SLUGS, provider_bans._BANNED, provider_bans._BAN_COUNTS): + for state in (provider_bans._SLUGS, provider_bans._BANNED): state.clear() response = httpx.Response( 200, @@ -44,7 +44,6 @@ def test_ban_maps_provider_names_to_slugs(endpoints: MagicMock) -> None: assert provider_bans.ban_provider(MODEL, "InferenceNet", "context_length") assert provider_bans.banned_providers(MODEL) == ["google-vertex", "inference-net"] - assert provider_bans.ban_count("z-ai/glm-5.3") == 2 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 @@ -66,7 +65,6 @@ def test_banning_every_provider_unbans_them_all() -> None: assert not provider_bans.ban_provider(MODEL, "Google", "provider_unavailable") assert provider_bans.banned_providers(MODEL) == [] - assert provider_bans.ban_count(MODEL) == 3 assert logger.warning.call_args.args[0].startswith("provider_bans_reset ") @@ -76,7 +74,6 @@ def test_ban_is_skipped_for_an_unknown_provider_or_failed_lookup(endpoints: Magi assert not provider_bans.ban_provider("openrouter/other/model", "Relace", "x") assert provider_bans.banned_providers(MODEL) == [] - assert provider_bans.ban_count(MODEL) == 0 def test_bans_can_be_turned_off(endpoints: MagicMock) -> None: