fix(router): scope alias-to-model resolution to the caller's team

The alias resolver picked the first deployment under a model group while
the credential resolver skipped deployments owned by other teams, so a
team-owned deployment listed before a shared one under the same alias
leaked its private model id into an outside caller's batch while the
request ran on the shared deployment's credentials

_resolve_unblocked_deployment now carries the full team-aware lookup
(team-usable name match, exact team public model name, team wildcard
before global wildcard) and both resolvers delegate to it, so model and
credentials always come from the same deployment by construction. The
batch endpoints pass user_api_key_dict.team_id through both paths
This commit is contained in:
Kent 2026-08-01 18:42:50 +08:00
parent b1c409b94b
commit adbb3e1bce
6 changed files with 126 additions and 60 deletions

View file

@ -52,6 +52,7 @@ def _swap_alias_for_deployment_model(
create_batch_data: LiteLLMBatchCreateRequest,
alias: str,
llm_router: Optional["Router"],
team_id: "str | None",
) -> None:
"""
Replace a proxy model-group alias on the batch request with the
@ -61,11 +62,14 @@ def _swap_alias_for_deployment_model(
cannot resolve a proxy alias, so a provider transform (e.g. Bedrock's, which
forwards ``model`` as the batch ``modelId``) would otherwise receive the
alias and the provider would reject it. Falls back to the alias when the
router is unavailable or the alias resolves to nothing.
router is unavailable or the alias resolves to nothing. ``team_id`` keeps
this lookup on the same team-usable deployment the credential resolver
picked, so a team-owned deployment sharing the alias can't leak its model
to callers outside that team.
"""
if llm_router is None:
return
resolved_model = llm_router.get_deployment_model_for_alias(model_id=alias)
resolved_model = llm_router.get_deployment_model_for_alias(model_id=alias, team_id=team_id)
if resolved_model is not None:
create_batch_data["model"] = resolved_model
@ -214,6 +218,7 @@ async def create_batch(
llm_router=llm_router,
model_id=model_from_file_id,
operation_context="batch creation (file created with model)",
team_id=user_api_key_dict.team_id,
)
original_file_id = get_original_file_id(input_file_id)
@ -226,6 +231,7 @@ async def create_batch(
create_batch_data=_create_batch_data,
alias=model_from_file_id,
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
)
# Create batch using model credentials
@ -308,6 +314,7 @@ async def create_batch(
llm_router=llm_router,
model_id=model_param,
operation_context="batch creation",
team_id=user_api_key_dict.team_id,
)
prepare_data_with_credentials(
@ -318,6 +325,7 @@ async def create_batch(
create_batch_data=_create_batch_data,
alias=model_param,
llm_router=llm_router,
team_id=user_api_key_dict.team_id,
)
# Create batch using model credentials
@ -518,6 +526,7 @@ async def retrieve_batch(
llm_router=llm_router,
model_id=model_from_id,
operation_context="batch retrieval (batch created with model)",
team_id=user_api_key_dict.team_id,
)
original_batch_id = get_original_file_id(batch_id)
@ -724,6 +733,7 @@ async def list_batches(
llm_router=llm_router,
model_id=model_param,
operation_context="batch listing",
team_id=user_api_key_dict.team_id,
)
data.update(credentials)
@ -905,6 +915,7 @@ async def cancel_batch(
llm_router=llm_router,
model_id=model_from_id,
operation_context="batch cancellation (batch created with model)",
team_id=user_api_key_dict.team_id,
)
original_batch_id = get_original_file_id(batch_id)

View file

@ -259,6 +259,7 @@ def get_credentials_for_model(
llm_router, # Router instance
model_id: str,
operation_context: str = "file operation",
team_id: "str | None" = None,
):
"""
Retrieve API credentials for a model from the LLM Router.
@ -267,6 +268,8 @@ def get_credentials_for_model(
llm_router: LiteLLM Router instance
model_id: Model name or deployment ID
operation_context: Description for error messages (e.g., "file upload", "batch creation")
team_id: Caller's team id; unlocks that team's own deployments and keeps
shared model names from resolving another team's credentials
Returns:
Dictionary with credentials (api_key, api_base, custom_llm_provider, etc.)
@ -282,7 +285,7 @@ def get_credentials_for_model(
detail={"error": "Router not initialized. Cannot use model-based routing."},
)
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 None:
raise HTTPException(

View file

@ -8639,19 +8639,54 @@ class Router:
raise Exception("Model Name invalid - {}".format(type(model)))
return None
def _resolve_unblocked_deployment(self, model_id: str) -> "Deployment | None":
def _resolve_unblocked_deployment(self, model_id: str, team_id: "str | None" = None) -> "Deployment | None":
"""
Resolve a model id, model-group alias, or wildcard pattern to a single
deployment, returning None when nothing matches or the match is paused
via ``LiteLLM_ProxyModelTable.blocked``.
Both the credential resolver and the alias-to-model resolver delegate
here so a mixed model group can never hand one caller the credentials
of one deployment and the model of another. Name and wildcard lookups
skip deployments owned by a team other than ``team_id``; passing
``team_id`` also unlocks that team's own deployments (exact team public
model name and team wildcard patterns).
"""
deployment = self.get_deployment(model_id=model_id)
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)
# Team-scoped deployments whose team public model name exactly matches
# model_id (wildcard team names are matched via team_pattern_routers
# below).
if deployment is None and team_id is not None:
team_indices = self.team_model_to_deployment_indices.get((team_id, model_id)) or ()
team_match = next((self.model_list[idx] for idx in team_indices), None)
if isinstance(team_match, dict):
deployment = Deployment(**team_match)
elif isinstance(team_match, Deployment):
deployment = team_match
# Wildcard pattern matches. Team wildcard matches take priority so a
# global pattern (e.g. "openai/*") doesn't shadow the team's own entry.
if deployment is None:
wildcard_match = next(iter(self.pattern_router.route(model_id) or ()), None)
team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None
team_wildcard_match = (
next(iter(team_pattern_router.route(model_id) or ()), None) if team_pattern_router else None
)
wildcard_match = (
team_wildcard_match
if team_wildcard_match is not None
else next(
(
wildcard_model
for wildcard_model in (self.pattern_router.route(model_id) or ())
if self._deployment_usable_by_team(wildcard_model, team_id)
),
None,
)
)
if isinstance(wildcard_match, dict):
deployment = Deployment(**wildcard_match)
elif isinstance(wildcard_match, Deployment):
@ -8661,7 +8696,7 @@ class Router:
return None
return deployment
def get_deployment_model_for_alias(self, model_id: str) -> "str | None":
def get_deployment_model_for_alias(self, model_id: str, team_id: "str | None" = None) -> "str | None":
"""
Resolve a model-group alias (or deployment id / wildcard) to the
deployment's underlying ``litellm_params.model``.
@ -8671,8 +8706,10 @@ class Router:
``get_llm_provider`` cannot resolve an alias, so passing it straight
through reaches the provider as an invalid model identifier. Returns
None when the alias resolves to nothing or to a paused deployment.
Pass the caller's ``team_id`` so the deployment picked here is the same
one the credential resolver picks for that caller.
"""
deployment = self._resolve_unblocked_deployment(model_id=model_id)
deployment = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id)
if deployment is None:
return None
return deployment.litellm_params.model
@ -8753,43 +8790,8 @@ class Router:
credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm")
# Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...}
"""
# Try to get deployment by model_id first
deployment = self.get_deployment(model_id=model_id)
# If not found, try by model_group_name
deployment = self._resolve_unblocked_deployment(model_id=model_id, team_id=team_id)
if deployment is None:
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
# team_pattern_routers below).
if deployment is None and team_id is not None:
team_indices = self.team_model_to_deployment_indices.get((team_id, model_id), [])
if team_indices:
team_model = self.model_list[team_indices[0]]
deployment = Deployment(**team_model) if isinstance(team_model, dict) else team_model
# If still not found, check for wildcard pattern matches. Team wildcard
# matches take priority so a global pattern (e.g. "openai/*") doesn't
# shadow the team's own entry.
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 []
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]
if isinstance(deployment_dict, dict):
deployment = Deployment(**deployment_dict)
elif isinstance(deployment_dict, Deployment):
deployment = deployment_dict
if deployment is None or self._is_deployment_blocked(deployment):
return None
# Get basic credentials

