fix(lm_studio): keep last limits on failed discovery, fall back to v0 on unknown-route bodies, refresh each interval

litellm/llms/lm_studio/model_info.py

@@ -3,16 +3,17 @@ import json
 from collections.abc import Mapping
 from itertools import chain
 from types import MappingProxyType
-from typing import Annotated, Final, Protocol, TypeAlias
+from typing import Annotated, Final, Protocol, TypeAlias, TypeVar
 
 import httpx
-from pydantic import BeforeValidator, ConfigDict
+from pydantic import BeforeValidator, ConfigDict, ValidationError
 
 from litellm._logging import verbose_logger
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
-from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
 from litellm.types.llms.base import LiteLLMBaseModel
 
+RESPONSE_REUSE_SECONDS: Final = 60
+
 
 class _MetadataCache(Protocol):
     def get_cache(self, key: str) -> object: ...
@@ -80,6 +81,13 @@ class _ModelsV0(LiteLLMBaseModel):
         return card.loaded_context_length if card is not None and card.state == "loaded" else None
 
 
+class _Unavailable(LiteLLMBaseModel):
+    model_config = ConfigDict(frozen=True)
+
+
+_Models: Final = TypeVar("_Models", _ModelsV1, _ModelsV0)
+
+
 async def get_lm_studio_model_info(
     *,
     model: str,
@@ -87,51 +95,70 @@ async def get_lm_studio_model_info(
     headers: Mapping[str, str],
     client: AsyncHTTPHandler,
     cache: _MetadataCache,
-) -> Mapping[str, int]:
+) -> Mapping[str, int] | None:
     try:
         base: Final = httpx.URL(api_base)
     except (ValueError, httpx.InvalidURL):
-        return MappingProxyType({})
+        return None
     if base.scheme not in ("http", "https") or not base.host or base.username or base.password:
-        return MappingProxyType({})
+        return None
     path: Final = base.path.rstrip("/")
     suffix: Final = next((suffix for suffix in ("/api/v1", "/api/v0", "/v1") if path.endswith(suffix)), "")
     prefix: Final = path[: -len(suffix)] if suffix else path
-    url: Final = str(base.copy_with(path=prefix + "/api/v1/models", fragment=None))
+    current_url: Final = base.copy_with(path=prefix + "/api/v1/models", fragment=None)
     cache_key: Final = (
-        "lm_studio_model_info:" + hashlib.sha256(json.dumps((url, sorted(headers.items()))).encode()).hexdigest()
+        "lm_studio_model_info:"
+        + hashlib.sha256(json.dumps((str(current_url), sorted(headers.items()))).encode()).hexdigest()
     )
     cached: Final[object] = cache.get_cache(cache_key)
+    if isinstance(cached, _Unavailable):
+        return None
     if isinstance(cached, (_ModelsV1, _ModelsV0)):
         return _token_limits(cached.context_length(model))
-    try:
-        response: Final = await client.get(  # pyright: ignore[reportUnknownMemberType]  # untyped options
-            url=url,
-            headers=dict(headers),
-            timeout=httpx.Timeout(5.0),
-            follow_redirects=False,
-            max_response_bytes=2 * 1024 * 1024,
-        )
-        if response.status_code in (404, 405):
-            legacy_response: Final = await client.get(  # pyright: ignore[reportUnknownMemberType]  # untyped options
-                url=str(base.copy_with(path=prefix + "/api/v0/models", fragment=None)),
-                headers=dict(headers),
-                timeout=httpx.Timeout(5.0),
-                follow_redirects=False,
-                max_response_bytes=2 * 1024 * 1024,
-            )
-            legacy_response.raise_for_status()
-            models: Final = _ModelsV0.model_validate_json(legacy_response.content)
-            cache.set_cache(cache_key, models, ttl=MODEL_INFO_REFRESH_SECONDS)
-            return _token_limits(models.context_length(model))
-        response.raise_for_status()
-        current_models: Final = _ModelsV1.model_validate_json(response.content)
-        cache.set_cache(cache_key, current_models, ttl=MODEL_INFO_REFRESH_SECONDS)
-        return _token_limits(current_models.context_length(model))
-    except Exception:  # noqa: BLE001  # optional metadata must not interrupt proxy refresh
+    models: Final = await _fetch_models(
+        client=client,
+        current_url=current_url,
+        legacy_url=base.copy_with(path=prefix + "/api/v0/models", fragment=None),
+        headers=headers,
+    )
+    if models is None:
         verbose_logger.debug("Could not discover LM Studio model token limits")
