diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index f419812947e..08afb2cbfb9 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -140,7 +140,7 @@ from litellm.router_rust.routing_stores import ( RoutingStores, service_target_of, ) -from litellm.router_rust.steps.access import caller_fields, selection_refusal +from litellm.router_rust.steps.access import caller_fields, pinned_deployment_wildcard_grant, selection_refusal from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks from litellm.router_rust.steps.plugins import SYNC_PATH_MESSAGE, candidate_models, run_plugins from litellm.router_rust.steps.precall import ( @@ -2900,6 +2900,11 @@ class RustCore( ) fields: dict[str, object] = dict(_start_fields(request)) # mutable-ok: built up below # rebind-ok: same fields.update(self._group_policy_fields(request)) + pinned_grant: Final = pinned_deployment_wildcard_grant( + request.get("model"), fields.get("allowed_models"), get_request_team_id(request), self.registry() + ) + if pinned_grant is not None: + fields["allowed_models"] = list(pinned_grant) if filter_tags: fields["filter_tags"] = True if request_carries_image(messages, request): diff --git a/litellm/router_rust/steps/access.py b/litellm/router_rust/steps/access.py index 087a87e2c6b..2f513b616b5 100644 --- a/litellm/router_rust/steps/access.py +++ b/litellm/router_rust/steps/access.py @@ -5,13 +5,16 @@ caller to its groups' deployments) and, for a proxy admin, the model the request from __future__ import annotations from collections.abc import Mapping -from typing import Final +from typing import TYPE_CHECKING, Final from pydantic import TypeAdapter from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request, request_user_api_key_auth +if TYPE_CHECKING: + from litellm.router_api import RouterRegistry + _NAMES: Final = TypeAdapter(tuple[object, ...]) @@ -42,6 +45,22 @@ def caller_fields(request: Mapping[str, object]) -> Mapping[str, object]: } +def pinned_deployment_wildcard_grant( + name: object, allowed_models: object, team_id: str | None, registry: RouterRegistry +) -> tuple[str, ...] | None: + granted: Final = _names(allowed_models) + if not isinstance(name, str) or not granted or name in granted or name in registry.index("model_name"): + return None + deployment: Final = registry.deployment(name) + if deployment is None: + return None + from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list + + access_groups: Final = registry.access_groups(model_name=deployment.model_name, team_id=team_id) + patterns: Final = tuple(model for model in granted if model not in access_groups) + return (*granted, name) if model_matches_any_wildcard_pattern_in_list(deployment.model_name, patterns) else None + + def _router_weights_error(raw: object) -> BaseException: """The error ``validate_router_weights`` raises for weights the Rust core refused.""" from pydantic import ValidationError diff --git a/litellm/rust/crates/router-core/src/resolve.rs b/litellm/rust/crates/router-core/src/resolve.rs index d897716b1e5..fdd87742aac 100644 --- a/litellm/rust/crates/router-core/src/resolve.rs +++ b/litellm/rust/crates/router-core/src/resolve.rs @@ -326,7 +326,17 @@ pub fn resolve_request( { let group = deployment.model_name.clone(); let litellm_model = deployment.params.model.clone(); - let kept = reserved(&litellm_model, vec![deployment], caller, now)?; + let accessible = if caller.allowed_models.iter().any(|m| m == name) { + vec![deployment] + } else { + filter_by_access_groups(&group, vec![deployment], caller.allowed_models, || { + model_access_groups(registry, &group, team_id) + }) + }; + if accessible.is_empty() { + return Ok(None); + } + let kept = reserved(&litellm_model, accessible, caller, now)?; return resolved(ResolveKind::Id, &group, kept); } let aliased = registry.alias_of(name); @@ -440,7 +450,7 @@ pub fn resolve_request( mod tests { use std::collections::BTreeMap; - use rstest::rstest; + use rstest::{fixture, rstest}; use serde_json::{Value, json}; use super::*; @@ -483,6 +493,47 @@ mod tests { assert_eq!(ids(by_model.unwrap()), ["b"]); } + #[fixture] + fn shared_group() -> Registry { + registry(json!([ + {"model_name": "shared", "litellm_params": {"model": "openai/pub"}, "model_info": {"id": "pub", "access_groups": ["public"]}}, + {"model_name": "shared", "litellm_params": {"model": "openai/priv"}, "model_info": {"id": "priv", "access_groups": ["private"]}}, + ])) + } + + #[rstest] + #[case::access_group_pins_its_own_deployment(&["public"], "pub", &["pub"])] + #[case::access_group_cannot_pin_another_groups_deployment(&["public"], "priv", &[])] + #[case::the_other_access_group_pins_its_deployment(&["private"], "priv", &["priv"])] + #[case::the_owning_group_pins_any_of_its_deployments(&["shared"], "priv", &["priv"])] + #[case::an_exact_id_grant_beside_another_access_group(&["public", "priv"], "priv", &["priv"])] + #[case::an_unrestricted_caller_pins_any_deployment(&[], "priv", &["priv"])] + #[case::every_model_pins_any_deployment(&["*"], "priv", &["priv"])] + #[case::every_proxy_model_pins_any_deployment(&["all-proxy-models"], "priv", &["priv"])] + #[case::the_group_name_stays_kept_to_the_callers_access_group(&["public"], "shared", &["pub"])] + fn pinned_deployment_respects_caller_access_groups( + shared_group: Registry, + #[case] allowed: &[&str], + #[case] name: &str, + #[case] expected: &[&str], + ) { + let allowed_models: Vec = allowed.iter().map(|m| (*m).to_owned()).collect(); + let caller = Caller { + allowed_models: &allowed_models, + ..Caller::default() + }; + let resolved = resolve_request( + &shared_group, + name, + false, + caller, + None, + None, + Timestamp::UNIX_EPOCH, + ); + assert_eq!(ids(resolved.unwrap()), expected); + } + #[test] fn a_proxy_admin_is_refused_a_public_name_several_teams_share() { let registry = registry(json!([ diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 6203c12be7c..fc9623c5ed9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -10,6 +10,7 @@ import json import re import unittest from collections.abc import AsyncIterator, Callable, Mapping +from types import MappingProxyType from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -1680,6 +1681,88 @@ async def test_a_plain_read_through_the_proxy_still_answers_json(monkeypatch: py assert [dict(request.url.params) for request in upstream_requests] == [{}] +_PUBLIC_UPSTREAM: Final = "https://public.example/v1" +_PRIVATE_UPSTREAM: Final = "https://private.example/v1" + + +def _auth_allowing(models: tuple[str, ...]) -> Callable[[], UserAPIKeyAuth]: + def auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-access-group", models=list(models)) + + return auth + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("models", "model_id", "status_code", "served_upstream"), + [ + pytest.param(("public",), "shared-public", 200, _PUBLIC_UPSTREAM, id="pinned-to-the-keys-access-group"), + pytest.param(("public",), "shared-private", 400, None, id="pinned-outside-the-keys-access-group"), + pytest.param( + ("shared-responses*", "public"), + "shared-private", + 200, + _PRIVATE_UPSTREAM, + id="a-wildcard-for-the-group-keeps-every-pin", + ), + pytest.param( + ("shared-other*", "public"), "shared-private", 400, None, id="an-unrelated-wildcard-grants-no-pin" + ), + ], +) +async def test_a_response_read_reaches_only_a_deployment_in_the_keys_access_group( + monkeypatch: pytest.MonkeyPatch, + models: tuple[str, ...], + model_id: str, + status_code: int, + served_upstream: str | None, +) -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.responses.utils import ResponsesAPIRequestUtils + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "shared-responses", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": upstream, "api_key": "fixture-key"}, + "model_info": {"id": deployment_id, "access_groups": [access_group]}, + } + for deployment_id, upstream, access_group in ( + ("shared-public", _PUBLIC_UPSTREAM, "public"), + ("shared-private", _PRIVATE_UPSTREAM, "private"), + ) + ], + num_retries=0, + ) + try: + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_allowing(models)) + response_id: Final = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id=model_id, response_id="resp_pinned" + ) + stored: Final = _background_response("resp_pinned", "in_progress", 100).model_dump(mode="json") + with respx.mock(assert_all_called=False) as transport: + reads: Final = MappingProxyType( + { + upstream: transport.get(f"{upstream}/responses/resp_pinned").mock( + return_value=httpx.Response(200, json=stored) + ) + for upstream in (_PUBLIC_UPSTREAM, _PRIVATE_UPSTREAM) + } + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.get(f"/v1/responses/{response_id}") + + assert result.status_code == status_code, result.text + assert {upstream: route.call_count for upstream, route in reads.items()} == { + upstream: 1 if upstream == served_upstream else 0 for upstream in reads + } + finally: + router.retire() + + @pytest.mark.asyncio @pytest.mark.parametrize( "query, bad_param",