View file

@ -146,12 +146,12 @@ class Harness:
return dict(self.router_acreate.call_args.kwargs)
def _creds_lookup(*, model_id: str) -> Dict[str, str]:
def _creds_lookup(*, model_id: str, team_id: Optional[str] = None) -> Dict[str, str]:
# KeyError on an unknown/hardcoded model_id - the bug cannot hide.
return dict(CREDS[model_id])
def _alias_lookup(*, model_id: str) -> str:
def _alias_lookup(*, model_id: str, team_id: Optional[str] = None) -> str:
# The endpoint swaps the request model for the deployment's real provider
# model before calling the provider; mirror that with the CREDS model so a
# wrong/hardcoded model_id KeyErrors instead of hiding.
@ -267,7 +267,7 @@ async def test_create__model_encoded_file_id(harness):
harness.router_acreate.assert_not_called()
# 2. CREDENTIALS - resolved for the model decoded FROM the file id.
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# 3. SEAM PAYLOAD - exact, whole dict. A new forwarded key breaks this.
assert harness.acreate_kwargs() == {
@ -323,7 +323,7 @@ async def test_create__model_encoded_file_id__resolver_gets_decoded_model(harnes
await call_create(harness)
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# =========================================================================== #
@ -347,7 +347,7 @@ async def test_create__model_from_body(harness):
assert harness.litellm_acreate.call_count == 1
harness.router_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None)
payload = harness.acreate_kwargs()
assert payload["custom_llm_provider"] == "vertex_ai"
assert payload["input_file_id"] == "file-plain"
@ -367,7 +367,7 @@ async def test_create__model_from_header(harness):
await call_create(harness, headers={"x-litellm-model": "vertex-model"})
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None)
harness.router_acreate.assert_not_called()
@ -384,7 +384,7 @@ async def test_create__model_from_query(harness):
await call_create(harness, query={"model": "vertex-model"})
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
harness.creds_resolver.assert_called_once_with(model_id="vertex-model", team_id=None)
harness.router_acreate.assert_not_called()
@ -407,7 +407,7 @@ async def test_create__body_model_beats_header_and_query(harness):
query={"model": "vertex-model"},
)
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# =========================================================================== #
@ -715,7 +715,7 @@ async def test_create__model_encoded_beats_unified(harness):
assert harness.litellm_acreate.call_count == 1
harness.router_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# =========================================================================== #
@ -761,7 +761,7 @@ async def test_create__model_encoded_beats_loadbalancing(harness):
assert harness.litellm_acreate.call_count == 1
harness.router_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# =========================================================================== #
@ -1072,7 +1072,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness):
retrieve_harness.router_aretrieve.assert_not_called()
# 2. CREDENTIALS - resolved for the model decoded FROM the batch id.
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# 3. SEAM PAYLOAD - exact, whole dict forwarded to the provider call.
# Note `model` is the DECODED model, not the deployment from creds: the
@ -1130,7 +1130,7 @@ async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness):
assert retrieve_harness.litellm_aretrieve.call_count == 1
retrieve_harness.router_aretrieve.assert_not_called()
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# --------------------------------------------------------------------------- #
@ -1546,7 +1546,7 @@ async def test_list__model_from_body_routes_and_encodes(list_harness):
assert list_harness.litellm_alist.call_count == 1
list_harness.router_alist.assert_not_called()
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch")
assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch")
@ -1836,7 +1836,7 @@ async def test_cancel__model_encoded_id(cancel_harness):
cancel_harness.router_acancel.assert_not_called()
# CREDENTIALS - resolved for the model decoded from the batch id.
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# SEAM PAYLOAD - exact dict. NOTE current behavior: `model` is the
# DEPLOYMENT name from creds, NOT the decoded model (cancel, unlike
@ -1874,7 +1874,7 @@ async def test_cancel__model_encoded_beats_unified(cancel_harness):
assert cancel_harness.litellm_acancel.call_count == 1
cancel_harness.router_acancel.assert_not_called()
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
# --------------------------------------------------------------------------- #

