scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #35028 from BerriAI/litellm_batch_provider_credentials

Browse source 
fix(proxy): resolve named credentials on provider-only batch and files calls
This commit is contained in:
Mateo Wang 2026-07-29 11:59:09 -07:00 • committed by GitHub
commit 8b03315ac6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 945 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -23,6 +23,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
apply_team_provider_credentials,
decode_model_from_file_id,
encode_batch_response_ids,
encode_file_id_with_model,
@ -295,6 +296,12 @@ async def create_batch(
verbose_proxy_logger.debug(f"Created batch using model: {model_param}")
else:
# SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
apply_team_provider_credentials(
data=cast(dict, _create_batch_data), # cast-ok: TypedDict is a dict at runtime
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
response = await litellm.acreate_batch(
custom_llm_provider=custom_llm_provider,
**_create_batch_data, # type: ignore
@ -525,6 +532,12 @@ async def retrieve_batch(
or get_custom_llm_provider_from_request_query(request=request)
or "openai"
)
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
response = await litellm.aretrieve_batch(
custom_llm_provider=custom_llm_provider,
**data, # type: ignore
@ -718,6 +731,12 @@ async def list_batches(
or get_custom_llm_provider_from_request_query(request=request)
or "openai"
)
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
response = await litellm.alist_batches(
custom_llm_provider=custom_llm_provider, # type: ignore
after=after,
@ -908,6 +927,12 @@ async def cancel_batch(
# Extract batch_id from data to avoid "multiple values for keyword argument" error
# data was cast from CancelBatchRequest which already contains batch_id
data.pop("batch_id", None)
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
_cancel_batch_data = CancelBatchRequest(batch_id=batch_id, **data)
response = await litellm.acancel_batch(
custom_llm_provider=custom_llm_provider, # type: ignore

102
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -14,6 +14,7 @@ from litellm.types.utils import SpecialEnums
if TYPE_CHECKING:
from fastapi import Request
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
@ -294,9 +295,8 @@ def get_credentials_for_model(
def get_team_provider_credentials(
llm_router: Optional["Router"],
team_models: List[str],
user_api_key_dict: "UserAPIKeyAuth",
custom_llm_provider: str,
team_id: Optional[str] = None,
) -> Optional[dict]:
"""
Resolve upstream credentials for a provider-scoped file operation
@ -304,21 +304,61 @@ def get_team_provider_credentials(
Priority:
1. The team's own (BYOK) deployment for this provider — a deployment whose
``model_info.team_id`` matches ``team_id``. This keeps team-scoped listings
on the team's own provider account/key instead of a shared global one.
2. Fallback: any deployment the team is granted access to for this provider,
expanding wildcard routes and the all-proxy-models sentinel.
``model_info.team_id`` matches the caller's team. This keeps team-scoped
listings on the team's own provider account/key instead of a shared
global one.
2. Fallback: any deployment the caller is granted access to for this
provider, expanding wildcard routes and the all-proxy-models sentinel.
Credential lookup is always scoped to the team's allowlist, so a team can
never resolve a provider key for a deployment it isn't authorized to use.
Credential lookup is scoped to both the team's allowlist and the key's own
model allowlist (``user_api_key_dict.models``), so neither a team nor a
restricted key within a team can resolve a provider key for a deployment
it isn't authorized to use. A key restricted to an explicit model list
only narrows the team scope; sentinel-bearing keys (all-proxy-models /
all-team-models) defer to the team scope instead of widening past it.
Returns None when the router is unavailable or no authorized deployment
matches, so the caller can fall back to default credential resolution.
"""
if llm_router is None:
return None
from litellm.proxy._types import SpecialModelNames
from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models
team_id = user_api_key_dict.team_id
team_models = user_api_key_dict.team_models or []
proxy_model_list = llm_router.get_model_names(team_id=team_id)
model_access_groups = llm_router.get_model_access_groups()
raw_key_models = user_api_key_dict.models or []
sentinel_values = {
SpecialModelNames.all_proxy_models.value,
SpecialModelNames.all_team_models.value,
}
key_is_restricted = bool(raw_key_models) and not (set(raw_key_models) & sentinel_values)
key_model_allowlist = (
tuple(
dict.fromkeys(
get_key_models(
user_api_key_dict=user_api_key_dict,
proxy_model_list=proxy_model_list,
model_access_groups=model_access_groups,
)
)
)
if key_is_restricted
else ()
)
key_model_allowlist_set = frozenset(key_model_allowlist)
def _key_may_use(public_model_name: Optional[str]) -> bool:
if not key_model_allowlist_set:
return True
return public_model_name is not None and public_model_name in key_model_allowlist_set
def _provider_credentials(model_id: str) -> Optional[dict]:
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id)
if credentials is not None and credentials.get("custom_llm_provider") == custom_llm_provider:
return credentials
return None
@ -332,27 +372,27 @@ def get_team_provider_credentials(
deployment_id = model_info.get("id")
if deployment_id is None:
continue
if not _key_may_use(model_info.get("team_public_model_name") or deployment.get("model_name")):
continue
credentials = _provider_credentials(deployment_id)
if credentials is not None:
return credentials
# 2. Fall back to deployments the team is allowed to access. The
# all-proxy-models sentinel isn't expanded by get_complete_model_list, so
# normalize it to an empty allowlist, which defers to the team-scoped
# proxy model list. A team with a restricted allowlist (e.g. anthropic
# only) therefore never resolves another provider's key.
from litellm.proxy._types import SpecialModelNames
from litellm.proxy.auth.model_checks import get_complete_model_list
# 2. Fall back to deployments the caller is allowed to access. The key's
# effective allowlist (sentinels and access groups already expanded by
# get_key_models) wins when set; otherwise the team's allowlist applies.
# The all-proxy-models sentinel isn't expanded by
# get_complete_model_list, so normalize it to an empty allowlist, which
# defers to the team-scoped proxy model list. A team or key with a
# restricted allowlist (e.g. anthropic only) therefore never resolves
# another provider's key.
grants_all_models = SpecialModelNames.all_proxy_models.value in team_models
effective_team_models = [] if grants_all_models else team_models
proxy_model_list = llm_router.get_model_names(team_id=team_id)
model_access_groups = llm_router.get_model_access_groups()
models_to_try = list(
dict.fromkeys(
get_complete_model_list(
key_models=[],
key_models=list(key_model_allowlist),
team_models=effective_team_models,
proxy_model_list=proxy_model_list,
user_model=None,
@ -373,6 +413,28 @@ def get_team_provider_credentials(
return None
def apply_team_provider_credentials(
data: dict, # mutable-ok: credentials are merged into the request payload in place, same contract as prepare_data_with_credentials
llm_router: Optional["Router"],
user_api_key_dict: "UserAPIKeyAuth",
custom_llm_provider: str,
) -> None:
"""
Resolve credentials for a provider-only request (no model pinned) via
``get_team_provider_credentials`` and merge them into ``data`` in-place.
Leaves ``data`` untouched when no authorized deployment matches, so the
caller falls back to environment-variable credentials exactly as before.
"""
credentials = get_team_provider_credentials(
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
if credentials is None:
return
prepare_data_with_credentials(data=data, credentials=credentials)
def prepare_data_with_credentials(
data: dict,
credentials: dict,

36
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -43,10 +43,10 @@ from litellm.litellm_core_utils.cloud_storage_security import (
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
apply_team_provider_credentials,
encode_file_id_with_model,
extract_file_creation_params,
get_credentials_for_model,
get_team_provider_credentials,
handle_model_based_routing,
prepare_data_with_credentials,
validate_managed_files_requirement,
@ -253,6 +253,12 @@ async def route_create_file(
_create_file_request=_create_file_request,
)
else:
apply_team_provider_credentials(
data=cast(dict, _create_file_request), # cast-ok: TypedDict is a plain dict at runtime; merged in place
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
# get configs for custom_llm_provider
llm_provider_config = get_files_provider_config(custom_llm_provider=custom_llm_provider)
if llm_provider_config is not None:
@ -735,6 +741,14 @@ async def get_file_content(
check_file_id_encoding=True,
)
if not should_route:
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
from litellm.proxy.openai_files_endpoints.file_content_streaming_handler import (
FileContentStreamingHandler,
)
@ -983,6 +997,12 @@ async def get_file(
# Remove file_id from data to avoid "multiple values for keyword argument" error
# data was initialized with {"file_id": file_id}
data.pop("file_id", None)
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
response = await litellm.afile_retrieve(
custom_llm_provider=custom_llm_provider,
file_id=file_id,
@ -1183,6 +1203,12 @@ async def delete_file(
)
else:
data.pop("file_id", None)
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
response = await litellm.afile_delete(
custom_llm_provider=custom_llm_provider,
file_id=file_id,
@ -1354,14 +1380,12 @@ async def list_files(
# No model/target_model_names pinned: resolve upstream credentials from
# the team's deployment for this provider so the call is authenticated
# against the team's own account (e.g. the team's openai deployment).
team_credentials = get_team_provider_credentials(
apply_team_provider_credentials(
data=data,
llm_router=llm_router,
team_models=user_api_key_dict.team_models or [],
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
team_id=user_api_key_dict.team_id,
)
if team_credentials is not None:
prepare_data_with_credentials(data=data, credentials=team_credentials)
response = await litellm.afile_list(
custom_llm_provider=custom_llm_provider,

42
litellm/router.py
View file

@ -31,6 +31,7 @@ from typing import (
Generator,
List,
Literal,
Mapping,
Optional,
Set,
Tuple,
@ -8636,6 +8637,33 @@ class Router:
raise Exception("Model Name invalid - {}".format(type(model)))
return None
@staticmethod
def _deployment_usable_by_team(model: Union[Mapping, Deployment], team_id: str | None) -> bool:
"""
A team-scoped deployment (``model_info.team_id`` set) is only usable by
callers from that same team; deployments without a team owner are shared.
"""
model_info = model.get("model_info") if isinstance(model, dict) else model.model_info
owner_team_id = model_info.get("team_id") if model_info is not None else None
return owner_team_id is None or owner_team_id == team_id
def _get_model_group_deployment_usable_by_team(
self, model_group_name: str, team_id: str | None
) -> Deployment | None:
"""
Like ``get_deployment_by_model_group_name``, but skips deployments owned
by other teams so a shared model name never resolves another team's
credentials.
"""
indices = self.model_name_to_deployment_indices.get(model_group_name) or ()
usable = (
self.model_list[idx] for idx in indices if self._deployment_usable_by_team(self.model_list[idx], team_id)
)
first_usable = next(usable, None)
if first_usable is None:
return None
return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable
def get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]":
"""
Return (max_input_tokens, max_output_tokens) explicitly configured in a concrete
@ -8670,7 +8698,10 @@ class Router:
model_id: Model ID or model name from model_list (e.g., "gpt-4o-litellm")
team_id: Optional team id of the caller. When set, team-scoped
deployments (indexed by team public model name, including team
wildcard models like "openai/*") are also considered.
wildcard models like "openai/*") are also considered. Name and
wildcard lookups never resolve a deployment owned by a
different team, so shared model names can't leak another
team's credentials.
Returns:
Dictionary containing api_key, api_base, custom_llm_provider, etc.
@ -8687,7 +8718,7 @@ class Router:
# If not found, try by model_group_name
if deployment is None:
deployment = self.get_deployment_by_model_group_name(model_group_name=model_id)
deployment = self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id)
# If not found, check team-scoped deployments whose team public model
# name exactly matches model_id (wildcard team names are matched via
@ -8704,7 +8735,12 @@ class Router:
if deployment is None:
team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None
team_wildcard_models = (team_pattern_router.route(model_id) or []) if team_pattern_router else []
potential_wildcard_models = team_wildcard_models or self.pattern_router.route(model_id) or []
global_wildcard_models = [
wildcard_model
for wildcard_model in (self.pattern_router.route(model_id) or [])
if self._deployment_usable_by_team(wildcard_model, team_id)
]
potential_wildcard_models = team_wildcard_models or global_wildcard_models
if potential_wildcard_models:
# Use the first matching wildcard deployment
deployment_dict = potential_wildcard_models[0]

153
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -51,7 +51,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.llms.openai import BatchJobStatus
from litellm.types.utils import LiteLLMBatch
from litellm.types.utils import CredentialItem, LiteLLMBatch
from fastapi import Response
@ -2091,3 +2091,154 @@ async def test_retrieve__unified_no_router_500(retrieve_harness):
assert exc.value.code == "500"
retrieve_harness.router_aretrieve.assert_not_called()
retrieve_harness.litellm_aretrieve.assert_not_called()
# =========================================================================== #
# SCENARIO 3 + configured deployments: a provider-only call (custom-llm-provider
# header, no model anywhere) must resolve the gateway/team deployment's named
# credential for that provider and attach it to the provider call kwargs,
# instead of silently falling through to the host environment's default
# credentials (regression: vertex batch jobs landing in the hosting env's GCP
# project because litellm_credential_name never reached the call).
# =========================================================================== #
VERTEX_NAMED_CREDENTIAL = CredentialItem(
credential_name="vertex-named-cred",
credential_info={},
credential_values={
"vertex_project": "customer-project",
"vertex_location": "us-central1",
"vertex_credentials": "/creds/customer-sa.json",
},
)
def vertex_named_credential_router() -> Router:
return Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"litellm_credential_name": "vertex-named-cred",
},
}
]
)
@pytest.mark.asyncio
async def test_create__provider_only_resolves_named_vertex_credentials(harness):
"""Provider-only create must attach the configured named credential, and must
NOT turn the call into a model-routed one (no model kwarg injected)."""
set_body(
harness,
{
"input_file_id": "file-plain",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
)
harness.provider_from_headers.return_value = "vertex_ai"
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
await call_create(harness)
assert harness.acreate_kwargs() == {
"custom_llm_provider": "vertex_ai",
"input_file_id": "file-plain",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"metadata": None,
"vertex_project": "customer-project",
"vertex_location": "us-central1",
"vertex_credentials": "/creds/customer-sa.json",
}
@pytest.mark.asyncio
async def test_create__provider_only_ignores_other_provider_deployments(harness):
"""A provider-only vertex call must not pick up credentials from deployments
of a different provider; with no vertex deployment the payload is exactly the
pre-fix env-var fallback."""
set_body(
harness,
{
"input_file_id": "file-plain",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
)
harness.provider_from_headers.return_value = "vertex_ai"
openai_only_router = Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-global-openai"},
}
]
)
with patch.object(proxy_server, "llm_router", openai_only_router):
await call_create(harness)
assert harness.acreate_kwargs() == {
"custom_llm_provider": "vertex_ai",
"input_file_id": "file-plain",
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"metadata": None,
}
@pytest.mark.asyncio
async def test_retrieve__provider_only_resolves_named_vertex_credentials(retrieve_harness):
retrieve_harness.provider_from_headers.return_value = "vertex_ai"
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
await call_retrieve(retrieve_harness, "batch-raw-xyz")
assert retrieve_harness.aretrieve_kwargs() == {
"custom_llm_provider": "vertex_ai",
"batch_id": "batch-raw-xyz",
"vertex_project": "customer-project",
"vertex_location": "us-central1",
"vertex_credentials": "/creds/customer-sa.json",
}
@pytest.mark.asyncio
async def test_list__provider_only_resolves_named_vertex_credentials(list_harness):
list_harness.provider_from_headers.return_value = "vertex_ai"
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
await call_list(list_harness)
assert list_harness.alist_kwargs() == {
"custom_llm_provider": "vertex_ai",
"after": None,
"limit": None,
"vertex_project": "customer-project",
"vertex_location": "us-central1",
"vertex_credentials": "/creds/customer-sa.json",
}
@pytest.mark.asyncio
async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_harness):
cancel_harness.provider_from_headers.return_value = "vertex_ai"
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
await call_cancel(cancel_harness, "batch-raw-xyz")
assert cancel_harness.acancel_kwargs() == {
"custom_llm_provider": "vertex_ai",
"batch_id": "batch-raw-xyz",
"vertex_project": "customer-project",
"vertex_location": "us-central1",
"vertex_credentials": "/creds/customer-sa.json",
}

441
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
View file

@ -2610,3 +2610,444 @@ def test_list_files_with_all_proxy_models_team_uses_openai_deployment(
assert captured_kwargs.get("api_key") == "team-openai-key"
assert captured_kwargs.get("custom_llm_provider") == "openai"
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def _setup_vertex_named_credential_router(monkeypatch) -> Router:
from litellm.types.utils import CredentialItem
monkeypatch.setattr(
litellm,
"credential_list",
[
CredentialItem(
credential_name="vertex-named-cred",
credential_info={},
credential_values={
"vertex_project": "customer-project",
"vertex_location": "us-central1",
"vertex_credentials": "/creds/customer-sa.json",
},
)
],
)
return Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"litellm_credential_name": "vertex-named-cred",
},
}
]
)
def _assert_vertex_named_credentials_attached(captured_kwargs: dict) -> None:
assert captured_kwargs.get("custom_llm_provider") == "vertex_ai"
assert captured_kwargs.get("vertex_project") == "customer-project"
assert captured_kwargs.get("vertex_location") == "us-central1"
assert captured_kwargs.get("vertex_credentials") == "/creds/customer-sa.json"
assert captured_kwargs.get("model") is None
def test_create_file_provider_only_resolves_named_vertex_credentials(
mocker: MockerFixture, monkeypatch
):
"""
POST /v1/files with only a custom-llm-provider header (no model, no
target_model_names) must attach the configured named vertex credential to
the upstream call instead of falling through to google.auth.default(),
which uploads into the hosting environment's GCP project.
"""
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
router = _setup_vertex_named_credential_router(monkeypatch)
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
proxy_logging_obj.update_request_status = mocker.AsyncMock()
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
captured_kwargs: dict = {}
async def _mock_acreate_file(**kwargs):
captured_kwargs.update(kwargs)
return OpenAIFileObject(
id="file-vertex-123",
object="file",
bytes=2,
created_at=1234567890,
filename="batch.jsonl",
purpose="batch",
status="uploaded",
)
monkeypatch.setattr(litellm, "acreate_file", _mock_acreate_file)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="test-user",
)
try:
response = client.post(
"/v1/files",
files={"file": ("batch.jsonl", b"{}", "application/jsonl")},
data={"purpose": "batch"},
headers={
"Authorization": "Bearer test-key",
"custom-llm-provider": "vertex_ai",
},
)
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
_assert_vertex_named_credentials_attached(captured_kwargs)
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def test_get_file_provider_only_resolves_named_vertex_credentials(
mocker: MockerFixture, monkeypatch
):
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
router = _setup_vertex_named_credential_router(monkeypatch)
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
proxy_logging_obj.update_request_status = mocker.AsyncMock()
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
captured_kwargs: dict = {}
async def _mock_afile_retrieve(**kwargs):
captured_kwargs.update(kwargs)
return OpenAIFileObject(
id="file-abc123",
object="file",
bytes=2,
created_at=1234567890,
filename="batch.jsonl",
purpose="batch",
status="uploaded",
)
monkeypatch.setattr(litellm, "afile_retrieve", _mock_afile_retrieve)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="test-user",
)
try:
response = client.get(
"/v1/files/file-abc123",
headers={
"Authorization": "Bearer test-key",
"custom-llm-provider": "vertex_ai",
},
)
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
assert captured_kwargs.get("file_id") == "file-abc123"
_assert_vertex_named_credentials_attached(captured_kwargs)
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def test_get_file_content_provider_only_resolves_named_vertex_credentials(
mocker: MockerFixture, monkeypatch
):
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
router = _setup_vertex_named_credential_router(monkeypatch)
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
proxy_logging_obj.update_request_status = mocker.AsyncMock()
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
captured_kwargs: dict = {}
async def _mock_afile_content(**kwargs):
captured_kwargs.update(kwargs)
return HttpxBinaryResponseContent(
response=httpx.Response(
status_code=200,
content=b"vertex-bytes",
headers={"content-type": "application/octet-stream"},
)
)
monkeypatch.setattr(litellm, "afile_content", _mock_afile_content)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="test-user",
)
try:
response = client.get(
"/v1/files/file-abc123/content",
headers={
"Authorization": "Bearer test-key",
"custom-llm-provider": "vertex_ai",
},
)
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
assert response.content == b"vertex-bytes"
assert captured_kwargs.get("file_id") == "file-abc123"
_assert_vertex_named_credentials_attached(captured_kwargs)
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def test_delete_file_provider_only_resolves_named_vertex_credentials(
mocker: MockerFixture, monkeypatch
):
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
router = _setup_vertex_named_credential_router(monkeypatch)
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
proxy_logging_obj.update_request_status = mocker.AsyncMock()
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
captured_kwargs: dict = {}
async def _mock_afile_delete(**kwargs):
captured_kwargs.update(kwargs)
return OpenAIFileObject(
id="file-abc123",
object="file",
bytes=2,
created_at=1234567890,
filename="batch.jsonl",
purpose="batch",
status="uploaded",
)
monkeypatch.setattr(litellm, "afile_delete", _mock_afile_delete)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="test-user",
)
try:
response = client.delete(
"/v1/files/file-abc123",
headers={
"Authorization": "Bearer test-key",
"custom-llm-provider": "vertex_ai",
},
)
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
assert captured_kwargs.get("file_id") == "file-abc123"
_assert_vertex_named_credentials_attached(captured_kwargs)
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def test_create_file_provider_only_skips_other_team_vertex_deployment(
mocker: MockerFixture, monkeypatch
):
"""
Regression: with a team-scoped vertex deployment indexed before a global
one under the same model name, a provider-only upload from a different
team must use the global deployment's credentials, never the other
team's.
"""
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
router = Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "team-b-project",
},
"model_info": {
"id": "team-b-vertex",
"team_id": "team-b",
"team_public_model_name": "gemini-2.5-pro",
},
},
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "shared-project",
},
},
]
)
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
proxy_logging_obj.update_request_status = mocker.AsyncMock()
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
captured_kwargs: dict = {}
async def _mock_acreate_file(**kwargs):
captured_kwargs.update(kwargs)
return OpenAIFileObject(
id="file-vertex-456",
object="file",
bytes=2,
created_at=1234567890,
filename="batch.jsonl",
purpose="batch",
status="uploaded",
)
monkeypatch.setattr(litellm, "acreate_file", _mock_acreate_file)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key",
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="test-user",
team_id="team-a",
team_models=["gemini-2.5-pro"],
)
try:
response = client.post(
"/v1/files",
files={"file": ("batch.jsonl", b"{}", "application/jsonl")},
data={"purpose": "batch"},
headers={
"Authorization": "Bearer test-key",
"custom-llm-provider": "vertex_ai",
},
)
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
assert captured_kwargs.get("vertex_project") == "shared-project"
proxy_logging_obj.post_call_failure_hook.assert_not_called()
def _team_openai_plus_global_anthropic_router() -> Router:
return Router(
model_list=[
{
"model_name": "team-gpt",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "team-openai-key",
},
"model_info": {
"id": "team-a-openai",
"team_id": "team-a",
"team_public_model_name": "team-gpt",
},
},
{
"model_name": "claude-opus-4-6",
"litellm_params": {
"model": "anthropic/claude-opus-4-6",
"api_key": "anthropic-key",
},
},
]
)
def _list_files_captured_kwargs(
mocker: MockerFixture, monkeypatch, router: Router, key_models: list
) -> dict:
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, router)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
proxy_logging_obj.update_request_status = mocker.AsyncMock()
proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=[])
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
captured_kwargs: dict = {}
async def _mock_afile_list(**kwargs):
captured_kwargs.update(kwargs)
return []
monkeypatch.setattr(litellm, "afile_list", _mock_afile_list)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
api_key="test-key",
user_role=LitellmUserRoles.INTERNAL_USER,
user_id="test-user",
team_id="team-a",
team_models=["team-gpt", "claude-opus-4-6"],
models=key_models,
)
try:
response = client.get(
"/v1/files",
headers={
"Authorization": "Bearer test-key",
"custom-llm-provider": "openai",
},
)
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
assert response.status_code == 200, response.text
return captured_kwargs
def test_list_files_key_restricted_to_other_provider_does_not_leak_team_openai_credentials(
mocker: MockerFixture, monkeypatch
):
"""
Regression: a key restricted to an anthropic model on a team that also has
an openai deployment must not attach the team's openai credentials to a
provider-only openai files call; key-level model restrictions apply to
credential resolution, not just completions.
"""
captured_kwargs = _list_files_captured_kwargs(
mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["claude-opus-4-6"]
)
assert captured_kwargs.get("api_key") != "team-openai-key"
def test_list_files_key_allowed_openai_model_still_resolves_team_credentials(
mocker: MockerFixture, monkeypatch
):
"""
A key whose allowlist includes the team's openai model keeps resolving that
deployment's credentials for provider-only openai files calls.
"""
captured_kwargs = _list_files_captured_kwargs(
mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"]
)
assert captured_kwargs.get("api_key") == "team-openai-key"

176
tests/test_litellm/test_router.py
View file

@ -3755,6 +3755,182 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority():
assert global_credentials["api_key"] == "global-key"
def test_get_deployment_credentials_with_provider_skips_other_team_deployment():
"""
Regression: a team-scoped deployment sharing a model_name with a global
deployment must never resolve for another team's (or an unscoped) caller,
even when it is indexed first; the shared global deployment wins instead.
"""
router = litellm.Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "team-b-project",
},
"model_info": {
"id": "team-b-vertex",
"team_id": "team-b",
"team_public_model_name": "gemini-2.5-pro",
},
},
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "shared-project",
},
},
],
)
other_team_credentials = router.get_deployment_credentials_with_provider(
model_id="gemini-2.5-pro", team_id="team-a"
)
assert other_team_credentials is not None
assert other_team_credentials["vertex_project"] == "shared-project"
unscoped_credentials = router.get_deployment_credentials_with_provider(
model_id="gemini-2.5-pro"
)
assert unscoped_credentials is not None
assert unscoped_credentials["vertex_project"] == "shared-project"
owner_credentials = router.get_deployment_credentials_with_provider(
model_id="gemini-2.5-pro", team_id="team-b"
)
assert owner_credentials is not None
assert owner_credentials["vertex_project"] == "team-b-project"
def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only_name():
"""
When the only deployments under a model name belong to another team, other
callers must get None (env fallback) instead of that team's credentials.
"""
router = litellm.Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "team-b-project",
},
"model_info": {
"id": "team-b-vertex",
"team_id": "team-b",
"team_public_model_name": "gemini-2.5-pro",
},
},
],
)
assert (
router.get_deployment_credentials_with_provider(
model_id="gemini-2.5-pro", team_id="team-a"
)
is None
)
assert (
router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro")
is None
)
def test_deployment_usable_by_team_helpers():
"""
Direct coverage of the team-ownership filter: a team-scoped deployment is
usable only by its owning team, shared deployments by anyone, and the
model-group picker returns the first usable deployment or None.
"""
router = litellm.Router(
model_list=[
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "team-b-project",
},
"model_info": {
"id": "team-b-vertex",
"team_id": "team-b",
"team_public_model_name": "gemini-2.5-pro",
},
},
{
"model_name": "gemini-2.5-pro",
"litellm_params": {
"model": "vertex_ai/gemini-2.5-pro",
"vertex_project": "shared-project",
},
},
],
)
team_owned, shared = router.model_list
assert router._deployment_usable_by_team(team_owned, "team-b") is True
assert router._deployment_usable_by_team(team_owned, "team-a") is False
assert router._deployment_usable_by_team(team_owned, None) is False
assert router._deployment_usable_by_team(shared, "team-a") is True
assert router._deployment_usable_by_team(shared, None) is True
picked = router._get_model_group_deployment_usable_by_team(
model_group_name="gemini-2.5-pro", team_id="team-a"
)
assert picked is not None
assert picked.litellm_params.vertex_project == "shared-project"
owner_picked = router._get_model_group_deployment_usable_by_team(
model_group_name="gemini-2.5-pro", team_id="team-b"
)
assert owner_picked is not None
assert owner_picked.litellm_params.vertex_project == "team-b-project"
assert (
router._get_model_group_deployment_usable_by_team(
model_group_name="unknown-model", team_id="team-a"
)
is None
)
def test_get_deployment_credentials_with_provider_skips_other_team_wildcard():
"""
Global wildcard resolution must skip a team-scoped wildcard deployment for
callers outside that team, falling through to the shared wildcard entry.
"""
router = litellm.Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*", "api_key": "team-b-key"},
"model_info": {
"id": "team-b-wildcard",
"team_id": "team-b",
"team_public_model_name": "openai/*",
},
},
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*", "api_key": "global-key"},
},
],
)
other_team_credentials = router.get_deployment_credentials_with_provider(
model_id="openai/gpt-5.2", team_id="team-a"
)
assert other_team_credentials is not None
assert other_team_credentials["api_key"] == "global-key"
owner_credentials = router.get_deployment_credentials_with_provider(
model_id="openai/gpt-5.2", team_id="team-b"
)
assert owner_credentials is not None
assert owner_credentials["api_key"] == "team-b-key"
def test_team_wildcard_credentials_not_usable_after_delete_deployment():
"""
Regression: team_pattern_routers retained deleted deployments, so a team