-        cache.set_cache(cache_key, _ModelsV1(models=()), ttl=60)
-        return MappingProxyType({})
+        cache.set_cache(cache_key, _Unavailable(), ttl=RESPONSE_REUSE_SECONDS)
+        return None
+    cache.set_cache(cache_key, models, ttl=RESPONSE_REUSE_SECONDS)
+    return _token_limits(models.context_length(model))
+
+
+async def _fetch_models(
+    *, client: AsyncHTTPHandler, current_url: httpx.URL, legacy_url: httpx.URL, headers: Mapping[str, str]
+) -> _ModelsV1 | _ModelsV0 | None:
+    try:
+        current: Final = await _get(client=client, url=current_url, headers=headers)
+        current_models: Final = _parse(current, _ModelsV1)
+        if current_models is not None or not (current.is_success or current.status_code in (404, 405)):
+            return current_models
+        return _parse(await _get(client=client, url=legacy_url, headers=headers), _ModelsV0)
+    except Exception:  # noqa: BLE001  # optional metadata must not interrupt proxy refresh
+        return None
+
+
+async def _get(*, client: AsyncHTTPHandler, url: httpx.URL, headers: Mapping[str, str]) -> httpx.Response:
+    return await client.get(  # pyright: ignore[reportUnknownMemberType]  # untyped options
+        url=str(url),
+        headers=dict(headers),
+        timeout=httpx.Timeout(5.0),
+        follow_redirects=False,
+        max_response_bytes=2 * 1024 * 1024,
+    )
+
+
+def _parse(response: httpx.Response, shape: type[_Models]) -> _Models | None:
+    if not response.is_success:
+        return None
+    try:
+        return shape.model_validate_json(response.content)
+    except ValidationError:
+        return None
 
 
 def _token_limits(context: int | None) -> Mapping[str, int]:

litellm/router_python_core.py

@@ -48,7 +48,20 @@ from datetime import datetime, timezone
 from functools import lru_cache, partial
 from itertools import chain
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Final,
+    Generic,
+    Literal,
+    NoReturn,
+    Optional,
+    Protocol,
+    TypeAlias,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import anyio
 import httpx
@@ -672,6 +685,29 @@ _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
 _RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None)
+
+
+class _ModelInfoDiscoverer(Protocol):
+    def __call__(
+        self, *, model: str, api_base: str, headers: Mapping[str, str], client: AsyncHTTPHandler, cache: InMemoryCache
+    ) -> Awaitable[Mapping[str, int] | None]: ...
+
+
+async def _discover_openai_compatible_model_info(
+    *, model: str, api_base: str, headers: Mapping[str, str], client: AsyncHTTPHandler, cache: InMemoryCache
+) -> Mapping[str, int] | None:
+    limits: Final = await get_openai_compatible_model_info(
+        model=model, api_base=api_base, headers=headers, client=client, cache=cache
+    )
+    return limits or None
+
+
+_MODEL_INFO_DISCOVERERS: Final[Mapping[str, _ModelInfoDiscoverer]] = MappingProxyType(
+    {
+        **dict.fromkeys(MODEL_INFO_DISCOVERY_PROVIDERS, _discover_openai_compatible_model_info),
+        "lm_studio": get_lm_studio_model_info,
+    }
+)
 _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
 _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