View file

@ -444,6 +444,7 @@ async def test_create_batch_swaps_alias_for_deployment_model_before_provider_cal
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.parent_otel_span = None
mock_user_api_key_dict.user_id = "test_user"
mock_user_api_key_dict.team_id = "team-caller"
mock_user_api_key_dict.team_metadata = {}
mock_router = MagicMock()
@ -510,7 +511,7 @@ async def test_create_batch_swaps_alias_for_deployment_model_before_provider_cal
user_api_key_dict=mock_user_api_key_dict,
)
mock_router.get_deployment_model_for_alias.assert_called_once_with(model_id=alias)
mock_router.get_deployment_model_for_alias.assert_called_once_with(model_id=alias, team_id="team-caller")
create_kwargs = mock_create_batch.call_args.kwargs
assert create_kwargs["model"] == real_model, (
"Bedrock batch transform receives modelId from this 'model'; it must be the "

View file

@ -5844,6 +5844,55 @@ def test_get_deployment_model_for_alias_returns_none_for_blocked_deployment():
assert router.get_deployment_model_for_alias(model_id="dep-1") == "openai/gpt-4o-1"
def test_get_deployment_model_for_alias_matches_credential_deployment_per_team():
"""
Model and credential resolution must pick the SAME deployment for a caller.
Regression: with a team-owned deployment listed before a shared one under
the same alias, an unscoped alias lookup returned the team deployment's
model while the team-aware credential resolver returned the shared
deployment's credentials, so an outside caller's batch reached the provider
with team A's private model id on the shared account.
"""
router = litellm.Router(
model_list=[
{
"model_name": "bedrock-batch",
"litellm_params": {
"model": "bedrock/team-a-private-model",
"aws_region_name": "team-a-region",
},
"model_info": {
"id": "team-a-dep",
"team_id": "team-a",
"team_public_model_name": "bedrock-batch",
},
},
{
"model_name": "bedrock-batch",
"litellm_params": {
"model": "bedrock/shared-model",
"aws_region_name": "shared-region",
},
"model_info": {"id": "shared-dep"},
},
]
)
for team_id, expected_model, expected_region in [
(None, "bedrock/shared-model", "shared-region"),
("team-b", "bedrock/shared-model", "shared-region"),
("team-a", "bedrock/team-a-private-model", "team-a-region"),
]:
resolved_model = router.get_deployment_model_for_alias(model_id="bedrock-batch", team_id=team_id)
credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch", team_id=team_id)
assert resolved_model == expected_model, f"team_id={team_id}"
assert credentials is not None
assert credentials["aws_region_name"] == expected_region, (
f"team_id={team_id}: credentials came from a different deployment than the model"
)
def test_resolve_unblocked_deployment_resolves_alias_id_and_wildcard():
"""
_resolve_unblocked_deployment underpins both the credential resolver and the