From e6f86d7a06352201a5c56b94cdc0b3e5b6ef1561 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 31 Aug 2026 20:53:56 -0700 Subject: [PATCH] test(router): encode credential scope regressions --- .../proxy/batches_endpoints/test_endpoints.py | 84 +++++- .../test_files_endpoint.py | 150 ++++++++++ tests/test_litellm/test_router.py | 259 ++++++++++++++---- 3 files changed, 438 insertions(+), 55 deletions(-) diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 2f6a5a3b0e0..4c38153952b 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -159,7 +159,7 @@ 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]) @@ -232,6 +232,10 @@ def set_body(harness: Harness, body: Dict[str, Any]) -> None: harness.body["body"] = body +def _team_a_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", team_id="team-a") + + async def call_create( harness: Harness, *, @@ -291,6 +295,42 @@ async def test_create__model_encoded_file_id(harness): assert resp.input_file_id == AZURE_FILE_ID +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "expected_model_id"), + [ + pytest.param( + { + "input_file_id": AZURE_FILE_ID, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + "azure/gpt-4o", + id="encoded-file-id", + ), + pytest.param( + { + "input_file_id": "file-plain", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": "vertex-model", + }, + "vertex-model", + id="body-model", + ), + ], +) +async def test_create__model_routing__scopes_credentials_to_callers_team(harness, body, expected_model_id): + set_body(harness, body) + + await call_create(harness, user=_team_a_user()) + + harness.creds_resolver.assert_called_once_with( + model_id=expected_model_id, + team_id="team-a", + ) + + @pytest.mark.asyncio async def test_create__model_encoded_file_id__encodes_output_and_error_ids(harness): set_body( @@ -1219,6 +1259,20 @@ async def test_retrieve__model_encoded_id(retrieve_harness): assert retrieve_harness.update_batch_in_db.call_args.kwargs["operation"] == "retrieve" +@pytest.mark.asyncio +async def test_retrieve__model_encoded_id__scopes_credentials_to_callers_team(retrieve_harness): + await call_retrieve( + retrieve_harness, + AZURE_BATCH_ID, + user=_team_a_user(), + ) + + retrieve_harness.creds_resolver.assert_called_once_with( + model_id="azure/gpt-4o", + team_id="team-a", + ) + + @pytest.mark.asyncio async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment( retrieve_harness, @@ -1758,6 +1812,20 @@ async def test_list__model_from_body_routes_and_encodes(list_harness): assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch") +@pytest.mark.asyncio +async def test_list__model_from_body__scopes_credentials_to_callers_team(list_harness): + await call_list( + list_harness, + body={"model": "azure/gpt-4o"}, + user=_team_a_user(), + ) + + list_harness.creds_resolver.assert_called_once_with( + model_id="azure/gpt-4o", + team_id="team-a", + ) + + # --------------------------------------------------------------------------- # # Branch 3 - target_model_names (function param or body) -> llm_router. Routes # to the FIRST model in the comma list; `model` is stripped from data first. @@ -2066,6 +2134,20 @@ async def test_cancel__model_encoded_id(cancel_harness): assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel" +@pytest.mark.asyncio +async def test_cancel__model_encoded_id__scopes_credentials_to_callers_team(cancel_harness): + await call_cancel( + cancel_harness, + AZURE_BATCH_ID, + user=_team_a_user(), + ) + + cancel_harness.creds_resolver.assert_called_once_with( + model_id="azure/gpt-4o", + team_id="team-a", + ) + + @pytest.mark.asyncio async def test_cancel__model_encoded_id_forwards_deployment_model(cancel_harness): """Pin the current contract: cancel forwards the creds' deployment model. diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 87e0319f6a1..94ae8913070 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1979,6 +1979,81 @@ def test_get_file_content_routed_provider_skips_streaming_when_resolved_provider proxy_logging_obj.post_call_failure_hook.assert_not_called() +def test_get_file_content_keeps_bedrock_metadata_for_empty_named_credential(mocker: MockerFixture, monkeypatch): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="bedrock-ambient-auth", + credential_info={"custom_llm_provider": "bedrock"}, + credential_values={}, + ) + ], + ) + router = Router( + model_list=[ + { + "model_name": "bedrock-batch-model", + "litellm_params": { + "model": "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0", + "litellm_credential_name": "bedrock-ambient-auth", + "aws_region_name": "us-east-1", + "s3_bucket_name": "configured-batch-bucket", + }, + "model_info": {"id": "bedrock-batch-deployment"}, + } + ] + ) + 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() + provider_call = mocker.patch.object( + litellm, + "afile_content", + new=mocker.AsyncMock( + return_value=HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=b"bedrock-bytes", + headers={"content-type": "application/octet-stream"}, + ) + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + ) + + try: + response = client.get( + "/v1/files/file-abc123/content?model=bedrock-batch-deployment", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert response.content == b"bedrock-bytes" + provider_call.assert_awaited_once() + call_kwargs = provider_call.await_args.kwargs + credentials = call_kwargs["_litellm_internal_model_credentials"] + assert call_kwargs["custom_llm_provider"] == "bedrock" + assert credentials["s3_bucket_name"] == "configured-batch-bucket" + assert credentials["aws_region_name"] == "us-east-1" + assert credentials["model"] == "bedrock/anthropic.claude-haiku-4-5-20251001-v1:0" + + def test_get_file_content_non_openai_provider_skips_streaming_handler( mocker: MockerFixture, monkeypatch, llm_router: Router ): @@ -2441,6 +2516,81 @@ def test_list_files_resolves_wildcard_deployment_credentials( proxy_logging_obj.post_call_failure_hook.assert_not_called() +@pytest.mark.parametrize( + ("caller_team_id", "expected_status", "expected_api_key"), + [ + pytest.param("team-b", 200, "team-b-key", id="owner"), + pytest.param("team-a", 400, None, id="other-team"), + pytest.param(None, 400, None, id="teamless"), + ], +) +def test_list_files_scopes_exact_deployment_id_to_owning_team( + mocker: MockerFixture, + monkeypatch, + caller_team_id: str | None, + expected_status: int, + expected_api_key: str | None, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + router = Router( + model_list=[ + { + "model_name": "team-b-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "team-b-key", + }, + "model_info": { + "id": "team-b-deployment", + "team_id": "team-b", + "team_public_model_name": "team-b-model", + }, + }, + { + "model_name": "*", + "litellm_params": { + "model": "openai/*", + "api_key": "shared-wildcard-key", + }, + }, + ] + ) + 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() + provider_call = mocker.patch.object(litellm, "afile_list", new=mocker.AsyncMock(return_value=[])) + + 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=caller_team_id, + team_models=[], + models=[], + ) + + try: + response = client.get( + "/v1/files?target_model_names=team-b-deployment", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == expected_status, response.text + if expected_api_key is None: + provider_call.assert_not_awaited() + else: + provider_call.assert_awaited_once() + assert provider_call.await_args.kwargs["api_key"] == expected_api_key + + def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice( mocker: MockerFixture, monkeypatch, llm_router: Router ): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index e436b4852b2..2cec72102b3 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -4491,73 +4491,93 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): assert credentials["custom_llm_provider"] == "vertex_ai" -def test_get_deployment_credentials_with_provider_resolves_credential_name(): - """ - Test that get_deployment_credentials_with_provider correctly resolves - litellm_credential_name to actual credential values (for UI-created models). - """ - from litellm.types.utils import CredentialItem - - # Setup credential list with a test credential - litellm.credential_list = [ - CredentialItem( - credential_name="test-azure-cred", - credential_info={"custom_llm_provider": "azure"}, - credential_values={ - "api_key": "resolved-api-key", - "api_base": "https://resolved.openai.azure.com", +@pytest.mark.parametrize( + ("credential_values_by_name", "deployment_auth", "expected_auth"), + [ + pytest.param( + { + "saved": { + "api_key": "saved-key", + "api_base": "https://saved.example.com", + "api_version": "2024-02-01", + } + }, + {"api_key": "inline-key"}, + { + "api_key": "saved-key", + "api_base": "https://saved.example.com", "api_version": "2024-02-01", }, - ) - ] + id="saved-name-overrides-inline", + ), + pytest.param( + {"other": {"api_key": "other-key"}}, + {}, + None, + id="missing-name-without-fallback-fails-closed", + ), + pytest.param( + {"other": {"api_key": "other-key"}}, + {"api_key": "inline-key"}, + {"api_key": "inline-key"}, + id="stale-name-keeps-inline-key", + ), + pytest.param( + {}, + {"aws_role_name": "deployment-role", "aws_region_name": "us-west-2"}, + {"aws_role_name": "deployment-role", "aws_region_name": "us-west-2"}, + id="empty-worker-keeps-deployment-auth", + ), + pytest.param( + {"saved": {}}, + {}, + {}, + id="empty-saved-credential-allows-ambient-auth", + ), + ], +) +def test_get_deployment_credentials_with_provider_named_credential_states( + credential_values_by_name: dict[str, dict[str, str]], + deployment_auth: dict[str, str], + expected_auth: dict[str, str] | None, +) -> None: + from litellm.types.utils import CredentialItem router = litellm.Router( model_list=[ { - "model_name": "azure-gpt-4", + "model_name": "configured-model", "litellm_params": { - "model": "azure/gpt-4", - "litellm_credential_name": "test-azure-cred", - }, - } - ], - ) - - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) - - assert credentials is not None - assert credentials["api_key"] == "resolved-api-key" - assert credentials["api_base"] == "https://resolved.openai.azure.com" - assert credentials["api_version"] == "2024-02-01" - assert credentials["custom_llm_provider"] == "azure" - # Ensure credential name is removed after resolution - assert "litellm_credential_name" not in credentials - - # Cleanup - litellm.credential_list = [] - - -def test_get_deployment_credentials_with_provider_fails_closed_for_missing_named_credential(): - router = litellm.Router( - model_list=[ - { - "model_name": "embedding-model", - "litellm_params": { - "model": "azure/text-embedding-3-small", - "litellm_credential_name": "deleted-credential", + "model": "bedrock/anthropic.claude-3-sonnet", + "litellm_credential_name": "saved", + "s3_bucket_name": "configured-bucket", + **deployment_auth, }, } ] ) + credential_list = [ + CredentialItem( + credential_name=name, + credential_info={}, + credential_values=values, + ) + for name, values in credential_values_by_name.items() + ] - with patch.object( # test-quality-ok: CredentialAccessor reads the process-global credential list - litellm, "credential_list", [] - ): - credentials = router.get_deployment_credentials_with_provider(model_id="embedding-model") + with patch.object(litellm, "credential_list", credential_list): + result = router.get_deployment_credentials_with_provider(model_id="configured-model") - assert credentials is None + if expected_auth is None: + assert result is None + return + + assert result == { + **expected_auth, + "s3_bucket_name": "configured-bucket", + "model": "bedrock/anthropic.claude-3-sonnet", + "custom_llm_provider": "bedrock", + } def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): @@ -4648,6 +4668,26 @@ def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> di } +def _router_with_foreign_specific_and_shared_wildcards() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "openai/gpt-*", + "litellm_params": {"model": "openai/gpt-*", "api_key": "team-b-key"}, + "model_info": { + "id": "team-b-specific-wildcard", + "team_id": "team-b", + "team_public_model_name": "openai/gpt-*", + }, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "shared-key"}, + }, + ] + ) + + def test_get_deployment_credentials_with_provider_team_wildcard_priority(): """ Regression: a global wildcard pattern (e.g. "openai/*") must not shadow a @@ -4795,6 +4835,64 @@ def test_get_deployment_credentials_with_provider_rejects_other_team_exact_id(): ) +def test_foreign_exact_deployment_id_is_not_reinterpreted_as_a_shared_wildcard(): + router = litellm.Router( + model_list=[ + { + "model_name": "team-b-model", + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "team-b-key"}, + "model_info": { + "id": "openai/team-b-deployment", + "team_id": "team-b", + }, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "shared-key"}, + }, + ] + ) + + assert ( + router.get_deployment_credentials_with_provider( + model_id="openai/team-b-deployment", team_id="team-a" + ) + is None + ) + + +def test_team_public_model_name_beats_colliding_shared_model_name_for_owner(): + router = litellm.Router( + model_list=[ + { + "model_name": "model_name_team-a_private-id", + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "team-a-key"}, + "model_info": { + "id": "private-id", + "team_id": "team-a", + "team_public_model_name": "friendly-model", + }, + }, + { + "model_name": "friendly-model", + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "shared-key"}, + }, + ] + ) + + owner_credentials = router.get_deployment_credentials_with_provider( + model_id="friendly-model", team_id="team-a" + ) + other_team_credentials = router.get_deployment_credentials_with_provider( + model_id="friendly-model", team_id="team-b" + ) + + assert owner_credentials is not None + assert owner_credentials["api_key"] == "team-a-key" + assert other_team_credentials is not None + assert other_team_credentials["api_key"] == "shared-key" + + def test_deployment_usable_by_team_helpers(): """ Direct coverage of the team-ownership filter: a team-scoped deployment is @@ -4888,6 +4986,15 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert owner_credentials["api_key"] == "team-b-key" +def test_foreign_specific_wildcard_falls_back_to_broader_shared_wildcard(): + credentials = _router_with_foreign_specific_and_shared_wildcards().get_deployment_credentials_with_provider( + model_id="openai/gpt-5.2", team_id="team-a" + ) + + assert credentials is not None + assert credentials["api_key"] == "shared-key" + + def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ Regression: team_pattern_routers retained deleted deployments, so a team @@ -6344,6 +6451,50 @@ def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallba ) +def test_common_checks_rejects_foreign_exact_deployment_id_for_team(): + from litellm.proxy._types import UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "team-b-model", + "litellm_params": {"model": "openai/gpt-5.2", "api_key": "team-b-key"}, + "model_info": { + "id": "team-b-deployment", + "team_id": "team-b", + }, + } + ] + ) + team_a_auth = UserAPIKeyAuth( + api_key="team-a-key", + team_id="team-a", + models=[], + team_models=[], + ) + + with pytest.raises(litellm.BadRequestError): + router._common_checks_available_deployment( + model="team-b-deployment", + request_kwargs={ + "metadata": { + "user_api_key_team_id": "team-a", + "user_api_key_auth": team_a_auth, + } + }, + ) + + +def test_common_checks_skips_foreign_specific_wildcard_for_team(): + _, deployments = _router_with_foreign_specific_and_shared_wildcards()._common_checks_available_deployment( + model="openai/gpt-5.2", + request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, + ) + + assert isinstance(deployments, list) + assert [deployment["litellm_params"]["api_key"] for deployment in deployments] == ["shared-key"] + + def test_try_early_resolve_deployments_for_model_not_in_names(): """ Direct coverage for ``_try_early_resolve_deployments_for_model_not_in_names``: