scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

fix(router)!: resolve trusted JEV connection environment references (#48717)

Browse source
This commit is contained in:
tin-berri 2026-10-28 14:28:55 -04:00 • committed by GitHub
parent 8c9a76859d
commit c40bdc475f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 301 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -49,6 +49,8 @@ from litellm.proxy.management_endpoints.common_daily_activity import daily_activ
from litellm.proxy.management_helpers.auto_router_permissions import (
authorize_member_auto_router_dependencies,
authorize_member_auto_router_team,
auto_router_environment_reference_error,
raise_auto_router_environment_reference_error,
validate_member_auto_router_config,
)
from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository
@ -402,6 +404,10 @@ async def validate_complexity_router_config(
"""
member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id)
raise_auto_router_environment_reference_error(
auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict)
)
from litellm.router_utils.auto_router_model_naming import (
validate_complexity_router_config_write,
)
@ -560,6 +566,12 @@ async def preview_auto_router_routing(
detail={"error": CommonProxyErrors.no_llm_router.value},
)
resolved: Final = await _resolve_saved_routing_test(data, user_api_key_dict, llm_router)
if data.saved_model_id is None:
raise_auto_router_environment_reference_error(
auto_router_environment_reference_error(
data.complexity_router_config.model_dump(exclude_none=True), user_api_key_dict
)
)
actor: Final = (
await _authorize_member_dry_run_config(
config=resolved.complexity_router_config.model_dump(exclude_none=True),

10
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -102,6 +102,8 @@ from litellm.proxy.management_helpers.auto_router_permissions import (
authorize_member_auto_router_dependencies,
authorize_member_auto_router_team,
authorize_member_auto_router_write,
auto_router_environment_reference_error,
raise_auto_router_environment_reference_error,
)
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
from litellm.proxy.management_helpers.model_quota_policy import (
@ -2627,6 +2629,14 @@ class ModelManagementAuthChecks:
ModelManagementAuthChecks._reject_team_inline_api_key_secret_reference(
model_params=model_params, incoming_params=incoming_params, incoming_model_params=incoming_model_params
)
raise_auto_router_environment_reference_error(
auto_router_environment_reference_error(
config=_OPTIONAL_JSON_OBJECT.validate_python(
incoming_params.get("complexity_router_config") if incoming_params is not None else None
),
user_api_key_dict=user_api_key_dict,
)
)
# Federation fields choose which server-side secret is read and where the org-scoped token
# it buys is sent, so only a proxy admin may point a federated deployment somewhere else.
# Evaluated on the RESULTING deployment: a patch attaching a credential by name inherits

35
litellm/proxy/management_helpers/auto_router_permissions.py
View file

@ -5,8 +5,8 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
from fastapi import HTTPException
from pydantic import ConfigDict, Field, ValidationError
from typing_extensions import ReadOnly, TypedDict
from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm.models.organization import LiteLLM_OrganizationTable
from litellm.models.project import LiteLLM_ProjectTable
@ -36,6 +36,7 @@ from litellm.repositories.table_repositories import TeamMembershipRepository
from litellm.router import Router
from litellm.router_strategy.complexity_router.config import (
ComplexityRouterConfigWrite,
normalize_classifier_config_aliases,
resolve_complexity_router_config_write,
)
from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig
@ -47,6 +48,36 @@ if TYPE_CHECKING:
from prisma import types as prisma_types
def auto_router_environment_reference_error(
config: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth
) -> Literal["environment_reference_denied"] | None:
if config is None or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return None
classifier: Final = normalize_classifier_config_aliases(config).get("opensource_classifier_config")
if not isinstance(classifier, Mapping):
return None
fields: Final = TypeAdapter(Mapping[str, object]).validate_python(classifier)
if any(
isinstance(value, str) and value.startswith("os.environ/")
for value in (fields.get("api_key"), fields.get("api_base"))
):
return "environment_reference_denied"
return None
def raise_auto_router_environment_reference_error(error: Literal["environment_reference_denied"] | None) -> None:
match error:
case None:
return
case "environment_reference_denied":
raise HTTPException(
status_code=403,
detail="Only proxy admins can submit classifier connection environment references",
)
case _:
assert_never(error)
class _MemberRouterThinking(LiteLLMBaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)

33
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -23,6 +23,7 @@ import random
import re
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from itertools import accumulate, chain, islice, takewhile
from threading import Lock
@ -1390,6 +1391,11 @@ def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None)
return {"model": model, "tier": tier_value}
@dataclass(frozen=True, slots=True)
class _JevConnectionError:
message: str
class ComplexityRouter(CustomLogger):
"""
Complexity router that classifies requests and routes to appropriate models.
@ -1405,7 +1411,7 @@ class ComplexityRouter(CustomLogger):
- Question complexity (multiple questions)
"""
def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient:
def _build_jev_client(self, config: OpenSourceClassifierConfig) -> JevClassifierClient | _JevConnectionError:
if config.deployment_name is not None:
return RouterJevClassifierClient(self.litellm_router_instance, config.deployment_name)
if config.provider in ("laya", "bespoke"):
@ -1426,12 +1432,22 @@ class ComplexityRouter(CustomLogger):
http_client=get_async_httpx_client(httpxSpecialProvider.PassThroughEndpoint),
provider="databricks",
)
api_key: Final = config.api_key or get_secret_str("TYPESAFE_API_KEY")
if not api_key:
raise ValueError(
api_key: Final = (
get_secret_str(config.api_key)
if config.api_key is not None and config.api_key.startswith("os.environ/")
else config.api_key or get_secret_str("TYPESAFE_API_KEY")
)
if not api_key or not api_key.strip():
return _JevConnectionError(
"opensource_classifier_config.api_key or TYPESAFE_API_KEY is required for classifier_type 'oss_classifier'"
)
api_base: Final = config.api_base or get_secret_str("TYPESAFE_API_BASE") or "https://api.typesafe.ai"
api_base: Final = (
get_secret_str(config.api_base)
if config.api_base is not None and config.api_base.startswith("os.environ/")
else config.api_base or get_secret_str("TYPESAFE_API_BASE") or "https://api.typesafe.ai"
)
if not api_base or not api_base.strip():
return _JevConnectionError("opensource_classifier_config.api_base must resolve to a non-empty value")
return HttpJevClassifierClient(
api_key=api_key,
api_base=api_base,
@ -1474,13 +1490,18 @@ class ComplexityRouter(CustomLogger):
self.config.default_model = default_model
jev_config: Final = self.config.opensource_classifier_config
self._jev_client: JevClassifierClient | None = (
jev_client_result: Final = (
jev_client
if jev_client is not None
else self._build_jev_client(jev_config)
if self.config.classifier_type == "oss_classifier" and jev_config is not None
else None
)
match jev_client_result:
case _JevConnectionError(message=message):
raise ValueError(message)
case _:
self._jev_client: JevClassifierClient | None = jev_client_result
self._tier_affinity_config = hashlib.sha256(
self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode()

75
tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py
View file

@ -2769,15 +2769,22 @@ async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typ
@pytest.mark.asyncio
@pytest.mark.parametrize(
"case", ["allowed", "credential-free", "missing", "blocked", "key", "budget", "team", "not-router"]
"case",
["allowed", "credential-free", "team-admin-environment", "missing", "blocked", "key", "budget", "team", "not-router"],
)
async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: pytest.MonkeyPatch, case: str) -> None:
router: Final = RecordingRouter("SIMPLE")
stored_key: Final = "synthetic-server-jev-key"
monkeypatch.setenv("SAVED_JEV_KEY", stored_key)
monkeypatch.setenv("SAVED_JEV_BASE", "https://saved-jev.test")
stored_config: Final = {
"classifier_type": "jev",
"tiers": TIERS,
"jev_classifier_config": {"api_key": stored_key, "api_base": "https://saved-jev.test"},
"jev_classifier_config": (
{"api_key": "os.environ/SAVED_JEV_KEY", "api_base": "os.environ/SAVED_JEV_BASE"}
if case == "team-admin-environment"
else {"api_key": stored_key, "api_base": "https://saved-jev.test"}
),
}
router.add_deployment(
Deployment.model_validate(
@ -2790,7 +2797,10 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch:
"model_info": {
"id": "saved-jev-id",
"blocked": case == "blocked",
"team_id": "owner-team" if case == "team" else None,
"team_id": (
"owner-team" if case == "team"
else "admin-preview-team" if case == "team-admin-environment" else None
),
},
}
)
@ -2799,6 +2809,8 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch:
actor: Final = (
_configure_member_preview(monkeypatch)
if case == "team"
else _configure_team_admin_preview(monkeypatch)
if case == "team-admin-environment"
else UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-probe",
@ -2812,12 +2824,17 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch:
{
"prompt": "what is 2+2",
"saved_model_id": "missing-id" if case == "missing" else "saved-jev-id",
"team_id": "member-preview-team" if case == "team" else None,
"team_id": (
"member-preview-team" if case == "team"
else "admin-preview-team" if case == "team-admin-environment" else None
),
},
classifier_type="jev",
jev_classifier_config=(
{"model": "jev-latest", "timeout_ms": 3000}
if case == "credential-free"
else {"api_key": "os.environ/IGNORED_BROWSER_KEY", "api_base": "https://browser-override.test"}
if case == "team-admin-environment"
else {"api_key": "masked-key", "api_base": "https://browser-override.test"}
),
)
@ -2856,7 +2873,7 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch:
assert result.routed_model == "cheap-model"
assert evaluation.calls.last.request.headers["authorization"] == f"Bearer {stored_key}"
assert stored_key not in result.model_dump_json()
assert evaluation.call_count == (1 if case in ("allowed", "credential-free") else 0)
assert evaluation.call_count == (1 if case in ("allowed", "credential-free", "team-admin-environment") else 0)
assert router.recorded_calls == []
await handler.client.aclose()
@ -3378,6 +3395,54 @@ async def test_validate_config_gates_like_the_write_it_rehearses(monkeypatch: py
assert not_their_team.value.status_code == 403
def _configure_team_admin_preview(monkeypatch: pytest.MonkeyPatch) -> UserAPIKeyAuth:
from litellm.proxy._types import LiteLLM_TeamTable
team: Final = LiteLLM_TeamTable(
team_id="admin-preview-team",
models=["saved-jev", "typesafe/jev-latest", *[TIERS[name][0] for name in TIERS]],
members_with_roles=[{"role": "admin", "user_id": "preview-admin"}],
)
database: Final = MagicMock()
database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team)
monkeypatch.setattr(proxy_server, "prisma_client", database)
monkeypatch.setattr(proxy_server, "premium_user", True)
return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="preview-admin", api_key="sk-preview")
@pytest.mark.asyncio
@pytest.mark.parametrize("legacy", [False, True])
@pytest.mark.parametrize("field", ["api_key", "api_base"])
async def test_team_admin_cannot_preview_or_validate_environment_references(
monkeypatch: pytest.MonkeyPatch, legacy: bool, field: str
) -> None:
from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest
actor: Final = _configure_team_admin_preview(monkeypatch)
monkeypatch.setattr(proxy_server, "llm_router", _router())
config: Final = {
"tiers": TIERS,
"classifier_type": "jev" if legacy else "oss_classifier",
"jev_classifier_config" if legacy else "opensource_classifier_config": {
"api_key": "caller-key", "api_base": "https://denied-jev.test", field: "os.environ/PRIVATE_JEV_VALUE",
},
}
validation: Final = ComplexityRouterConfigValidationRequest(
team_id="admin-preview-team", complexity_router_config=config
)
with pytest.raises(HTTPException) as denied_validation:
await auto_router_endpoints.validate_complexity_router_config(validation, actor)
assert denied_validation.value.status_code == 403
with respx.mock(assert_all_called=False) as http, pytest.raises(HTTPException) as denied_preview:
await preview_auto_router_routing(
_request_from({"prompt": "what is 2+2", "team_id": "admin-preview-team"}, **config),
actor,
ROUTING_HTTP_REQUEST,
)
assert denied_preview.value.status_code == 403
assert len(http.calls) == 0
def _configure_member_preview(monkeypatch: pytest.MonkeyPatch, *, allowed: bool = True) -> UserAPIKeyAuth:
from litellm.proxy import proxy_server
from litellm.proxy._types import UI_TEAM_ID, LiteLLM_TeamTable

70
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -8475,6 +8475,76 @@ class TestTeamMemberAutoRouterWrites:
"jev_classifier_config" if legacy else "opensource_classifier_config": classifier,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("endpoint", ["create", "patch", "legacy"])
@pytest.mark.parametrize("legacy", [False, True])
@pytest.mark.parametrize("field", ["api_key", "api_base"])
async def test_team_admin_classifier_writes_reject_environment_references(
self, endpoint: str, legacy: bool, field: str
) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
stored_config: Final = self._classifier_config(
{"api_key": "os.environ/PRIVATE_JEV_VALUE", "api_base": "https://trusted-jev.test"}, legacy
)
row: Final = self._row().model_copy(update={"litellm_params": {
"model": "auto_router/complexity_router", "complexity_router_config": stored_config,
}})
team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]})
database: Final = self._database(team, row)
actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER)
config: Final = self._classifier_config(
{"api_key": "caller-key", "api_base": "https://untrusted-jev.test", field: "os.environ/PRIVATE_JEV_VALUE"},
legacy,
)
request: Final = updateDeployment(
litellm_params=updateLiteLLMParams(complexity_router_config=config),
model_info=ModelInfo(id=row.model_id, team_id="member-team"),
)
operation: Final = (
add_new_model(Deployment(
model_name="new-jev-router",
litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=config),
model_info=ModelInfo(id=row.model_id, team_id="member-team"),
), actor)
if endpoint == "create"
else patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor)
)
with self._environment(database, row), pytest.raises((HTTPException, ProxyException)) as denied:
await operation
assert (denied.value.status_code if isinstance(denied.value, HTTPException) else int(denied.value.code)) == 403
database.db.litellm_proxymodeltable.create.assert_not_awaited()
database.db.litellm_proxymodeltable.update.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("endpoint", ["patch", "legacy"])
@pytest.mark.parametrize("change_config", [False, True])
async def test_team_admin_unrelated_updates_preserve_stored_environment_references(
self, endpoint: str, change_config: bool
) -> None:
connection: Final = {"api_key": "os.environ/SAVED_JEV_KEY", "api_base": "os.environ/SAVED_JEV_BASE"}
stored_config: Final = self._classifier_config(connection, True)
row: Final = self._row().model_copy(update={"litellm_params": {
"model": "auto_router/complexity_router", "complexity_router_config": stored_config,
}})
team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]})
database: Final = self._database(team, row)
actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER)
params: Final = (
{"complexity_router_config": self._classifier_config({"timeout_ms": 8100}, False)}
if change_config else {"max_tokens": 100}
)
request: Final = updateDeployment(
litellm_params=updateLiteLLMParams.model_validate(params), model_info=ModelInfo(id=row.model_id),
)
with self._environment(database, row):
await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor))
written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"]
saved: Final = json.loads(written["litellm_params"])["complexity_router_config"]
assert saved == (
self._classifier_config({**connection, "timeout_ms": 8100}, False) if change_config else stored_config
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"stored_params,removed,denied_param",

79
tests/unit/router_strategy/complexity_router/test_jev_classifier.py
View file

@ -732,6 +732,85 @@ def test_jev_api_base_without_its_own_key_is_rejected_so_the_environment_key_sta
assert JevClassifierConfig(api_key="sk-own").api_base is None
@pytest.mark.asyncio
@pytest.mark.parametrize("legacy", [False, True])
@pytest.mark.parametrize("connection", ["literal", "key-reference", "base-reference", "references", "defaults"])
async def test_jev_connection_resolves_only_explicit_environment_references(
monkeypatch: pytest.MonkeyPatch, legacy: bool, connection: str
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("JEV_CONNECTION_KEY", "resolved-key")
monkeypatch.setenv("JEV_CONNECTION_BASE", "https://jev.test")
monkeypatch.setenv("TYPESAFE_API_KEY", "default-key")
monkeypatch.setenv("TYPESAFE_API_BASE", "https://default-jev.test")
monkeypatch.setenv("literal-key", "must-not-replace-literal")
fields: Final = (
{}
if connection == "defaults"
else {
"api_key": "os.environ/JEV_CONNECTION_KEY"
if connection in ("key-reference", "references")
else "literal-key",
"api_base": "os.environ/JEV_CONNECTION_BASE"
if connection in ("base-reference", "references")
else "https://jev.test",
}
)
config: Final = {
"classifier_type": "jev" if legacy else "oss_classifier",
"jev_classifier_config" if legacy else "opensource_classifier_config": {"model": "openjev-latest", **fields},
"tiers": {"SIMPLE": "cheap", "COMPLEX": "strong"},
}
classifier: Final = ComplexityRouter("env-router", litellm.Router(), config, derive_savings_baseline=False)
with respx.mock(assert_all_called=True) as upstream:
route: Final = upstream.post(
"https://default-jev.test/v1/systemone" if connection == "defaults" else "https://jev.test/v1/systemone"
).respond(json={"answers": {"tier": _answer("SIMPLE").model_dump()}})
outcome: Final = await classifier.aclassify("What is two plus two?")
expected_key: Final = (
"default-key"
if connection == "defaults"
else "resolved-key" if connection in ("key-reference", "references") else "literal-key"
)
assert outcome.cause == "jev_classifier"
assert route.calls[0].request.headers["authorization"] == f"Bearer {expected_key}"
assert classifier.config.opensource_classifier_config is not None
assert classifier.config.opensource_classifier_config.api_key == fields.get("api_key")
assert classifier.config.opensource_classifier_config.api_base == fields.get("api_base")
@pytest.mark.parametrize("field", ["api_key", "api_base"])
@pytest.mark.parametrize("value", [None, "", " "])
def test_jev_missing_explicit_environment_reference_never_uses_provider_defaults(
monkeypatch: pytest.MonkeyPatch, field: str, value: str | None
) -> None:
monkeypatch.setenv("TYPESAFE_API_KEY", "never-send-default-key")
monkeypatch.setenv("TYPESAFE_API_BASE", "https://never-use-default.test")
if value is None:
monkeypatch.delenv("JEV_MISSING_CONNECTION", raising=False)
else:
monkeypatch.setenv("JEV_MISSING_CONNECTION", value)
with (
respx.mock(assert_all_called=False) as upstream,
pytest.raises(ValueError, match=rf"opensource_classifier_config\.{field}"),
):
ComplexityRouter(
"env-router",
litellm.Router(),
{
"classifier_type": "jev",
"jev_classifier_config": {
"api_key": "literal-key",
"api_base": "https://jev.test",
field: "os.environ/JEV_MISSING_CONNECTION",
},
"tiers": {"SIMPLE": "cheap"},
},
derive_savings_baseline=False,
)
assert len(upstream.calls) == 0
def test_databricks_api_base_without_its_own_key_is_rejected_so_the_workspace_token_stays_home() -> None:
with pytest.raises(ValueError, match="DATABRICKS_API_KEY or DATABRICKS_TOKEN is only sent to DATABRICKS_API_BASE"):
JevClassifierConfig.model_validate(