@@ -12467,7 +12503,8 @@ class PythonCore(DeprecatedRouterMembers):
         model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")(
             model=params.model, litellm_params=params
         )
-        if provider not in MODEL_INFO_DISCOVERY_PROVIDERS and provider != "lm_studio":
+        discover: Final = _MODEL_INFO_DISCOVERERS.get(provider)
+        if discover is None:
             return
         if api_base is None or "*" in model or params.get("use_clientside_credentials"):
             return
@@ -12483,7 +12520,6 @@ class PythonCore(DeprecatedRouterMembers):
         auth_headers: Final = (
             MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({})
         )
-        discover: Final = get_lm_studio_model_info if provider == "lm_studio" else get_openai_compatible_model_info
         limits: Final = await discover(
             model=model,
             api_base=api_base,
@@ -12497,11 +12533,7 @@ class PythonCore(DeprecatedRouterMembers):
             cache=self._core_cache.in_memory_cache,
         )
         model_id: Final = deployment.model_info.id
-        if (
-            (not limits and provider != "lm_studio")
-            or model_id is None
-            or self._core_get_model_info(model_id) is not raw_deployment
-        ):
+        if limits is None or model_id is None or self._core_get_model_info(model_id) is not raw_deployment:
             return
         self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1)
         self._discovered_model_info_cache.delete_cache(model_id)

tests/unit/llms/lm_studio/test_model_info.py

@@ -66,13 +66,22 @@ async def test_matches_loaded_instances_without_using_theoretical_limits(model:
 
 
 @pytest.mark.parametrize("base_path", ("", "/v1/", "/api/v1", "/api/v0"))
-@pytest.mark.parametrize("status", (404, 405))
-async def test_legacy_fallback_preserves_origin_and_proxy_prefix(base_path: str, status: int) -> None:
+@pytest.mark.parametrize(
+    "current_api_response",
+    (
+        httpx.Response(404),
+        httpx.Response(405),
+        httpx.Response(200, json={"error": "Unexpected endpoint or method. (GET /api/v1/models)"}),
+    ),
+)
+async def test_legacy_fallback_preserves_origin_and_proxy_prefix(
+    base_path: str, current_api_response: httpx.Response
+) -> None:
     def respond(request: httpx.Request) -> httpx.Response:
         assert request.url.host == "backend.test"
         assert request.headers["authorization"] == "Bearer local-key"
         if request.url.path == "/tenant/api/v1/models":
-            return httpx.Response(status)
+            return current_api_response
         assert request.url.path == "/tenant/api/v0/models"
         return httpx.Response(
             200,
@@ -146,14 +155,13 @@ async def test_invalid_context_is_not_advertised(context: object) -> None:
         (httpx.Response(404), 2),
         (httpx.Response(405), 2),
         (httpx.Response(302, headers={"location": "https://elsewhere.test"}), 1),
-        (httpx.Response(200, content=b"not json"), 1),
-        (httpx.Response(200, json={"models": None}), 1),
-        (httpx.Response(200, json={"data": []}), 1),
+        (httpx.Response(200, content=b"not json"), 2),
+        (httpx.Response(200, json={"models": None}), 2),
         (httpx.ReadTimeout("offline"), 1),
         (httpx.ConnectError("offline"), 1),
     ),
 )
-async def test_failures_are_negative_cached_without_redirects_or_unbounded_fallbacks(
+async def test_failures_are_reported_as_unknown_and_negative_cached_without_redirects_or_unbounded_fallbacks(
     response: httpx.Response | Exception,
     requests_per_attempt: int,
 ) -> None:
@@ -164,30 +172,37 @@ async def test_failures_are_negative_cached_without_redirects_or_unbounded_fallb
     async with httpx.AsyncClient(transport=httpx.MockTransport(responder), follow_redirects=True) as client:
         handler.client = client
         cache: Final = InMemoryCache(clock=clock)
-        for _ in range(2):
-            assert (
-                await get_lm_studio_model_info(
-                    model="model",
-                    api_base="https://backend.test",
-                    headers={"authorization": "secret"},
-                    client=handler,
-                    cache=cache,
-                )
-                == {}
-            )
-        assert responder.call_count == requests_per_attempt
-        clock.return_value = 61
-        assert (
-            await get_lm_studio_model_info(
+
+        async def lookup() -> Mapping[str, int] | None:
+            return await get_lm_studio_model_info(
                 model="model",
                 api_base="https://backend.test",
                 headers={"authorization": "secret"},
                 client=handler,
                 cache=cache,
             )
+
+        assert await lookup() is None
+        assert await lookup() is None
+        assert responder.call_count == requests_per_attempt
+        clock.return_value = 61
+        assert await lookup() is None
+        assert responder.call_count == 2 * requests_per_attempt
+
+
+async def test_server_without_the_model_loaded_is_known_to_have_no_limits() -> None:
+    responder: Final = Mock(return_value=httpx.Response(200, json={"models": []}))
+    handler: Final = AsyncHTTPHandler()
+    await handler.client.aclose()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(responder)) as client:
+        handler.client = client
+        assert (
+            await get_lm_studio_model_info(
+                model="model", api_base="https://backend.test", headers={}, client=handler, cache=InMemoryCache()
+            )
             == {}
         )
-        assert responder.call_count == 2 * requests_per_attempt
+        assert responder.call_count == 1
 
 
 async def test_cache_is_scoped_to_server_and_credentials_and_refreshes_context() -> None:
@@ -215,16 +230,18 @@ async def test_cache_is_scoped_to_server_and_credentials_and_refreshes_context()
         cache: Final = InMemoryCache(clock=clock)
 
         async def lookup(model: str = "model", host: str = "one.test", key: str = "one") -> Mapping[str, int]:
-            return await get_lm_studio_model_info(
+            limits: Final = await get_lm_studio_model_info(
                 model=model, api_base=f"https://{host}/v1", headers={"authorization": key}, client=handler, cache=cache
             )
+            assert limits is not None
+            return limits
 
         assert (await lookup())["max_input_tokens"] == 8192
         assert (await lookup("instance"))["max_input_tokens"] == 8192
         assert responder.call_count == 1
         assert (await lookup(key="two"))["max_input_tokens"] == 4096
         assert (await lookup(host="two.test"))["max_input_tokens"] == 2048
-        clock.return_value = MODEL_INFO_REFRESH_SECONDS + 1
+        clock.return_value = MODEL_INFO_REFRESH_SECONDS // 2
         assert (await lookup())["max_input_tokens"] == 1024
         assert responder.call_count == 4
 
@@ -244,6 +261,6 @@ async def test_invalid_or_missing_server_never_sends_credentials(api_base: str)
                 client=handler,
                 cache=InMemoryCache(),
             )
-            == {}
+            is None
         )
         assert responder.call_count == 0

tests/unit/test_router_model_cost_isolation.py

@@ -174,20 +174,26 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch
     _invalidate_model_cost_lowercase_map()
 
 
-async def test_discovery_sends_the_secret_behind_an_oidc_api_key(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+async def test_discovery_sends_the_secret_behind_an_oidc_api_key(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
     monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
     monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path))
     token: Final = tmp_path / "token"
     token.write_text("local-secret")
-    router: Final = Router(model_list=[{
-        "model_name": "local",
-        "litellm_params": {
-            "model": "hosted_vllm/local-model",
-            "api_base": "https://discovery.test/v1",
-            "api_key": f"oidc/file/{token}",
-        },
-        "model_info": {"id": "oidc-deployment"},
-    }])
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "local",
+                "litellm_params": {
+                    "model": "hosted_vllm/local-model",
+                    "api_base": "https://discovery.test/v1",
+                    "api_key": f"oidc/file/{token}",
+                },
+                "model_info": {"id": "oidc-deployment"},
+            }
+        ]
+    )
     authorizations: Final[list[str | None]] = []
 
     def respond(request: httpx.Request) -> httpx.Response:
@@ -293,7 +299,28 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa
     _invalidate_model_cost_lowercase_map()
 
 
-async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
+def _upstream_models_response(provider: str, context: int) -> httpx.Response:
+    return httpx.Response(
+        200,
+        json=(
+            {
+                "models": [
+                    {
+                        "key": "local-model",
+                        "loaded_instances": [{"id": "local-model", "config": {"context_length": context}}],
+                    }
+                ]
+            }
+            if provider == "lm_studio"
+            else {"data": [{"id": "local-model", "max_model_len": context}]}
+        ),
+    )
+
+
+@pytest.mark.parametrize("provider", ("hosted_vllm", "lm_studio"))
+async def test_discovered_limits_expire_after_the_last_successful_refresh(
+    provider: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
     monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
     clock: Final = Mock(return_value=0.0)
     router: Final = Router(
@@ -301,7 +328,7 @@ async def test_discovered_limits_expire_after_the_last_successful_refresh(monkey
             {
                 "model_name": "local",
                 "litellm_params": {
-                    "model": "hosted_vllm/local-model",
+                    "model": f"{provider}/local-model",
                     "api_base": "https://expiry.test/v1",
                     "api_key": "local-key",
                 },
@@ -312,8 +339,8 @@ async def test_discovered_limits_expire_after_the_last_successful_refresh(monkey
     router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS)
     responses: Final = iter(
         (
-            httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
-            httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}),
+            _upstream_models_response(provider, 4096),
+            _upstream_models_response(provider, 8192),
             httpx.Response(503),
         )
     )
@@ -601,7 +628,9 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing():
     )
 
 
-@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"])
+@pytest.mark.parametrize(
+    "flag", ["supports_regex_lookaround", "supports_assistant_prefill", "supports_sampling_params"]
+)
 def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None:
     """A deployment's capability override must not land on the shared
     ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it."""
@@ -1083,8 +1112,12 @@ def test_bedrock_deployment_keeps_custom_prices_when_profile_lookup_fails(
     tier_rates: Final = inherit_configured and service_tier != "default"
     input_rate: Final = configured_prices[f"input_cost_per_token_{service_tier}"] if tier_rates else 0.01
     output_rate: Final = configured_prices[f"output_cost_per_token_{service_tier}"] if tier_rates else 0.02
-    effective_input: Final = 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate
-    cache_rate: Final = 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input)
+    effective_input: Final = (
+        0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate
+    )
+    cache_rate: Final = (
+        0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input)
+    )
     assert completion_cost(
         model=model,
         completion_response=ModelResponse(
@@ -2859,7 +2892,8 @@ def test_a_github_copilot_chat_pin_opts_a_dual_endpoint_model_out_of_the_respons
 
 @pytest.mark.parametrize("registered_mode, pinned_mode", [("chat", "responses"), ("responses", "chat")])
 @pytest.mark.parametrize(
-    "model, custom_llm_provider", [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")]
+    "model, custom_llm_provider",
+    [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")],
 )
 def test_a_github_copilot_responses_pin_still_opts_the_model_into_the_responses_api(
     monkeypatch: pytest.MonkeyPatch,
@@ -4037,7 +4071,9 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info(
     assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price
 
 
-def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
     catalog: Final = litellm.get_model_cost_map(url="")
     monkeypatch.setattr(litellm, "model_cost", dict(catalog))
@@ -4525,7 +4561,6 @@ def test_pricing_provider_fills_unset_rates_from_that_providers_catalog_row() ->
     (
         httpx.Response(200, json={"models": [{"key": "org/model", "loaded_instances": []}]}),
         httpx.Response(200, json={"models": []}),
-        httpx.Response(503),
     ),
 )
 @pytest.mark.parametrize("input_limit", (None, 1000))