From f07ea5921b361471198aced3c941d83fe7727ffb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 16 Jul 2026 08:36:46 -0700 Subject: [PATCH 001/167] feat(proxy): enforce team isolation for provider-format batch ids and output files Managed batch ownership rows were already written for every create path, but retrieve/cancel/file-content checks only fired for unified (litellm_proxy-prefixed) ids, so model-encoded and raw provider batch ids bypassed isolation entirely. - enforce can_access_resource on retrieve/cancel for provider-format batch ids when an ownership row exists; ids with no row stay accessible so pass-through reads keep working - make batch sync operations update-only so a retrieve/cancel can never mint an ownership row attributed to the first caller - write ownership rows for a synced batch's provider-format output/error file ids, inherited from the owning batch row, and enforce them on file content/retrieve/delete --- .../proxy/hooks/managed_files.py | 142 ++++++ .../proxy/hooks/test_managed_files.py | 412 ++++++++++++++++++ 2 files changed, 554 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 3f42867d90e..27363f45505 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -30,10 +30,12 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + decode_model_from_file_id, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, get_models_from_unified_file_id, + get_original_file_id, normalize_mime_type_for_provider, ) from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue] @@ -165,7 +167,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_object_id: str, file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, + update_only: bool = False, ) -> None: + """Persist a managed object row. + + With ``update_only=True`` an existing row is refreshed but a missing + row is NOT created: sync operations (retrieve/cancel) must never mint + an ownership row attributed to whoever happened to call them first. + """ verbose_logger.info( f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" ) @@ -175,6 +184,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose=file_purpose, file_object=file_object, ) + + if update_only: + updated_count = ( + await self.prisma_client.db.litellm_managedobjecttable.update_many( + where={"unified_object_id": unified_object_id}, + data={ + "file_object": file_object.model_dump_json(), + "status": file_object.status, + "updated_by": user_api_key_dict.user_id, + }, + ) + ) + if updated_count == 0: + return + await self.internal_usage_cache.async_set_cache( + key=unified_object_id, + value=litellm_managed_object.model_dump(), + litellm_parent_otel_span=litellm_parent_otel_span, + ) + return + await self.internal_usage_cache.async_set_cache( key=unified_object_id, value=litellm_managed_object.model_dump(), @@ -284,6 +314,103 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"Object not found: {unified_object_id}", ) + async def enforce_batch_object_access( + self, object_id: str, user_api_key_dict: UserAPIKeyAuth + ) -> None: + """Deny access to a provider-format batch id owned by another caller. + + Ids with no ownership row (batches created before ownership tracking, + or directly on the provider account) stay accessible so pass-through + reads keep working. + """ + if self.prisma_client is None: + return + managed_object = ( + await self.prisma_client.db.litellm_managedobjecttable.find_first( + where={"unified_object_id": object_id} + ) + ) + if managed_object is None: + return + if not can_access_resource( + user_api_key_dict=user_api_key_dict, + created_by=managed_object.created_by, + resource_team_id=managed_object.team_id, + ): + raise HTTPException( + status_code=403, + detail=f"User {user_api_key_dict.user_id} does not have access to the object {object_id}", + ) + + async def enforce_provider_file_access( + self, file_id: str, user_api_key_dict: UserAPIKeyAuth + ) -> None: + """Deny access to a provider-format file id owned by another caller. + + Ownership rows for provider-format ids are written when a managed + batch's output/error files are first synced; ids with no row stay + accessible so pass-through reads keep working. + """ + if self.prisma_client is None: + return + managed_file = ( + await self.prisma_client.db.litellm_managedfiletable.find_first( + where={"unified_file_id": file_id} + ) + ) + if managed_file is None: + return + if not can_access_resource( + user_api_key_dict=user_api_key_dict, + created_by=managed_file.created_by, + resource_team_id=managed_file.team_id, + ): + raise HTTPException( + status_code=403, + detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", + ) + + async def store_batch_output_file_ownership( + self, response: LiteLLMBatch, litellm_parent_otel_span: Optional[Span] + ) -> None: + """Record ownership rows for a batch's provider-format output/error + file ids, inherited from the owning batch row (never the caller), so + file reads can be isolation-checked.""" + provider_file_ids = tuple( + file_id + for file_id in ( + getattr(response, "output_file_id", None), + getattr(response, "error_file_id", None), + ) + if file_id and not _is_base64_encoded_unified_file_id(file_id) + ) + if not provider_file_ids: + return + if self.prisma_client is None: + return + batch_row = ( + await self.prisma_client.db.litellm_managedobjecttable.find_first( + where={"unified_object_id": response.id} + ) + ) + if batch_row is None or ( + batch_row.created_by is None and batch_row.team_id is None + ): + return + owner_identity = UserAPIKeyAuth( + user_id=batch_row.created_by, team_id=batch_row.team_id + ) + for file_id in provider_file_ids: + model_name = decode_model_from_file_id(file_id) + raw_file_id = get_original_file_id(file_id) + await self.store_unified_file_id( + file_id=file_id, + file_object=None, + litellm_parent_otel_span=litellm_parent_otel_span, + model_mappings={model_name: raw_file_id} if model_name else {}, + user_api_key_dict=owner_identity, + ) + async def list_user_batches( self, user_api_key_dict: UserAPIKeyAuth, @@ -397,6 +524,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the file {retrieve_file_id}", ) + if retrieve_file_id: + await self.enforce_provider_file_access( + retrieve_file_id, user_api_key_dict + ) return False async def check_file_ids_access( @@ -580,6 +711,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): data[accessor_key] = get_batch_id_from_unified_batch_id( potential_llm_object_id ) + elif retrieve_object_id and accessor_key == "batch_id": + await self.enforce_batch_object_access( + retrieve_object_id, user_api_key_dict + ) elif call_type == CallTypes.acreate_fine_tuning_job.value: input_file_id = cast(Optional[str], data.get("training_file")) if input_file_id: @@ -1183,6 +1318,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings={model_id: provider_file_id}, user_api_key_dict=user_api_key_dict, ) + is_batch_create = "completion_window" in data or "input_file_id" in data await self.store_unified_object_id( unified_object_id=response.id, file_object=response, @@ -1190,7 +1326,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_object_id=original_response_id, file_purpose="batch", user_api_key_dict=user_api_key_dict, + update_only=not is_batch_create, ) + if not is_batch_create: + await self.store_batch_output_file_ownership( + response=response, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + ) # Only record batch creation metric on actual create (not retrieve/cancel). # unified_file_id in _hidden_params is only set by the create_batch endpoint. diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 68dc3269f34..49d7a53d40a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -11,6 +11,7 @@ from litellm.caching import DualCache from litellm.proxy._types import CallTypes from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + encode_file_id_with_model, ) @@ -2367,3 +2368,414 @@ async def test_same_user_different_keys_can_access_batch(): assert "batch_id" in result2 # Both keys should get the same result assert result1["batch_id"] == result2["batch_id"] + + +MODEL_ENCODED_BATCH_ID = encode_file_id_with_model( + "batch_provider123", "gpt-4o-team-alias", id_type="batch" +) +MODEL_ENCODED_OUTPUT_FILE_ID = encode_file_id_with_model( + "file-output456", "gpt-4o-team-alias", id_type="file" +) +RAW_PROVIDER_BATCH_ID = "batch_provider123" +RAW_PROVIDER_FILE_ID = "file-output456" + + +def _owned_record(created_by, team_id): + record = MagicMock() + record.created_by = created_by + record.team_id = team_id + return record + + +def _batch_response(batch_id, output_file_id=None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1700000000, + endpoint="/v1/chat/completions", + input_file_id="file-input789", + object="batch", + status="completed", + output_file_id=output_file_id, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["aretrieve_batch", "acancel_batch"]) +@pytest.mark.parametrize( + "batch_id", [MODEL_ENCODED_BATCH_ID, RAW_PROVIDER_BATCH_ID] +) +async def test_team_b_cannot_access_team_a_provider_format_batch( + call_type, batch_id +): + """ + Cross-team retrieve/cancel of a model-encoded or raw provider batch id + must 403 when an ownership row exists for another team. + + Regression test: before this check only unified (litellm_proxy-prefixed) + batch ids were enforced, so any key could read any model-encoded or raw + provider batch. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + with pytest.raises(HTTPException) as exc_info: + await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() + ), + cache=MagicMock(), + data={"batch_id": batch_id}, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403 + prisma_client.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( + where={"unified_object_id": batch_id} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller_kwargs", + [ + {"user_id": "user_a", "team_id": "team_a"}, + {"user_id": "teammate_of_a", "team_id": "team_a"}, + {"user_id": "admin_user", "user_role": "proxy_admin"}, + ], +) +async def test_authorized_callers_can_access_provider_format_batch(caller_kwargs): + """ + The creator, a same-team member, and a proxy admin can all retrieve a + model-encoded batch owned by team_a. Data must pass through unmodified so + the endpoint's own routing still applies. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + parent_otel_span=MagicMock(), **caller_kwargs + ), + cache=MagicMock(), + data={"batch_id": MODEL_ENCODED_BATCH_ID}, + call_type="aretrieve_batch", + ) + + assert result["batch_id"] == MODEL_ENCODED_BATCH_ID + assert "model" not in result + + +@pytest.mark.asyncio +async def test_provider_format_batch_without_ownership_row_stays_accessible(): + """ + A provider-format batch id with no ownership row (created before ownership + tracking, or directly on the provider account) must stay retrievable. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_first.return_value = None + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() + ), + cache=MagicMock(), + data={"batch_id": RAW_PROVIDER_BATCH_ID}, + call_type="aretrieve_batch", + ) + + assert result["batch_id"] == RAW_PROVIDER_BATCH_ID + + +@pytest.mark.asyncio +async def test_fine_tuning_provider_format_id_not_enforced(): + """ + Provider-format fine-tuning job ids are deliberately out of scope for + ownership enforcement; only unified fine-tuning ids are checked. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() + ), + cache=MagicMock(), + data={"fine_tuning_job_id": "ftjob-abc123"}, + call_type="aretrieve_fine_tuning_job", + ) + + assert result["fine_tuning_job_id"] == "ftjob-abc123" + prisma_client.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type", ["afile_content", "afile_retrieve", "afile_delete"] +) +@pytest.mark.parametrize( + "file_id", [MODEL_ENCODED_OUTPUT_FILE_ID, RAW_PROVIDER_FILE_ID] +) +async def test_team_b_cannot_access_team_a_provider_format_file( + call_type, file_id +): + """ + Cross-team content/retrieve/delete of a model-encoded or raw provider + file id must 403 when an ownership row exists for another team. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(), prisma_client=prisma_client + ) + + with pytest.raises(HTTPException) as exc_info: + await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() + ), + cache=MagicMock(), + data={"file_id": file_id}, + call_type=call_type, + ) + + assert exc_info.value.status_code == 403 + prisma_client.db.litellm_managedfiletable.find_first.assert_awaited_once_with( + where={"unified_file_id": file_id} + ) + + +@pytest.mark.asyncio +async def test_same_team_can_access_provider_format_file(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id="teammate_of_a", team_id="team_a", parent_otel_span=MagicMock() + ), + cache=MagicMock(), + data={"file_id": MODEL_ENCODED_OUTPUT_FILE_ID}, + call_type="afile_content", + ) + + assert result["file_id"] == MODEL_ENCODED_OUTPUT_FILE_ID + + +@pytest.mark.asyncio +async def test_provider_format_file_without_ownership_row_stays_accessible(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_first.return_value = None + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() + ), + cache=MagicMock(), + data={"file_id": RAW_PROVIDER_FILE_ID}, + call_type="afile_content", + ) + + assert result["file_id"] == RAW_PROVIDER_FILE_ID + + +@pytest.mark.asyncio +async def test_post_call_batch_create_stores_ownership_row(): + """ + Batch creation (request data carries completion_window/input_file_id) + must write an ownership row attributed to the creating key. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client + ) + + await proxy_managed_files.async_post_call_success_hook( + data={ + "input_file_id": "file-input789", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + user_api_key_dict=UserAPIKeyAuth( + user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() + ), + response=_batch_response(MODEL_ENCODED_BATCH_ID), + ) + + upsert_call = prisma_client.db.litellm_managedobjecttable.upsert.await_args + assert upsert_call.kwargs["where"] == { + "unified_object_id": MODEL_ENCODED_BATCH_ID + } + create_data = upsert_call.kwargs["data"]["create"] + assert create_data["created_by"] == "user_a" + assert create_data["team_id"] == "team_a" + prisma_client.db.litellm_managedobjecttable.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_post_call_batch_sync_does_not_claim_ownership(): + """ + Retrieve/cancel of a batch with no ownership row must NOT create one: + otherwise the first foreign key to touch a legacy batch would become its + owner and lock out the real creator once enforcement is on. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update_many.return_value = 0 + internal_usage_cache = MagicMock(async_set_cache=AsyncMock()) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + internal_usage_cache, prisma_client=prisma_client + ) + + await proxy_managed_files.async_post_call_success_hook( + data={"batch_id": MODEL_ENCODED_BATCH_ID}, + user_api_key_dict=UserAPIKeyAuth( + user_id="user_b", team_id="team_b", parent_otel_span=MagicMock() + ), + response=_batch_response(MODEL_ENCODED_BATCH_ID), + ) + + prisma_client.db.litellm_managedobjecttable.upsert.assert_not_awaited() + internal_usage_cache.async_set_cache.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_post_call_batch_sync_updates_existing_row(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update_many.return_value = 1 + prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client + ) + + await proxy_managed_files.async_post_call_success_hook( + data={"batch_id": MODEL_ENCODED_BATCH_ID}, + user_api_key_dict=UserAPIKeyAuth( + user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() + ), + response=_batch_response(MODEL_ENCODED_BATCH_ID), + ) + + update_call = prisma_client.db.litellm_managedobjecttable.update_many.await_args + assert update_call.kwargs["where"] == { + "unified_object_id": MODEL_ENCODED_BATCH_ID + } + assert update_call.kwargs["data"]["status"] == "completed" + prisma_client.db.litellm_managedobjecttable.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_post_call_batch_sync_stores_output_file_ownership_from_batch_row(): + """ + When a synced batch reports a provider-format output file id, an + ownership row for that file must be written with the BATCH row's + created_by/team_id, not the caller's identity. + """ + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.update_many.return_value = 1 + prisma_client.db.litellm_managedobjecttable.find_first.return_value = ( + _owned_record(created_by="user_a", team_id="team_a") + ) + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client + ) + + await proxy_managed_files.async_post_call_success_hook( + data={"batch_id": MODEL_ENCODED_BATCH_ID}, + user_api_key_dict=UserAPIKeyAuth( + user_id="admin_user", + user_role="proxy_admin", + parent_otel_span=MagicMock(), + ), + response=_batch_response( + MODEL_ENCODED_BATCH_ID, output_file_id=MODEL_ENCODED_OUTPUT_FILE_ID + ), + ) + + file_upsert = prisma_client.db.litellm_managedfiletable.upsert.await_args + assert file_upsert.kwargs["where"] == { + "unified_file_id": MODEL_ENCODED_OUTPUT_FILE_ID + } + create_data = file_upsert.kwargs["data"]["create"] + assert create_data["created_by"] == "user_a" + assert create_data["team_id"] == "team_a" + assert create_data["flat_model_file_ids"] == ["file-output456"] + + +@pytest.mark.asyncio +async def test_post_call_batch_create_does_not_store_output_file_ownership(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client + ) + + await proxy_managed_files.async_post_call_success_hook( + data={ + "input_file_id": "file-input789", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + user_api_key_dict=UserAPIKeyAuth( + user_id="user_a", team_id="team_a", parent_otel_span=MagicMock() + ), + response=_batch_response( + MODEL_ENCODED_BATCH_ID, output_file_id=MODEL_ENCODED_OUTPUT_FILE_ID + ), + ) + + prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() From 26805f75fe44170f36b54d4e8e57d9ce7ad19b5b Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 5 Aug 2026 19:36:37 +0000 Subject: [PATCH 002/167] fix(bedrock_mantle): stop dropping the web_search tool on /v1/responses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../responses/transformation.py | 4 +- ...odel_prices_and_context_window_backup.json | 15 ++-- model_prices_and_context_window.json | 15 ++-- ...bedrock_mantle_responses_transformation.py | 90 +++++++++++++++++-- 4 files changed, 107 insertions(+), 17 deletions(-) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 92da5835b2d..3b01c7dbad0 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -44,7 +44,9 @@ _BASE_SUFFIXES_TO_STRIP: Final = ( ) # Per Bedrock Mantle Responses API validation errors. -_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES = frozenset({"function", "mcp", "custom", "namespace", "tool_search"}) +_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset( + {"function", "mcp", "custom", "namespace", "tool_search", "web_search"} +) _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02b3cde217a..e85d49cfd12 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45247,7 +45247,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -45275,7 +45276,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -45303,7 +45305,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, @@ -45330,7 +45333,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.4": { "input_cost_per_token": 2.75e-06, @@ -45357,7 +45361,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bc7330ec99c..7514f1596b0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45368,7 +45368,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -45396,7 +45397,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.6-luna": { "input_cost_per_token": 2.2e-07, @@ -45424,7 +45426,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, @@ -45451,7 +45454,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/openai.gpt-5.4": { "input_cost_per_token": 2.75e-06, @@ -45478,7 +45482,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_web_search": true }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index bea979aec64..3810ef39062 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -339,7 +339,7 @@ class TestBedrockMantleResponsesTools: params = cfg.map_openai_params( response_api_optional_params={ "tools": [ - {"type": "web_search"}, + {"type": "file_search", "vector_store_ids": ["vs_123"]}, {"type": "function", "name": "exec_command"}, ] }, @@ -351,7 +351,7 @@ class TestBedrockMantleResponsesTools: def test_map_openai_params_removes_tools_when_all_unsupported(self): cfg = BedrockMantleResponsesAPIConfig() params = cfg.map_openai_params( - response_api_optional_params={"tools": [{"type": "web_search"}]}, + response_api_optional_params={"tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}]}, model="openai.gpt-5.5", drop_params=False, ) @@ -365,12 +365,86 @@ class TestBedrockMantleResponsesTools: "litellm.llms.bedrock_mantle.responses.transformation.verbose_logger.warning" ) as mock_warning: cfg.map_openai_params( - response_api_optional_params={"tools": [{"type": "web_search"}]}, + response_api_optional_params={"tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}]}, model="openai.gpt-5.5", drop_params=False, ) assert mock_warning.call_count == 1 - assert "web_search" in str(mock_warning.call_args) + assert "file_search" in str(mock_warning.call_args) + + +class TestBedrockMantleResponsesWebSearch: + """Web Search on Amazon Bedrock is a server-side built-in tool that Mantle runs + itself when the caller passes {"type": "web_search"} on the Responses path, so + the config must forward the tool and its options untouched instead of filtering + it out and returning an ungrounded answer.""" + + _WEB_SEARCH_TOOL = {"type": "web_search", "external_web_access": False} + + def test_web_search_survives_map_openai_params_with_its_options(self): + cfg = BedrockMantleResponsesAPIConfig() + params = cfg.map_openai_params( + response_api_optional_params={"tools": [self._WEB_SEARCH_TOOL]}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert params["tools"] == [self._WEB_SEARCH_TOOL] + + def test_web_search_reaches_outbound_body_alongside_function_tools(self): + cfg = BedrockMantleResponsesAPIConfig() + function_tool = {"type": "function", "name": "exec_command"} + params = cfg.map_openai_params( + response_api_optional_params={"tools": [self._WEB_SEARCH_TOOL, function_tool]}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + body = cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input="What did AWS announce today?", + response_api_optional_request_params=params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["tools"] == [self._WEB_SEARCH_TOOL, function_tool] + + def test_web_search_is_not_logged_as_dropped(self): + from unittest.mock import patch + + cfg = BedrockMantleResponsesAPIConfig() + with patch("litellm.llms.bedrock_mantle.responses.transformation.verbose_logger.warning") as mock_warning: + cfg.map_openai_params( + response_api_optional_params={"tools": [self._WEB_SEARCH_TOOL]}, + model="openai.gpt-5.6-sol", + drop_params=False, + ) + assert mock_warning.call_count == 0 + + def test_hoisted_web_search_tool_survives(self): + cfg = BedrockMantleResponsesAPIConfig() + body = cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input=[ + {"type": "additional_tools", "role": "developer", "tools": [self._WEB_SEARCH_TOOL]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + ], + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["tools"] == [self._WEB_SEARCH_TOOL] + + @pytest.mark.parametrize( + "model", + [ + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + "bedrock_mantle/openai.gpt-5.5", + "bedrock_mantle/openai.gpt-5.4", + ], + ) + def test_cost_map_advertises_web_search_support(self, model): + assert litellm.supports_web_search(model=model) is True def _codex_exec_tool(): @@ -558,7 +632,7 @@ class TestBedrockMantleCodexAdditionalTools: "type": "additional_tools", "role": "developer", "tools": [ - {"type": "web_search"}, + {"type": "file_search", "vector_store_ids": ["vs_123"]}, {"type": "function", "name": "wait"}, ], }, @@ -570,7 +644,11 @@ class TestBedrockMantleCodexAdditionalTools: def test_item_stripped_even_when_no_hoisted_tool_survives(self): body = self._transform( input=[ - {"type": "additional_tools", "role": "developer", "tools": [{"type": "web_search"}]}, + { + "type": "additional_tools", + "role": "developer", + "tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}], + }, self._USER_MESSAGE, ] ) From 7bac4a41af549070d7084174f9e6e52fea282300 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:17:22 -0700 Subject: [PATCH 003/167] fix(prompts): key the in-memory prompt registry by environment LiteLLM_PromptTable is unique on (prompt_id, version, environment) and version numbering restarts at 1 per environment, but the in-memory registry keyed prompts as {prompt_id}.v{version} with no environment, so environments sharing a prompt id shadowed each other and only one environment's template ever served. Registry entries are now keyed {versioned_id}::{environment}, and serve time resolution goes through resolve_prompt_spec(base_id, version, environment): production > staging > development when no environment is requested, latest version within the chosen environment when no version is requested. Chat requests can pin an environment with a new optional prompt_environment body param, filtered from provider-bound params like prompt_id and prompt_version. The newest-updated_at dedupe in _init_prompts_in_db is dropped since registry keys can no longer collide, and the key-parsing serve helpers plus dead registry getters are removed --- litellm/proxy/prompts/prompt_endpoints.py | 247 +++--------------- litellm/proxy/prompts/prompt_registry.py | 174 +++++++++--- litellm/proxy/proxy_server.py | 11 +- litellm/proxy/utils.py | 28 +- litellm/types/utils.py | 1 + .../proxy/prompts/test_prompt_endpoints.py | 96 +------ .../prompts/test_prompt_endpoints_crud.py | 68 ++--- .../proxy/prompts/test_prompt_environment.py | 12 +- .../proxy/prompts/test_prompt_registry.py | 127 ++++++++- tests/test_litellm/proxy/test_proxy_server.py | 63 +++-- .../proxy_logging/test_guardrail_pipeline.py | 54 +++- 11 files changed, 413 insertions(+), 468 deletions(-) diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 425ff7572d0..32fd1d53777 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -29,6 +29,12 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.path_utils import safe_filename +from litellm.proxy.prompts.prompt_registry import ( + DEFAULT_PROMPT_ENVIRONMENT, + get_base_prompt_id, + get_version_number, + prompt_environment_or_default, +) from litellm.repositories.table_repositories import PromptRepository from litellm.types.prompts.init_prompts import ( ListPromptsResponse, @@ -102,165 +108,20 @@ def _prompt_table(prisma_client: "PrismaClient") -> _PromptTableActions: return PromptRepository(prisma_client).table -def get_base_prompt_id(prompt_id: str) -> str: - """ - Extract the base prompt ID by stripping the version suffix if present. - - Args: - prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v1" or "jack_success_v1") - - Returns: - Base prompt ID without version suffix (e.g., "jack_success") - - Examples: - >>> get_base_prompt_id("jack_success.v1") - "jack_success" - >>> get_base_prompt_id("jack_success_v1") - "jack_success" - >>> get_base_prompt_id("jack_success") - "jack_success" - """ - # Try dot separator first (.v) - if ".v" in prompt_id: - return prompt_id.split(".v")[0] - # Try underscore separator (_v) - if "_v" in prompt_id: - return prompt_id.split("_v")[0] - return prompt_id - - -def get_version_number(prompt_id: str) -> int: - """ - Extract the version number from a versioned prompt ID. - - Args: - prompt_id: Prompt ID that may include version suffix (e.g., "jack_success.v2" or "jack_success_v2") - - Returns: - Version number (defaults to 1 if no version suffix or invalid format) - - Examples: - >>> get_version_number("jack_success.v2") - 2 - >>> get_version_number("jack_success_v2") - 2 - >>> get_version_number("jack_success") - 1 - """ - # Try dot separator first (.v) - if ".v" in prompt_id: - version_str = prompt_id.split(".v")[1] - try: - return int(version_str) - except ValueError: - pass - - # Try underscore separator (_v) - if "_v" in prompt_id: - version_str = prompt_id.split("_v")[1] - try: - return int(version_str) - except ValueError: - pass - - return 1 - - -def construct_versioned_prompt_id(prompt_id: str, version: int | None = None) -> str: - """ - Construct a versioned prompt ID from a base prompt_id and version number. - - Args: - prompt_id: Base prompt ID (e.g., "jack_success") - version: Version number (if None, returns the base prompt_id unchanged) - - Returns: - Versioned prompt ID (e.g., "jack_success.v4") - - Examples: - >>> construct_versioned_prompt_id("jack_success", 4) - "jack_success.v4" - >>> construct_versioned_prompt_id("jack_success", None) - "jack_success" - >>> construct_versioned_prompt_id("jack_success.v2", 4) - "jack_success.v4" - """ - if version is None: - return prompt_id - - # Strip any existing version suffix first - base_id: Final = get_base_prompt_id(prompt_id) - return f"{base_id}.v{version}" - - -def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Mapping[str, object]) -> str: - """ - Find the latest version of a prompt from available prompt IDs. - - Args: - prompt_id: Base prompt ID or versioned prompt ID (e.g., "jack_success" or "jack_success.v2") - all_prompt_ids: Dictionary of all available prompt IDs (keys are prompt IDs) - - Returns: - The prompt ID with the highest version number, or the original prompt_id if no versions exist - - Examples: - >>> all_ids = {"jack.v1": {}, "jack.v2": {}, "jack.v3": {}} - >>> get_latest_version_prompt_id("jack", all_ids) - "jack.v3" - >>> get_latest_version_prompt_id("jack.v1", all_ids) - "jack.v3" - >>> all_ids = {"simple": {}} - >>> get_latest_version_prompt_id("simple", all_ids) - "simple" - """ - base_id: Final = get_base_prompt_id(prompt_id=prompt_id) - - # Find all versions of this prompt - matching_versions: Final = [] - for stored_prompt_id in all_prompt_ids: - if get_base_prompt_id(prompt_id=stored_prompt_id) == base_id: - version_num = get_version_number(prompt_id=stored_prompt_id) - matching_versions.append((version_num, stored_prompt_id)) - - # Use the highest version number - if matching_versions: - matching_versions.sort(reverse=True) - return matching_versions[0][1] - else: - # No versioned prompts found, use the base ID as-is - return prompt_id - - def get_latest_prompt_versions(prompts: list[PromptSpec]) -> list[PromptSpec]: """ - Filter a list of prompts to return only the latest version of each unique prompt. - - Args: - prompts: List of PromptSpec objects - - Returns: - List of PromptSpec objects with only the latest version of each prompt + Filter prompts down to the latest version per (base prompt id, environment). """ - latest_prompts: Final[dict[str, PromptSpec]] = {} - - for prompt in prompts: - base_id = get_base_prompt_id(prompt_id=prompt.prompt_id) - version = get_version_number(prompt_id=prompt.prompt_id) - - # Keep the prompt with the highest version number - if base_id not in latest_prompts: - latest_prompts[base_id] = prompt - else: - existing_version = get_version_number(prompt_id=latest_prompts[base_id].prompt_id) - if version > existing_version: - latest_prompts[base_id] = prompt - + sorted_prompts: Final = sorted(prompts, key=lambda prompt: get_version_number(prompt_id=prompt.prompt_id)) + latest_prompts: Final = { + (get_base_prompt_id(prompt_id=prompt.prompt_id), prompt_environment_or_default(prompt.environment)): prompt + for prompt in sorted_prompts + } return list(latest_prompts.values()) async def get_next_version_for_prompt( - prisma_client: "PrismaClient", prompt_id: str, environment: str = "development" + prisma_client: "PrismaClient", prompt_id: str, environment: str = DEFAULT_PROMPT_ENVIRONMENT ) -> int: """ Get the next version number for a prompt in a specific environment. @@ -403,11 +264,14 @@ async def list_prompts( if key_metadata is not None: prompts: Final = cast(list[str] | None, key_metadata.get("prompts", None)) if prompts is not None: - all_prompts = [ - IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[prompt_id] - for prompt_id in prompts - if prompt_id in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS + allowed_prompt_ids: Final = frozenset(prompts) + allowed_prompts: Final = [ + spec + for spec in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.values() + if spec.prompt_id in allowed_prompt_ids + or get_base_prompt_id(prompt_id=spec.prompt_id) in allowed_prompt_ids ] + all_prompts = get_latest_prompt_versions(prompts=allowed_prompts) if environment: all_prompts = [p for p in all_prompts if p.environment == environment] prompt_list: Final = [] @@ -576,7 +440,7 @@ def _get_prompt_template(prompt_spec: PromptSpec, base_prompt_id: str) -> Prompt metadata=parsed.get("metadata"), ) else: - prompt_callback: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(prompt_spec.prompt_id) + prompt_callback: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=prompt_spec) if prompt_callback is not None: integration_name: Final = prompt_callback.integration_name if integration_name == "dotprompt": @@ -690,15 +554,8 @@ async def get_prompt_info( if env_prompts: prompt_spec = create_versioned_prompt_spec(db_prompt=env_prompts[0]) - # Fallback: use in-memory registry (no environment filter) if prompt_spec is None and environment is None: - prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if prompt_spec is None: - latest_prompt_id: Final = get_latest_version_prompt_id( - prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, - ) - prompt_spec = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) + prompt_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec(prompt_id, version=requested_version) if prompt_spec is None: raise HTTPException( @@ -785,7 +642,7 @@ async def create_prompt( environment: Final = ( request.prompt_info.environment if request.prompt_info and request.prompt_info.environment - else "development" + else DEFAULT_PROMPT_ENVIRONMENT ) # Get next version number @@ -885,7 +742,7 @@ async def update_prompt( environment: Final = ( request.prompt_info.environment if request.prompt_info and request.prompt_info.environment - else "development" + else DEFAULT_PROMPT_ENVIRONMENT ) # Check if any version of this prompt exists (in any environment) @@ -897,9 +754,7 @@ async def update_prompt( detail=f"Prompt with ID {base_prompt_id} not found", ) - # Check if it's a config prompt - existing_in_memory: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - if existing_in_memory and existing_in_memory.prompt_info.prompt_type == "config": + if IN_MEMORY_PROMPT_REGISTRY.has_config_prompt(base_prompt_id=base_prompt_id): raise HTTPException( status_code=400, detail="Cannot update config prompts.", @@ -988,52 +843,24 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - # Try to get prompt directly first - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(prompt_id) - - # If not found, try to find the latest version - if existing_prompt is None: - latest_prompt_id: Final = get_latest_version_prompt_id( - prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, - ) - existing_prompt = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(latest_prompt_id) - # Use the resolved prompt_id for deletion - prompt_id = latest_prompt_id + base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) + existing_prompt: Final = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec(prompt_id, environment=environment) if existing_prompt is None: raise HTTPException(status_code=404, detail=f"Prompt with ID {prompt_id} not found") - if existing_prompt.prompt_info.prompt_type == "config": + if IN_MEMORY_PROMPT_REGISTRY.has_config_prompt(base_prompt_id=base_prompt_id): raise HTTPException( status_code=400, detail="Cannot delete config prompts.", ) - # Get the base prompt ID (without version suffix) for database deletion - base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) - - # Build delete filter; scope to environment if provided - delete_where: Final[dict[str, str]] = {"prompt_id": base_prompt_id} - if environment: - delete_where["environment"] = environment - - # Delete versions from the database (scoped to environment if provided) + delete_where: Final[dict[str, str]] = { + "prompt_id": base_prompt_id, + **({"environment": environment} if environment else {}), + } await _prompt_table(prisma_client).delete_many(where=delete_where) - - # Remove matching prompts from memory — scope to environment if provided - if environment: - prompts_to_delete: Final = [ - pid - for pid, prompt in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.items() - if get_base_prompt_id(prompt_id=pid) == base_prompt_id and prompt.environment == environment - ] - for pid in prompts_to_delete: - del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[pid] - if pid in IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt: - del IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt[pid] - else: - IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id) + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id=base_prompt_id, environment=environment) env_msg: Final = f" from {environment}" if environment else "" return {"message": f"Prompt {base_prompt_id} deleted successfully{env_msg}"} @@ -1105,7 +932,7 @@ async def patch_prompt( try: # Resolve the target row: find the latest version in the given environment base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) - env: Final = environment or "development" + env: Final = prompt_environment_or_default(environment) requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None # Build query to find the exact row by composite unique key @@ -1129,11 +956,7 @@ async def patch_prompt( target_row: Final = db_rows[0] - # Check if prompt exists in memory - versioned_id: Final = f"{base_prompt_id}.v{target_row.version}" - existing_prompt: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(versioned_id) - - if existing_prompt and existing_prompt.prompt_info.prompt_type == "config": + if IN_MEMORY_PROMPT_REGISTRY.has_config_prompt(base_prompt_id=base_prompt_id): raise HTTPException( status_code=400, detail="Cannot update config prompts.", diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index d4342773a85..b575184a229 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -1,6 +1,6 @@ import importlib import os -from collections.abc import Callable +from collections.abc import Callable, Sequence from pathlib import Path from typing import Final @@ -14,6 +14,77 @@ from litellm.types.prompts.init_prompts import ( prompt_initializer_registry = {} +DEFAULT_PROMPT_ENVIRONMENT: Final = "development" +PROMPT_ENVIRONMENT_SERVE_PRECEDENCE: Final = ("production", "staging", "development") + + +def get_base_prompt_id(prompt_id: str) -> str: + """ + Extract the base prompt ID by stripping the version suffix if present. + + Examples: + >>> get_base_prompt_id("jack_success.v1") + "jack_success" + >>> get_base_prompt_id("jack_success_v1") + "jack_success" + >>> get_base_prompt_id("jack_success") + "jack_success" + """ + if ".v" in prompt_id: + return prompt_id.split(".v")[0] + if "_v" in prompt_id: + return prompt_id.split("_v")[0] + return prompt_id + + +def get_version_number(prompt_id: str) -> int: + """ + Extract the version number from a versioned prompt ID (defaults to 1). + + Examples: + >>> get_version_number("jack_success.v2") + 2 + >>> get_version_number("jack_success_v2") + 2 + >>> get_version_number("jack_success") + 1 + """ + if ".v" in prompt_id: + version_str = prompt_id.split(".v")[1] + try: + return int(version_str) + except ValueError: + pass + + if "_v" in prompt_id: + version_str = prompt_id.split("_v")[1] + try: + return int(version_str) + except ValueError: + pass + + return 1 + + +def prompt_environment_or_default(environment: str | None) -> str: + return environment or DEFAULT_PROMPT_ENVIRONMENT + + +def registry_key_for_prompt(prompt: PromptSpec) -> str: + return f"{prompt.prompt_id}::{prompt_environment_or_default(prompt.environment)}" + + +def _spec_version(prompt: PromptSpec) -> int: + return prompt.version if prompt.version is not None else get_version_number(prompt_id=prompt.prompt_id) + + +def _default_serve_environment(prompts: Sequence[PromptSpec]) -> str: + present: Final = frozenset(prompt_environment_or_default(prompt.environment) for prompt in prompts) + ladder_pick: Final = next((env for env in PROMPT_ENVIRONMENT_SERVE_PRECEDENCE if env in present), None) + if ladder_pick is not None: + return ladder_pick + return min(present) if present else DEFAULT_PROMPT_ENVIRONMENT + def get_prompt_initializer_from_integrations(): """ @@ -113,17 +184,16 @@ class InMemoryPromptRegistry: """ import litellm - prompt_id: Final = prompt.prompt_id - if prompt_id in self.IN_MEMORY_PROMPTS: - verbose_proxy_logger.debug("prompt_id already exists in IN_MEMORY_PROMPTS") - return self.IN_MEMORY_PROMPTS[prompt_id] + registry_key: Final = registry_key_for_prompt(prompt) + if registry_key in self.IN_MEMORY_PROMPTS: + verbose_proxy_logger.debug("prompt already exists in IN_MEMORY_PROMPTS") + return self.IN_MEMORY_PROMPTS[registry_key] parsed_prompt, custom_prompt_callback = self._build_prompt_callback(prompt=prompt) litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback) - # store references to the prompt in memory - self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt - self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback + self.IN_MEMORY_PROMPTS[registry_key] = parsed_prompt + self.prompt_id_to_custom_prompt[registry_key] = custom_prompt_callback return parsed_prompt @@ -166,57 +236,85 @@ class InMemoryPromptRegistry: import litellm parsed_prompt, new_callback = self._build_prompt_callback(prompt=prompt) - stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None) - self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None) + registry_key: Final = registry_key_for_prompt(parsed_prompt) + stale_callback: Final = self.prompt_id_to_custom_prompt.pop(registry_key, None) + self.IN_MEMORY_PROMPTS.pop(registry_key, None) if stale_callback is not None: litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback) litellm.logging_callback_manager.add_litellm_callback(new_callback) - self.IN_MEMORY_PROMPTS[prompt.prompt_id] = parsed_prompt - self.prompt_id_to_custom_prompt[prompt.prompt_id] = new_callback + self.IN_MEMORY_PROMPTS[registry_key] = parsed_prompt + self.prompt_id_to_custom_prompt[registry_key] = new_callback return parsed_prompt def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None: - existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id) + existing: Final = self.IN_MEMORY_PROMPTS.get(registry_key_for_prompt(prompt)) if existing is None: return self.initialize_prompt(prompt=prompt) if existing.litellm_params == prompt.litellm_params and existing.prompt_info == prompt.prompt_info: return existing return self.reload_prompt(prompt=prompt) - def get_prompt_by_id(self, prompt_id: str) -> PromptSpec | None: + def resolve_prompt_spec( + self, + prompt_id: str, + version: int | None = None, + environment: str | None = None, + ) -> PromptSpec | None: """ - Get a prompt by its ID from memory - """ - return self.IN_MEMORY_PROMPTS.get(prompt_id) + Resolve a prompt spec by base prompt id, optional version, and optional environment. - def get_prompt_callback_by_id(self, prompt_id: str) -> CustomPromptManagement | None: + With no environment, resolves within the default serve environment + (production > staging > development > alphabetical first present). + With no version, resolves to the highest version in the chosen environment. """ - Get a prompt callback by its ID from memory + base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) + base_matches: Final = tuple( + spec + for spec in self.IN_MEMORY_PROMPTS.values() + if get_base_prompt_id(prompt_id=spec.prompt_id) == base_prompt_id + ) + if not base_matches: + return None + resolved_environment: Final = ( + environment if environment is not None else _default_serve_environment(base_matches) + ) + env_matches: Final = tuple( + spec for spec in base_matches if prompt_environment_or_default(spec.environment) == resolved_environment + ) + if not env_matches: + return None + if version is not None: + return next((spec for spec in env_matches if _spec_version(spec) == version), None) + return max(env_matches, key=_spec_version) + + def get_prompt_callback_for_prompt(self, prompt: PromptSpec) -> CustomPromptManagement | None: + return self.prompt_id_to_custom_prompt.get(registry_key_for_prompt(prompt)) + + def has_config_prompt(self, base_prompt_id: str) -> bool: + return any( + spec.prompt_info.prompt_type == "config" + for spec in self.IN_MEMORY_PROMPTS.values() + if get_base_prompt_id(prompt_id=spec.prompt_id) == base_prompt_id + ) + + def delete_prompts_by_base_id(self, base_prompt_id: str, environment: str | None = None) -> list[str]: """ - return self.prompt_id_to_custom_prompt.get(prompt_id) + Delete matching prompts from memory, scoped to one environment when given. - def delete_prompts_by_base_id(self, base_prompt_id: str) -> list[str]: + Returns the registry keys that were deleted. """ - Delete all prompts matching the given base prompt ID from memory. - - Args: - base_prompt_id: The base prompt ID (without version suffix) - - Returns: - List of prompt IDs that were deleted - """ - from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id - - prompts_to_delete: Final = [ - pid for pid in self.IN_MEMORY_PROMPTS if get_base_prompt_id(prompt_id=pid) == base_prompt_id + keys_to_delete: Final = [ + key + for key, spec in self.IN_MEMORY_PROMPTS.items() + if get_base_prompt_id(prompt_id=spec.prompt_id) == base_prompt_id + and (environment is None or prompt_environment_or_default(spec.environment) == environment) ] - for pid in prompts_to_delete: - del self.IN_MEMORY_PROMPTS[pid] - if pid in self.prompt_id_to_custom_prompt: - del self.prompt_id_to_custom_prompt[pid] + for key in keys_to_delete: + del self.IN_MEMORY_PROMPTS[key] + self.prompt_id_to_custom_prompt.pop(key, None) - return prompts_to_delete + return keys_to_delete IN_MEMORY_PROMPT_REGISTRY: Final = InMemoryPromptRegistry() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2cfe08fe332..bc4eaaab761 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7273,16 +7273,7 @@ class ProxyConfig: parsed_specs: Final[tuple[PromptSpec, ...]] = tuple( spec for row in prompts_in_db if (spec := parse_row(row)) is not None ) - newest_spec_per_id: Final[Mapping[str, PromptSpec]] = MappingProxyType( - { - spec.prompt_id: spec - for spec in sorted( - parsed_specs, - key=lambda s: s.updated_at.timestamp() if s.updated_at else float("-inf"), - ) - } - ) - for prompt_spec in newest_spec_per_id.values(): + for prompt_spec in parsed_specs: try: IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8cbf5b685fd..95f02279f19 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1397,27 +1397,26 @@ class ProxyLogging: ) -> None: """Process prompt template if applicable.""" - from litellm.proxy.prompts.prompt_endpoints import ( - construct_versioned_prompt_id, - get_latest_version_prompt_id, - ) from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.utils import get_non_default_completion_params - if prompt_version is None: - lookup_prompt_id = get_latest_version_prompt_id( - prompt_id=prompt_id, - all_prompt_ids=IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS, - ) - else: - lookup_prompt_id = construct_versioned_prompt_id(prompt_id=prompt_id, version=prompt_version) - - custom_logger: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id(lookup_prompt_id) - prompt_spec: Final = IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id(lookup_prompt_id) + raw_prompt_environment: Final = data.get("prompt_environment", None) + prompt_environment: Final = raw_prompt_environment if isinstance(raw_prompt_environment, str) else None + prompt_spec: Final = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec( + prompt_id, + version=prompt_version, + environment=prompt_environment, + ) + custom_logger: Final = ( + IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=prompt_spec) + if prompt_spec is not None + else None + ) litellm_prompt_id: str | None = None if prompt_spec is not None: litellm_prompt_id = prompt_spec.litellm_params.prompt_id data.pop("prompt_id", None) + data.pop("prompt_environment", None) if custom_logger and prompt_spec is not None: ( @@ -1444,6 +1443,7 @@ class ProxyLogging: data.pop("prompt_variables", None) data.pop("prompt_label", None) data.pop("prompt_version", None) + data.pop("prompt_environment", None) def _process_guardrail_metadata(self, data: dict) -> None: """Process guardrails from metadata and add to applied_guardrails.""" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ef3586f2559..27437b74251 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3522,6 +3522,7 @@ all_litellm_params = ( "litellm_system_prompt", "provider_specific_header", "prompt_version", + "prompt_environment", "api_base", "force_timeout", "logger_fn", diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py index 6ebb10eff76..35402db219f 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py @@ -104,97 +104,6 @@ class TestPromptVersioning: assert get_base_prompt_id(prompt_id="jack") == "jack" assert get_base_prompt_id(prompt_id="my_prompt.v10") == "my_prompt" - def test_get_latest_version_prompt_id(self): - """ - Test that get_latest_version_prompt_id returns the highest version - """ - from litellm.proxy.prompts.prompt_endpoints import get_latest_version_prompt_id - - # Mock prompt IDs dictionary - all_prompt_ids = { - "jack.v1": {}, - "jack.v2": {}, - "jack.v3": {}, - "jane.v1": {}, - "simple_prompt": {}, - } - - # Test with base prompt ID - should return latest version - assert ( - get_latest_version_prompt_id( - prompt_id="jack", all_prompt_ids=all_prompt_ids - ) - == "jack.v3" - ) - - # Test with versioned prompt ID - should still return latest version - assert ( - get_latest_version_prompt_id( - prompt_id="jack.v1", all_prompt_ids=all_prompt_ids - ) - == "jack.v3" - ) - - # Test with single version - assert ( - get_latest_version_prompt_id( - prompt_id="jane", all_prompt_ids=all_prompt_ids - ) - == "jane.v1" - ) - - # Test with non-versioned prompt - assert ( - get_latest_version_prompt_id( - prompt_id="simple_prompt", all_prompt_ids=all_prompt_ids - ) - == "simple_prompt" - ) - - # Test with non-existent prompt - assert ( - get_latest_version_prompt_id( - prompt_id="nonexistent", all_prompt_ids=all_prompt_ids - ) - == "nonexistent" - ) - - def test_construct_versioned_prompt_id(self): - """ - Test that construct_versioned_prompt_id correctly builds versioned IDs - """ - from litellm.proxy.prompts.prompt_endpoints import construct_versioned_prompt_id - - # Test with base prompt ID and version - assert ( - construct_versioned_prompt_id(prompt_id="jack_success", version=4) - == "jack_success.v4" - ) - - # Test with None version - should return base ID unchanged - assert ( - construct_versioned_prompt_id(prompt_id="jack_success", version=None) - == "jack_success" - ) - - # Test with existing versioned ID - should replace version - assert ( - construct_versioned_prompt_id(prompt_id="jack_success.v2", version=4) - == "jack_success.v4" - ) - - # Test with hyphenated prompt ID - assert ( - construct_versioned_prompt_id(prompt_id="my-prompt", version=1) - == "my-prompt.v1" - ) - - # Test with double-digit version - assert ( - construct_versioned_prompt_id(prompt_id="test_prompt", version=10) - == "test_prompt.v10" - ) - class TestPromptVersionsEndpoint: """ @@ -444,7 +353,7 @@ class TestAdminViewerReadAccess: "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = PromptSpec( + mock_registry.resolve_prompt_spec.return_value = PromptSpec( prompt_id="jack.v2", litellm_params=PromptLiteLLMParams( prompt_id="jack", @@ -453,8 +362,7 @@ class TestAdminViewerReadAccess: ), prompt_info=PromptInfo(prompt_type="db"), ) - mock_registry.IN_MEMORY_PROMPTS = {"jack.v1": {}, "jack.v2": {}} - mock_registry.get_prompt_callback_by_id.return_value = None + mock_registry.get_prompt_callback_for_prompt.return_value = None response = await get_prompt_info(prompt_id="jack", user_api_key_dict=viewer) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py index 688b739fb5a..6916a7163f8 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -52,8 +52,6 @@ async def test_delete_prompt_success(): with patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - # User passes "test_prompt.v2" - # We simulate that get_prompt_by_id returns the prompt spec for v2 prompt_spec = PromptSpec( prompt_id="test_prompt.v2", litellm_params=PromptLiteLLMParams( @@ -61,7 +59,8 @@ async def test_delete_prompt_success(): ), prompt_info=PromptInfo(prompt_type="db"), ) - mock_registry.get_prompt_by_id.return_value = prompt_spec + mock_registry.resolve_prompt_spec.return_value = prompt_spec + mock_registry.has_config_prompt.return_value = False # Patch the prisma client in the endpoint module with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): @@ -79,7 +78,7 @@ async def test_delete_prompt_success(): # 2. Memory deletion should use base ID mock_registry.delete_prompts_by_base_id.assert_called_once_with( - expected_base_id + base_prompt_id=expected_base_id, environment=None ) assert response == { @@ -108,31 +107,14 @@ async def test_delete_prompt_by_base_id_success(): with patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - # User passes "test_prompt" (base ID) - # 1. get_prompt_by_id("test_prompt") -> None (if it's not registered as base) - # 2. It calls get_latest_version_prompt_id -> returns "test_prompt.v3" - # 3. get_prompt_by_id("test_prompt.v3") -> returns Spec - - # Setup mocks behavior - def get_prompt_side_effect(prompt_id): - if prompt_id == "test_prompt": - return None - if prompt_id == "test_prompt.v3": - return PromptSpec( - prompt_id="test_prompt.v3", - litellm_params=PromptLiteLLMParams( - prompt_id="test_prompt", prompt_integration="dotprompt" - ), - prompt_info=PromptInfo(prompt_type="db"), - ) - return None - - mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect - mock_registry.IN_MEMORY_PROMPTS = { - "test_prompt.v1": {}, - "test_prompt.v2": {}, - "test_prompt.v3": {}, - } + mock_registry.resolve_prompt_spec.return_value = PromptSpec( + prompt_id="test_prompt.v3", + litellm_params=PromptLiteLLMParams( + prompt_id="test_prompt", prompt_integration="dotprompt" + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + mock_registry.has_config_prompt.return_value = False # Patch the prisma client in the endpoint module with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): @@ -150,7 +132,7 @@ async def test_delete_prompt_by_base_id_success(): # 2. Memory deletion should use base ID mock_registry.delete_prompts_by_base_id.assert_called_once_with( - expected_base_id + base_prompt_id=expected_base_id, environment=None ) assert response == { @@ -187,24 +169,8 @@ async def test_get_prompt_info_by_base_id(): prompt_info=PromptInfo(prompt_type="db"), ) - # When get_prompt_by_id is called with "test_prompt", return None (so it searches versions) - # When called with "test_prompt.v3", return the spec - def get_prompt_side_effect(prompt_id): - if prompt_id == "test_prompt": - return None - if prompt_id == "test_prompt.v3": - return prompt_spec_v3 - return None - - mock_registry.get_prompt_by_id.side_effect = get_prompt_side_effect - mock_registry.IN_MEMORY_PROMPTS = { - "test_prompt.v1": {}, - "test_prompt.v2": {}, - "test_prompt.v3": {}, - } - - # We also need to mock get_prompt_callback_by_id to avoid content extraction errors/logic - mock_registry.get_prompt_callback_by_id.return_value = None + mock_registry.resolve_prompt_spec.return_value = prompt_spec_v3 + mock_registry.get_prompt_callback_for_prompt.return_value = None response = await get_prompt_info( prompt_id="test_prompt", user_api_key_dict=mock_user_auth @@ -253,7 +219,7 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404(): "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = existing_prompt + mock_registry.has_config_prompt.return_value = False with pytest.raises(HTTPException) as exc_info: await patch_prompt( @@ -294,7 +260,7 @@ async def test_patch_prompt_merges_unsent_fields_from_db_row_not_stale_memory(): "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = stale_in_memory + mock_registry.has_config_prompt.return_value = False mock_registry.reload_prompt.side_effect = lambda prompt: prompt response = await patch_prompt( @@ -456,7 +422,7 @@ async def test_patch_prompt_info_only_keeps_legacy_keyed_row_patchable(): "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry, ): - mock_registry.get_prompt_by_id.return_value = existing_prompt + mock_registry.has_config_prompt.return_value = False await patch_prompt( prompt_id="agent-prompt", diff --git a/tests/test_litellm/proxy/prompts/test_prompt_environment.py b/tests/test_litellm/proxy/prompts/test_prompt_environment.py index ecd89afefbe..3cb647dea13 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_environment.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_environment.py @@ -191,11 +191,7 @@ async def test_update_prompt_stores_environment_and_created_by(): with patch( "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" ) as mock_registry: - mock_registry.get_prompt_by_id.return_value = PromptSpec( - prompt_id="my_prompt.v1", - litellm_params=request.litellm_params, - prompt_info=PromptInfo(prompt_type="db"), - ) + mock_registry.has_config_prompt.return_value = False mock_registry.initialize_prompt.return_value = PromptSpec( prompt_id="my_prompt.v2", litellm_params=request.litellm_params, @@ -239,7 +235,8 @@ async def test_delete_prompt_scoped_to_environment(): prompt_info=PromptInfo(prompt_type="db"), environment="staging", ) - mock_registry.get_prompt_by_id.return_value = prompt_spec + mock_registry.resolve_prompt_spec.return_value = prompt_spec + mock_registry.has_config_prompt.return_value = False with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): await delete_prompt( @@ -251,3 +248,6 @@ async def test_delete_prompt_scoped_to_environment(): mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( where={"prompt_id": "test_prompt", "environment": "staging"} ) + mock_registry.delete_prompts_by_base_id.assert_called_once_with( + base_prompt_id="test_prompt", environment="staging" + ) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py index 47f1ba13627..689cbd79873 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_registry.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -1,26 +1,35 @@ import pytest import litellm +from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec -def _db_prompt_spec(content: str) -> PromptSpec: +def _db_prompt_spec(content: str, environment: str = "development", version: int = 1) -> PromptSpec: return PromptSpec( - prompt_id="greeting.v1", + prompt_id=f"greeting.v{version}", litellm_params=PromptLiteLLMParams( prompt_id="greeting", prompt_integration="dotprompt", prompt_data={"content": content, "metadata": {}}, ), prompt_info=PromptInfo(prompt_type="db"), + version=version, + environment=environment, ) -def _served_content(registry: InMemoryPromptRegistry) -> str: - callback = registry.get_prompt_callback_by_id("greeting.v1") +def _resolved_callback(registry: InMemoryPromptRegistry, environment: str | None = None) -> CustomPromptManagement: + spec = registry.resolve_prompt_spec("greeting", environment=environment) + assert spec is not None + callback = registry.get_prompt_callback_for_prompt(prompt=spec) assert callback is not None - return callback.prompt_manager.get_prompt("greeting").content + return callback + + +def _served_content(registry: InMemoryPromptRegistry, environment: str | None = None) -> str: + return _resolved_callback(registry, environment=environment).prompt_manager.get_prompt("greeting").content @pytest.fixture @@ -32,32 +41,34 @@ def isolated_callbacks(monkeypatch: pytest.MonkeyPatch) -> list: def test_sync_prompt_from_db_reloads_row_edited_elsewhere(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) - stale_callback = registry.get_prompt_callback_by_id("greeting.v1") + stale_callback = _resolved_callback(registry) assert _served_content(registry) == "begin every reply with AHOY" registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY")) assert _served_content(registry) == "begin every reply with HOWDY" - assert registry.get_prompt_by_id("greeting.v1").litellm_params.prompt_data["content"] == "begin every reply with HOWDY" + reloaded_spec = registry.resolve_prompt_spec("greeting", environment="development") + assert reloaded_spec is not None + assert reloaded_spec.litellm_params.prompt_data["content"] == "begin every reply with HOWDY" assert stale_callback not in isolated_callbacks - assert isolated_callbacks == [registry.get_prompt_callback_by_id("greeting.v1")] + assert isolated_callbacks == [_resolved_callback(registry)] def test_sync_prompt_from_db_keeps_unchanged_row_in_place(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) - first_callback = registry.get_prompt_callback_by_id("greeting.v1") + first_callback = _resolved_callback(registry) registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY")) - assert registry.get_prompt_callback_by_id("greeting.v1") is first_callback + assert _resolved_callback(registry) is first_callback assert isolated_callbacks == [first_callback] def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) - stale_callback = registry.get_prompt_callback_by_id("greeting.v1") + stale_callback = _resolved_callback(registry) reloaded = registry.reload_prompt(prompt=_db_prompt_spec("begin every reply with HOWDY")) @@ -70,7 +81,7 @@ def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_ca def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolated_callbacks: list) -> None: registry = InMemoryPromptRegistry() registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY")) - old_callback = registry.get_prompt_callback_by_id("greeting.v1") + old_callback = _resolved_callback(registry) broken = PromptSpec( prompt_id="greeting.v1", @@ -80,11 +91,101 @@ def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolate prompt_data={"content": "begin every reply with HOWDY", "metadata": {}}, ), prompt_info=PromptInfo(prompt_type="db"), + version=1, + environment="development", ) with pytest.raises(ValueError, match="Unsupported prompt"): registry.reload_prompt(prompt=broken) - assert registry.get_prompt_callback_by_id("greeting.v1") is old_callback + assert _resolved_callback(registry) is old_callback assert _served_content(registry) == "begin every reply with AHOY" assert isolated_callbacks == [old_callback] + + +def test_environments_sharing_a_prompt_id_keep_separate_templates(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + assert _served_content(registry, environment="development") == "begin every reply with AHOY" + assert _served_content(registry, environment="production") == "begin every reply with HOWDY" + assert _resolved_callback(registry, environment="development") is not _resolved_callback( + registry, environment="production" + ) + + +def test_default_resolution_prefers_production(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + assert _served_content(registry) == "begin every reply with HOWDY" + + +@pytest.mark.parametrize("environment", ["staging", "qa"]) +def test_default_resolution_serves_the_only_environment_present(isolated_callbacks: list, environment: str) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment=environment)) + + assert _served_content(registry) == "begin every reply with AHOY" + + +def test_resolution_picks_exact_version_and_latest_within_an_environment(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development", version=1)) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with YO", environment="development", version=2)) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production", version=1)) + + exact = registry.resolve_prompt_spec("greeting", version=1, environment="development") + assert exact is not None + assert exact.litellm_params.prompt_data["content"] == "begin every reply with AHOY" + + latest = registry.resolve_prompt_spec("greeting", environment="development") + assert latest is not None + assert latest.litellm_params.prompt_data["content"] == "begin every reply with YO" + + assert registry.resolve_prompt_spec("greeting", version=3, environment="development") is None + + +def test_resolution_returns_none_for_unknown_environment_or_prompt(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + + assert registry.resolve_prompt_spec("greeting", environment="production") is None + assert registry.resolve_prompt_spec("no_such_prompt") is None + + +def test_delete_prompts_by_base_id_scoped_to_one_environment(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY", environment="development")) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + deleted = registry.delete_prompts_by_base_id(base_prompt_id="greeting", environment="development") + + assert deleted == ["greeting.v1::development"] + assert registry.resolve_prompt_spec("greeting", environment="development") is None + assert _served_content(registry, environment="production") == "begin every reply with HOWDY" + + deleted_rest = registry.delete_prompts_by_base_id(base_prompt_id="greeting") + + assert deleted_rest == ["greeting.v1::production"] + assert registry.resolve_prompt_spec("greeting") is None + + +def test_has_config_prompt_matches_any_version_of_the_base_id(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + config_spec = PromptSpec( + prompt_id="greeting", + litellm_params=PromptLiteLLMParams( + prompt_id="greeting", + prompt_integration="dotprompt", + prompt_data={"content": "begin every reply with AHOY", "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="config"), + ) + registry.initialize_prompt(prompt=config_spec) + registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY", environment="production")) + + assert registry.has_config_prompt(base_prompt_id="greeting") is True + assert registry.has_config_prompt(base_prompt_id="other_prompt") is False diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 4266a13bf11..e2692a99732 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11379,10 +11379,15 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp } return row - def served_content() -> str: - callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1") + def served_callback(): + spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_sync") + assert spec is not None + callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=spec) assert callback is not None - return callback.prompt_manager.get_prompt("greeting_sync").content + return callback + + def served_content() -> str: + return served_callback().prompt_manager.get_prompt("greeting_sync").content prisma_client = MagicMock() try: @@ -11394,7 +11399,7 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) assert served_content() == "Begin every reply with HOWDY" - assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")] + assert litellm.callbacks == [served_callback()] finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_sync") @@ -11433,22 +11438,25 @@ async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkey ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("broken_sync.v1") is None - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1") is not None - assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1")] + assert IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("broken_sync") is None + healthy_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("healthy_sync") + assert healthy_spec is not None + healthy_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=healthy_spec) + assert healthy_callback is not None + assert litellm.callbacks == [healthy_callback] finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("healthy_sync") IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("broken_sync") @pytest.mark.asyncio -async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collide_on_a_versioned_id(monkeypatch): +async def test_init_prompts_in_db_syncs_every_environment_sharing_a_versioned_id(monkeypatch): from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.proxy.proxy_server import ProxyConfig monkeypatch.setattr(litellm, "callbacks", []) - def db_row(environment: str, content: str, updated_at: datetime) -> MagicMock: + def db_row(environment: str, content: str) -> MagicMock: row = MagicMock() row.model_dump.return_value = { "prompt_id": "greeting_env", @@ -11464,30 +11472,41 @@ async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collid ), "prompt_info": json.dumps({"prompt_type": "db"}), "created_at": None, - "updated_at": updated_at, + "updated_at": None, } return row - freshly_patched = db_row( - "production", "Begin every reply with HOWDY", datetime(2026, 8, 26, 12, 0, tzinfo=timezone.utc) - ) - stale_sibling = db_row( - "development", "Begin every reply with AHOY", datetime(2026, 8, 26, 11, 0, tzinfo=timezone.utc) - ) + def served_content(environment: str | None) -> str: + spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec("greeting_env", environment=environment) + assert spec is not None + callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_for_prompt(prompt=spec) + assert callback is not None + return callback.prompt_manager.get_prompt("greeting_env").content prisma_client = MagicMock() try: - prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[freshly_patched, stale_sibling]) + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[ + db_row("development", "Begin every reply with AHOY"), + db_row("production", "Begin every reply with HOWDY"), + ] + ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - first_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") - assert first_callback is not None - assert first_callback.prompt_manager.get_prompt("greeting_env").content == "Begin every reply with HOWDY" + assert served_content("development") == "Begin every reply with AHOY" + assert served_content("production") == "Begin every reply with HOWDY" + assert served_content(None) == "Begin every reply with HOWDY" + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[ + db_row("development", "Begin every reply with YO"), + db_row("production", "Begin every reply with HOWDY"), + ] + ) await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) - assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") is first_callback - assert litellm.callbacks == [first_callback] + assert served_content("development") == "Begin every reply with YO" + assert served_content("production") == "Begin every reply with HOWDY" finally: IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env") diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 7df39b0ef82..69310124ed3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -726,10 +726,7 @@ async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, from litellm.proxy.prompts import prompt_registry monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_by_id", lambda *a, **kw: None - ) - monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: None + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None ) data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} await proxy_logging._process_prompt_template( @@ -752,11 +749,11 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging, monkeypatch.setattr( prompt_registry.IN_MEMORY_PROMPT_REGISTRY, - "get_prompt_callback_by_id", + "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: prompt_spec + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec ) logging_obj = MagicMock() @@ -802,11 +799,11 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi prompt_spec.litellm_params = MagicMock(prompt_id="x") monkeypatch.setattr( prompt_registry.IN_MEMORY_PROMPT_REGISTRY, - "get_prompt_callback_by_id", + "get_prompt_callback_for_prompt", lambda *a, **kw: custom_logger, ) monkeypatch.setattr( - prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_by_id", lambda *a, **kw: prompt_spec + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec ) logging_obj = MagicMock() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt")) @@ -818,3 +815,44 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi prompt_version=None, call_type="completion", ) + + +@pytest.mark.asyncio +async def test_process_prompt_template_resolves_the_requested_environment(proxy_logging, monkeypatch): + from litellm.proxy.prompts import prompt_registry + + prompt_spec = MagicMock() + prompt_spec.litellm_params = MagicMock(prompt_id="greeting") + resolve_calls: list[dict] = [] + + def fake_resolve(prompt_id, version=None, environment=None): + resolve_calls.append({"prompt_id": prompt_id, "version": version, "environment": environment}) + return prompt_spec + + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", fake_resolve) + monkeypatch.setattr( + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_for_prompt", lambda *a, **kw: MagicMock() + ) + logging_obj = MagicMock() + logging_obj.async_get_chat_completion_prompt = AsyncMock( + return_value=("m", [{"role": "user", "content": "rendered"}], {}) + ) + data: Dict[str, Any] = { + "messages": [{"role": "user", "content": "orig"}], + "model": "m", + "prompt_id": "greeting", + "prompt_version": 1, + "prompt_environment": "development", + } + await proxy_logging._process_prompt_template( + data=data, + litellm_logging_obj=logging_obj, + prompt_id="greeting", + prompt_version=1, + call_type="completion", + ) + + assert resolve_calls == [{"prompt_id": "greeting", "version": 1, "environment": "development"}] + assert "prompt_environment" not in data + assert "prompt_id" not in data + assert data["messages"] == [{"role": "user", "content": "rendered"}] From 6f94554713ea69f108b356ab6691ee2bc72896ac Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 20:36:16 -0700 Subject: [PATCH 004/167] fix(ui): key prompt table rows by environment and dedupe key prompt options --- .../prompts/_components/PromptTable.tsx | 4 +++- .../components/organisms/create_key_button.tsx | 2 +- .../components/templates/key_edit_view.test.tsx | 16 ++++++++++++++++ .../src/components/templates/key_edit_view.tsx | 2 +- 4 files changed, 21 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx index 47d4f64f254..a64d7dbfe0c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx @@ -74,7 +74,9 @@ const PromptTable: React.FC = ({ prompt.prompt_id || String(index)} + getRowId={(prompt, index) => + prompt.prompt_id ? `${prompt.prompt_id}::${prompt.environment || "development"}` : String(index) + } sortingMode="client" sorting={sorting} onSortingChange={setSorting} diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index e1e6dcfa442..d0da4ddfb2b 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -349,7 +349,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const fetchPrompts = async () => { try { const response = await getPromptsList(accessToken); - setPromptsList(response.prompts.map((prompt) => prompt.prompt_id)); + setPromptsList(Array.from(new Set(response.prompts.map((prompt) => prompt.prompt_id)))); } catch (error) { console.error("Failed to fetch prompts:", error); } diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index bf8f43b8b5f..08edce6eda1 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -425,6 +425,22 @@ describe("KeyEditView", () => { expect(screen.getByText("Policies")).toBeInTheDocument(); }); + it("lists a prompt existing in several environments once in the dropdown", async () => { + vi.mocked(getPromptsList).mockResolvedValueOnce({ + prompts: [ + { prompt_id: "envgreet", litellm_params: {}, prompt_info: { prompt_type: "db" }, environment: "development" }, + { prompt_id: "envgreet", litellm_params: {}, prompt_info: { prompt_type: "db" }, environment: "production" }, + ], + }); + + renderAs("Admin"); + + const prompts = await screen.findByLabelText(/Prompts/); + await userEvent.type(prompts, "envgreet"); + + expect(await screen.findAllByRole("option", { name: "envgreet" })).toHaveLength(1); + }); + it("should omit both fields and fire neither admin-only request for an internal user", async () => { renderAs("Internal User"); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index df1af2ca8e9..d4a8a804f29 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -165,7 +165,7 @@ export function KeyEditView({ if (!accessToken) return; try { const response = await getPromptsList(accessToken); - setPromptsList(response.prompts.map((prompt) => prompt.prompt_id)); + setPromptsList(Array.from(new Set(response.prompts.map((prompt) => prompt.prompt_id)))); } catch (error) { console.error("Failed to fetch prompts:", error); } From dcba64ab72a259a14dcd1fd8b631973741338bda Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:16:16 -0700 Subject: [PATCH 005/167] fix(ui): scope prompt row delete and info view to the row's environment --- .../prompts/_components/PromptTable.test.tsx | 8 ++-- .../prompts/_components/PromptTable.tsx | 4 +- .../_components/PromptTableColumns.tsx | 20 ++++++-- .../prompts/_components/index.test.tsx | 48 ++++++++++++++----- .../(dashboard)/prompts/_components/index.tsx | 18 ++++--- .../prompts/_components/prompt_info.test.tsx | 29 +++++++++++ .../prompts/_components/prompt_info.tsx | 17 +++++-- .../src/components/networking.tsx | 7 ++- 8 files changed, 117 insertions(+), 34 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx index 52efd6407f8..edbb897fb2f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.test.tsx @@ -65,11 +65,13 @@ describe("PromptTable", () => { expect(within(rows[1]).getByText("prompt-older")).toBeInTheDocument(); }); - it("should call onPromptClick when the prompt ID is clicked", async () => { + it("should call onPromptClick with the row's environment, defaulting to development", async () => { const user = userEvent.setup(); render(); await user.click(screen.getByRole("button", { name: "prompt-newer" })); - expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-newer"); + expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-newer", "production"); + await user.click(screen.getByRole("button", { name: "prompt-older" })); + expect(mockOnPromptClick).toHaveBeenCalledWith("prompt-older", "development"); }); it("should label the environment and default missing environments to development", () => { @@ -83,7 +85,7 @@ describe("PromptTable", () => { render(); await user.click(screen.getByTestId("prompt-actions-prompt-newer")); await user.click(await screen.findByTestId("prompt-action-delete")); - expect(mockOnDeleteClick).toHaveBeenCalledWith("prompt-newer", "prompt-newer"); + expect(mockOnDeleteClick).toHaveBeenCalledWith("prompt-newer", "prompt-newer", "production"); }); it("should copy the prompt ID through the actions menu", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx index a64d7dbfe0c..c766042ac44 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx @@ -13,8 +13,8 @@ import { ModelGroupInfo } from "./prompt_utils"; interface PromptTableProps { promptsList: PromptSpec[]; isLoading: boolean; - onPromptClick?: (id: string) => void; - onDeleteClick?: (id: string, name: string) => void; + onPromptClick?: (id: string, environment: string) => void; + onDeleteClick?: (id: string, name: string, environment: string) => void; accessToken: string | null; isAdmin: boolean; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx index ae584ef6df6..f927a6d1486 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTableColumns.tsx @@ -64,7 +64,7 @@ function PromptModelCell({ prompt, modelHubData }: { prompt: PromptSpec; modelHu interface PromptRowActionsProps { prompt: PromptSpec; isAdmin: boolean; - onDeleteClick?: (id: string, name: string) => void; + onDeleteClick?: (id: string, name: string, environment: string) => void; } function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsProps) { @@ -91,7 +91,13 @@ function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsPr onDeleteClick?.(prompt.prompt_id, prompt.prompt_id || "Unknown Prompt")} + onClick={() => + onDeleteClick?.( + prompt.prompt_id, + prompt.prompt_id || "Unknown Prompt", + prompt.environment || "development", + ) + } > Delete @@ -106,8 +112,8 @@ function PromptRowActions({ prompt, isAdmin, onDeleteClick }: PromptRowActionsPr interface PromptTableColumnsDeps { modelHubData: Map; isAdmin: boolean; - onPromptClick?: (id: string) => void; - onDeleteClick?: (id: string, name: string) => void; + onPromptClick?: (id: string, environment: string) => void; + onDeleteClick?: (id: string, name: string, environment: string) => void; } export const getPromptTableColumns = ({ @@ -128,7 +134,11 @@ export const getPromptTableColumns = ({ title={row.original.prompt_id} titleClassName="font-mono text-xs font-normal" className="max-w-60" - onClick={onPromptClick ? () => onPromptClick(row.original.prompt_id) : undefined} + onClick={ + onPromptClick + ? () => onPromptClick(row.original.prompt_id, row.original.environment || "development") + : undefined + } /> ), }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx index d6a4aaea2e1..334339e6f1d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.test.tsx @@ -15,21 +15,31 @@ vi.mock("./PromptTable", () => ({ __esModule: true, default: ({ isLoading, + onPromptClick, onDeleteClick, }: { isLoading: boolean; - onDeleteClick: (id: string, name: string) => void; + onPromptClick: (id: string, environment: string) => void; + onDeleteClick: (id: string, name: string, environment: string) => void; }) => (
{isLoading ? "table-loading" : "table-loaded"} - +
), })); -vi.mock("./prompt_info", () => ({ __esModule: true, default: () =>
prompt-info-view
})); +vi.mock("./prompt_info", () => ({ + __esModule: true, + default: ({ initialEnvironment }: { initialEnvironment?: string }) => ( +
prompt-info-view:{initialEnvironment ?? "none"}
+ ), +})); vi.mock("./add_prompt_form", () => ({ __esModule: true, default: ({ visible }: { visible: boolean }) => (visible ?
add-prompt-form
: null), @@ -143,6 +153,22 @@ describe("PromptsPanel toolbar", () => { }); }); +describe("PromptsPanel row navigation", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockGetPromptsList.mockResolvedValue({ prompts: [] } as never); + }); + + it("should open the info view preselected to the clicked row's environment", async () => { + const user = userEvent.setup(); + renderPanel("Admin"); + + await user.click(await screen.findByRole("button", { name: "row-open" })); + + expect(screen.getByText("prompt-info-view:staging")).toBeInTheDocument(); + }); +}); + describe("PromptsPanel delete confirmation", () => { beforeEach(() => { vi.clearAllMocks(); @@ -156,13 +182,13 @@ describe("PromptsPanel delete confirmation", () => { await user.click(await screen.findByRole("button", { name: "row-delete" })); - expect(await screen.findByText(/delete prompt: my-prompt/i)).toBeInTheDocument(); + expect(await screen.findByText(/the staging copy of prompt: my-prompt/i)).toBeInTheDocument(); expect(screen.getByText(/cannot be undone/i)).toBeInTheDocument(); expect(mockDeletePromptCall).not.toHaveBeenCalled(); await user.click(screen.getByRole("button", { name: /^delete$/i })); - await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1")); + await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1", "staging")); }); it("should abandon the delete when the confirmation is dismissed", async () => { @@ -170,11 +196,11 @@ describe("PromptsPanel delete confirmation", () => { renderPanel("Admin"); await user.click(await screen.findByRole("button", { name: "row-delete" })); - await screen.findByText(/delete prompt: my-prompt/i); + await screen.findByText(/the staging copy of prompt: my-prompt/i); await user.click(screen.getByRole("button", { name: /cancel/i })); - await waitFor(() => expect(screen.queryByText(/delete prompt: my-prompt/i)).not.toBeInTheDocument()); + await waitFor(() => expect(screen.queryByText(/the staging copy of prompt: my-prompt/i)).not.toBeInTheDocument()); expect(mockDeletePromptCall).not.toHaveBeenCalled(); }); @@ -189,14 +215,14 @@ describe("PromptsPanel delete confirmation", () => { renderPanel("Admin"); await user.click(await screen.findByRole("button", { name: "row-delete" })); - await screen.findByText(/delete prompt: my-prompt/i); + await screen.findByText(/the staging copy of prompt: my-prompt/i); await user.click(screen.getByRole("button", { name: /^delete$/i })); - await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1")); + await waitFor(() => expect(mockDeletePromptCall).toHaveBeenCalledWith("sk-test", "prompt-1", "staging")); await user.keyboard("{Escape}"); - expect(screen.getByText(/delete prompt: my-prompt/i)).toBeInTheDocument(); + expect(screen.getByText(/the staging copy of prompt: my-prompt/i)).toBeInTheDocument(); finishDelete(); - await waitFor(() => expect(screen.queryByText(/delete prompt: my-prompt/i)).not.toBeInTheDocument()); + await waitFor(() => expect(screen.queryByText(/the staging copy of prompt: my-prompt/i)).not.toBeInTheDocument()); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx index 3594e933704..2d8d905c480 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/index.tsx @@ -41,11 +41,12 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { const [isLoading, setIsLoading] = useState(true); const [selectedEnvironment, setSelectedEnvironment] = useState(undefined); const [selectedPromptId, setSelectedPromptId] = useState(null); + const [selectedPromptEnvironment, setSelectedPromptEnvironment] = useState(undefined); const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [showEditorView, setShowEditorView] = useState(false); const [editPromptData, setEditPromptData] = useState(null); const [isDeleting, setIsDeleting] = useState(false); - const [promptToDelete, setPromptToDelete] = useState<{ id: string; name: string } | null>(null); + const [promptToDelete, setPromptToDelete] = useState<{ id: string; name: string; environment: string } | null>(null); // Admin Viewer follows the read-parity rule: see prompts, no writes. const canModify = userRole ? isProxyAdminRole(userRole) : false; @@ -71,8 +72,9 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { fetchPrompts(); }, [accessToken, selectedEnvironment]); - const handlePromptClick = (promptId: string) => { + const handlePromptClick = (promptId: string, environment: string) => { setSelectedPromptId(promptId); + setSelectedPromptEnvironment(environment); }; const handleAddPrompt = () => { @@ -111,8 +113,8 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { setSelectedPromptId(null); }; - const handleDeleteClick = (promptId: string, promptName: string) => { - setPromptToDelete({ id: promptId, name: promptName }); + const handleDeleteClick = (promptId: string, promptName: string, environment: string) => { + setPromptToDelete({ id: promptId, name: promptName, environment }); }; const handleDeleteConfirm = async () => { @@ -120,8 +122,8 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { setIsDeleting(true); try { - await deletePromptCall(accessToken, promptToDelete.id); - toast.success(`Prompt "${promptToDelete.name}" deleted successfully`); + await deletePromptCall(accessToken, promptToDelete.id, promptToDelete.environment); + toast.success(`Prompt "${promptToDelete.name}" deleted successfully from ${promptToDelete.environment}`); fetchPrompts(); // Refresh the list } catch (error) { console.error("Error deleting prompt:", error); @@ -148,6 +150,7 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { ) : selectedPromptId ? ( setSelectedPromptId(null)} accessToken={accessToken} isAdmin={canModify} @@ -219,7 +222,8 @@ const PromptsPanel: React.FC = ({ accessToken, userRole }) => { Delete Prompt - Are you sure you want to delete prompt: {promptToDelete.name} ? This action cannot be undone. + Are you sure you want to delete the {promptToDelete.environment} copy of prompt: {promptToDelete.name}? + This action cannot be undone. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx index bb29f12ac42..a5b195e5057 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx @@ -29,6 +29,35 @@ const promptWithoutTemplate = { environments: [], }; +describe("PromptInfoView environment scoping", () => { + beforeEach(() => { + vi.mocked(networking.getPromptInfo).mockReset().mockResolvedValue(promptWithoutTemplate); + vi.mocked(networking.getPromptVersions).mockReset().mockResolvedValue({ prompts: [] }); + }); + + it("fetches the initial environment it was opened with", async () => { + render( + , + ); + + await screen.findByRole("tab", { name: "Raw JSON" }); + expect(networking.getPromptInfo).toHaveBeenCalledWith("sk-test", "support-reply", "staging"); + }); + + it("fetches the serve default when opened without an environment", async () => { + render(); + + await screen.findByRole("tab", { name: "Raw JSON" }); + expect(networking.getPromptInfo).toHaveBeenCalledWith("sk-test", "support-reply", undefined); + }); +}); + describe("PromptInfoView tabs", () => { beforeEach(() => { vi.mocked(networking.getPromptInfo).mockReset().mockResolvedValue(promptWithoutTemplate); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx index c6a6f09fcaa..af4e3bf2121 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx @@ -20,6 +20,7 @@ import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from " export interface PromptInfoProps { promptId: string; + initialEnvironment?: string; onClose: () => void; accessToken: string | null; isAdmin: boolean; @@ -27,7 +28,15 @@ export interface PromptInfoProps { onEdit?: (promptData: any) => void; } -const PromptInfoView: React.FC = ({ promptId, onClose, accessToken, isAdmin, onDelete, onEdit }) => { +const PromptInfoView: React.FC = ({ + promptId, + initialEnvironment, + onClose, + accessToken, + isAdmin, + onDelete, + onEdit, +}) => { const [promptData, setPromptData] = useState(null); const [promptTemplate, setPromptTemplate] = useState(null); const [rawApiResponse, setRawApiResponse] = useState(null); @@ -43,7 +52,7 @@ const PromptInfoView: React.FC = ({ promptId, onClose, accessTo const [selectedVersion, setSelectedVersion] = useState(null); const [loadingVersions, setLoadingVersions] = useState(false); - // Initial fetch — no environment filter, gets default + all environments list + // Fetches the requested environment (or the serve-time default when omitted) plus the environments list const fetchPromptInfo = async (environment?: string) => { try { setLoading(true); @@ -89,7 +98,7 @@ const PromptInfoView: React.FC = ({ promptId, onClose, accessTo setSelectedEnv(null); setEnvironments([]); setVersionHistory([]); - fetchPromptInfo(); + fetchPromptInfo(initialEnvironment); }, [promptId, accessToken]); // When environment changes (user clicks tab), re-fetch — skip initial mount @@ -493,7 +502,7 @@ const PromptInfoView: React.FC = ({ promptId, onClose, accessTo Delete Prompt

- Are you sure you want to delete prompt: {basePromptId}? + Are you sure you want to delete prompt: {basePromptId} from every environment?

This action cannot be undone.

diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 5b6d70b4771..e297bcc5270 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4625,9 +4625,12 @@ export const updatePromptCall = async (accessToken: string, promptId: string, pr } }; -export const deletePromptCall = async (accessToken: string, promptId: string) => { +export const deletePromptCall = async (accessToken: string, promptId: string, environment?: string) => { try { - const data = await apiClient.delete(`/prompts/${promptId}`, { accessToken }); + const data = await apiClient.delete(`/prompts/${promptId}`, { + accessToken, + query: { environment: environment || undefined }, + }); return data; } catch (error) { console.error("Failed to delete prompt:", error); From adfa42096d99e320e1413c37578158e6ec2d6465 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 23:07:18 -0700 Subject: [PATCH 006/167] fix(prompts): resolve config prompts in /prompts/{id}/info when environment is set --- litellm/proxy/prompts/prompt_endpoints.py | 6 +- .../proxy/prompts/test_prompt_endpoints.py | 92 +++++++++++++++++++ 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 1b7932cdb18..b16611ac157 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -554,8 +554,10 @@ async def get_prompt_info( if env_prompts: prompt_spec = create_versioned_prompt_spec(db_prompt=env_prompts[0]) - if prompt_spec is None and environment is None: - prompt_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec(prompt_id, version=requested_version) + if prompt_spec is None: + prompt_spec = IN_MEMORY_PROMPT_REGISTRY.resolve_prompt_spec( + prompt_id, version=requested_version, environment=environment + ) if prompt_spec is None: raise HTTPException( diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py index 35402db219f..41c3003af34 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py @@ -368,3 +368,95 @@ class TestAdminViewerReadAccess: assert response.prompt_spec.prompt_id == "jack" assert response.prompt_spec.version == 2 + + +class TestConfigPromptInfoWithEnvironment: + """ + Regression: /prompts/{id}/info with an environment param must still resolve + config-file (in-memory) prompts on a DB-backed proxy instead of 400ing. + """ + + def _registry_with_config_prompt(self): + from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry + + registry = InMemoryPromptRegistry() + registry.IN_MEMORY_PROMPTS["envgreet::development"] = PromptSpec( + prompt_id="envgreet", + litellm_params=PromptLiteLLMParams( + prompt_id="envgreet", + prompt_integration="dotprompt", + dotprompt_content="AHOY {{user_message}}", + ), + prompt_info=PromptInfo(prompt_type="config"), + ) + return registry + + def _prisma_client_with_empty_prompt_table(self): + from unittest.mock import AsyncMock + + mock_prisma = MagicMock() + mock_prisma.db.litellm_prompttable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + @pytest.mark.asyncio + async def test_get_prompt_info_with_environment_falls_back_to_registry(self): + from unittest.mock import patch + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.prompts.prompt_endpoints import get_prompt_info + + admin = UserAPIKeyAuth( + api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.proxy_server.prisma_client", + self._prisma_client_with_empty_prompt_table(), + ), + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY", + self._registry_with_config_prompt(), + ), + ): + response = await get_prompt_info( + prompt_id="envgreet", + environment="development", + user_api_key_dict=admin, + ) + + assert response.prompt_spec.prompt_id == "envgreet" + assert response.prompt_spec.litellm_params.dotprompt_content == "AHOY {{user_message}}" + + @pytest.mark.asyncio + async def test_get_prompt_info_with_wrong_environment_still_400s(self): + from unittest.mock import patch + + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.prompts.prompt_endpoints import get_prompt_info + + admin = UserAPIKeyAuth( + api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + with ( + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.proxy_server.prisma_client", + self._prisma_client_with_empty_prompt_table(), + ), + patch( # test-quality-ok: endpoint reads prisma_client and the registry from module globals at call time; no injection seam + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY", + self._registry_with_config_prompt(), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await get_prompt_info( + prompt_id="envgreet", + environment="production", + user_api_key_dict=admin, + ) + + assert exc_info.value.status_code == 400 + assert "environment production" in exc_info.value.detail From 3088db3f0ec77fd8256eff2b70dee882755f9a4c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:55:34 -0700 Subject: [PATCH 007/167] fix(prompts): accept a string prompt_version and carry the viewed environment into code snippets --- litellm/proxy/prompts/prompt_registry.py | 10 ++++++ litellm/proxy/utils.py | 5 +-- .../proxy/prompts/test_prompt_registry.py | 10 +++++- .../proxy_logging/test_guardrail_pipeline.py | 35 ++++++++++++++++++ .../PromptCodeSnippets.test.tsx | 24 +++++++++++++ .../prompt_editor_view/PromptCodeSnippets.tsx | 29 ++++++++------- .../PromptEditorHeader.test.tsx | 5 ++- .../prompt_editor_view/PromptEditorHeader.tsx | 1 + .../prompts/_components/prompt_info.test.tsx | 36 ++++++++++++++++++- .../prompts/_components/prompt_info.tsx | 1 + 10 files changed, 139 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index ba53809690d..7803352e9e7 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -74,6 +74,16 @@ def registry_key_for_prompt(prompt: PromptSpec) -> str: return f"{prompt.prompt_id}::{prompt_environment_or_default(prompt.environment)}" +def parse_prompt_version(raw_version: object) -> int | None: + if isinstance(raw_version, bool): + return None + if isinstance(raw_version, int): + return raw_version + if isinstance(raw_version, str) and raw_version.isdigit(): + return int(raw_version) + return None + + def _spec_version(prompt: PromptSpec) -> int: return prompt.version if prompt.version is not None else get_version_number(prompt_id=prompt.prompt_id) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 09efa0cc72d..0263ce7598a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1749,7 +1749,6 @@ class ProxyLogging: litellm_logging_obj: Final = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)) prompt_id: Final[str | None] = data.get("prompt_id", None) - prompt_version: Final[int | None] = data.get("prompt_version", None) ## PROMPT TEMPLATE CHECK ## @@ -1759,11 +1758,13 @@ class ProxyLogging: and prompt_id is not None and (call_type == "completion" or call_type == "acompletion" or call_type == "aresponses") ): + from litellm.proxy.prompts.prompt_registry import parse_prompt_version + await self._process_prompt_template( data=data, litellm_logging_obj=litellm_logging_obj, prompt_id=prompt_id, - prompt_version=prompt_version, + prompt_version=parse_prompt_version(data.get("prompt_version", None)), call_type=call_type, ) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py index 32930fbbdff..a0743b65dd7 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_registry.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -2,7 +2,7 @@ import pytest import litellm from litellm.integrations.custom_prompt_management import CustomPromptManagement -from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry +from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry, parse_prompt_version from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec @@ -214,3 +214,11 @@ def test_remove_prompt_is_a_no_op_for_an_unknown_registry_key(isolated_callbacks assert registry.resolve_prompt_spec("greeting") is not None assert len(isolated_callbacks) == 1 + + +@pytest.mark.parametrize( + ("raw_version", "expected"), + [(2, 2), ("2", 2), (None, None), ("v2", None), (True, None), (2.0, None)], +) +def test_parse_prompt_version_accepts_integers_and_json_strings(raw_version: object, expected: int | None) -> None: + assert parse_prompt_version(raw_version) == expected diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 6589b14e53a..2ba58bd5644 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -858,6 +858,41 @@ async def test_process_prompt_template_resolves_the_requested_environment(proxy_ assert data["messages"] == [{"role": "user", "content": "rendered"}] +@pytest.mark.asyncio +async def test_pre_call_hook_matches_a_prompt_version_sent_as_a_json_string(proxy_logging, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.prompts import prompt_registry + + prompt_spec = MagicMock() + prompt_spec.litellm_params = MagicMock(prompt_id="greeting") + resolve_calls: list[dict] = [] + + def fake_resolve(prompt_id, version=None, environment=None): + resolve_calls.append({"prompt_id": prompt_id, "version": version, "environment": environment}) + return prompt_spec + + monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", fake_resolve) + monkeypatch.setattr( + prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_for_prompt", lambda *a, **kw: MagicMock() + ) + logging_obj = MagicMock() + logging_obj.async_get_chat_completion_prompt = AsyncMock( + return_value=("m", [{"role": "user", "content": "rendered"}], {}) + ) + data: Dict[str, Any] = { + "messages": [{"role": "user", "content": "orig"}], + "model": "m", + "prompt_id": "greeting", + "prompt_version": "2", + "litellm_logging_obj": logging_obj, + } + + result = await proxy_logging.pre_call_hook(user_api_key_dict=UserAPIKeyAuth(), data=data, call_type="completion") + + assert resolve_calls == [{"prompt_id": "greeting", "version": 2, "environment": None}] + assert result["messages"] == [{"role": "user", "content": "rendered"}] + + @pytest.mark.asyncio async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(proxy_logging, monkeypatch): from litellm.proxy.prompts import prompt_registry diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx index a1b4ad52634..7fa44a4dfe5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx @@ -44,4 +44,28 @@ describe("PromptCodeSnippets", () => { expect(screen.getByRole("combobox", { name: "Language" })).toHaveTextContent("Python (OpenAI SDK)"); }); + + it("includes the viewed environment in every generated request", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + render( + , + ); + await user.click(screen.getByRole("button", { name: /get code/i })); + await screen.findByText("Generated Code"); + + await user.click(screen.getByRole("button", { name: /copy to clipboard/i })); + expect(await navigator.clipboard.readText()).toContain('"prompt_environment": "development"'); + + await user.click(screen.getByRole("tab", { name: "With Version" })); + await user.click(screen.getByRole("button", { name: /copy to clipboard/i })); + const versionSnippet = await navigator.clipboard.readText(); + expect(versionSnippet).toContain('"prompt_environment": "development"'); + expect(versionSnippet).toContain('"prompt_version": 2'); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx index af7d6421265..a6adc160674 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx @@ -22,6 +22,7 @@ interface PromptCodeSnippetsProps { promptVariables?: Record; accessToken: string | null; version?: string; + environment?: string; proxySettings?: { PROXY_BASE_URL?: string; LITELLM_UI_API_DOC_BASE_URL?: string | null; @@ -34,6 +35,7 @@ const PromptCodeSnippets: React.FC = ({ promptVariables = {}, accessToken, version = "1", + environment, proxySettings, }) => { const syntaxTheme = useSyntaxTheme(coy); @@ -64,6 +66,9 @@ const PromptCodeSnippets: React.FC = ({ // Generate code based on selected language and tab const generateCode = () => { const hasVariables = Object.keys(promptVariables).length > 0; + const curlEnvironment = environment ? `,\n "prompt_environment": "${environment}"` : ""; + const pythonEnvironment = environment ? `,\n "prompt_environment": "${environment}"` : ""; + const jsEnvironment = environment ? `,\n prompt_environment: "${environment}"` : ""; if (selectedLanguage === "curl") { if (selectedTab === "basic") { @@ -72,7 +77,7 @@ const PromptCodeSnippets: React.FC = ({ -H 'Authorization: Bearer ${effectiveApiKey}' \\ -d '{ "model": "${model}", - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${curlEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 6).replace(/\n/g, "\n ")}` @@ -85,7 +90,7 @@ const PromptCodeSnippets: React.FC = ({ -H 'Authorization: Bearer ${effectiveApiKey}' \\ -d '{ "model": "${model}", - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${curlEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 6).replace(/\n/g, "\n ")}` @@ -104,7 +109,7 @@ const PromptCodeSnippets: React.FC = ({ -H 'Authorization: Bearer ${effectiveApiKey}' \\ -d '{ "model": "${model}", - "prompt_id": "${promptId}", + "prompt_id": "${promptId}"${curlEnvironment}, "prompt_version": ${version}, "messages": [ { @@ -127,7 +132,7 @@ client = openai.OpenAI( response = client.chat.completions.create( model="${model}", extra_body={ - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${pythonEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` @@ -145,7 +150,7 @@ response = client.chat.completions.create( {"role": "user", "content": "hi"} ], extra_body={ - "prompt_id": "${promptId}"${ + "prompt_id": "${promptId}"${pythonEnvironment}${ hasVariables ? `, "prompt_variables": ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` @@ -163,7 +168,7 @@ response = client.chat.completions.create( {"role": "user", "content": "Who are u"} ], extra_body={ - "prompt_id": "${promptId}", + "prompt_id": "${promptId}"${pythonEnvironment}, "prompt_version": ${version} } ) @@ -186,9 +191,9 @@ async function main() { model: "${model}", ${ hasVariables - ? `prompt_id: "${promptId}", + ? `prompt_id: "${promptId}"${jsEnvironment}, prompt_variables: ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` - : `prompt_id: "${promptId}"` + : `prompt_id: "${promptId}"${jsEnvironment}` } }); @@ -206,9 +211,9 @@ async function main() { ], ${ hasVariables - ? `prompt_id: "${promptId}", + ? `prompt_id: "${promptId}"${jsEnvironment}, prompt_variables: ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}` - : `prompt_id: "${promptId}"` + : `prompt_id: "${promptId}"${jsEnvironment}` } }); @@ -224,7 +229,7 @@ async function main() { messages: [ { role: "user", content: "Who are u" } ], - prompt_id: "${promptId}", + prompt_id: "${promptId}"${jsEnvironment}, prompt_version: ${version} }); @@ -241,7 +246,7 @@ main();`; if (isModalVisible) { setGeneratedCode(generateCode()); } - }, [isModalVisible, selectedLanguage, selectedTab, promptId, model, promptVariables]); + }, [isModalVisible, selectedLanguage, selectedTab, promptId, model, promptVariables, version, environment]); return ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx index 5194cdd4e63..3afc00e37a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx @@ -2,7 +2,9 @@ import { fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; import PromptEditorHeader from "./PromptEditorHeader"; -vi.mock("./PromptCodeSnippets", () => ({ default: () => })); +vi.mock("./PromptCodeSnippets", () => ({ + default: ({ environment }: { environment?: string }) => , +})); describe("PromptEditorHeader", () => { it("preserves navigation, naming, and save actions", () => { @@ -48,5 +50,6 @@ describe("PromptEditorHeader", () => { ); expect(screen.getByRole("combobox", { name: "Environment" })).toHaveTextContent(label); + expect(screen.getByRole("button", { name: "Get Code" })).toHaveAttribute("data-environment", environment); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx index eea9755054f..04cac01365a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx @@ -89,6 +89,7 @@ const PromptEditorHeader: React.FC = ({ promptVariables={promptVariables} accessToken={accessToken} version={version?.replace("v", "") || "1"} + environment={environment} proxySettings={proxySettings} /> {editMode && onShowHistory && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx index a5b195e5057..b14d5c3d91f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx @@ -12,7 +12,9 @@ vi.mock("@/components/networking", () => ({ })); vi.mock("./prompt_editor_view/PromptCodeSnippets", () => ({ - default: () =>
, + default: ({ environment }: { environment?: string }) => ( +
+ ), })); const promptWithoutTemplate = { @@ -58,6 +60,38 @@ describe("PromptInfoView environment scoping", () => { }); }); +describe("PromptInfoView code snippets", () => { + beforeEach(() => { + vi.mocked(networking.getPromptVersions).mockReset().mockResolvedValue({ prompts: [] }); + }); + + it.each([ + ["a prompt with several environments", "staging", ["development", "staging"]], + ["a config prompt with no environment list", "development", []], + ])("hands the viewed environment of %s to the code snippets", async (_label, environment, environments) => { + vi.mocked(networking.getPromptInfo) + .mockReset() + .mockResolvedValue({ + ...promptWithoutTemplate, + prompt_spec: { ...promptWithoutTemplate.prompt_spec, environment }, + environments, + }); + + render( + , + ); + + await screen.findByRole("tab", { name: "Raw JSON" }); + expect(screen.getByTestId("prompt-code-snippets")).toHaveAttribute("data-environment", environment); + }); +}); + describe("PromptInfoView tabs", () => { beforeEach(() => { vi.mocked(networking.getPromptInfo).mockReset().mockResolvedValue(promptWithoutTemplate); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx index af4e3bf2121..062e1d84a0a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx @@ -221,6 +221,7 @@ const PromptInfoView: React.FC = ({ promptVariables={extractTemplateVariables(promptTemplate?.content)} accessToken={accessToken} version={currentVersion} + environment={selectedEnv ?? promptData.environment} /> diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index c368a43dd1e..93066c106ee 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -139,19 +139,20 @@ describe("RequestLogsPanel", () => { }); describe("server-grouped session pagination (#38060)", () => { - it("requests session-grouped pages of 10 rows by default", async () => { + it("requests session-grouped pages of 10 rows by default without a cursor", async () => { renderPanel(); await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); expect(lastCall()?.params?.group_by_session).toBe(true); + expect(lastCall()?.params?.session_cursor).toBeUndefined(); expect(lastCall()?.page_size).toBe(10); }); it("renders every row the server returns without client-side collapsing", async () => { respondWith([ - logEntry({ request_id: "req-a", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-b", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-c", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-a", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-b", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-c", session_id: "sess-1", session_total_count: 3 }), ]); renderPanel(); @@ -177,6 +178,138 @@ describe("RequestLogsPanel", () => { expect(within(row("req-llm") as HTMLElement).getByText("3")).toBeInTheDocument(); }); + it("passes the server keyset cursor when navigating to the next page", async () => { + const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: firstPage, + total: 80, + page: 1, + page_size: 50, + total_pages: 2, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + fireEvent.click(screen.getByTestId("pagination-next")); + + await waitFor(() => { + const call = lastCall(); + expect(call?.params?.session_cursor).toBe("2026-07-07 09:50:13|key-1|sess-1"); + expect(call?.page).toBe(2); + }); + }); + + it("drops the cursor and returns to the first page when a filter changes", async () => { + const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: firstPage, + total: 80, + page: 1, + page_size: 50, + total_pages: 2, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + fireEvent.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastCall()?.page).toBe(2)); + + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "req-elsewhere" } }); + + await waitFor(() => { + const call = lastCall(); + expect(call?.page).toBe(1); + expect(call?.params?.session_cursor).toBeUndefined(); + }); + }); + + it("ignores another next click while the next page is still fetching", async () => { + const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); + const firstResponse = { + data: firstPage, + total: 150, + page: 1, + page_size: 50, + total_pages: 3, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }; + vi.mocked(uiSpendLogsCall) + .mockResolvedValueOnce(firstResponse) + .mockImplementation(() => new Promise(() => {})); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + fireEvent.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastCall()?.page).toBe(2)); + + fireEvent.click(screen.getByTestId("pagination-next")); + + expect(lastCall()?.page).toBe(2); + expect(vi.mocked(uiSpendLogsCall).mock.calls.filter(([options]) => options.page === 3)).toHaveLength(0); + }); + + it("still moves to the next page while a live-tail refetch of the current page is in flight", async () => { + const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); + const firstResponse = { + data: firstPage, + total: 150, + page: 1, + page_size: 50, + total_pages: 3, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }; + vi.mocked(uiSpendLogsCall) + .mockResolvedValueOnce(firstResponse) + .mockImplementation(() => new Promise(() => {})); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + void testQueryClient.refetchQueries({ queryKey: ["logs", "table"] }); + await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalledTimes(2)); + + fireEvent.click(screen.getByTestId("pagination-next")); + + await waitFor(() => { + const call = lastCall(); + expect(call?.page).toBe(2); + expect(call?.params?.session_cursor).toBe("2026-07-07 09:50:13|key-1|sess-1"); + }); + }); + + it("drops the cursor and returns to the first page when Custom Range is toggled", async () => { + const user = userEvent.setup(); + const firstPage = Array.from({ length: 50 }, (_, index) => logEntry({ request_id: `req-${index}` })); + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: firstPage, + total: 80, + page: 1, + page_size: 50, + total_pages: 2, + next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1", + has_more: true, + }); + renderPanel(); + + await waitFor(() => expect(row("req-0")).not.toBeNull()); + fireEvent.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastCall()?.page).toBe(2)); + + await user.click(screen.getByRole("button", { name: /Last 24 Hours/i })); + await user.click(await screen.findByRole("button", { name: "Custom Range" })); + + await waitFor(() => { + const call = lastCall(); + expect(call?.page).toBe(1); + expect(call?.params?.session_cursor).toBeUndefined(); + }); + }); + it("leaves single-call rows untouched", async () => { respondWith([ logEntry({ request_id: "req-solo-a", session_id: "sess-a", session_total_count: 1 }), @@ -395,9 +528,7 @@ describe("RequestLogsPanel", () => { }); it("opens a deep-linked multi-call session log in session mode", async () => { - respondWith([ - logEntry({ request_id: "req-llm", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - ]); + respondWith([logEntry({ request_id: "req-llm", session_id: "sess-1", session_total_count: 3 })]); renderPanel("?log_id=req-llm"); await waitFor(() => { @@ -410,8 +541,8 @@ describe("RequestLogsPanel", () => { it("clicking a multi-call session's row writes ?session_id= alongside ?log_id=", async () => { const user = userEvent.setup(); respondWith([ - logEntry({ request_id: "req-llm", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), - logEntry({ request_id: "req-llm-2", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-llm", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-llm-2", session_id: "sess-1", session_total_count: 3 }), ]); renderPanel(); @@ -426,7 +557,7 @@ describe("RequestLogsPanel", () => { it("selecting another log while a session view is open keeps the session open", async () => { const user = userEvent.setup(); respondWith([ - logEntry({ request_id: "req-llm", call_type: "acompletion", session_id: "sess-1", session_total_count: 3 }), + logEntry({ request_id: "req-llm", session_id: "sess-1", session_total_count: 3 }), logEntry({ request_id: "req-unenriched" }), ]); renderPanel("?log_id=req-llm"); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 2ebabf64ab0..96cf2bd5dde 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -39,6 +39,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); const [sorting, setSorting] = useState(DEFAULT_LOGS_SORTING); const [columnFilters, setColumnFilters] = useState([]); + const [sessionCursors, setSessionCursors] = useState>({}); const [startTime, setStartTime] = useState(moment().subtract(24, "hours").format("YYYY-MM-DDTHH:mm")); const [endTime, setEndTime] = useState(moment().format("YYYY-MM-DDTHH:mm")); @@ -74,7 +75,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, sessionStorage.setItem("excludeInternalHealthChecks", JSON.stringify(excludeInternalHealthChecks)); }, [excludeInternalHealthChecks]); - const { logsQuery, filteredLogs, allTeams } = useLogFilterLogic({ + const { logsQuery, filteredLogs, allTeams, usesSessionCursor } = useLogFilterLogic({ accessToken, token, userRole, @@ -88,6 +89,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, pagination, isCustomDate, sorting, + sessionCursors, }); // Follow the table's own last fetch so a live-tail refresh carries the filter @@ -163,23 +165,52 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.REQUEST_ID); return value === "" ? others : [...others, { id: LOG_FILTER_IDS.REQUEST_ID, value }]; }); + setSessionCursors({}); setPagination((previous) => ({ ...previous, pageIndex: 0 })); }, []); const handleSortingChange = useCallback>((updaterOrValue) => { setSorting(updaterOrValue); + setSessionCursors({}); setPagination((previous) => ({ ...previous, pageIndex: 0 })); }, []); const handleColumnFiltersChange = useCallback>((updaterOrValue) => { setColumnFilters(updaterOrValue); + setSessionCursors({}); setPagination((previous) => ({ ...previous, pageIndex: 0 })); }, []); const resetToFirstPage = useCallback(() => { + setSessionCursors({}); setPagination((previous) => ({ ...previous, pageIndex: 0 })); }, []); + const handlePaginationChange = useCallback>( + (updaterOrValue) => { + const requested = typeof updaterOrValue === "function" ? updaterOrValue(pagination) : updaterOrValue; + if (!usesSessionCursor) { + setPagination(requested); + return; + } + if (requested.pageSize !== pagination.pageSize) { + setSessionCursors({}); + setPagination({ ...requested, pageIndex: 0 }); + return; + } + if (requested.pageIndex <= pagination.pageIndex) { + setPagination(requested); + return; + } + const nextCursor = filteredLogs.next_session_cursor; + if (!nextCursor || logsQuery.isPlaceholderData) return; + const nextPageIndex = pagination.pageIndex + 1; + setSessionCursors((previous) => ({ ...previous, [nextPageIndex]: nextCursor })); + setPagination({ ...requested, pageIndex: nextPageIndex }); + }, + [usesSessionCursor, pagination, filteredLogs.next_session_cursor, logsQuery.isPlaceholderData], + ); + const handleExcludeInternalHealthChecksChange = useCallback( (value: boolean) => { setExcludeInternalHealthChecks(value); @@ -256,7 +287,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, isLoading={logsQuery.isLoading} isRefreshing={logsQuery.isFetching} pagination={pagination} - onPaginationChange={setPagination} + onPaginationChange={handlePaginationChange} sorting={sorting} onSortingChange={handleSortingChange} columnFilters={columnFilters} diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index acd5be06593..90f0f0a60f1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -15,6 +15,8 @@ export interface PaginatedResponse { page_size: number; total_pages: number; total_is_capped?: boolean; + next_session_cursor?: string | null; + has_more?: boolean; } export const LOG_FILTER_IDS = { @@ -109,6 +111,7 @@ export function useLogFilterLogic({ pagination, isCustomDate, sorting, + sessionCursors = {}, }: { accessToken: string | null; token: string | null; @@ -123,11 +126,14 @@ export function useLogFilterLogic({ pagination: PaginationState; isCustomDate: boolean; sorting: SortingState; + sessionCursors?: Record; }) { const pageSize = pagination.pageSize || defaultPageSize; const activeSort = sorting[0] ?? DEFAULT_LOGS_SORTING[0]; const sortBy: LogsSortField = isSortField(activeSort.id) ? activeSort.id : "startTime"; const sortOrder: "asc" | "desc" = activeSort.desc ? "desc" : "asc"; + const usesSessionCursor = sortBy === "startTime"; + const sessionCursor = usesSessionCursor ? sessionCursors[pagination.pageIndex] : undefined; const logsQueryOptions: UseQueryOptions = { queryKey: [ @@ -142,6 +148,7 @@ export function useLogFilterLogic({ sortBy, sortOrder, excludeInternalHealthChecks, + sessionCursor, ], queryFn: async () => { if (!accessToken || !token || !userRole || !userID) { @@ -182,6 +189,7 @@ export function useLogFilterLogic({ sort_order: sortOrder, exclude_internal_health_checks: excludeInternalHealthChecks, group_by_session: true, + session_cursor: sessionCursor, }, }); }, @@ -219,5 +227,6 @@ export function useLogFilterLogic({ logsQuery, filteredLogs, allTeams, + usesSessionCursor, }; } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3a4e013b684..098b43f6433 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -56863,6 +56863,8 @@ export interface operations { exclude_internal_health_checks?: boolean; /** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */ group_by_session?: boolean; + /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ + session_cursor?: string | null; }; header?: never; path?: never; @@ -56977,6 +56979,8 @@ export interface operations { exclude_internal_health_checks?: boolean; /** @description Paginate over sessions instead of raw logs: one representative row per session, total counts sessions */ group_by_session?: boolean; + /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ + session_cursor?: string | null; }; header?: never; path?: never; From a0f44af838bba4dc51be8b8addd11f98dedf8c9b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:27:46 -0700 Subject: [PATCH 098/167] fix(proxy/db): translate libpq sslrootcert and verify-* into Prisma's strict TLS params (#39563) * fix(proxy/db): translate libpq sslrootcert and verify-* into Prisma's strict TLS params Prisma silently drops sslrootcert and treats sslmode=verify-ca/verify-full as prefer, so a DATABASE_URL copied from the RDS docs connected over TLS without checking the server certificate. The URL handed to Prisma (writer, DIRECT_URL, read replica, componentized entrypoints) now carries sslmode=require, sslcert= and sslaccept=strict instead. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy/db): ruff format translate_libpq_ssl_params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_url_settings.py | 43 +++++++++- litellm/proxy/proxy_cli.py | 25 ++++-- .../proxy/db/test_db_url_settings.py | 82 +++++++++++++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 46 +++++++++++ 4 files changed, 186 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/db/db_url_settings.py b/litellm/proxy/db/db_url_settings.py index 01f66e4f3c5..f28e505246a 100644 --- a/litellm/proxy/db/db_url_settings.py +++ b/litellm/proxy/db/db_url_settings.py @@ -124,6 +124,42 @@ def add_missing_query_params(url: str, params: Mapping[str, str | int | float]) return urllib.parse.urlunsplit(parsed._replace(query=query)) +LIBPQ_VERIFY_SSLMODES: Final[frozenset[str]] = frozenset({"verify-ca", "verify-full"}) + + +def translate_libpq_ssl_params(url: str) -> str: + """Rewrite libpq's certificate-verification params into Prisma's dialect. + + Prisma's engine only knows ``sslmode=disable|prefer|require``, ``sslcert`` + (the CA bundle) and ``sslaccept=strict``. It silently discards + ``sslrootcert`` and downgrades ``sslmode=verify-ca`` / ``verify-full`` to + ``prefer``, so a URL copied from libpq / RDS docs connects over TLS with no + certificate check at all. ``verify-ca`` and ``verify-full`` both become + ``require`` (Prisma has no CA-only mode), ``sslrootcert`` becomes + ``sslcert``, and either one turns on ``sslaccept=strict`` (chain and + hostname), matching libpq where a root cert makes ``require`` verify. + Prisma params the operator pinned themselves win; anything else is left + untouched. + """ + parsed: Final = urllib.parse.urlsplit(url) + pairs: Final = tuple(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True)) + keys: Final = frozenset(key for key, _ in pairs) + wants_verify: Final = any(key == "sslmode" and value in LIBPQ_VERIFY_SSLMODES for key, value in pairs) + if not wants_verify and "sslrootcert" not in keys: + return url + translated: Final = tuple( + ("sslmode", "require") if key == "sslmode" and value in LIBPQ_VERIFY_SSLMODES else (key, value) + for key, value in pairs + if key != "sslrootcert" + ) + root_cert: Final = tuple( + ("sslcert", value) for key, value in pairs if key == "sslrootcert" and "sslcert" not in keys + ) + strict: Final = () if "sslaccept" in keys else (("sslaccept", "strict"),) + query: Final = urllib.parse.urlencode(translated + root_cert + strict) + return urllib.parse.urlunsplit(parsed._replace(query=query)) + + def reader_shareable_params(params: Mapping[str, str | int | float]) -> Mapping[str, str | int | float]: """Return the subset of ``params`` the read replica is allowed to inherit.""" return MappingProxyType({key: value for key, value in params.items() if key in CONNECTION_PARAM_KEYS}) @@ -403,6 +439,11 @@ class DatabaseURLSettings(BaseSettings): self._raise_for_unsupported_scheme() wrote_writer: Final = self.apply_writer_url_to_env() + for env_var in ("DATABASE_URL", "DIRECT_URL"): + url = os.environ.get(env_var) + if url: + os.environ[env_var] = translate_libpq_ssl_params(url) + # DATABASE_DISABLE_PREPARED_STATEMENTS maps to Prisma's `pgbouncer=true` # URL param, same as the CLI's `database_disable_prepared_statements` # config key. An explicit `pgbouncer` value already on the URL wins. @@ -418,7 +459,7 @@ class DatabaseURLSettings(BaseSettings): reader_url: Final = self.build_reader_url() or self.database_url_read_replica if reader_url is not None: os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params( - reader_url, + translate_libpq_ssl_params(reader_url), connection_params_from_url(os.environ.get("DATABASE_URL", "")), ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index ed247d52ce2..e780beb4410 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1228,6 +1228,7 @@ def run_server( add_missing_query_params, idle_lifetime_params, reader_shareable_params, + translate_libpq_ssl_params, unsupported_db_scheme, unsupported_db_scheme_message, ) @@ -1275,11 +1276,15 @@ def run_server( writer_url, connection_url_params, ) - os.environ["DATABASE_URL"] = add_missing_query_params(modified_url, lifetime_params) + os.environ["DATABASE_URL"] = translate_libpq_ssl_params( + add_missing_query_params(modified_url, lifetime_params) + ) if os.getenv("DIRECT_URL", None) is not None: database_url = os.getenv("DIRECT_URL") modified_url = append_query_params(database_url, connection_url_params) - os.environ["DIRECT_URL"] = add_missing_query_params(modified_url, lifetime_params) + os.environ["DIRECT_URL"] = translate_libpq_ssl_params( + add_missing_query_params(modified_url, lifetime_params) + ) # The reader pool is a real pool against the same configured cap, so it # gets the allowlisted pool params. Schema-affecting ones, including any # the operator smuggled in through database_extra_connection_params, stay @@ -1292,14 +1297,16 @@ def run_server( db_statement_timeout, db_lock_timeout, ) - os.environ["DATABASE_URL_READ_REPLICA"] = add_missing_query_params( + os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( add_missing_query_params( - _with_query_value(read_replica_url, "options", reader_options) - if reader_options - else read_replica_url, - reader_shareable_params(connection_url_params), - ), - lifetime_params, + add_missing_query_params( + _with_query_value(read_replica_url, "options", reader_options) + if reader_options + else read_replica_url, + reader_shareable_params(connection_url_params), + ), + lifetime_params, + ) ) subprocess.run(["prisma"], capture_output=True) is_prisma_runnable = True diff --git a/tests/test_litellm/proxy/db/test_db_url_settings.py b/tests/test_litellm/proxy/db/test_db_url_settings.py index 2552e52fb77..db524625a93 100644 --- a/tests/test_litellm/proxy/db/test_db_url_settings.py +++ b/tests/test_litellm/proxy/db/test_db_url_settings.py @@ -739,3 +739,85 @@ def test_unsupported_db_scheme_message_names_var_and_scheme(): assert "DIRECT_URL" in msg assert "sqlite" in msg assert "postgresql://" in msg + + +def _query(url: str) -> dict[str, list[str]]: + return urllib.parse.parse_qs(urllib.parse.urlsplit(url).query, keep_blank_values=True) + + +def test_libpq_verify_full_and_sslrootcert_become_prisma_strict_sslcert(monkeypatch): + """Prisma drops ``sslrootcert`` and treats ``verify-full`` as ``prefer``, so a + libpq-style URL (the form the RDS docs give) connects with no certificate + check. The URL Prisma actually receives must carry its own strict dialect.""" + monkeypatch.setenv( + "DATABASE_URL", + "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=verify-full&sslrootcert=/certs/rds-bundle.pem", + ) + + assert _apply() is False + + assert _query(os.environ["DATABASE_URL"]) == { + "sslmode": ["require"], + "sslcert": ["/certs/rds-bundle.pem"], + "sslaccept": ["strict"], + } + + +def test_libpq_verify_ca_becomes_prisma_strict(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=verify-ca") + + _apply() + + assert _query(os.environ["DATABASE_URL"]) == {"sslmode": ["require"], "sslaccept": ["strict"]} + + +def test_sslrootcert_alone_turns_on_strict_verification(monkeypatch): + """libpq verifies the chain whenever a root cert is supplied under + ``sslmode=require``; Prisma needs ``sslaccept=strict`` to do the same.""" + monkeypatch.setenv( + "DATABASE_URL", "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=require&sslrootcert=/certs/ca.pem" + ) + + _apply() + + assert _query(os.environ["DATABASE_URL"]) == { + "sslmode": ["require"], + "sslcert": ["/certs/ca.pem"], + "sslaccept": ["strict"], + } + + +def test_pinned_prisma_ssl_params_win_over_libpq_translation(monkeypatch): + monkeypatch.setenv( + "DATABASE_URL", + "postgresql://u:p@db.example.com:5432/litellm_db" + "?sslmode=verify-full&sslrootcert=/ignored.pem&sslcert=/pinned.pem&sslaccept=accept_invalid_certs", + ) + + _apply() + + assert _query(os.environ["DATABASE_URL"]) == { + "sslmode": ["require"], + "sslcert": ["/pinned.pem"], + "sslaccept": ["accept_invalid_certs"], + } + + +def test_prisma_native_ssl_url_is_left_untouched(monkeypatch): + url = "postgresql://u:p@db.example.com:5432/litellm_db?sslmode=require&sslcert=/certs/ca.pem&sslaccept=strict" + monkeypatch.setenv("DATABASE_URL", url) + + _apply() + + assert os.environ["DATABASE_URL"] == url + + +def test_libpq_ssl_translation_covers_direct_url_and_read_replica(monkeypatch): + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@writer.example.com:5432/db?sslmode=verify-full") + monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@direct.example.com:5432/db?sslmode=verify-full") + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://u:p@reader.example.com:5432/db?sslmode=verify-full") + + _apply() + + for env_var in ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA"): + assert _query(os.environ[env_var]) == {"sslmode": ["require"], "sslaccept": ["strict"]}, env_var diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 7b3528f3a68..ff2bd1114de 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -2623,3 +2623,49 @@ class TestTokenAuthCliFlags: assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "ENTRA_TOKEN" not in (database_url or "") assert toggle is None + + +class TestLibpqSslParamTranslation: + """Prisma ignores ``sslrootcert`` and treats ``sslmode=verify-full`` as + ``prefer``, so the URL handed to it must carry Prisma's own strict dialect.""" + + @staticmethod + def _config(tmp_path, general_settings): + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml.dump({"model_list": [], "general_settings": general_settings})) + return str(config_path) + + def test_libpq_url_is_translated_on_every_prisma_url(self, tmp_path): + libpq = "?sslmode=verify-full&sslrootcert=/certs/rds-bundle.pem" + captured = _run_server_and_capture_urls( + self._config(tmp_path, {}), + database_url=f"postgresql://t:t@localhost:5432/t{libpq}", + direct_url=f"postgresql://t:t@direct:5432/t{libpq}", + read_replica_url=f"postgresql://t:t@reader:5432/t{libpq}", + ) + + for env_var in ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA"): + query = urlparse.parse_qs(urlparse.urlparse(captured[env_var]).query) + assert "sslrootcert" not in query, env_var + assert query["sslmode"] == ["require"], env_var + assert query["sslcert"] == ["/certs/rds-bundle.pem"], env_var + assert query["sslaccept"] == ["strict"], env_var + + def test_libpq_params_from_extra_connection_params_are_translated(self, tmp_path): + captured = _run_server_and_capture_urls( + self._config( + tmp_path, + { + "database_extra_connection_params": { + "sslmode": "verify-full", + "sslrootcert": "/certs/rds-bundle.pem", + } + }, + ), + ) + + query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) + assert "sslrootcert" not in query + assert query["sslmode"] == ["require"] + assert query["sslcert"] == ["/certs/rds-bundle.pem"] + assert query["sslaccept"] == ["strict"] From 7256bd307a2cd5df177abec073673bfc7d53a64f Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 3 Sep 2026 10:32:03 -0700 Subject: [PATCH 099/167] fix(mcp): scope allow-all servers to virtual keys (#39531) --- .../mcp_server/auth/user_api_key_auth_mcp.py | 57 +++++++++--- .../mcp_server/mcp_server_manager.py | 37 +++++++- tests/mcp_tests/test_mcp_server.py | 16 ++-- .../auth/test_user_api_key_auth_mcp.py | 38 ++++++++ .../mcp_server/test_mcp_server.py | 5 +- .../mcp_server/test_mcp_server_manager.py | 93 +++++++++++++++++-- 6 files changed, 214 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 66d4aedba06..b9bfb062ec7 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -3,7 +3,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, Final, Literal, cast from fastapi import HTTPException from starlette.datastructures import Headers @@ -305,6 +305,12 @@ def _admission_failure_fallback( raise exc +@dataclass(frozen=True, slots=True) +class MCPServerAccess: + server_ids: tuple[str, ...] + scope: Literal["unscoped", "scoped", "unresolved"] = "unscoped" + + @dataclass(frozen=True, slots=True) class DcrBridgeTarget: """The single DCR-bridge server a request targets, paired with the exact name the caller @@ -1456,6 +1462,18 @@ class MCPRequestHandler: *, keyless_source: bool = False, ) -> list[str]: + access: Final = await MCPRequestHandler.get_mcp_server_access( + user_api_key_auth, + keyless_source=keyless_source, + ) + return list(access.server_ids) + + @staticmethod + async def get_mcp_server_access( + user_api_key_auth: UserAPIKeyAuth | None = None, + *, + keyless_source: bool = False, + ) -> MCPServerAccess: """ Get list of allowed MCP servers for the given user/key based on permissions. @@ -1478,13 +1496,17 @@ class MCPRequestHandler: """ from litellm.proxy.proxy_server import general_settings + key_object_permission: Final = MCPRequestHandler._get_key_object_permission(user_api_key_auth) + try: # A keyless admitted subject resolves per source BEFORE any single-source rule here. Ordering # matters: the no_mcp_servers opt-out below reads the caller's own object_permission, so above # this branch a user's own opt-out would wrongly zero their TEAMS' grants too (each source is # independent; an opt-out silences only its own source, inside the recursive call). if _is_mcp_admitted_user_subject(user_api_key_auth) and user_api_key_auth is not None: - return await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth) + return MCPServerAccess( + server_ids=tuple(await MCPRequestHandler._resolve_admitted_subject_servers(user_api_key_auth)), + ) # Get allowed servers from key and team allowed_mcp_servers_for_key = await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) @@ -1492,7 +1514,7 @@ class MCPRequestHandler: # The key explicitly opted out of every MCP server. This overrides # team inheritance and additive grants (mirrors no-default-models). if SpecialMCPServerNames.no_mcp_servers.value in allowed_mcp_servers_for_key: - return [] + return MCPServerAccess(server_ids=(), scope="scoped") allowed_mcp_servers_for_team = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth) @@ -1572,7 +1594,7 @@ class MCPRequestHandler: "require_end_user_mcp_access_defined=True and end_user %s has no MCP permissions - blocking MCP access", user_api_key_auth.end_user_id, ) - return [] + return MCPServerAccess(server_ids=(), scope="scoped") ######################################################### # Check agent permissions if agent_id is set on the key @@ -1601,14 +1623,22 @@ class MCPRequestHandler: ######################################################### # Apply org-level ceiling if org_id is set ######################################################### - allowed_mcp_servers = await MCPRequestHandler._apply_primary_org_ceiling( + allowed_mcp_servers, org_restricts = await MCPRequestHandler._apply_primary_org_ceiling( allowed_mcp_servers, user_api_key_auth, has_lower_level_mcp_restrictions, keyless_source=keyless_source, ) - return list(set(allowed_mcp_servers)) + declares_key_mcp_scope: Final = getattr(key_object_permission, "mcp_servers", None) is not None + return MCPServerAccess( + server_ids=tuple(set(allowed_mcp_servers)), + scope=( + "scoped" + if has_lower_level_mcp_restrictions or org_restricts or declares_key_mcp_scope + else "unscoped" + ), + ) except Exception as e: if isinstance(e, UnloadableEntitlementError): # A ceiling we KNOW exists and cannot read. Denying is the only answer that does not @@ -1616,7 +1646,10 @@ class MCPRequestHandler: verbose_logger.warning("Denying MCP access, entitlement unreadable: %s", e) else: verbose_logger.warning("Failed to get allowed MCP servers: %s", e) - return [] + return MCPServerAccess( + server_ids=(), + scope="scoped" if getattr(key_object_permission, "mcp_servers", None) is not None else "unresolved", + ) @staticmethod async def _apply_primary_org_ceiling( @@ -1624,7 +1657,7 @@ class MCPRequestHandler: user_api_key_auth: UserAPIKeyAuth | None, has_lower_level_mcp_restrictions: bool, keyless_source: bool = False, - ) -> list[str]: + ) -> tuple[list[str], bool]: """Cap the resolved server list by this caller's org ceiling: an explicit org list intersects lower-level restrictions (else becomes the ceiling); no org or an empty list leaves it unchanged. @@ -1638,7 +1671,7 @@ class MCPRequestHandler: cannot be read raises out of ``_get_allowed_mcp_servers_for_org`` and never arrives here as ``None``, so key auth cannot silently shed a ceiling an operator did configure.""" if not (user_api_key_auth and user_api_key_auth.org_id): - return allowed_mcp_servers + return allowed_mcp_servers, False allowed_mcp_servers_for_org: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_org(user_api_key_auth) if allowed_mcp_servers_for_org is None: verbose_logger.warning( @@ -1646,9 +1679,9 @@ class MCPRequestHandler: user_api_key_auth.org_id, "denying (keyless admitted subject)" if keyless_source else "leaving uncapped (key auth)", ) - return [] if keyless_source else allowed_mcp_servers + return ([] if keyless_source else allowed_mcp_servers), False if len(allowed_mcp_servers_for_org) == 0: - return allowed_mcp_servers + return allowed_mcp_servers, False if has_lower_level_mcp_restrictions or keyless_source: # Org can only cap lower-level restrictions. A keyless admitted source ALWAYS takes this # arm: its model unions GRANTS, so an org list may only narrow a source, never become one. @@ -1657,7 +1690,7 @@ class MCPRequestHandler: # No lower-level restrictions → org list becomes the ceiling. capped = allowed_mcp_servers_for_org verbose_logger.debug("Applied org ceiling filter. Final allowed servers: %s", capped) - return capped + return capped, True @staticmethod def _scoped_source_auth( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1434fa5bfea..a772c569bfa 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -55,6 +55,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + MCPServerAccess, _is_mcp_admitted_user_subject, ) from litellm.proxy._experimental.mcp_server.elicitation_handler import ( @@ -2958,7 +2959,13 @@ class MCPServerManager: return None return user_api_key_auth.mcp_session_resource_server_id - async def get_allowed_mcp_servers(self, user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str]: + async def get_allowed_mcp_servers( + self, + user_api_key_auth: UserAPIKeyAuth | None = None, + *, + access: MCPServerAccess | None = None, + general_settings: Mapping[str, object] | None = None, + ) -> list[str]: """ Get the allowed MCP Servers for the user. @@ -2967,6 +2974,9 @@ class MCPServerManager: 2. If admin and no object_permission, return all servers 3. Otherwise, use standard permission checks """ + from litellm.proxy.proxy_server import general_settings as proxy_general_settings + + resolved_general_settings: Final = proxy_general_settings if general_settings is None else general_settings allow_all_server_ids: Final = self.get_allow_all_keys_server_ids() # A keyless admitted subject is resolved per grant source, and channel decisions that are @@ -3007,11 +3017,16 @@ class MCPServerManager: # whole registry, for keys AND admitted session subjects alike (one predicate owns the # question). Seeded into the union rather than returned early so the session resource # scope below still bounds a per-server envelope held by an admin. - combined_servers: Final = ( - set(self.get_registry().keys()) - if await MCPRequestHandler.admin_view_unscoped(user_api_key_auth) - else set(await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)) + admin_unscoped: Final = await MCPRequestHandler.admin_view_unscoped(user_api_key_auth) + resolved_access: Final = ( + MCPServerAccess(server_ids=()) + if admin_unscoped + else access or await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) ) + resolved_server_ids: Final = ( + set(self.get_registry().keys()) if admin_unscoped else set(resolved_access.server_ids) + ) + combined_servers: Final = set(resolved_server_ids) verbose_logger.debug("Allowed MCP Servers for user api key auth: %s", combined_servers) combined_servers.update( await self.operator_open_server_ids( @@ -3052,6 +3067,18 @@ class MCPServerManager: ] combined_servers.update(delegate_server_ids) + restrict_allow_all: Final = ( + resolved_general_settings.get("mcp_allow_all_keys_respects_mcp_scope", False) + and user_api_key_auth is not None + and user_api_key_auth.via_virtual_key + and resolved_access.scope != "unscoped" + ) + if restrict_allow_all: + combined_servers.difference_update( + set(allow_all_server_ids) + - resolved_server_ids + - (set(submitted_server_ids) if resolved_access.scope != "unresolved" else set()) + ) if len(combined_servers) == 0: verbose_logger.debug("No allowed MCP Servers found for user api key auth.") scope = MCPServerManager._admitted_session_resource_scope(user_api_key_auth) diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index e06c33263fb..1781dfe2fc2 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -2815,6 +2815,7 @@ async def test_mcp_server_manager_with_access_groups_integration(): """Integration test for MCPServerManager with access group filtering""" from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + MCPServerAccess, ) from litellm.proxy._types import UserAPIKeyAuth @@ -2848,11 +2849,11 @@ async def test_mcp_server_manager_with_access_groups_integration(): ) # Mock the permission lookup to return staff access group - with patch.object(MCPRequestHandler, "get_allowed_mcp_servers") as mock_get_allowed: - mock_get_allowed.return_value = [ - "staff-server-id", - "ops-server-id", - ] # User has access to staff and ops + with patch.object(MCPRequestHandler, "get_mcp_server_access") as mock_get_allowed: # test-quality-ok: manager resolver seam + mock_get_allowed.return_value = MCPServerAccess( + server_ids=("staff-server-id", "ops-server-id"), + scope="scoped", + ) allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth) @@ -2901,6 +2902,7 @@ async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permi from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + MCPServerAccess, ) test_manager = MCPServerManager() @@ -2923,9 +2925,9 @@ async def test_get_allowed_mcp_servers_returns_empty_for_non_admin_without_permi ) with patch.object( - MCPRequestHandler, "get_allowed_mcp_servers", new_callable=AsyncMock + MCPRequestHandler, "get_mcp_server_access", new_callable=AsyncMock ) as mock_permission_lookup: - mock_permission_lookup.return_value = [] + mock_permission_lookup.return_value = MCPServerAccess(server_ids=()) allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth) assert allowed_servers == [] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index c8ea4867f2c..f1e299802fb 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -667,6 +667,44 @@ class TestMCPRequestHandler: assert result == [] + async def test_db_default_empty_key_scope_keeps_org_substitution(self): + """A key whose object_permission row carries only the DB-default empty mcp_servers + list (e.g. a vector-stores-only key) places no lower-level MCP restriction: the org + list still substitutes with the flag off, while the access result stays scoped so + the opt-in allow-all ceiling can still bind""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", org_id="org-1") + key_object_permission = self._toolset_only_object_permission([]) + key_object_permission.mcp_toolsets = None + mock_manager = self._mock_manager_with_toolsets({}) + + with ( + patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_key_object_permission", return_value=key_object_permission + ), + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here + MCPRequestHandler, "_get_key_access_group_mcp_server_extras", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here + MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[]) + ), + patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling tests + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object( # test-quality-ok: stub the level's perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, + "_get_allowed_mcp_servers_for_org", + AsyncMock(return_value=["server-x", "server-y"]), + ), + ): + access = await MCPRequestHandler.get_mcp_server_access(user_api_key_auth) + + assert sorted(access.server_ids) == ["server-x", "server-y"] + assert access.scope == "scoped" + async def test_team_dangling_toolset_denies_key_own_grants(self): """A team toolset that cannot be resolved must deny on the SERVER axis too, not silently drop the team ceiling and pass the key's own grants through""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ff0007f47aa..0fd35e674b7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -3367,6 +3367,7 @@ async def test_mcp_manager_returns_public_when_permission_lookup_fails(): @pytest.mark.asyncio async def test_mcp_manager_merges_public_and_restricted_servers(): try: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, ) @@ -3398,8 +3399,8 @@ async def test_mcp_manager_merges_public_and_restricted_servers(): return_value=False, ), patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_allowed_mcp_servers", - AsyncMock(return_value=["restricted"]), + "litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPRequestHandler.get_mcp_server_access", + AsyncMock(return_value=MCPServerAccess(server_ids=("restricted",), scope="scoped")), ), ): allowed = await manager.get_allowed_mcp_servers(UserAPIKeyAuth()) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 5cde2e83f62..91e870d2d95 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5884,6 +5884,7 @@ class TestMCPServerManager: """ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + MCPServerAccess, ) from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth @@ -5903,19 +5904,22 @@ class TestMCPServerManager: object_permission_id="perm_123", ) - # Mock MCPRequestHandler.get_allowed_mcp_servers to verify it receives user_api_key_auth + # Mock MCPRequestHandler.get_mcp_server_access to verify it receives user_api_key_auth with patch.object( MCPRequestHandler, - "get_allowed_mcp_servers", + "get_mcp_server_access", new_callable=AsyncMock, ) as mock_get_allowed: # Configure mock to return servers from object_permission - mock_get_allowed.return_value = ["test_server_1", "test_server_2"] + mock_get_allowed.return_value = MCPServerAccess( + server_ids=("test_server_1", "test_server_2"), + scope="scoped", + ) # Call get_allowed_mcp_servers with user_api_key_auth result = await manager.get_allowed_mcp_servers(user_api_key_auth) - # Verify MCPRequestHandler.get_allowed_mcp_servers was called with user_api_key_auth + # Verify MCPRequestHandler.get_mcp_server_access was called with user_api_key_auth mock_get_allowed.assert_called_once() call_args = mock_get_allowed.call_args assert call_args[0][0] is user_api_key_auth # First positional arg should be user_api_key_auth @@ -6072,6 +6076,7 @@ class TestMCPServerManager: from litellm.proxy import proxy_server as proxy_server_module from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + MCPServerAccess, ) from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_active_toolset_id, @@ -6104,9 +6109,12 @@ class TestMCPServerManager: patch.object(manager, "get_allow_all_keys_server_ids", return_value=["global-server"]), patch.object( MCPRequestHandler, - "get_allowed_mcp_servers", + "get_mcp_server_access", new_callable=AsyncMock, - return_value=["toolset-server"], + return_value=MCPServerAccess( + server_ids=("toolset-server",), + scope="scoped", + ), ), ): result = await manager.get_allowed_mcp_servers(user_api_key_auth) @@ -6205,6 +6213,79 @@ class TestMCPServerManager: assert set(result) == {"global-server", "submitted-server"} + @pytest.mark.asyncio + @pytest.mark.parametrize( + "flag_enabled, via_virtual_key, resolved_server_ids, scope, submitted_server_ids, expected_server_ids", + [ + (False, True, ("granted",), "scoped", (), {"granted", "public"}), + (True, True, ("granted",), "scoped", (), {"granted"}), + (True, True, ("granted", "public"), "scoped", (), {"granted", "public"}), + (True, True, (), "scoped", (), set()), + (True, True, (), "unscoped", (), {"public"}), + ( + True, + True, + ("team-granted",), + "scoped", + ("submitted",), + {"team-granted", "submitted"}, + ), + ( + True, + True, + ("team-granted",), + "scoped", + ("public",), + {"team-granted", "public"}, + ), + (True, False, ("granted",), "scoped", (), {"granted", "public"}), + ], + ids=( + "flag_off_preserves_allow_all", + "flag_on_scoped_key_excludes_allow_all", + "flag_on_keeps_allow_all_when_granted", + "flag_on_restricted_empty_excludes_allow_all", + "flag_on_unscoped_key_preserves_allow_all", + "flag_on_preserves_submitted_byom", + "flag_on_preserves_submitted_byom_when_it_is_allow_all", + "flag_on_non_virtual_key_preserves_allow_all", + ), + ) + async def test_allow_all_keys_scope_flag( + self, + flag_enabled, + via_virtual_key, + resolved_server_ids, + scope, + submitted_server_ids, + expected_server_ids, + ): # test-quality-ok: parameterized matrix covers the scope state machine + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPServerAccess + + manager = MCPServerManager() + auth = UserAPIKeyAuth(api_key="sk-test", user_id="user-123") + auth.via_virtual_key = via_virtual_key + access = MCPServerAccess(server_ids=resolved_server_ids, scope=scope) + + with ( + patch.object(manager, "get_allow_all_keys_server_ids", return_value=["public"]), + patch.object( + manager, + "_get_active_submitted_mcp_server_ids_for_user", + new=AsyncMock(return_value=list(submitted_server_ids)), + ), + ): + assert ( + set( + await manager.get_allowed_mcp_servers( + auth, + access=access, + general_settings={"mcp_allow_all_keys_respects_mcp_scope": flag_enabled}, + ) + ) + == expected_server_ids + ) + @pytest.mark.asyncio async def test_get_allowed_mcp_servers_anonymous_delegate_requires_oauth2(self): """Anonymous delegated auth listing should only include oauth2 servers.""" From bb7d787425ee35c9530296c8aee10e2a30e08d7d Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 3 Sep 2026 10:35:01 -0700 Subject: [PATCH 100/167] Merge pull request #39571 from BerriAI/codex/team-id-empty-field fix(team): generate team IDs for blank input --- litellm/proxy/_types.py | 7 +++++++ .../proxy/management_endpoints/test_team_endpoints.py | 8 ++++++++ 2 files changed, 15 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c6ecb0be8b9..97f5f59d2dc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1930,6 +1930,13 @@ class NewTeamRequest(TeamBase): model_config = ConfigDict(protected_namespaces=()) + @field_validator("team_id", mode="before") + @classmethod + def treat_blank_team_id_as_unset(cls, v: object) -> object: + if isinstance(v, str) and not v.strip(): + return None + return v + class GlobalEndUsersSpend(LiteLLMPydanticObjectBase): api_key: str | None = None diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 30b2ab86b9a..088e82b370f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11083,6 +11083,14 @@ async def test_new_team_rejects_reserved_ui_session_team_id(): mock_prisma.get_data.assert_not_called() +@pytest.mark.parametrize("team_id", ["", " "]) +def test_new_team_request_blank_team_id_is_unset(team_id: str) -> None: + from litellm.proxy._types import NewTeamRequest + + assert NewTeamRequest(team_alias="t", team_id=team_id).team_id is None + assert NewTeamRequest(team_id="custom").team_id == "custom" + + # --------------------------------------------------------------------------- # PATCH /team/{team_id} — RFC 7386 JSON Merge Patch # From b9e030ddd662cb7abb98b8fa5139efb14d57cb6c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:38:25 -0700 Subject: [PATCH 101/167] fix(cost): apply off_peak_pricing in the dashscope cost calculator --- .../litellm_core_utils/llm_cost_calc/utils.py | 4 +- litellm/llms/dashscope/cost_calculator.py | 120 ++++++++------- .../test_dashscope_cost_calculator.py | 138 +++++++++++++++++- 3 files changed, 209 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index b34c416cd40..21587af73aa 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -428,7 +428,7 @@ def _coerce_off_peak_rate(value: object, default: float) -> float: return default -def _apply_off_peak_pricing( +def apply_off_peak_pricing( model_info: ModelInfo, current_time: datetime | None, prompt_base_cost: float, @@ -462,7 +462,7 @@ def _apply_off_peak_to_base_costs( has no field for them. """ prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs - off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing( + off_peak_prompt, off_peak_completion, off_peak_cache_read = apply_off_peak_pricing( model_info, current_time, prompt, completion, cache_read ) return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read) diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index dd5bee1fe8b..d8eb1f9f8d7 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,11 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ -from dataclasses import dataclass +from dataclasses import dataclass, replace +from datetime import datetime from typing import Final from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.litellm_core_utils.llm_cost_calc.utils import ( + apply_off_peak_pricing, parse_completion_tokens_details, parse_prompt_tokens_details, ) @@ -32,6 +34,19 @@ class TokenBreakdown: return self.text_tokens + self.cached_tokens + self.cache_creation_tokens +@dataclass(frozen=True, slots=True) +class TokenRates: + input_rate: float + cache_read_rate: float + cache_creation_rate: float + output_rate: float + reasoning_rate: float | None + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: prompt_details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = prompt_details["cache_hit_tokens"] @@ -57,69 +72,75 @@ def _flat_rate(model_info: ModelInfo, cost_key: str, fallback_cost_key: str) -> return float(value) -def _calculate_prompt_cost( - breakdown: TokenBreakdown, - model_info: ModelInfo, - tier: dict | None, -) -> float: - if tier is not None: - return ( - (breakdown.text_tokens * tier_rate(tier, "input_cost_per_token")) - + (breakdown.cached_tokens * tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token")) - + ( - breakdown.cache_creation_tokens - * tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token") - ) - ) - - input_cost: Final = float(model_info.get("input_cost_per_token") or 0.0) - cache_read_cost: Final = _flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token") - cache_creation_cost: Final = _flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token") - - return ( - (breakdown.text_tokens * input_cost) - + (breakdown.cached_tokens * cache_read_cost) - + (breakdown.cache_creation_tokens * cache_creation_cost) +def _flat_rates(model_info: ModelInfo) -> TokenRates: + reasoning_rate: Final = model_info.get("output_cost_per_reasoning_token") + return TokenRates( + input_rate=float(model_info.get("input_cost_per_token") or 0.0), + cache_read_rate=_flat_rate(model_info, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=_flat_rate(model_info, "cache_creation_input_token_cost", "input_cost_per_token"), + output_rate=float(model_info.get("output_cost_per_token") or 0.0), + reasoning_rate=None if reasoning_rate is None else float(reasoning_rate), ) -def _calculate_completion_cost( - breakdown: TokenBreakdown, - model_info: ModelInfo, - tier: dict | None, -) -> float: +def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: # A tier that declares output rates keeps the request on them, all-or-nothing. A tier table # spelling out only input rates would serve every completion for free, so there the model's # own output rates stand in - tier_declares_output: Final = tier is not None and "output_cost_per_token" in tier - output_cost: Final = ( - tier_rate(tier, "output_cost_per_token") - if tier_declares_output - else float(model_info.get("output_cost_per_token") or 0.0) - ) - tier_declares_reasoning: Final = tier is not None and "output_cost_per_reasoning_token" in tier - model_reasoning_rate: Final = None if tier_declares_output else model_info.get("output_cost_per_reasoning_token") - reasoning_cost: Final = ( - tier_rate(tier, "output_cost_per_reasoning_token", "output_cost_per_token") - if tier_declares_reasoning - else float(model_reasoning_rate) - if model_reasoning_rate is not None - else output_cost + flat_rates: Final = _flat_rates(model_info) + tier_declares_output: Final = "output_cost_per_token" in tier + tier_declares_reasoning: Final = "output_cost_per_reasoning_token" in tier + return TokenRates( + input_rate=tier_rate(tier, "input_cost_per_token"), + cache_read_rate=tier_rate(tier, "cache_read_input_token_cost", "input_cost_per_token"), + cache_creation_rate=tier_rate(tier, "cache_creation_input_token_cost", "input_cost_per_token"), + output_rate=tier_rate(tier, "output_cost_per_token") if tier_declares_output else flat_rates.output_rate, + reasoning_rate=( + tier_rate(tier, "output_cost_per_reasoning_token") + if tier_declares_reasoning + else None + if tier_declares_output + else flat_rates.reasoning_rate + ), ) - return (breakdown.completion_tokens * output_cost) + (breakdown.reasoning_tokens * reasoning_cost) + +def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: + input_rate, output_rate, cache_read_rate = apply_off_peak_pricing( + model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate + ) + return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate) -def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashscope") -> tuple[float, float]: +def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]: + prompt_cost: Final = ( + (breakdown.text_tokens * rates.input_rate) + + (breakdown.cached_tokens * rates.cache_read_rate) + + (breakdown.cache_creation_tokens * rates.cache_creation_rate) + ) + completion_cost: Final = (breakdown.completion_tokens * rates.output_rate) + ( + breakdown.reasoning_tokens * rates.billed_reasoning_rate + ) + return prompt_cost, completion_cost + + +def cost_per_token( + model: str, + usage: Usage, + custom_llm_provider: str = "dashscope", + current_time: datetime | None = None, +) -> tuple[float, float]: """ Calculate cost per token for Dashscope models. - Supports both tiered and flat pricing with cached and reasoning tokens. + Supports both tiered and flat pricing with cached and reasoning tokens, and swaps in the + model's off_peak_pricing rates while one of its windows is open. Args: model: Model name without provider prefix usage: LiteLLM Usage block custom_llm_provider: The provider id the request resolved to; dashscope or one of its brand aliases + current_time: The moment the request is billed at; defaults to now, UTC Returns: Tuple[float, float] - (prompt_cost_in_usd, completion_cost_in_usd) @@ -133,8 +154,7 @@ def cost_per_token(model: str, usage: Usage, custom_llm_provider: str = "dashsco if tiered_pricing else None ) + standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier) + rates: Final = _off_peak_rates(model_info, current_time, standard_rates) - prompt_cost: Final = _calculate_prompt_cost(breakdown=breakdown, model_info=model_info, tier=tier) - completion_cost: Final = _calculate_completion_cost(breakdown=breakdown, model_info=model_info, tier=tier) - - return prompt_cost, completion_cost + return _bill(breakdown, rates) diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index 8dc4620dd1b..b6281834f24 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -10,11 +10,11 @@ Tests the cost calculation for Dashscope models including: import math import os +from datetime import datetime, timezone import pytest # Add the project root to Python path - import litellm from litellm.llms.dashscope.cost_calculator import ( cost_per_token as dashscope_cost_per_token, @@ -526,3 +526,139 @@ class TestDashscopeCostCalculator: assert prompt_cost == 0.0 assert math.isclose(completion_cost, 500 * 1.6e-06, rel_tol=1e-10) + + OFF_PEAK_WINDOW = "14:00-00:00" + INSIDE_WINDOW = datetime(2026, 9, 3, 17, 25, tzinfo=timezone.utc) + OUTSIDE_WINDOW = datetime(2026, 9, 3, 9, 0, tzinfo=timezone.utc) + + def _register_off_peak_flat_model(self, model_key: str, off_peak_pricing: dict) -> None: + litellm.model_cost[model_key] = { + "litellm_provider": "dashscope", + "mode": "chat", + "input_cost_per_token": 2.4e-06, + "output_cost_per_token": 4.8e-06, + "cache_read_input_token_cost": 2e-07, + "cache_creation_input_token_cost": 3e-06, + "off_peak_pricing": off_peak_pricing, + } + + def test_dashscope_off_peak_window_swaps_in_the_off_peak_rates(self): + """ + Regression (LIT-6782): a deployment configured with off_peak_pricing kept billing the + standard dashscope rates inside its window, while the same block on a deepseek + deployment billed the off-peak rates. + """ + self._register_off_peak_flat_model( + "dashscope/deepseek-off-peak-test", + { + "hours_utc": self.OFF_PEAK_WINDOW, + "input_cost_per_token": 1.2e-06, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1e-07, + }, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=200, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="deepseek-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, (600 * 1.2e-06) + (300 * 1e-07) + (100 * 3e-06), rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="deepseek-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + + assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10) + assert math.isclose(peak_completion_cost, 200 * 4.8e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_window_overrides_the_selected_tier(self): + """An open off-peak window bills the whole request at the flat off-peak rates, whichever tier + the input volume selected.""" + self._register_tiered_model( + "dashscope/qwen-tiered-off-peak-test", + [ + {"range": [0, 1000], "input_cost_per_token": 4e-07, "output_cost_per_token": 1.6e-06}, + {"range": [1000, 2000], "input_cost_per_token": 8e-07, "output_cost_per_token": 3.2e-06}, + ], + ) + litellm.model_cost["dashscope/qwen-tiered-off-peak-test"]["off_peak_pricing"] = { + "hours_utc": self.OFF_PEAK_WINDOW, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + } + usage = Usage(prompt_tokens=1500, completion_tokens=300) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-tiered-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, 1500 * 1e-07, rel_tol=1e-10) + assert math.isclose(completion_cost, 300 * 4e-07, rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="qwen-tiered-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + + assert math.isclose(peak_prompt_cost, 1500 * 8e-07, rel_tol=1e-10) + assert math.isclose(peak_completion_cost, 300 * 3.2e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_rates_left_unset_keep_the_standard_rates(self): + """A block that only overrides the input rate leaves output and cache reads on the standard + rates, and an explicit reasoning rate is never swapped out.""" + self._register_off_peak_flat_model( + "dashscope/qwen-partial-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "input_cost_per_token": 1.2e-06}, + ) + litellm.model_cost["dashscope/qwen-partial-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06 + usage = Usage( + prompt_tokens=1000, + completion_tokens=200, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-partial-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(prompt_cost, (700 * 1.2e-06) + (300 * 2e-07), rel_tol=1e-10) + assert math.isclose(completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_output_rate_covers_reasoning_without_a_dedicated_rate(self): + """Reasoning tokens on a model with no dedicated reasoning rate follow the off-peak output + rate, the same way they follow the standard output rate outside the window.""" + self._register_off_peak_flat_model( + "dashscope/qwen-reasoning-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "output_cost_per_token": 2.4e-06}, + ) + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + _, completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + + def test_dashscope_off_peak_defaults_to_the_current_time(self): + """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the + default current time.""" + self._register_off_peak_flat_model( + "dashscope/qwen-all-day-off-peak-test", + {"hours_utc": "00:00-00:00", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 2.4e-06}, + ) + usage = Usage(prompt_tokens=1000, completion_tokens=200) + + prompt_cost, completion_cost = dashscope_cost_per_token(model="qwen-all-day-off-peak-test", usage=usage) + + assert math.isclose(prompt_cost, 1000 * 1.2e-06, rel_tol=1e-10) + assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) From 68ffa1db230ee2a03e851f567ce03ec48ce83e9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:43:41 -0700 Subject: [PATCH 102/167] fix(router): pin JWT-authenticated callers by user id in deployment_affinity --- .../deployment_affinity_check.py | 36 +-- .../test_deployment_affinity_check.py | 207 ++++++++++++++++++ 2 files changed, 225 insertions(+), 18 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 39d3e25aacb..ea450864604 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -82,6 +82,7 @@ class DeploymentAffinityCheck(CustomLogger): """ CACHE_KEY_PREFIX = "deployment_affinity:v1" + USER_ID_AFFINITY_PREFIX: Final = "user_id:" def __init__( self, @@ -253,15 +254,6 @@ class DeploymentAffinityCheck(CustomLogger): hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" - @staticmethod - def _get_user_key_from_metadata_dict(metadata: dict) -> str | None: - # NOTE: affinity is keyed on the *API key hash* provided by the proxy (not the - # OpenAI `user` parameter, which is an end-user identifier). - user_key: Final = metadata.get("user_api_key_hash") - if user_key is None: - return None - return str(user_key) - @staticmethod def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: session_id: Final = metadata.get("session_id") @@ -285,22 +277,30 @@ class DeploymentAffinityCheck(CustomLogger): return metadata_dicts @staticmethod - def _get_user_key_from_request_kwargs(request_kwargs: dict) -> str | None: + def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None: + value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None) + return None if value is None else str(value) + + @classmethod + def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None: """ Extract a stable affinity key from request kwargs. - Source (proxy): `metadata.user_api_key_hash` + Source (proxy): `metadata.user_api_key_hash` for virtual-key callers. JWT-authenticated + callers carry no key hash, so their `metadata.user_api_key_user_id` stands in for it, + namespaced under `USER_ID_AFFINITY_PREFIX` so a user id can never alias a key hash. Note: the OpenAI `user` parameter is an end-user identifier and is intentionally not used for deployment affinity. """ - # Check metadata dicts (Proxy usage) - for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): - user_key = DeploymentAffinityCheck._get_user_key_from_metadata_dict(metadata=metadata) - if user_key is not None: - return user_key - - return None + metadata_dicts: Final = cls._iter_metadata_dicts(request_kwargs) + user_api_key_hash: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_hash") + if user_api_key_hash is not None: + return user_api_key_hash + user_id: Final = cls._first_metadata_value(metadata_dicts, "user_api_key_user_id") + if user_id is None: + return None + return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}" @staticmethod def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 60433921de6..1852d641f0a 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -998,3 +998,210 @@ async def test_model_group_affinity_config_overrides_global(): ) # All deployments returned (user-key affinity disabled for this group) assert len(filtered) == 2 + + +def _jwt_metadata(user_id: str) -> dict: + return {"user_api_key_hash": None, "user_api_key_user_id": user_id} + + +def _two_deployments(model_group: str) -> list[dict]: + return [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "openai-deployment-a"}, + }, + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini"}, + "model_info": {"id": "openai-deployment-b"}, + }, + ] + + +@pytest.mark.asyncio +async def test_async_jwt_user_affinity_routes_to_same_deployment(): + """ + JWT-authenticated proxy requests carry no `user_api_key_hash`, only `user_api_key_user_id`. + They must still pin to one deployment per user. + """ + model_group = "gpt-5.4-mini" + router = litellm.Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-a"}, + "model_info": {"id": "openai-deployment-a"}, + }, + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock-api-key-b"}, + "model_info": {"id": "openai-deployment-b"}, + }, + ], + optional_pre_call_checks=["deployment_affinity"], + ) + + choice_calls = {"count": 0} + + def deterministic_choice(seq): + choice_calls["count"] += 1 + if choice_calls["count"] == 1: + return seq[0] + return seq[1] if len(seq) > 1 else seq[0] + + with patch( # test-quality-ok: simple-shuffle has no injectable RNG; forcing the other pick is what proves the pin overrides the strategy + "litellm.router_strategy.simple_shuffle.random.choice", + side_effect=deterministic_choice, + ): + first_response = await router.acompletion( + model=model_group, + messages=[{"role": "user", "content": "Reply with the single word ok"}], + mock_response="ok", + metadata=_jwt_metadata("jwt-user-alice"), + ) + second_response = await router.acompletion( + model=model_group, + messages=[{"role": "user", "content": "Reply with the single word ok"}], + mock_response="ok", + metadata=_jwt_metadata("jwt-user-alice"), + ) + + first_model_id = first_response._hidden_params["model_id"] + assert first_model_id in ("openai-deployment-a", "openai-deployment-b") + assert second_response._hidden_params["model_id"] == first_model_id + + +@pytest.mark.asyncio +async def test_proxy_jwt_auth_metadata_pins_per_user(): + """ + The metadata the proxy stamps for a JWT caller (`UserAPIKeyAuth(api_key=None, user_id=)`) + must claim a pin and be read back by the filter, and another JWT user must not inherit it. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + def proxy_request(user_id: str) -> dict: + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id), + _metadata_variable_name="metadata", + ) + + alice_request = proxy_request("jwt-user-alice") + assert alice_request["metadata"]["user_api_key_hash"] is None + + await callback.async_pre_call_deployment_hook( + kwargs={ + **alice_request, + "metadata": {**alice_request["metadata"], "deployment_model_name": model_group}, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + alice_pinned = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=alice_request, + parent_otel_span=None, + ) + assert [deployment["model_info"]["id"] for deployment in alice_pinned] == ["openai-deployment-b"] + + bob_filtered = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=proxy_request("jwt-user-bob"), + parent_otel_span=None, + ) + assert bob_filtered == healthy_deployments + + +@pytest.mark.asyncio +async def test_jwt_user_id_never_reads_a_virtual_key_pin(): + """ + A JWT user id that happens to equal a virtual key's 64-hex hash must not read that key's pin. + """ + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + key_hash = "a" * 64 + + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": {"user_api_key_hash": key_hash, "deployment_model_name": model_group}, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + key_pinned = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"user_api_key_hash": key_hash}}, + parent_otel_span=None, + ) + assert [deployment["model_info"]["id"] for deployment in key_pinned] == ["openai-deployment-b"] + + lookalike_jwt_user = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": _jwt_metadata(key_hash)}, + parent_otel_span=None, + ) + assert lookalike_jwt_user == healthy_deployments + + +@pytest.mark.asyncio +async def test_virtual_key_hash_wins_over_user_id_for_affinity(): + """ + A virtual-key caller with a user id pins on the key hash, so two keys owned by one user + keep independent pins. + """ + model_group = "gpt-5.4-mini" + healthy_deployments = _two_deployments(model_group) + callback = DeploymentAffinityCheck( + cache=DualCache(), + ttl_seconds=60, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + + await callback.async_pre_call_deployment_hook( + kwargs={ + "metadata": { + "user_api_key_hash": "key-one", + "user_api_key_user_id": "shared-user", + "deployment_model_name": model_group, + }, + "model_info": {"id": "openai-deployment-b"}, + }, + call_type=None, + ) + + other_key_same_user = await callback.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"metadata": {"user_api_key_hash": "key-two", "user_api_key_user_id": "shared-user"}}, + parent_otel_span=None, + ) + assert other_key_same_user == healthy_deployments From 1b42b81f4e048db9403c967365f15165f07a0dd4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:50:21 -0700 Subject: [PATCH 103/167] fix(router): log the hashed affinity key so JWT callers stay distinguishable --- .../pre_call_checks/deployment_affinity_check.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index ea450864604..6f3ea8eb78a 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -533,9 +533,9 @@ class DeploymentAffinityCheck(CustomLogger): return typed_healthy_deployments verbose_router_logger.debug( - "DeploymentAffinityCheck: api-key affinity hit -> deployment=%s user_key=%s", + "DeploymentAffinityCheck: caller affinity hit -> deployment=%s user_key=%s", model_id, - self._shorten_for_logs(user_key), + self._shorten_for_logs(self._hash_user_key(user_key)), ) return [deployment] @@ -626,7 +626,7 @@ class DeploymentAffinityCheck(CustomLogger): deployment_model_name, model_id, self.ttl_seconds, - self._shorten_for_logs(user_key), + self._shorten_for_logs(self._hash_user_key(user_key)), ) else: verbose_router_logger.debug( From 4d3c1998affa9249fbfdcd0c8157acaa991f9186 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:50:23 -0700 Subject: [PATCH 104/167] fix(image_gen): report the requested output_format on gpt-image responses --- .../image_generation/gpt_transformation.py | 2 +- .../test_gpt_transformation.py | 38 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 05494c497ca..64244cfaeda 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -84,6 +84,6 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): # set optional params image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 image_response.quality = optional_params.get("quality", "high") # always hd for dall-e-3 - image_response.output_format = optional_params.get("response_format", "png") # always png for dall-e-3 + image_response.output_format = optional_params.get("output_format", "png") return image_response diff --git a/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py b/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py new file mode 100644 index 00000000000..de54713a570 --- /dev/null +++ b/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py @@ -0,0 +1,38 @@ +from unittest.mock import MagicMock + +import httpx +import pytest + +from litellm.llms.azure.image_generation.gpt_transformation import AzureGPTImageGenerationConfig +from litellm.llms.openai.image_generation.gpt_transformation import GPTImageGenerationConfig +from litellm.types.utils import ImageResponse + + +@pytest.mark.parametrize("config", [GPTImageGenerationConfig(), AzureGPTImageGenerationConfig()]) +def test_transform_image_generation_response_reports_requested_output_format(config): + raw_response = httpx.Response( + status_code=200, + json={ + "created": 1788457009, + "data": [{"b64_json": "/9j/4AAQSkZJRg=="}], + "output_format": "jpeg", + "background": "opaque", + "quality": "low", + "size": "1024x1024", + }, + request=httpx.Request("POST", "https://api.openai.com/v1/images/generations"), + ) + + image_response = config.transform_image_generation_response( + model="gpt-image-2", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={"prompt": "a red apple", "output_format": "jpeg"}, + optional_params={"output_format": "jpeg"}, + litellm_params={}, + encoding=None, + ) + + assert image_response.output_format == "jpeg" + assert image_response.background == "opaque" From ba9bb752980ea18d8b92830b11c7040bc48b5675 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 3 Sep 2026 10:53:50 -0700 Subject: [PATCH 105/167] bump: litellm-enterprise 0.1.63 -> 0.1.64, litellm-proxy-extras 0.4.92 -> 0.4.93 --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 4 ++-- uv.lock | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 8360c0a077d..b6f482ccd86 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.63" +version = "0.1.64" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.63" +version = "0.1.64" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 0944f99ad54..97e9eb66bf2 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.92" +version = "0.4.93" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.92" +version = "0.4.93" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index 161994635b8..d3038a60c42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.92", - "litellm-enterprise==0.1.63", + "litellm-proxy-extras==0.4.93", + "litellm-enterprise==0.1.64", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/uv.lock b/uv.lock index de3181edd5a..dfa77c66dfe 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-08-31T17:52:45.782441Z" exclude-newer-span = "P3D" [manifest] @@ -4765,12 +4765,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.63" +version = "0.1.64" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.92" +version = "0.4.93" source = { editable = "litellm-proxy-extras" } [[package]] From ec2e35b6796b75231f8ff54686baa1604e7f3697 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:07:08 -0700 Subject: [PATCH 106/167] fix(image_gen): keep the provider's echoed size, quality, and output_format on gpt-image responses --- litellm/llms/openai/image_generation/gpt_transformation.py | 6 +++--- .../llms/openai/image_generation/test_gpt_transformation.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 64244cfaeda..090b2eba387 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -82,8 +82,8 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): ) # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "high") # always hd for dall-e-3 - image_response.output_format = optional_params.get("output_format", "png") + image_response.size = image_response.size or optional_params.get("size", "1024x1024") + image_response.quality = image_response.quality or optional_params.get("quality", "high") + image_response.output_format = image_response.output_format or optional_params.get("output_format", "png") return image_response diff --git a/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py b/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py index de54713a570..d9b87627b58 100644 --- a/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/image_generation/test_gpt_transformation.py @@ -9,7 +9,7 @@ from litellm.types.utils import ImageResponse @pytest.mark.parametrize("config", [GPTImageGenerationConfig(), AzureGPTImageGenerationConfig()]) -def test_transform_image_generation_response_reports_requested_output_format(config): +def test_transform_image_generation_response_keeps_provider_echo(config): raw_response = httpx.Response( status_code=200, json={ @@ -35,4 +35,5 @@ def test_transform_image_generation_response_reports_requested_output_format(con ) assert image_response.output_format == "jpeg" + assert image_response.quality == "low" assert image_response.background == "opaque" From aaef5d219aa3ed1ba262302705df8805570a6a45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:07:29 -0700 Subject: [PATCH 107/167] fix(proxy): drop anthropic-beta on the Vertex passthrough count-tokens route --- .../llm_passthrough_endpoints.py | 12 ++- .../test_vertex_passthrough_load_balancing.py | 88 +++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b48b8d81494..29f216fd450 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1730,6 +1730,16 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: return headers +def _is_vertex_anthropic_count_tokens_route(endpoint: str) -> bool: + return endpoint.rsplit("/", 1)[-1].split(":", 1)[0] == "count-tokens" + + +def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]: + if not _is_vertex_anthropic_count_tokens_route(endpoint): + return headers + return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"}) + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: @@ -2128,7 +2138,7 @@ async def _base_vertex_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=target, - custom_headers=headers, + custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 961479c0393..6735f2a3780 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -5,6 +5,7 @@ import pytest from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _base_vertex_proxy_route, + _upstream_headers_for_vertex_route, ) from litellm.types.router import DeploymentTypedDict @@ -348,6 +349,93 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): assert headers_passed_through is False +VERTEX_ANTHROPIC_MODELS_PREFIX = "v1/projects/test-project/locations/global/publishers/anthropic/models/" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_segment", "expects_anthropic_beta"), + [ + ("count-tokens:rawPredict", False), + ("claude-sonnet-4-6:streamRawPredict", True), + ], +) +async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( + model_segment: str, expects_anthropic_beta: bool +): + with ( + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.proxy_server.llm_router", None + ), + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router" + ) as mock_pt_router, + patch( # test-quality-ok: the route offers no injection point for its header preparation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", + new_callable=AsyncMock, + ) as mock_prep_headers, + patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route, + patch( # test-quality-ok: the route calls auth directly rather than through Depends + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", + new_callable=AsyncMock, + return_value=False, + ), + ): + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ( + { + "anthropic-beta": "tool-search-tool-2025-10-19,web-search-2025-03-05", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + }, + "https://aiplatform.googleapis.com", + False, + "test-project", + "global", + ) + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key") + + await _base_vertex_proxy_route( + endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", + request=MagicMock(), + fastapi_response=MagicMock(), + get_vertex_pass_through_handler=MagicMock(), + ) + + upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] + assert ("anthropic-beta" in upstream_headers) is expects_anthropic_beta + assert upstream_headers["Authorization"] == "Bearer vertex-access-token" + assert upstream_headers["content-type"] == "application/json" + + +def test_upstream_headers_for_vertex_route_filters_anthropic_beta_by_route(): + headers = { + "Anthropic-Beta": "effort-2025-11-24", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + + count_tokens_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}count-tokens:rawPredict", headers + ) + model_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}claude-sonnet-4-6:rawPredict", headers + ) + + assert dict(count_tokens_headers) == { + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + assert dict(model_headers) == headers + + @pytest.mark.asyncio async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): """ From b3325750ae4d5b462c7162084c5f74261f458c42 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 11:09:44 -0700 Subject: [PATCH 108/167] fix(ui): aggregate session token usage in the logs table The logs table already rolled up cost per session but the Tokens column only showed the representative call's usage. The per-session aggregate query now also sums prompt, completion and total tokens, and the Tokens cell switches to those sums for multi-call sessions the same way the Cost cell does. Claude-Session: https://claude.ai/code/session_01CNasFqyjnLN3Rqman25vde --- .../spend_management_endpoints.py | 21 ++++- .../test_spend_management_endpoints.py | 85 +++++++++++++++++++ .../RequestLogsTableColumns.test.tsx | 43 +++++++++- .../view_logs/RequestLogsTableColumns.tsx | 17 ++-- .../src/components/view_logs/columns.tsx | 3 + 5 files changed, 160 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 295faaa980a..2a50d5170f0 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -173,6 +173,9 @@ class _SessionSpendRow(TypedDict): session_cache_hit_count: ReadOnly[int] session_llm_count: ReadOnly[int] session_agent_count: ReadOnly[int] + session_total_prompt_tokens: ReadOnly[int] + session_total_completion_tokens: ReadOnly[int] + session_total_tokens: ReadOnly[int] session_models: ReadOnly[Sequence[str]] @@ -188,6 +191,9 @@ class _SessionSpendStats(NamedTuple): session_cache_hit_count: int session_llm_count: int session_agent_count: int + session_total_prompt_tokens: int + session_total_completion_tokens: int + session_total_tokens: int session_models: Sequence[str] session_models_truncated: bool @@ -4287,8 +4293,8 @@ async def _build_ui_spend_logs_response( Build the paginated response for the UI spend-logs endpoint. When ``enrich_session_counts`` is ``True`` (the default for the v1/UI - endpoint), each row is enriched with ``session_total_count`` plus spend - and call-type aggregates so the frontend knows which sessions are + endpoint), each row is enriched with ``session_total_count`` plus spend, + token and call-type aggregates so the frontend knows which sessions are expandable (multi-call sessions). One ``GROUP BY (session_id, api_key)`` query serves every referenced session, keyed per api key so two callers reusing a session id never see each other's totals. Rows without a @@ -4356,7 +4362,10 @@ async def _build_ui_spend_logs_response( COUNT(*) FILTER ( WHERE call_type NOT IN {_MCP_CALL_TYPES_SQL} AND call_type != {_AGENT_CALL_TYPE_SQL} )::int AS session_llm_count, - COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count + COUNT(*) FILTER (WHERE call_type = {_AGENT_CALL_TYPE_SQL})::int AS session_agent_count, + COALESCE(SUM(prompt_tokens), 0)::bigint AS session_total_prompt_tokens, + COALESCE(SUM(completion_tokens), 0)::bigint AS session_total_completion_tokens, + COALESCE(SUM(total_tokens), 0)::bigint AS session_total_tokens FROM "LiteLLM_SpendLogs" WHERE session_id = ANY($1::text[]) AND api_key = ANY($2::text[]) @@ -4389,6 +4398,9 @@ async def _build_ui_spend_logs_response( session_cache_hit_count=int(row.get("session_cache_hit_count") or 0), session_llm_count=int(row.get("session_llm_count") or 0), session_agent_count=int(row.get("session_agent_count") or 0), + session_total_prompt_tokens=int(row.get("session_total_prompt_tokens") or 0), + session_total_completion_tokens=int(row.get("session_total_completion_tokens") or 0), + session_total_tokens=int(row.get("session_total_tokens") or 0), session_models=models[:_SESSION_MODELS_LIMIT], session_models_truncated=len(models) > _SESSION_MODELS_LIMIT, ) @@ -4418,6 +4430,9 @@ async def _build_ui_spend_logs_response( row_dict["session_cache_hit_count"] = session_stats.session_cache_hit_count row_dict["session_llm_count"] = session_stats.session_llm_count row_dict["session_agent_count"] = session_stats.session_agent_count + row_dict["session_total_prompt_tokens"] = session_stats.session_total_prompt_tokens + row_dict["session_total_completion_tokens"] = session_stats.session_total_completion_tokens + row_dict["session_total_tokens"] = session_stats.session_total_tokens row_dict["session_models"] = session_stats.session_models row_dict["session_models_truncated"] = session_stats.session_models_truncated enriched.append(row_dict) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 60a32102946..f4cd8814bc1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -4234,6 +4234,91 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): assert call_args[2] == [api_key] +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens(): + """ + Regression test for LIT-4929: the logs table showed the summed session cost but + only the last call's token usage. Every row of a multi-round session must carry + the session-wide prompt, completion and total token sums from the aggregate + query, while rows outside a session carry none of them. + """ + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + session_id = "sess-multi-round-tokens" + api_key = "hashed-key-xyz" + dict_rows = [ + { + "request_id": "req-1", + "session_id": session_id, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 10, + "prompt_tokens": 7, + "completion_tokens": 3, + }, + { + "request_id": "req-2", + "session_id": session_id, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 50, + "prompt_tokens": 35, + "completion_tokens": 15, + }, + { + "request_id": "req-3", + "session_id": None, + "call_type": "completion", + "api_key": api_key, + "total_tokens": 5, + "prompt_tokens": 4, + "completion_tokens": 1, + }, + ] + + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "session_id": session_id, + "api_key": api_key, + "session_total_count": 2, + "session_total_spend": 0.06, + "mcp_tool_call_count": 0, + "mcp_tool_call_spend": 0.0, + "session_total_prompt_tokens": 42, + "session_total_completion_tokens": 18, + "session_total_tokens": 60, + } + ] + ) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=3, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + rows = result["data"] + session_rows = rows[:2] + assert [row["session_total_tokens"] for row in session_rows] == [60, 60] + assert [row["session_total_prompt_tokens"] for row in session_rows] == [42, 42] + assert [row["session_total_completion_tokens"] for row in session_rows] == [18, 18] + assert [(row["total_tokens"], row["prompt_tokens"], row["completion_tokens"]) for row in session_rows] == [ + (10, 7, 3), + (50, 35, 15), + ] + + token_keys = ("session_total_tokens", "session_total_prompt_tokens", "session_total_completion_tokens") + assert all(key not in rows[2] for key in token_keys) + + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_session_cache_hit_count(): """ diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index e3bacc0908a..d2c84173fd1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -1,4 +1,4 @@ -import { render, screen } from "@testing-library/react"; +import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; @@ -75,6 +75,47 @@ describe("Cost column", () => { }); }); +describe("Tokens column", () => { + it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => { + renderRows([ + logEntry({ + request_id: "req-session-tokens", + total_tokens: 10, + prompt_tokens: 7, + completion_tokens: 3, + session_id: "sess-1", + session_total_count: 3, + session_total_tokens: 60, + session_total_prompt_tokens: 42, + session_total_completion_tokens: 18, + }), + ]); + + const tokensCell = screen.getByText("60").closest("td")!; + expect(within(tokensCell).getByText("(42+18)")).toBeInTheDocument(); + expect(within(tokensCell).getByText("session total")).toBeInTheDocument(); + expect(screen.queryByText("10")).not.toBeInTheDocument(); + expect(screen.queryByText("(7+3)")).not.toBeInTheDocument(); + }); + + it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => { + renderRows([ + logEntry({ + request_id: "req-no-token-aggregate", + total_tokens: 10, + prompt_tokens: 7, + completion_tokens: 3, + session_id: "sess-2", + session_total_count: 3, + }), + ]); + + const tokensCell = screen.getByText("10").closest("td")!; + expect(within(tokensCell).getByText("(7+3)")).toBeInTheDocument(); + expect(within(tokensCell).queryByText("session total")).not.toBeInTheDocument(); + }); +}); + describe("Type column", () => { it("shows the conversation badge and composition even when an MCP call represents the conversation", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index 8db0b106851..9d4dc4f7898 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -263,13 +263,20 @@ export const getRequestLogsTableColumns = ({ meta: { numeric: true }, cell: ({ row }) => { const log = row.original; + const showSessionTotal = (log.session_total_count || 1) > 1 && log.session_total_tokens != null; + const total = showSessionTotal ? log.session_total_tokens : log.total_tokens; + const prompt = showSessionTotal ? log.session_total_prompt_tokens : log.prompt_tokens; + const completion = showSessionTotal ? log.session_total_completion_tokens : log.completion_tokens; return ( - - {String(log.total_tokens || "0")} - - ({String(log.prompt_tokens || "0")}+{String(log.completion_tokens || "0")}) +
+ + {String(total || "0")} + + ({String(prompt || "0")}+{String(completion || "0")}) + - + {showSessionTotal && session total} +
); }, }, diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index 21e09faf454..b2e29c3a0c1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -42,6 +42,9 @@ export type LogEntry = { request_duration_ms?: number; session_total_count?: number; session_total_spend?: number; + session_total_tokens?: number; + session_total_prompt_tokens?: number; + session_total_completion_tokens?: number; session_cache_hit_count?: number; mcp_tool_call_count?: number; mcp_tool_call_spend?: number; From 07c5908b18dd4e4d4d686f6e750e0db029dde738 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 11:27:56 -0700 Subject: [PATCH 109/167] fix(ui): let the Internal Users search box match user_id as well as email MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GET /user/list gains an optional search query param that ORs a case-insensitive contains match over user_id and user_email. The Users page search box now sends that param and reads "Search by email or ID…", the way the Teams page already searches by name or ID. Every existing /user/list param keeps its meaning and the Filters drawer is untouched Claude-Session: https://claude.ai/code/session_018yW93iDaEMhoQUXcYjus7D --- .../internal_user_endpoints.py | 16 +++++ .../internal_user_endpoints.py | 15 ++++- .../test_internal_user_endpoints.py | 61 +++++++++++++++++++ .../users/_components/view_users.test.tsx | 21 +++++++ .../users/_components/view_users.tsx | 9 +-- .../_components/view_users/UsersTable.tsx | 2 +- .../src/components/networking.test.ts | 38 ++++++++++++ .../src/components/networking.tsx | 2 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++ 9 files changed, 162 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 5326074ad3c..73d6761e17f 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -77,6 +77,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, UserListResponse, + UserSearchWhere, UserUpdateResult, ) from litellm.types.proxy.management_endpoints.scim_v2 import ( @@ -2079,6 +2080,10 @@ async def get_users( user_ids: str | None = fastapi.Query(default=None, description="Get list of users by user_ids"), sso_user_ids: str | None = fastapi.Query(default=None, description="Get list of users by sso_user_id"), user_email: str | None = fastapi.Query(default=None, description="Filter users by partial email match"), + search: str | None = fastapi.Query( + default=None, + description="Combined search: matches users whose 'user_id' or 'user_email' contains the value (case-insensitive).", + ), team: str | None = fastapi.Query(default=None, description="Filter users by team id"), page: int = fastapi.Query(default=1, ge=1, description="Page number"), page_size: int = fastapi.Query(default=25, ge=1, le=100, description="Number of items per page"), @@ -2109,6 +2114,8 @@ async def get_users( Get list of users by sso_ids. Comma separated list of sso_ids. user_email: Optional[str] Filter users by partial email match + search: Optional[str] + Combined search: matches users whose user_id or user_email contains the value (case-insensitive) team: Optional[str] Filter users by team id. Will match if user has this team in their teams array. page: int @@ -2168,6 +2175,15 @@ async def get_users( "mode": "insensitive", # Case-insensitive search } + if search: + search_where: Final[UserSearchWhere] = { + "OR": ( + {"user_id": {"contains": search, "mode": "insensitive"}}, + {"user_email": {"contains": search, "mode": "insensitive"}}, + ) + } + where_conditions["OR"] = search_where["OR"] + if team is not None and isinstance(team, str): where_conditions["teams"] = { "has": team # Array contains for string arrays in Prisma diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index df0a090cdb0..6973f1d1f12 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,6 +1,8 @@ -from typing import Any, Final +from collections.abc import Mapping +from typing import Any, Final, Literal from pydantic import BaseModel, field_validator +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._types import ( LiteLLM_UserTableWithKeyCount, @@ -9,6 +11,17 @@ from litellm.proxy._types import ( ) +class InsensitiveContains(TypedDict): + contains: ReadOnly[str] + mode: ReadOnly[Literal["insensitive"]] + + +class UserSearchWhere(TypedDict): + """Prisma filter behind `/user/list?search=`: user_id or user_email contains the term, case-insensitive.""" + + OR: ReadOnly[tuple[Mapping[Literal["user_id", "user_email"], InsensitiveContains], ...]] + + class UserListResponse(BaseModel): """ Response model for the user list endpoint diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index f231eb66a50..d02b42eb0bf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2007,6 +2007,67 @@ async def test_get_users_user_id_partial_match(mocker): assert captured_where_conditions["user_id"]["in"] == ["user1", "user2", "user3"] +def test_get_users_search_matches_user_id_or_email(mocker): + """ + `search` ORs a case-insensitive contains match over user_id and user_email on both the rows + query and the count, while the legacy `user_email` param keeps filtering only user_email. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + searched_user_id = "a6f5c02b-0163-45ce-815f-f88d10e95686" + mock_user_row = mocker.MagicMock() + mock_user_row.user_id = searched_user_id + mock_user_row.model_dump.return_value = { + "user_id": searched_user_id, + "user_email": "search@example.com", + "user_role": "internal_user", + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + find_many_wheres = [] + count_wheres = [] + + async def mock_find_many(*args, **kwargs): + find_many_wheres.append(kwargs["where"]) + return [mock_user_row] + + async def mock_count(*args, **kwargs): + count_wheres.append(kwargs["where"]) + return 1 + + async def mock_key_count(*args, **kwargs): + return 0 + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many + mock_prisma_client.db.litellm_usertable.count = mock_count + mock_prisma_client.db.litellm_verificationtoken.count = mock_key_count + mocker.patch( # test-quality-ok: /user/list reads prisma_client off proxy_server at call time + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + search_response = client.get("/user/list", params={"search": "A6F5C02B-0163"}) + assert search_response.status_code == 200, search_response.text + expected_or = ( + {"user_id": {"contains": "A6F5C02B-0163", "mode": "insensitive"}}, + {"user_email": {"contains": "A6F5C02B-0163", "mode": "insensitive"}}, + ) + assert find_many_wheres == [{"OR": expected_or}] + assert count_wheres == [{"OR": expected_or}] + assert [user["user_id"] for user in search_response.json()["users"]] == [searched_user_id] + assert search_response.json()["total"] == 1 + + legacy_response = client.get("/user/list", params={"user_email": "search@example.com"}) + assert legacy_response.status_code == 200, legacy_response.text + assert find_many_wheres[-1] == {"user_email": {"contains": "search@example.com", "mode": "insensitive"}} + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + def test_update_internal_user_params_reset_max_budget_with_none(): """ Test that _update_internal_user_params allows setting max_budget to None. diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 095bdbf7250..6b0423067fd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -323,5 +323,26 @@ describe("ViewUserDashboard", () => { expect(latest[2]).toBe(1); }); }); + + it("sends the toolbar search as the combined search param instead of user_email", async () => { + const user = userEvent.setup(); + renderDashboard(); + + await waitFor(() => { + expect(screen.getByText("test@example.com")).toBeInTheDocument(); + }); + + const searchedUserId = "a6f5c02b-0163-45ce-815f-f88d10e95686"; + await user.type(screen.getByPlaceholderText("Search by email or ID…"), searchedUserId); + + await waitFor(() => { + const latest = userListCall.mock.calls[userListCall.mock.calls.length - 1]; + expect(latest[11]).toBe(searchedUserId); + }); + const latest = userListCall.mock.calls[userListCall.mock.calls.length - 1]; + expect(latest[1]).toBeNull(); + expect(latest[4]).toBeNull(); + expect(latest[2]).toBe(1); + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 7a3133840be..1ed23e7f523 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -65,7 +65,7 @@ const ViewUserDashboard: React.FC = ({ const [sorting, setSorting] = useState(DEFAULT_SORTING); const [columnFilters, setColumnFilters] = useState([]); const [searchInput, setSearchInput] = useState(""); - const [searchEmail] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); + const [searchQuery] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); const [rowSelection, setRowSelection] = useState({}); const [selectionMode, setSelectionMode] = useState(false); @@ -222,12 +222,12 @@ const ViewUserDashboard: React.FC = ({ const ssoUserIdFilter = getFilterValue("sso_user_id"); const userRoleFilter = getFilterValue("user_role"); const teamFilter = getFilterValue("team"); - const emailFilter = searchEmail.trim() || null; + const searchFilter = searchQuery.trim() || null; const userListQueryFilters = { page: pagination.pageIndex + 1, pageSize: pagination.pageSize, - email: emailFilter, + search: searchFilter, userId: userIdFilter, ssoUserId: ssoUserIdFilter, role: userRoleFilter, @@ -247,13 +247,14 @@ const ViewUserDashboard: React.FC = ({ userIdFilter ? [userIdFilter] : null, pagination.pageIndex + 1, pagination.pageSize, - emailFilter, + null, userRoleFilter ?? null, teamFilter ?? null, ssoUserIdFilter ?? null, sortBy, sortOrder, orgAdminOrgIds ? orgAdminOrgIds.map((o) => o.organization_id) : null, + searchFilter, ); }, enabled: Boolean(accessToken && token && userRole && userID), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx index 26663f5d9e2..875df74652c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/UsersTable.tsx @@ -158,7 +158,7 @@ export function UsersTable({ table={table} searchValue={searchValue} onSearchChange={onSearchChange} - searchPlaceholder="Search by email…" + searchPlaceholder="Search by email or ID…" onOpenFilters={() => setFiltersOpen(true)} filterLabels={FILTER_LABELS} formatFilterValue={formatFilterValue} diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 335a5f8b816..8dcd8c39d9d 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -815,3 +815,41 @@ describe("daily activity api_key filter", () => { expect(requestedUrl(mockFetch)).toContain("user_id="); }); }); + +describe("userListCall search serialization", () => { + const originalFetch = global.fetch; + + afterEach(() => { + global.fetch = originalFetch; + }); + + const mockOkFetch = () => { + const body = JSON.stringify({ users: [], total: 0, page: 1, page_size: 25, total_pages: 0 }); + const mockFetch = vi.fn().mockResolvedValue({ ok: true, text: vi.fn().mockResolvedValue(body) } as any); + global.fetch = mockFetch as any; + return mockFetch; + }; + + const lastParams = (mockFetch: ReturnType) => { + const [url] = mockFetch.mock.calls.at(-1) ?? []; + return new URL(url as string, "http://example.com").searchParams; + }; + + it("sends the combined search term as search, not user_email", async () => { + const mockFetch = mockOkFetch(); + + await Networking.userListCall("token", null, 1, 25, null, null, null, null, null, null, null, "a6f5c02b"); + + expect(lastParams(mockFetch).get("search")).toBe("a6f5c02b"); + expect(lastParams(mockFetch).has("user_email")).toBe(false); + }); + + it("omits search when no search term is given and keeps user_email as before", async () => { + const mockFetch = mockOkFetch(); + + await Networking.userListCall("token", null, 1, 25, "ada@example.com"); + + expect(lastParams(mockFetch).has("search")).toBe(false); + expect(lastParams(mockFetch).get("user_email")).toBe("ada@example.com"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 7e0f6c7e4f5..a40b7704674 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1033,6 +1033,7 @@ export const userListCall = async ( sortBy: string | null = null, sortOrder: "asc" | "desc" | null = null, organizationIds: string[] | null = null, + search: string | null = null, ) => { /** * Get all available teams on proxy @@ -1051,6 +1052,7 @@ export const userListCall = async ( sort_by: sortBy || undefined, sort_order: sortOrder || undefined, organization_ids: organizationIds && organizationIds.length > 0 ? organizationIds.join(",") : undefined, + search: search || undefined, }, })) as UserListResponse; return data; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 098b43f6433..8714570d5c4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16606,6 +16606,8 @@ export interface paths { * Get list of users by sso_ids. Comma separated list of sso_ids. * user_email: Optional[str] * Filter users by partial email match + * search: Optional[str] + * Combined search: matches users whose user_id or user_email contains the value (case-insensitive) * team: Optional[str] * Filter users by team id. Will match if user has this team in their teams array. * page: int @@ -59835,6 +59837,8 @@ export interface operations { sso_user_ids?: string | null; /** @description Filter users by partial email match */ user_email?: string | null; + /** @description Combined search: matches users whose 'user_id' or 'user_email' contains the value (case-insensitive). */ + search?: string | null; /** @description Filter users by team id */ team?: string | null; /** @description Page number */ From 3b13a5fda968620d699d5f8cc15d7485f118bc5e Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 18:31:43 +0000 Subject: [PATCH 110/167] test(ui): query the tokens cell by role instead of walking the DOM Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../RequestLogsTableColumns.test.tsx | 56 +++++++++---------- 1 file changed, 25 insertions(+), 31 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index d2c84173fd1..3c9e6543c1c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -1,4 +1,4 @@ -import { render, screen, within } from "@testing-library/react"; +import { render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; @@ -76,43 +76,37 @@ describe("Cost column", () => { }); describe("Tokens column", () => { - it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => { - renderRows([ - logEntry({ - request_id: "req-session-tokens", - total_tokens: 10, - prompt_tokens: 7, - completion_tokens: 3, - session_id: "sess-1", - session_total_count: 3, - session_total_tokens: 60, - session_total_prompt_tokens: 42, - session_total_completion_tokens: 18, - }), - ]); + const sessionRow: Partial = { + request_id: "req-session-tokens", + total_tokens: 10, + prompt_tokens: 7, + completion_tokens: 3, + session_id: "sess-1", + session_total_count: 3, + }; - const tokensCell = screen.getByText("60").closest("td")!; - expect(within(tokensCell).getByText("(42+18)")).toBeInTheDocument(); - expect(within(tokensCell).getByText("session total")).toBeInTheDocument(); + it("shows the summed session token usage, not the representative call's tokens, for a multi-round session", () => { + const aggregatedRow: Partial = { + ...sessionRow, + session_total_tokens: 60, + session_total_prompt_tokens: 42, + session_total_completion_tokens: 18, + }; + renderRows([logEntry(aggregatedRow)]); + + const tokensCell = screen.getByRole("cell", { name: /\(42\+18\)/ }); + expect(tokensCell).toHaveTextContent("60"); + expect(tokensCell).toHaveTextContent("session total"); expect(screen.queryByText("10")).not.toBeInTheDocument(); expect(screen.queryByText("(7+3)")).not.toBeInTheDocument(); }); it("falls back to the call's own tokens with no session label when the backend sent no session token sums", () => { - renderRows([ - logEntry({ - request_id: "req-no-token-aggregate", - total_tokens: 10, - prompt_tokens: 7, - completion_tokens: 3, - session_id: "sess-2", - session_total_count: 3, - }), - ]); + renderRows([logEntry(sessionRow)]); - const tokensCell = screen.getByText("10").closest("td")!; - expect(within(tokensCell).getByText("(7+3)")).toBeInTheDocument(); - expect(within(tokensCell).queryByText("session total")).not.toBeInTheDocument(); + const tokensCell = screen.getByRole("cell", { name: /\(7\+3\)/ }); + expect(tokensCell).toHaveTextContent("10"); + expect(tokensCell).not.toHaveTextContent("session total"); }); }); From 4d2ffe2e8e8dd1769a8fd2f1a3ebc857a4a3701c Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 18:38:22 +0000 Subject: [PATCH 111/167] test(ui): hoist inherited-grant fixture out of the inline createMockTeamData arg Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/components/team/TeamInfo.test.tsx | 32 +++++++++---------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index ae78ac06a9c..a9c1077e96b 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -313,23 +313,21 @@ describe("TeamInfoView", () => { vi.mocked(networking.getAgentsList).mockResolvedValue({ agents: [{ agent_id: "agent-support-5678", agent_name: "support_agent" }], }); - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - object_permission: null, - access_group_ids: ["ag-1"], - access_group_mcp_server_ids: ["mcp-github-1234"], - access_group_agent_ids: ["agent-support-5678"], - access_group_details: [ - { - access_group_id: "ag-1", - access_group_name: "platform-tools", - models: [], - mcp_server_ids: ["mcp-github-1234"], - agent_ids: ["agent-support-5678"], - }, - ], - }), - ); + const platformToolsGroup = { + access_group_id: "ag-1", + access_group_name: "platform-tools", + models: [], + mcp_server_ids: ["mcp-github-1234"], + agent_ids: ["agent-support-5678"], + }; + const inheritedGrants = { + object_permission: null, + access_group_ids: ["ag-1"], + access_group_mcp_server_ids: ["mcp-github-1234"], + access_group_agent_ids: ["agent-support-5678"], + access_group_details: [platformToolsGroup], + }; + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData(inheritedGrants)); renderWithProviders(); From f87b9097eaffdde471df332c323a1a59be9c2014 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:50:51 -0700 Subject: [PATCH 112/167] test(bedrock): drop EOL cohere.command-r-plus-v1:0 from local_testing Bedrock retired cohere.command-r-plus-v1:0 on 2026-08-19 and lists no Cohere command chat model anymore, so the three local_testing cases that pinned it fail with a 404 end-of-life error on every pipeline. Drop the case from test_completion_bedrock_httpx_models and move the parallel-streaming Bedrock entry to mistral.mistral-7b-instruct-v0:2, which still takes the invoke route and is ACTIVE in the CI account. --- tests/local_testing/test_completion.py | 1 - tests/local_testing/test_streaming.py | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index ef8d6c55148..f8f23ea015a 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -2879,7 +2879,6 @@ def response_format_tests(response: litellm.ModelResponse): "model", [ "bedrock/mistral.mistral-large-2407-v1:0", - "bedrock/cohere.command-r-plus-v1:0", "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "mistral.mistral-7b-instruct-v0:2", "meta.llama3-8b-instruct-v1:0", diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 07d693af447..bf39d3155b7 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -1168,7 +1168,6 @@ async def test_completion_replicate_llama3_streaming(sync_mode): "model, region", [ # ["bedrock/ai21.jamba-instruct-v1:0", "us-east-1"], - # ["bedrock/cohere.command-r-plus-v1:0", None], ["us.anthropic.claude-sonnet-4-5-20250929-v1:0", None], # ["mistral.mistral-7b-instruct-v0:2", None], # ["meta.llama3-8b-instruct-v1:0", None], @@ -1271,7 +1270,7 @@ def test_bedrock_claude_3_streaming(): "model", [ "claude-haiku-4-5-20251001", - "cohere.command-r-plus-v1:0", # bedrock + "bedrock/mistral.mistral-7b-instruct-v0:2", "gpt-3.5-turbo", ], ) From c2265b0ef3569b72767359c33516a91aef7db7fb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:55:59 -0700 Subject: [PATCH 113/167] fix(proxy): return persisted team memberships from /user/new so first CLI login gets the default team (#39545) * fix(proxy): return persisted team memberships from /user/new new_user attached default teams after building its response from the pre-membership snapshot, so NewUserResponse.teams was always empty for users created with default_internal_user_params.teams. The CLI SSO flow reads that response on a user's first login and minted a teamless JWT, which skipped the default team's model allowlist. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): return team ids as a tuple to satisfy LIT001 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/internal_user_endpoints.py | 12 ++++++++++++ .../test_internal_user_endpoints.py | 7 +++++++ 2 files changed, 19 insertions(+) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 5326074ad3c..93423ca5a1a 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -426,6 +426,11 @@ async def add_new_user_to_default_team( await asyncio.gather(*tasks, return_exceptions=True) +async def _fetch_user_team_ids(user_id: str, prisma_client: "PrismaClient") -> tuple[str, ...]: + user_row: Final = await _user_table(prisma_client).find_unique(where={"user_id": user_id}) + return tuple(user_row.teams) if user_row is not None else () + + @router.post( "/user/new", tags=["Internal User management"], @@ -580,6 +585,11 @@ async def new_user( ) user_id: Final = cast(str | None, response.get("user_id", None)) + attached_team_ids: Final = ( + await _fetch_user_team_ids(user_id=user_id, prisma_client=prisma_client) + if user_id is not None and (_team_id is not None or teams is not None) + else None + ) if organization_ids is not None and user_id is not None: await _add_user_to_organizations( @@ -596,6 +606,8 @@ async def new_user( response_dict[key] = value response_dict["key"] = response.get("token", "") + if attached_team_ids is not None: + response_dict["teams"] = list(attached_team_ids) new_user_response: Final = NewUserResponse.model_validate(response_dict) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index f231eb66a50..022aeff4e20 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1449,6 +1449,11 @@ async def test_new_user_default_teams_flow(mocker): return 5 # Low user count, under limit mock_prisma_client.db.litellm_usertable.count = mock_count + persisted_user_row = mocker.MagicMock() + persisted_user_row.teams = ["96fed65b-0182-4ff4-8429-2721cd7d42af"] + mock_prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock( + return_value=persisted_user_row + ) # Mock duplicate checks to pass async def mock_check_duplicate_user_email(*args, **kwargs): @@ -1477,6 +1482,7 @@ async def test_new_user_default_teams_flow(mocker): "token": "sk-test-token-123", "expires": None, "max_budget": 100, + "teams": [], } # Mock _add_user_to_team @@ -1551,6 +1557,7 @@ async def test_new_user_default_teams_flow(mocker): # Verify response structure assert response.user_id == "test-user-123" assert response.key == "sk-test-token-123" + assert response.teams == ["96fed65b-0182-4ff4-8429-2721cd7d42af"] finally: # Restore original default params (always assign, never delattr — the attribute From e046aee3d52e2308d94399b033893fe77674ee50 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:57:56 -0700 Subject: [PATCH 114/167] fix(spend_tracking): add missing_session_id: omit to leave SpendLogs.session_id null without a client session (#39458) * fix(spend_tracking): leave SpendLogs.session_id null when no client session id was established Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(lint): ratchet basedpyright budget after session_id fix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): ignore trace ids as session ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): gate null SpendLogs.session_id behind missing_session_id: omit Unset, generate and reject keep the legacy trace id fallback. omit records only metadata.session_id, the key Langfuse reads, so a trace id copied into litellm_session_id by get_litellm_params never becomes a session. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): stamp the omit decision on the request so a config reload cannot fabricate a session Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): keep omit covering requests the pre-call stamp never reaches Router-model provider pass-through calls allm_passthrough_route directly and skips add_litellm_data_to_request, so those requests never run the pre-call helper and carry no omit stamp. Reading only the stamp made POST /anthropic/v1/messages write a fabricated uuid into SpendLogs.session_id under missing_session_id: omit while its Langfuse trace had no session, the exact divergence the policy exists to remove. The stamp now only pins omit on, and an unstamped request falls back to the configured policy, so a config reload still cannot fabricate a session for a request that was decided pre-call. * fix(spend_tracking): make the session-omission marker proxy-owned so clients cannot forge it * fix(spend_tracking): strip the client-sent omission marker from both metadata buckets before they merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): strip the session-omission marker from both metadata buckets The pre-call policy ran before litellm_metadata is merged into metadata, so a client that planted the marker in litellm_metadata had it copied back into the route's own bucket after the strip and still got a null SpendLogs.session_id. --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- basedpyright-code-budget.json | 6 +- litellm/constants.py | 1 + litellm/proxy/_types.py | 4 +- .../proxy/hooks/proxy_track_cost_callback.py | 7 +- litellm/proxy/litellm_pre_call_utils.py | 12 +- .../pass_through_endpoints.py | 6 +- .../spend_tracking/spend_tracking_utils.py | 38 +- .../test_pass_through_endpoints.py | 35 + .../test_spend_tracking_utils.py | 477 ++++++------- .../proxy/test_litellm_pre_call_utils.py | 649 ++++++------------ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 11 files changed, 519 insertions(+), 720 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 2967fc2a505..a37cb194757 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15288 + "limit": 15287 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38324 + "limit": 38323 }, "reportUnknownParameterType": { - "limit": 19625 + "limit": 19624 }, "reportUnknownVariableType": { "limit": 29861 diff --git a/litellm/constants.py b/litellm/constants.py index ef9329b9dfc..be13d9aac5f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1450,6 +1450,7 @@ SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affin CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" +SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 97f5f59d2dc..0aea72be1e2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2608,9 +2608,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", ) - missing_session_id: Literal["generate", "reject"] | None = Field( + missing_session_id: Literal["generate", "reject", "omit"] | None = Field( None, - description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", + description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.", ) enable_public_model_hub: bool = Field( default=False, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 47aafda2337..7254b05db2e 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -168,11 +168,8 @@ class _ProxyDBLogger(CustomLogger): "custom_llm_provider" ) or request_data.get("custom_llm_provider", "") - # Propagate standard_logging_object and litellm_trace_id from the - # Logging instance so that _get_session_id_for_spend_log uses the same - # trace_id that Langfuse received (via async_failure_handler). - # Without this, the DB session_id would be a random UUID that doesn't - # match the Langfuse trace_id, making failed requests unsearchable. + # Propagate standard_logging_object and litellm_trace_id from the Logging + # instance so the failure row carries the same trace_id Langfuse received. _litellm_logging_obj: Final = request_data.get("litellm_logging_obj") if _litellm_logging_obj is not None: if not request_data.get("standard_logging_object"): diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1d440448c2f..f752d7cfa89 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -733,12 +734,18 @@ def apply_missing_session_id_policy( general_settings: Mapping[str, object] | None, request: Request, ) -> None: + for metadata_key in ("metadata", "litellm_metadata"): + if isinstance(client_metadata := data.get(metadata_key), dict): + client_metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None) + metadata: Final = data.get(_metadata_variable_name) policy: Final = general_settings.get("missing_session_id") if general_settings else None if policy is None or not _is_llm_inference_route(request): return - metadata: Final = data.get(_metadata_variable_name) if not isinstance(metadata, dict): return + if policy == "omit": + metadata[SESSION_ID_OMITTED_METADATA_KEY] = True + return if data.get("litellm_session_id") or metadata.get("session_id"): return match policy: @@ -760,7 +767,8 @@ def apply_missing_session_id_policy( ) case _: verbose_proxy_logger.warning( - "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy + "Ignoring unknown general_settings.missing_session_id=%r; expected 'generate', 'reject' or 'omit'", + policy, ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 79d5d0a016f..323756bf204 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -40,6 +40,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( MAXIMUM_TRACEBACK_LINES_TO_LOG, + SESSION_ID_OMITTED_METADATA_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) from litellm.integrations.custom_guardrail import CustomGuardrail @@ -581,8 +582,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ) # Set internal keys after merging client-supplied metadata so a request - # body that mirrors them cannot clobber the authenticated key or the - # real parent span. + # body that mirrors them cannot clobber the authenticated key, the real + # parent span, or the proxy's own session-id decision. + _metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None) _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 7442d71bd96..a37c3ba4405 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,6 +15,7 @@ from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, REDACTED_BY_LITELM_STRING, + SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, @@ -578,7 +579,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs ), session_id=_get_session_id_for_spend_log( kwargs=kwargs, + metadata=metadata, standard_logging_payload=standard_logging_payload, + omit_when_missing=_omits_session_id_when_missing(metadata), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -602,26 +605,39 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs raise e +def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> bool: + """The pre-call stamp pins `omit` on for the requests that carry it, so a config reload between pre-call and spend + logging cannot fabricate a session. `apply_missing_session_id_policy` drops any client-supplied copy of the key + from both metadata buckets before stamping, which the merge of `litellm_metadata` into `metadata` makes + necessary, so a caller cannot forge it. Requests that never reach the pre-call helper, router-model + passthrough among them, carry no stamp, so they fall back to the configured policy and `omit` still covers their + spend logs.""" + if metadata is not None and metadata.get(SESSION_ID_OMITTED_METADATA_KEY): + return True + + from litellm.proxy.proxy_server import general_settings + + return general_settings.get("missing_session_id") == "omit" + + def _get_session_id_for_spend_log( - kwargs: dict, + kwargs: Mapping[str, object], + metadata: Mapping[str, object] | None, standard_logging_payload: StandardLoggingPayload | None, -) -> str: - """ - Get the session id for the spend log. + omit_when_missing: bool, +) -> str | None: + """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may + be a copied trace id.""" + if omit_when_missing: + session_id: Final = metadata.get("session_id") if metadata else None + return str(session_id) if session_id else None - This ensures each spend log is associated with a unique session id. - - """ from litellm._uuid import uuid if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: return str(standard_logging_payload.get("trace_id")) - - # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs if kwargs.get("litellm_trace_id") is not None: return str(kwargs.get("litellm_trace_id")) - - # Ensure we always have a session id, if none is provided return str(uuid.uuid4()) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d3f17c73499..91367507247 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5730,3 +5730,38 @@ async def test_pass_through_request_leaves_cost_router_logger_working(): verbose_logger.removeHandler(recorder) assert not raised, f"cost router logger raised on the passthrough logging params: {raised[0].exc_info}" + + +@pytest.mark.parametrize("client_metadata_key", ["litellm_metadata", "metadata"]) +def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key: str): + """The omit marker is proxy-owned: only the pre-call policy may set it. A pass-through body that carries + it in its own metadata must not null out SpendLogs.session_id on a request the proxy never omitted.""" + from litellm.constants import SESSION_ID_OMITTED_METADATA_KEY + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://0.0.0.0:4000/gemini/v1beta/models/gemini-2.5-flash:generateContent" + mock_request.headers = Headers({}) + mock_request.scope = {} + + kwargs = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + passthrough_logging_payload=MagicMock(), + logging_obj=MagicMock(), + _parsed_body={client_metadata_key: {SESSION_ID_OMITTED_METADATA_KEY: True}}, + litellm_call_id="lit-6694-call-id", + ) + + metadata = kwargs["litellm_params"]["metadata"] + assert SESSION_ID_OMITTED_METADATA_KEY not in metadata + assert ( + _get_session_id_for_spend_log( + kwargs={}, + metadata=metadata, + standard_logging_payload={"trace_id": "per-call-random-trace-id"}, + omit_when_missing=bool(metadata.get(SESSION_ID_OMITTED_METADATA_KEY)), + ) + == "per-call-random-trace-id" + ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9e5917637a8..323930eee60 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,6 +14,7 @@ from litellm.constants import ( LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, REDACTED_BY_LITELM_STRING, + SESSION_ID_OMITTED_METADATA_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.spend_tracking.spend_tracking_utils import ( @@ -21,6 +22,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_proxy_server_request_for_spend_logs_payload, _get_request_duration_ms, _get_response_for_spend_logs_payload, + _get_session_id_for_spend_log, _get_spend_logs_metadata, _get_vector_store_request_for_spend_logs_payload, _is_master_key, @@ -33,6 +35,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, get_spend_logs_id, ) +from litellm.proxy._types import SpendLogsPayload from litellm.proxy.utils import hash_token from litellm.types.utils import ( StandardLoggingHiddenParams, @@ -74,6 +77,110 @@ def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_token assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123 +_TRACE_ONLY_STANDARD_LOGGING: Final = cast( + StandardLoggingPayload, + { + "trace_id": "trace-abc", + "session_id": "trace-abc", + "metadata": {}, + "model_map_information": None, + "request_tags": [], + }, +) + + +def _trace_only_session_id(omit_when_missing: bool) -> str | None: + """get_litellm_params copies metadata.trace_id into litellm_session_id, so every field echoes the trace id.""" + return _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc", "litellm_session_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=omit_when_missing, + ) + + +def test_omit_leaves_session_id_none_when_only_a_trace_id_exists(): + assert _trace_only_session_id(omit_when_missing=True) is None + + +def test_omit_leaves_session_id_none_without_any_ids(): + assert ( + _get_session_id_for_spend_log(kwargs={}, metadata=None, standard_logging_payload=None, omit_when_missing=True) + is None + ) + + +def test_omit_records_metadata_session_id(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_session_id": "chain-1"}, + metadata={"trace_id": "chain-1", "session_id": "chain-1"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=True, + ) + assert session_id == "chain-1" + + +def test_legacy_policy_keeps_trace_id_fallback(): + assert _trace_only_session_id(omit_when_missing=False) == "trace-abc" + generated: Final = _get_session_id_for_spend_log( + kwargs={}, metadata=None, standard_logging_payload=None, omit_when_missing=False + ) + assert len(str(generated)) == 36 + + +@pytest.mark.parametrize( + ("request_metadata", "expected"), + [ + ({"trace_id": "trace-abc"}, "trace-abc"), + ({"trace_id": "trace-abc", SESSION_ID_OMITTED_METADATA_KEY: True}, None), + ({"trace_id": "trace-abc", "session_id": "chain-1", SESSION_ID_OMITTED_METADATA_KEY: True}, "chain-1"), + ], +) +def test_get_logging_payload_reads_omit_decision_stamped_on_request( + request_metadata: dict[str, object], expected: str | None +): + """The pre-call stamp, not the live general_settings, decides the policy, so a config reload between + pre-call and spend logging cannot fabricate a session for a request accepted under `omit`.""" + with patch( # test-quality-ok: proves log time ignores proxy config; general_settings is yaml, not an HTTP boundary + "litellm.proxy.proxy_server.general_settings", {"missing_session_id": "generate"} + ): + payload: SpendLogsPayload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_trace_id": "trace-abc", + "litellm_params": {"litellm_session_id": "trace-abc", "metadata": request_metadata}, + "standard_logging_object": _TRACE_ONLY_STANDARD_LOGGING, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[]), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["session_id"] == expected + + +@pytest.mark.parametrize("policy", ["omit", "generate", None]) +def test_get_logging_payload_applies_omit_to_requests_that_carry_no_stamp(policy: str | None): + """Router-model passthrough calls `allm_passthrough_route` directly and never reaches the pre-call helper that + stamps the omit decision, so an unstamped request falls back to the configured policy. Without that fallback + `missing_session_id: omit` would fabricate a uuid session id on every passthrough spend log while its Langfuse + trace has none, which is the divergence the policy exists to remove.""" + with patch( # test-quality-ok: general_settings is proxy config, loaded from yaml, not an HTTP boundary + "litellm.proxy.proxy_server.general_settings", {} if policy is None else {"missing_session_id": policy} + ): + payload: SpendLogsPayload = get_logging_payload( + kwargs={ + "model": "claude-opus-4", + "litellm_trace_id": "trace-abc", + "litellm_params": {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc"}}, + "standard_logging_object": _TRACE_ONLY_STANDARD_LOGGING, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-test", choices=[]), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["session_id"] == (None if policy == "omit" else "trace-abc") + + def test_get_logging_payload_preserves_anthropic_cache_read_input_tokens(): additional_usage_values = _get_additional_usage_values_for_usage( litellm.Usage( @@ -277,9 +384,7 @@ def test_sanitize_request_body_for_spend_logs_payload_long_string(): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB (2048) - long_string = ( - "a" * 3000 - ) # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB + long_string = "a" * 3000 # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB request_body = {"text": long_string, "normal_text": "short text"} sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) @@ -329,9 +434,7 @@ def test_sanitize_request_body_for_spend_logs_payload_nested_list(): # Create a string longer than MAX_STRING_LENGTH_PROMPT_IN_DB long_string = "a" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500) - request_body = { - "items": [{"text": long_string}, {"text": "short"}, [{"text": long_string}]] - } + request_body = {"items": [{"text": long_string}, {"text": "short"}, [{"text": long_string}]]} sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) # Calculate expected lengths based on actual MAX_STRING_LENGTH_PROMPT_IN_DB @@ -415,14 +518,10 @@ def test_sanitize_request_body_for_spend_logs_payload_circular_reference(): # Test that it handles circular reference without infinite recursion sanitized = _sanitize_request_body_for_spend_logs_payload(a) - assert sanitized == { - "b": {"a": {}} - } # Should return empty dict for circular reference + assert sanitized == {"b": {"a": {}}} # Should return empty dict for circular reference -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( mock_should_store, ): @@ -431,27 +530,16 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_true( # Sample vector store request metadata vector_store_request = [ - { - "vector_store_search_response": { - "data": [ - {"content": [{"text": "sensitive information", "type": "text"}]} - ] - } - } + {"vector_store_search_response": {"data": [{"content": [{"text": "sensitive information", "type": "text"}]}]}} ] # When store_prompts is True, the original data should be returned unchanged result = _get_vector_store_request_for_spend_logs_payload(vector_store_request) assert result == vector_store_request - assert ( - result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] - == "sensitive information" - ) + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == "sensitive information" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( mock_should_store, ): @@ -460,32 +548,18 @@ def test_get_vector_store_request_for_spend_logs_payload_store_prompts_false( # Sample vector store request metadata vector_store_request = [ - { - "vector_store_search_response": { - "data": [ - {"content": [{"text": "sensitive information", "type": "text"}]} - ] - } - } + {"vector_store_search_response": {"data": [{"content": [{"text": "sensitive information", "type": "text"}]}]}} ] # When store_prompts is False, text should be redacted result = _get_vector_store_request_for_spend_logs_payload(vector_store_request) assert result is not None - assert ( - result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] - == REDACTED_BY_LITELM_STRING - ) + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["text"] == REDACTED_BY_LITELM_STRING # Ensure other fields are unchanged - assert ( - result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] - == "text" - ) + assert result[0]["vector_store_search_response"]["data"][0]["content"][0]["type"] == "text" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_store): # When input is None mock_should_store.return_value = False @@ -493,9 +567,7 @@ def test_get_vector_store_request_for_spend_logs_payload_null_input(mock_should_ assert result is None -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns messages @@ -522,9 +594,7 @@ def test_get_messages_for_spend_logs_realtime_returns_messages(mock_should_store assert parsed[1]["content"] == "What is the weather today?" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from messages.""" mock_should_store.return_value = True @@ -541,9 +611,7 @@ def test_get_messages_for_spend_logs_strips_null_bytes(mock_should_store): assert parsed[0]["content"] == "helloworld" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_realtime_empty_when_disabled(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns '{}' for realtime calls @@ -561,9 +629,7 @@ def test_get_messages_for_spend_logs_realtime_empty_when_disabled(mock_should_st assert result == "{}" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_messages_for_spend_logs_non_realtime_returns_empty(mock_should_store): """ Test that _get_messages_for_spend_logs_payload returns '{}' for non-realtime @@ -581,9 +647,7 @@ def test_get_messages_for_spend_logs_non_realtime_returns_empty(mock_should_stor assert result == "{}" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_store): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB @@ -611,9 +675,7 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" mock_should_store.return_value = True @@ -626,18 +688,14 @@ def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store assert json.loads(response_json)["content"] == "answerhere" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_truncates_large_embedding( mock_should_store, ): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB mock_should_store.return_value = True - embedding_values = [ - round(i * 0.0001, 6) for i in range(MAX_STRING_LENGTH_PROMPT_IN_DB + 500) - ] + embedding_values = [round(i * 0.0001, 6) for i in range(MAX_STRING_LENGTH_PROMPT_IN_DB + 500)] large_embedding = json.dumps(embedding_values) payload = cast( StandardLoggingPayload, @@ -685,9 +743,7 @@ def test_truncation_includes_db_safeguard_note(): ) -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_response_truncation_logs_info_message(mock_should_store): """ Test that when response is truncated before DB storage, an info log is emitted @@ -702,18 +758,14 @@ def test_response_truncation_logs_info_message(mock_should_store): {"response": {"data": [{"content": large_text}]}}, ) - with patch( - "litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger" - ) as mock_logger: + with patch("litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger") as mock_logger: _get_response_for_spend_logs_payload(payload) mock_logger.info.assert_called_once() log_msg = mock_logger.info.call_args[0][0] assert "response was truncated" in log_msg -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_request_body_truncation_logs_info_message(mock_should_store): """ Test that when request body is truncated before DB storage, an info log is emitted. @@ -722,18 +774,10 @@ def test_request_body_truncation_logs_info_message(mock_should_store): mock_should_store.return_value = True large_prompt = "C" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500) - litellm_params = { - "proxy_server_request": { - "body": {"messages": [{"role": "user", "content": large_prompt}]} - } - } + litellm_params = {"proxy_server_request": {"body": {"messages": [{"role": "user", "content": large_prompt}]}}} - with patch( - "litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger" - ) as mock_logger: - _get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params=litellm_params, kwargs={} - ) + with patch("litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger") as mock_logger: + _get_proxy_server_request_for_spend_logs_payload(metadata={}, litellm_params=litellm_params, kwargs={}) mock_logger.info.assert_called_once() log_msg = mock_logger.info.call_args[0][0] assert "request body was truncated" in log_msg @@ -870,14 +914,10 @@ def test_get_logging_payload_api_key_preserved_when_standard_logging_payload_is_ ) # The api_key should be hashed (not the raw key) - assert ( - payload["api_key"] != test_api_key - ), "api_key should be hashed, not the raw key" + assert payload["api_key"] != test_api_key, "api_key should be hashed, not the raw key" # The api_key should be a valid hash (64 character hex string for SHA256) - assert ( - len(payload["api_key"]) == 64 - ), f"Expected 64 character hash, got {len(payload['api_key'])} characters" + assert len(payload["api_key"]) == 64, f"Expected 64 character hash, got {len(payload['api_key'])} characters" # Verify other fields are set correctly assert payload["model"] == "openai/gpt-4.1" @@ -1019,9 +1059,7 @@ async def test_api_key_preserved_through_failure_hook_to_database(): assert payload_api_key is not None, "🚨 CRITICAL: payload['api_key'] is None!" - assert ( - payload_api_key == hashed_key - ), f"🚨 CRITICAL: Expected api_key={hashed_key}, got {payload_api_key}" + assert payload_api_key == hashed_key, f"🚨 CRITICAL: Expected api_key={hashed_key}, got {payload_api_key}" # Verify token parameter matches assert data["token"] == hashed_key, f"Token parameter should be {hashed_key}" @@ -1066,9 +1104,7 @@ def test_get_logging_payload_includes_agent_id_from_kwargs(): end_time=end_time, ) - assert ( - payload["agent_id"] == test_agent_id - ), f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" + assert payload["agent_id"] == test_agent_id, f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" @patch("litellm.proxy.proxy_server.master_key", None) @@ -1093,9 +1129,7 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1173,9 +1207,9 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): metadata = json.loads(metadata_json) # Verify overhead is stored directly in metadata - assert ( - metadata.get("litellm_overhead_time_ms") == test_overhead_ms - ), f"Expected overhead '{test_overhead_ms}', got '{metadata.get('litellm_overhead_time_ms')}'" + assert metadata.get("litellm_overhead_time_ms") == test_overhead_ms, ( + f"Expected overhead '{test_overhead_ms}', got '{metadata.get('litellm_overhead_time_ms')}'" + ) @patch("litellm.proxy.proxy_server.master_key", None) @@ -1228,9 +1262,7 @@ def test_get_logging_payload_handles_missing_overhead_gracefully(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1309,14 +1341,12 @@ def test_get_logging_payload_handles_missing_overhead_gracefully(): metadata = json.loads(metadata_json) # When overhead is None, litellm_overhead_time_ms should be None or not present - assert ( - metadata.get("litellm_overhead_time_ms") is None - ), "litellm_overhead_time_ms should be None when overhead is not provided" + assert metadata.get("litellm_overhead_time_ms") is None, ( + "litellm_overhead_time_ms should be None when overhead is not provided" + ) -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_enabled( mock_should_store, ): @@ -1347,9 +1377,7 @@ def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_e ) parsed_request = json.loads(request_result) - assert parsed_request["messages"] == [ - {"role": "user", "content": "redacted-by-litellm"} - ] + assert parsed_request["messages"] == [{"role": "user", "content": "redacted-by-litellm"}] assert parsed_request["model"] == "gpt-4" # Test response redaction - use dict response to verify redaction @@ -1368,9 +1396,7 @@ def test_spend_logs_redacts_request_and_response_when_turn_off_message_logging_e {"response": response_dict}, ) - response_result = _get_response_for_spend_logs_payload( - payload=payload, kwargs=kwargs - ) + response_result = _get_response_for_spend_logs_payload(payload=payload, kwargs=kwargs) # When redaction is enabled and response is a dict (not ModelResponse), # perform_redaction redacts content in-place within the choices structure @@ -1415,30 +1441,22 @@ def test_should_store_prompts_and_responses_in_spend_logs_case_insensitive_strin # When env var is True, should return True mock_get_secret_bool.return_value = True result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is True - ), f"Expected True (from env var) for '{false_value}', got {result}" + assert result is True, f"Expected True (from env var) for '{false_value}', got {result}" # When env var is False, should return False mock_get_secret_bool.return_value = False result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is False - ), f"Expected False (from env var) for '{false_value}', got {result}" + assert result is False, f"Expected False (from env var) for '{false_value}', got {result}" # Test when general_settings doesn't have the key at all with patch("litellm.proxy.proxy_server.general_settings", {}): mock_get_secret_bool.return_value = True result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is True - ), "Expected True (from env var) when key missing, got False" + assert result is True, "Expected True (from env var) when key missing, got False" mock_get_secret_bool.return_value = False result = _should_store_prompts_and_responses_in_spend_logs() - assert ( - result is False - ), "Expected False (from env var) when key missing, got True" + assert result is False, "Expected False (from env var) when key missing, got True" def test_get_spend_logs_metadata_guardrail_info_fallback_from_metadata(): @@ -1831,9 +1849,7 @@ def test_get_logging_payload_includes_retry_info_in_spend_logs_metadata(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1897,12 +1913,10 @@ def test_get_logging_payload_includes_retry_info_in_spend_logs_metadata(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_retries") == 2 - ), f"Expected attempted_retries=2, got {metadata.get('attempted_retries')}" - assert ( - metadata.get("max_retries") == 3 - ), f"Expected max_retries=3, got {metadata.get('max_retries')}" + assert metadata.get("attempted_retries") == 2, ( + f"Expected attempted_retries=2, got {metadata.get('attempted_retries')}" + ) + assert metadata.get("max_retries") == 3, f"Expected max_retries=3, got {metadata.get('max_retries')}" @patch("litellm.proxy.proxy_server.master_key", None) @@ -1930,9 +1944,7 @@ def test_get_logging_payload_handles_missing_retry_info_gracefully(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -1996,20 +2008,14 @@ def test_get_logging_payload_handles_missing_retry_info_gracefully(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_retries") is None - ), "attempted_retries should be None when not provided" - assert ( - metadata.get("max_retries") is None - ), "max_retries should be None when not provided" + assert metadata.get("attempted_retries") is None, "attempted_retries should be None when not provided" + assert metadata.get("max_retries") is None, "max_retries should be None when not provided" def test_get_request_duration_ms_normal(): """Test that request duration is correctly computed in milliseconds.""" start = datetime.datetime(2025, 1, 1, 0, 0, 0, tzinfo=timezone.utc) - end = datetime.datetime( - 2025, 1, 1, 0, 0, 2, 500000, tzinfo=timezone.utc - ) # 2.5s later + end = datetime.datetime(2025, 1, 1, 0, 0, 2, 500000, tzinfo=timezone.utc) # 2.5s later result = _get_request_duration_ms(start, end) assert result == 2500 @@ -2039,9 +2045,7 @@ def test_get_logging_payload_includes_request_duration_ms(): "litellm_params": {"api_base": "https://api.openai.com"}, "standard_logging_object": None, } - response_obj = { - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} - } + response_obj = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}} with ( patch("litellm.proxy.proxy_server.master_key", None), @@ -2107,16 +2111,12 @@ def test_sanitize_request_body_strips_secret_fields(): } sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) - assert ( - "secret_fields" not in sanitized - ), "secret_fields must be stripped from the sanitized request body" + assert "secret_fields" not in sanitized, "secret_fields must be stripped from the sanitized request body" assert sanitized["model"] == "gpt-4" assert sanitized["messages"] == [{"role": "user", "content": "hi"}] -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store): """ End-to-end test: when the proxy_server_request body contains @@ -2140,14 +2140,10 @@ def test_proxy_server_request_payload_excludes_secret_fields(mock_should_store): } } - result = _get_proxy_server_request_for_spend_logs_payload( - metadata={}, litellm_params=litellm_params, kwargs={} - ) + result = _get_proxy_server_request_for_spend_logs_payload(metadata={}, litellm_params=litellm_params, kwargs={}) parsed = json.loads(result) - assert ( - "secret_fields" not in parsed - ), "secret_fields must never appear in the spend-log proxy_server_request column" + assert "secret_fields" not in parsed, "secret_fields must never appear in the spend-log proxy_server_request column" assert parsed["model"] == "gpt-4" assert parsed["messages"] == [{"role": "user", "content": "hello"}] @@ -2176,10 +2172,7 @@ def test_redact_prompt_leaks_strips_input_value_python_repr(): def test_redact_prompt_leaks_strips_input_value_json(): - error_text = ( - '{"error":{"message":"validation failed",' - '"input":[{"role":"user","content":"top-secret-content"}]}}' - ) + error_text = '{"error":{"message":"validation failed","input":[{"role":"user","content":"top-secret-content"}]}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "top-secret-content" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2203,9 +2196,7 @@ def test_redact_prompt_leaks_empty_string(): assert _redact_prompt_leaks_in_error_string("") == "" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_when_not_storing_prompts( mock_should_store, ): @@ -2233,9 +2224,7 @@ def test_sanitize_error_information_redacts_when_not_storing_prompts( assert sanitized["llm_provider"] == "openai" -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_skips_redaction_when_storing_prompts( mock_should_store, ): @@ -2246,9 +2235,7 @@ def test_sanitize_error_information_skips_redaction_when_storing_prompts( "error_class": "RateLimitError", "llm_provider": "openai", "traceback": "", - "error_message": ( - 'OpenAIException - {"error":{"input":[{"role":"user","content":"kept"}]}}' - ), + "error_message": ('OpenAIException - {"error":{"input":[{"role":"user","content":"kept"}]}}'), } sanitized = _sanitize_error_information_for_spend_logs(error_info) @@ -2259,9 +2246,7 @@ def test_sanitize_error_information_skips_redaction_when_storing_prompts( assert REDACTED_BY_LITELM_STRING not in sanitized["error_message"] -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_caps_size_regardless_of_prompt_flag( mock_should_store, ): @@ -2292,9 +2277,7 @@ def test_sanitize_error_information_none_passthrough(): assert _sanitize_error_information_for_spend_logs(None) is None -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_reproduces_lit_2992(mock_should_store): # Mirrors the reproduced row body from LIT-2992 — a RateLimitError whose # message embeds 178 pydantic validation errors, each carrying a full @@ -2335,10 +2318,7 @@ def test_redact_prompt_leaks_handles_nested_multimodal_content(): # Multi-modal payload: 'content' is itself a list. The depth-1 regex # would stop at the inner '['; the parser-based scanner must walk # through balanced nested brackets. - error_text = ( - '{"error":{"messages":[{"role":"user",' - '"content":[{"type":"text","text":"top-secret-multimodal"}]}]}}' - ) + error_text = '{"error":{"messages":[{"role":"user","content":[{"type":"text","text":"top-secret-multimodal"}]}]}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "top-secret-multimodal" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2347,9 +2327,7 @@ def test_redact_prompt_leaks_handles_nested_multimodal_content(): def test_redact_prompt_leaks_handles_bracket_in_prompt_text(): # Prompt text contains a literal '[' — the depth-1 regex would close # the outer ']' prematurely. The parser must respect string quoting. - error_text = ( - '{"error":{"input":[{"role":"user","content":"secret[123 still secret"}]}}' - ) + error_text = '{"error":{"input":[{"role":"user","content":"secret[123 still secret"}]}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "secret[123" not in redacted assert "still secret" not in redacted @@ -2368,8 +2346,7 @@ def test_redact_prompt_leaks_handles_escaped_quote_in_prompt_text(): def test_redact_prompt_leaks_handles_nested_input_python_repr(): # Python dict-repr with nested list inside 'input' — single quotes. error_text = ( - "validation error: {'input': [{'role': 'user', " - "'content': [{'type': 'text', 'text': 'leaked-nested-text'}]}]}" + "validation error: {'input': [{'role': 'user', 'content': [{'type': 'text', 'text': 'leaked-nested-text'}]}]}" ) redacted = _redact_prompt_leaks_in_error_string(error_text) assert "leaked-nested-text" not in redacted @@ -2385,9 +2362,7 @@ def test_redact_prompt_leaks_handles_unterminated_value(): assert REDACTED_BY_LITELM_STRING in redacted -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_traceback_when_not_storing_prompts( mock_should_store, ): @@ -2419,9 +2394,7 @@ def test_sanitize_error_information_redacts_traceback_when_not_storing_prompts( assert "ValueError: invalid request" in sanitized["traceback"] -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_skips_traceback_redaction_when_storing_prompts( mock_should_store, ): @@ -2431,9 +2404,7 @@ def test_sanitize_error_information_skips_traceback_redaction_when_storing_promp "error_code": "500", "error_class": "ValueError", "llm_provider": "", - "traceback": ( - 'raise ValueError({"input":[{"role":"user","content":"tb-kept"}]})' - ), + "traceback": ('raise ValueError({"input":[{"role":"user","content":"tb-kept"}]})'), "error_message": "invalid request", } @@ -2448,20 +2419,14 @@ def test_redact_prompt_leaks_strips_prompt_key_completions_payload(): # /v1/completions echoes the user input under the top-level 'prompt' key # rather than 'messages'. Without 'prompt' coverage the body would survive # the redactor when store_prompts_in_spend_logs is False. - error_text = ( - '{"error":{"message":"validation failed",' - '"prompt":"super-secret-completion-text"}}' - ) + error_text = '{"error":{"message":"validation failed","prompt":"super-secret-completion-text"}}' redacted = _redact_prompt_leaks_in_error_string(error_text) assert "super-secret-completion-text" not in redacted assert REDACTED_BY_LITELM_STRING in redacted def test_redact_prompt_leaks_strips_prompt_key_python_repr(): - error_text = ( - "{'model': 'gpt-3.5-turbo-instruct', " - "'prompt': 'leaked-completion-prompt-body'}" - ) + error_text = "{'model': 'gpt-3.5-turbo-instruct', 'prompt': 'leaked-completion-prompt-body'}" redacted = _redact_prompt_leaks_in_error_string(error_text) assert "leaked-completion-prompt-body" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2495,11 +2460,7 @@ def test_redact_prompt_leaks_strips_pydantic_input_value_list(): def test_redact_prompt_leaks_strips_pydantic_input_value_dict(): - error_text = ( - "[type=dict_type, " - "input_value={'role': 'user', 'content': 'leaked-dict-content'}, " - "input_type=dict]" - ) + error_text = "[type=dict_type, input_value={'role': 'user', 'content': 'leaked-dict-content'}, input_type=dict]" redacted = _redact_prompt_leaks_in_error_string(error_text) assert "leaked-dict-content" not in redacted assert REDACTED_BY_LITELM_STRING in redacted @@ -2541,9 +2502,7 @@ def test_redact_prompt_leaks_combined_quoted_key_and_pydantic_assignment(): assert redacted.count(REDACTED_BY_LITELM_STRING) >= 2 -@patch( - "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" -) +@patch("litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs") def test_sanitize_error_information_redacts_pydantic_assignment_form( mock_should_store, ): @@ -2741,9 +2700,7 @@ def test_get_spend_logs_metadata_non_sk_raw_key_hashed(): def test_get_spend_logs_metadata_already_hashed_unchanged_with_provenance(): already_hashed = hash_token("sk-some-key") - meta = _get_spend_logs_metadata( - {"user_api_key": already_hashed, "user_api_key_hash": already_hashed} - ) + meta = _get_spend_logs_metadata({"user_api_key": already_hashed, "user_api_key_hash": already_hashed}) assert meta["user_api_key"] == already_hashed assert hash_token(already_hashed) != meta["user_api_key"] # no double-hash @@ -2758,9 +2715,7 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed(): def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match(): already_hashed = hash_token("sk-some-key") different_hash = hash_token("sk-other-key") - meta = _get_spend_logs_metadata( - {"user_api_key": already_hashed, "user_api_key_hash": different_hash} - ) + meta = _get_spend_logs_metadata({"user_api_key": already_hashed, "user_api_key_hash": different_hash}) assert meta["user_api_key"] == hash_token(already_hashed) @@ -2797,16 +2752,12 @@ def test_get_logging_payload_uses_recovered_combined_usage_on_failure(): "model": "anthropic/claude-haiku-4-5", "call_type": "acompletion", "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } response_obj = Exception("MidStreamFallbackError: read timeout") now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert payload["prompt_tokens"] == 30 assert payload["completion_tokens"] == 1 @@ -2825,9 +2776,7 @@ def test_get_logging_payload_failure_without_recovered_usage_is_zero(): response_obj = Exception("BadRequestError") now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert payload["total_tokens"] == 0 @@ -2853,9 +2802,7 @@ def test_get_logging_payload_sets_litellm_call_id_for_correlation(): } now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) metadata = json.loads(payload["metadata"]) assert payload["request_id"] == provider_response_id @@ -2882,9 +2829,7 @@ def test_get_logging_payload_litellm_call_id_falls_back_to_litellm_params(): } now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id @@ -2901,14 +2846,10 @@ def test_get_logging_payload_litellm_call_id_when_response_has_no_id(): "litellm_call_id": trace_call_id, "litellm_params": {"metadata": {"user_api_key": "sk-test"}}, } - response_obj = { - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} - } + response_obj = {"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}} now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id assert payload["request_id"] == trace_call_id @@ -2932,9 +2873,7 @@ def test_get_logging_payload_cache_hit_keeps_raw_litellm_call_id(): } now = datetime.datetime.now(timezone.utc) - payload = get_logging_payload( - kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now - ) + payload = get_logging_payload(kwargs=kwargs, response_obj=response_obj, start_time=now, end_time=now) assert json.loads(payload["metadata"])["litellm_call_id"] == trace_call_id assert "_cache_hit" in payload["request_id"] @@ -3074,9 +3013,7 @@ def test_get_logging_payload_hashes_bearer_prefixed_api_key(): assert not payload["api_key"].startswith("Bearer"), ( f"api_key column contains plaintext Bearer key: {payload['api_key']}" ) - assert not payload["api_key"].startswith("sk-"), ( - f"api_key column contains unhashed key: {payload['api_key']}" - ) + assert not payload["api_key"].startswith("sk-"), f"api_key column contains unhashed key: {payload['api_key']}" metadata_dict = json.loads(payload["metadata"]) assert not metadata_dict["user_api_key"].startswith("Bearer"), ( @@ -3747,9 +3684,7 @@ def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -3813,12 +3748,12 @@ def test_get_logging_payload_includes_fallback_info_in_spend_logs_metadata(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_fallbacks") == 2 - ), f"Expected attempted_fallbacks=2, got {metadata.get('attempted_fallbacks')}" - assert ( - metadata.get("original_model_group") == "azure-gpt-fallback" - ), f"Expected original_model_group=azure-gpt-fallback, got {metadata.get('original_model_group')}" + assert metadata.get("attempted_fallbacks") == 2, ( + f"Expected attempted_fallbacks=2, got {metadata.get('attempted_fallbacks')}" + ) + assert metadata.get("original_model_group") == "azure-gpt-fallback", ( + f"Expected original_model_group=azure-gpt-fallback, got {metadata.get('original_model_group')}" + ) def test_get_logging_payload_handles_missing_fallback_info_gracefully(): @@ -3844,9 +3779,7 @@ def test_get_logging_payload_handles_missing_fallback_info_gracefully(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=None, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", model_id="model-123", model_group="openai", @@ -3910,12 +3843,10 @@ def test_get_logging_payload_handles_missing_fallback_info_gracefully(): metadata = json.loads(payload["metadata"]) - assert ( - metadata.get("attempted_fallbacks") is None - ), "attempted_fallbacks should be None when not provided" - assert ( - metadata.get("original_model_group") is None - ), "original_model_group should be None when not provided" + assert metadata.get("attempted_fallbacks") is None, "attempted_fallbacks should be None when not provided" + assert metadata.get("original_model_group") is None, "original_model_group should be None when not provided" + + @pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) def test_injected_cache_breakpoints_survive_into_spend_log_metadata(bucket): """The injection marker only gates savings if it reaches the spend-log row. diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 8366e5546a9..72d37650963 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -41,21 +41,18 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem - def test_check_if_token_is_service_account(): """ Test that only keys with `service_account_id` in metadata are considered service accounts """ # Test case 1: Service account token - service_account_token = UserAPIKeyAuth( - api_key="test-key", metadata={"service_account_id": "test-service-account"} - ) + service_account_token = UserAPIKeyAuth(api_key="test-key", metadata={"service_account_id": "test-service-account"}) assert check_if_token_is_service_account(service_account_token) == True # Test case 2: Regular user token @@ -63,9 +60,7 @@ def test_check_if_token_is_service_account(): assert check_if_token_is_service_account(regular_token) == False # Test case 3: Token with other metadata - other_metadata_token = UserAPIKeyAuth( - api_key="test-key", metadata={"user_id": "test-user"} - ) + other_metadata_token = UserAPIKeyAuth(api_key="test-key", metadata={"user_id": "test-user"}) assert check_if_token_is_service_account(other_metadata_token) == False @@ -112,15 +107,11 @@ class TestGetMetadataVariableName: def test_returns_litellm_metadata_for_bedrock_invoke(self): # GH#30629: bedrock passthrough must use litellm_metadata # to prevent key-level tags from leaking into provider body - request = self._make_request( - "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke" - ) + request = self._make_request("/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke") assert _get_metadata_variable_name(request) == "litellm_metadata" def test_returns_litellm_metadata_for_bedrock_converse(self): - request = self._make_request( - "/bedrock/model/us.anthropic.claude-sonnet-4-6/converse" - ) + request = self._make_request("/bedrock/model/us.anthropic.claude-sonnet-4-6/converse") assert _get_metadata_variable_name(request) == "litellm_metadata" @@ -128,9 +119,7 @@ def test_get_enforced_params_for_service_account_settings(): """ Test that service account enforced params are only added to service account keys """ - service_account_token = UserAPIKeyAuth( - api_key="test-key", metadata={"service_account_id": "test-service-account"} - ) + service_account_token = UserAPIKeyAuth(api_key="test-key", metadata={"service_account_id": "test-service-account"}) general_settings_with_service_account_settings = { "service_account_settings": {"enforced_params": ["metadata.service"]}, } @@ -140,9 +129,7 @@ def test_get_enforced_params_for_service_account_settings(): ) assert result == ["metadata.service"] - regular_token = UserAPIKeyAuth( - api_key="test-key", metadata={"enforced_params": ["user"]} - ) + regular_token = UserAPIKeyAuth(api_key="test-key", metadata={"enforced_params": ["user"]}) result = _get_enforced_params( general_settings=general_settings_with_service_account_settings, user_api_key_dict=regular_token, @@ -155,9 +142,7 @@ def test_get_enforced_params_for_service_account_settings(): [ ( {"enforced_params": ["param1", "param2"]}, - UserAPIKeyAuth( - api_key="test_api_key", user_id="test_user_id", org_id="test_org_id" - ), + UserAPIKeyAuth(api_key="test_api_key", user_id="test_user_id", org_id="test_org_id"), ["param1", "param2"], ), ( @@ -183,9 +168,7 @@ def test_get_enforced_params_for_service_account_settings(): ), ], ) -def test_get_enforced_params( - general_settings, user_api_key_dict, expected_enforced_params -): +def test_get_enforced_params(general_settings, user_api_key_dict, expected_enforced_params): from litellm.proxy.litellm_pre_call_utils import _get_enforced_params enforced_params = _get_enforced_params(general_settings, user_api_key_dict) @@ -441,9 +424,7 @@ async def test_add_litellm_data_to_request_strips_admin_injection_slots(): populated = updated["metadata"] assert populated["user_api_key_metadata"] == real_admin_metadata assert populated["user_api_key_team_metadata"] == real_admin_metadata - assert "_pipeline_managed_guardrails" not in populated or populated[ - "_pipeline_managed_guardrails" - ] != ["evaded"] + assert "_pipeline_managed_guardrails" not in populated or populated["_pipeline_managed_guardrails"] != ["evaded"] other = updated.get("litellm_metadata") or {} assert other.get("user_api_key_metadata") in (None, {}, real_admin_metadata) @@ -697,9 +678,7 @@ async def test_add_litellm_data_to_request_proxy_server_request_body_is_post_str snapshot_body = updated["proxy_server_request"]["body"] assert snapshot_body is not None snapshot_metadata = snapshot_body.get("metadata") or {} - assert "user_api_key_user_id" not in snapshot_metadata or ( - snapshot_metadata["user_api_key_user_id"] != "victim" - ) + assert "user_api_key_user_id" not in snapshot_metadata or (snapshot_metadata["user_api_key_user_id"] != "victim") @pytest.mark.asyncio @@ -754,9 +733,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_secret_fields( ) # secret_fields must exist on the live data dict - assert ( - "secret_fields" in updated - ), "secret_fields must still be present on the live data dict" + assert "secret_fields" in updated, "secret_fields must still be present on the live data dict" assert "raw_headers" in updated["secret_fields"] # But the body snapshot must NOT contain secret_fields @@ -815,8 +792,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r snapshot_body = updated["proxy_server_request"]["body"] assert "proxy_server_request" not in snapshot_body, ( - "proxy_server_request must be excluded from its own body snapshot " - "to prevent the body from self-referencing" + "proxy_server_request must be excluded from its own body snapshot to prevent the body from self-referencing" ) @@ -1344,23 +1320,18 @@ async def test_add_litellm_data_to_request_strips_client_redaction_bypass_contro assert "turn_off_message_logging" not in (updated.get("litellm_params") or {}).get("metadata", {}) assert "turn_off_message_logging" not in updated["metadata"] assert "turn_off_message_logging" not in (updated.get("litellm_metadata") or {}) + assert "litellm-disable-message-redaction" not in {header.lower() for header in updated["metadata"]["headers"]} assert "litellm-disable-message-redaction" not in { - header.lower() for header in updated["metadata"]["headers"] - } - assert "litellm-disable-message-redaction" not in { - header.lower() - for header in updated["metadata"]["requester_metadata"].get("headers", {}) + header.lower() for header in updated["metadata"]["requester_metadata"].get("headers", {}) } assert "litellm-disable-message-redaction" not in { header.lower() for header in updated["proxy_server_request"]["headers"] } assert "litellm-disable-message-redaction" not in { - header.lower() - for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] + header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } assert "litellm-disable-message-redaction" not in { - header.lower() - for header in (updated.get("litellm_metadata") or {}).get("headers", {}) + header.lower() for header in (updated.get("litellm_metadata") or {}).get("headers", {}) } @@ -1430,12 +1401,7 @@ async def test_add_litellm_data_to_request_admin_callback_vars_turn_off_message_ dynamic_params = initialize_standard_callback_dynamic_params(updated) assert dynamic_params.get("turn_off_message_logging") == "False" - assert ( - should_redact_message_logging( - {"standard_callback_dynamic_params": dynamic_params} - ) - is False - ) + assert should_redact_message_logging({"standard_callback_dynamic_params": dynamic_params}) is False finally: litellm.turn_off_message_logging = original_turn_off_message_logging @@ -1506,12 +1472,7 @@ async def test_add_litellm_data_to_request_admin_callback_vars_turn_off_message_ dynamic_params = initialize_standard_callback_dynamic_params(updated) assert dynamic_params.get("turn_off_message_logging") == "True" - assert ( - should_redact_message_logging( - {"standard_callback_dynamic_params": dynamic_params} - ) - is True - ) + assert should_redact_message_logging({"standard_callback_dynamic_params": dynamic_params}) is True finally: litellm.turn_off_message_logging = original_turn_off_message_logging @@ -1552,9 +1513,7 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o "headers": {"litellm-disable-message-redaction": "true"}, "turn_off_message_logging": False, }, - "litellm_metadata": json.dumps( - {"headers": {"LiteLLM-Disable-Message-Redaction": "true"}} - ), + "litellm_metadata": json.dumps({"headers": {"LiteLLM-Disable-Message-Redaction": "true"}}), }, request=request_mock, user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", **auth_kwargs), @@ -1567,19 +1526,15 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o assert updated["turn_off_message_logging"] is False assert updated["metadata"]["turn_off_message_logging"] is False + assert "litellm-disable-message-redaction" in {header.lower() for header in updated["metadata"]["headers"]} assert "litellm-disable-message-redaction" in { - header.lower() for header in updated["metadata"]["headers"] - } - assert "litellm-disable-message-redaction" in { - header.lower() - for header in updated["metadata"]["requester_metadata"].get("headers", {}) + header.lower() for header in updated["metadata"]["requester_metadata"].get("headers", {}) } assert "litellm-disable-message-redaction" in { header.lower() for header in updated["proxy_server_request"]["headers"] } assert "litellm-disable-message-redaction" in { - header.lower() - for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] + header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } assert "litellm_metadata" not in updated @@ -1870,9 +1825,7 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart(): request_mock.client.host = "127.0.0.1" # Simulate multipart data (metadata as string) - metadata_dict = { - "tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"] - } + metadata_dict = {"tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"]} stringified_metadata = json.dumps(metadata_dict) data = { @@ -2200,23 +2153,15 @@ def test_key_dynamic_logging_settings(): # Test with langfuse logging key_with_langfuse = UserAPIKeyAuth( api_key="test-key", - metadata={ - "logging": [{"callback_name": "langfuse", "callback_type": "success"}] - }, + metadata={"logging": [{"callback_name": "langfuse", "callback_type": "success"}]}, team_metadata={}, ) - result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - key_with_langfuse - ) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) assert result == [{"callback_name": "langfuse", "callback_type": "success"}] # Test with no logging metadata - key_without_logging = UserAPIKeyAuth( - api_key="test-key", metadata={}, team_metadata={} - ) - result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings( - key_without_logging - ) + key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) + result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_without_logging) assert result is None @@ -2228,35 +2173,23 @@ def test_team_dynamic_logging_settings(): key_with_team_arize = UserAPIKeyAuth( api_key="test-key", metadata={}, - team_metadata={ - "logging": [{"callback_name": "arize", "callback_type": "failure"}] - }, - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_with_team_arize + team_metadata={"logging": [{"callback_name": "arize", "callback_type": "failure"}]}, ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_arize) assert result == [{"callback_name": "arize", "callback_type": "failure"}] # Test with langfuse team logging key_with_team_langfuse = UserAPIKeyAuth( api_key="test-key", metadata={}, - team_metadata={ - "logging": [{"callback_name": "langfuse", "callback_type": "success"}] - }, - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_with_team_langfuse + team_metadata={"logging": [{"callback_name": "langfuse", "callback_type": "success"}]}, ) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_with_team_langfuse) assert result == [{"callback_name": "langfuse", "callback_type": "success"}] # Test with no team logging metadata - key_without_team_logging = UserAPIKeyAuth( - api_key="test-key", metadata={}, team_metadata={} - ) - result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings( - key_without_team_logging - ) + key_without_team_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) + result = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(key_without_team_logging) assert result is None @@ -2337,9 +2270,7 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): mock_proxy_config = MagicMock() # Call the function - result = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config - ) + result = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=mock_proxy_config) # Verify the result assert result is not None @@ -2355,9 +2286,7 @@ def test_add_team_callback_rejects_env_reference(): AddTeamCallback( callback_name="langfuse", callback_type="success", - callback_vars={ - "langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY_TEMP" - }, + callback_vars={"langfuse_secret_key": "os.environ/LANGFUSE_SECRET_KEY_TEMP"}, ) assert "os.environ/" in str(exc_info.value) @@ -2388,9 +2317,7 @@ def test_get_dynamic_logging_metadata_ignores_env_reference_from_key_metadata( team_metadata={}, ) - result = _get_dynamic_logging_metadata( - user_api_key_dict=user_api_key_dict, proxy_config=MagicMock() - ) + result = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) assert result is None @@ -2401,16 +2328,12 @@ def test_get_num_retries_from_request(): """ # Test case 1: Header is present with valid integer string headers_with_retries = {"x-litellm-num-retries": "3"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_retries - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_retries) assert result == 3 # Test case 2: Header is not present headers_without_retries = {"Content-Type": "application/json"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_without_retries - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_without_retries) assert result is None # Test case 3: Empty headers dictionary @@ -2425,9 +2348,7 @@ def test_get_num_retries_from_request(): # Test case 5: Header present with large number headers_with_large_number = {"x-litellm-num-retries": "100"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_large_number - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_large_number) assert result == 100 # Test case 6: Multiple headers with num retries header @@ -2441,19 +2362,17 @@ def test_get_num_retries_from_request(): # Test case 7: Header present with invalid value (should raise ValueError when int() is called) headers_with_invalid = {"x-litellm-num-retries": "invalid"} - with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): + with pytest.raises(ValueError, match="invalid literal for int\\(\\) with base"): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_invalid) # Test case 8: Header present with float string (should raise ValueError when int() is called) headers_with_float = {"x-litellm-num-retries": "3.5"} - with pytest.raises(ValueError, match='invalid literal for int\\(\\) with base'): + with pytest.raises(ValueError, match="invalid literal for int\\(\\) with base"): LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_float) # Test case 9: Header present with negative number headers_with_negative = {"x-litellm-num-retries": "-1"} - result = LiteLLMProxyRequestSetup._get_num_retries_from_request( - headers_with_negative - ) + result = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers_with_negative) assert result == -1 @@ -2463,15 +2382,11 @@ def test_get_keepalive_seconds_from_request(): """ # Header present with valid float string headers_with_keepalive = {"x-litellm-keepalive-seconds": "15"} - result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - headers_with_keepalive - ) + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request(headers_with_keepalive) assert result == 15.0 # Header not present - result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - {"Content-Type": "application/json"} - ) + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({"Content-Type": "application/json"}) assert result is None # Empty headers dictionary @@ -2479,17 +2394,13 @@ def test_get_keepalive_seconds_from_request(): assert result is None # Header present with a fractional value - result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - {"x-litellm-keepalive-seconds": "1.5"} - ) + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({"x-litellm-keepalive-seconds": "1.5"}) assert result == 1.5 # Header present with invalid value raises ValueError, matching the other # x-litellm-* numeric header helpers (_get_timeout_from_request, etc.) with pytest.raises(ValueError, match="could not convert string to float: 'not-a-number"): - LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( - {"x-litellm-keepalive-seconds": "not-a-number"} - ) + LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({"x-litellm-keepalive-seconds": "not-a-number"}) def test_add_litellm_data_for_backend_llm_call_merges_keepalive_seconds_header(): @@ -2728,9 +2639,7 @@ def test_management_endpoint_metadata_drops_callback_credentials(): ), ], ) -def test_add_headers_to_llm_call_by_model_group( - data, model_group_settings, expected_headers_added -): +def test_add_headers_to_llm_call_by_model_group(data, model_group_settings, expected_headers_added): """ Test LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group method @@ -2751,9 +2660,7 @@ def test_add_headers_to_llm_call_by_model_group( "X-Custom-Header": "custom-value", } - user_api_key_dict = UserAPIKeyAuth( - api_key="test-key", user_id="test-user", org_id="test-org" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="test-user", org_id="test-org") # Mock the model_group_settings original_model_group_settings = getattr(litellm, "model_group_settings", None) @@ -2771,7 +2678,6 @@ def test_add_headers_to_llm_call_by_model_group( "add_headers_to_llm_call", return_value=expected_returned_headers if expected_headers_added else {}, ) as mock_add_headers: - # Make a copy of original data to verify it's not mutated unexpectedly original_data = copy.deepcopy(data) @@ -2828,7 +2734,6 @@ def test_add_headers_to_llm_call_by_model_group_empty_headers_returned(): "add_headers_to_llm_call", return_value={}, # Return empty dict ) as mock_add_headers: - result = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=headers, user_api_key_dict=user_api_key_dict ) @@ -2876,7 +2781,6 @@ def test_add_headers_to_llm_call_by_model_group_existing_headers_in_data(): "add_headers_to_llm_call", return_value=new_headers, ) as mock_add_headers: - result = LiteLLMProxyRequestSetup.add_headers_to_llm_call_by_model_group( data=data, headers=headers, user_api_key_dict=user_api_key_dict ) @@ -2990,13 +2894,9 @@ async def test_add_litellm_metadata_from_request_headers(): general_settings = {} # Create mock select_data_generator with correct signature - def mock_select_data_generator( - response=None, user_api_key_dict=None, request_data=None - ): + def mock_select_data_generator(response=None, user_api_key_dict=None, request_data=None): async def mock_generator(): - yield "data: " + json.dumps( - {"choices": [{"delta": {"content": "Hello"}}]} - ) + "\n\n" + yield "data: " + json.dumps({"choices": [{"delta": {"content": "Hello"}}]}) + "\n\n" yield "data: [DONE]\n\n" return mock_generator() @@ -3023,21 +2923,19 @@ async def test_add_litellm_metadata_from_request_headers(): await asyncio.sleep(3) # Check if standard_logging_object was set - assert ( - test_logger.standard_logging_object is not None - ), "standard_logging_object should be populated after LLM request" + assert test_logger.standard_logging_object is not None, ( + "standard_logging_object should be populated after LLM request" + ) # Verify the logging object contains expected metadata standard_logging_obj = test_logger.standard_logging_object - print( - f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}" - ) + print(f"Standard logging object captured: {json.dumps(standard_logging_obj, indent=4, default=str)}") SPEND_LOGS_METADATA = standard_logging_obj["metadata"]["spend_logs_metadata"] - assert SPEND_LOGS_METADATA == dict( - json.loads(headers["x-litellm-spend-logs-metadata"]) - ), "spend_logs_metadata should be the same as the headers" + assert SPEND_LOGS_METADATA == dict(json.loads(headers["x-litellm-spend-logs-metadata"])), ( + "spend_logs_metadata should be the same as the headers" + ) finally: litellm.callbacks = original_callbacks @@ -3188,11 +3086,7 @@ def test_add_litellm_metadata_from_request_headers_generic_session_id_header(): def test_add_litellm_metadata_from_anthropic_user_id_sets_session_id(): - data = { - "metadata": { - "user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01" - } - } + data = {"metadata": {"user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01"}} LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( headers={}, data=data, _metadata_variable_name="metadata" ) @@ -3308,9 +3202,7 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers assert ( - get_chain_id_from_headers( - {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"} - ) + get_chain_id_from_headers({"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"}) == "e96634a3-fa28-4083-b354-55542e2dca01" ) # Short / non-alphanumeric values should be ignored @@ -3600,19 +3492,13 @@ def test_get_internal_user_header_from_mapping_returns_expected_header(): {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - mappings - ) + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert header_name == "X-OpenWebUI-User-Id" def test_get_internal_user_header_from_mapping_none_when_absent(): - mappings = [ - {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} - ] - header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - mappings - ) + mappings = [{"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}] + header_name = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert header_name is None single = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} @@ -3633,9 +3519,7 @@ def test_add_internal_user_from_user_mapping_sets_user_id_when_header_present(): ] } - result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( - general_settings, user_api_key_dict, headers - ) + result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping(general_settings, user_api_key_dict, headers) assert result is user_api_key_dict assert user_api_key_dict.user_id == "internal-user-123" @@ -3651,9 +3535,7 @@ def test_add_internal_user_from_user_mapping_no_header_or_mapping_returns_unchan assert user_api_key_dict.user_id is None general_settings = { - "user_header_mappings": [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} - ] + "user_header_mappings": [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] } result = LiteLLMProxyRequestSetup.add_internal_user_from_user_mapping( general_settings, user_api_key_dict, {"Other": "value"} @@ -3673,9 +3555,7 @@ def test_get_sanitized_user_information_from_key_includes_guardrails_metadata(): metadata={"guardrails": ["presidio", "aporia"], "other_field": "value"}, ) - result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) assert result["user_api_key_auth_metadata"] is not None assert "guardrails" in result["user_api_key_auth_metadata"] @@ -3704,9 +3584,7 @@ def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): team_max_budget=1000.0, ) - sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) assert sanitized["user_api_key_spend"] == 1.5 assert sanitized["user_api_key_max_budget"] == 10.0 @@ -3715,9 +3593,7 @@ def test_user_and_team_spend_and_budget_flow_to_standard_logging_metadata(): assert sanitized["user_api_key_team_spend"] == 250.75 assert sanitized["user_api_key_team_max_budget"] == 1000.0 - logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - dict(sanitized) - ) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(dict(sanitized)) assert logging_metadata["user_api_key_user_spend"] == 25.5 assert logging_metadata["user_api_key_user_max_budget"] == 100.0 @@ -3734,12 +3610,8 @@ def test_user_and_team_spend_and_budget_default_to_none_in_standard_logging_meta user_api_key_dict = UserAPIKeyAuth(api_key="test-key-hash") - sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) - logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - dict(sanitized) - ) + sanitized = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) + logging_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata(dict(sanitized)) assert logging_metadata["user_api_key_user_spend"] is None assert logging_metadata["user_api_key_user_max_budget"] is None @@ -4071,22 +3943,16 @@ async def test_embedding_header_forwarding_with_model_group(): # Verify that only x- prefixed headers (except x-stainless) were forwarded forwarded_headers = updated_data["headers"] - assert ( - "X-Custom-Header" in forwarded_headers - ), "X-Custom-Header should be forwarded" + assert "X-Custom-Header" in forwarded_headers, "X-Custom-Header should be forwarded" assert forwarded_headers["X-Custom-Header"] == "custom-value" assert "X-Request-ID" in forwarded_headers, "X-Request-ID should be forwarded" assert forwarded_headers["X-Request-ID"] == "test-request-123" # Verify that authorization header was NOT forwarded (sensitive header) - assert ( - "Authorization" not in forwarded_headers - ), "Authorization header should not be forwarded" + assert "Authorization" not in forwarded_headers, "Authorization header should not be forwarded" # Verify that Content-Type was NOT forwarded (doesn't start with x-) - assert ( - "Content-Type" not in forwarded_headers - ), "Content-Type should not be forwarded" + assert "Content-Type" not in forwarded_headers, "Content-Type should not be forwarded" # Verify original data fields are preserved assert updated_data["model"] == "local-openai/text-embedding-3-small" @@ -4142,9 +4008,9 @@ async def test_embedding_header_forwarding_without_model_group_config(): ) # Verify that headers were NOT added since model is not in forward list - assert ( - "headers" not in updated_data or updated_data.get("headers") is None - ), "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + assert "headers" not in updated_data or updated_data.get("headers") is None, ( + "Headers should not be forwarded for models not in forward_client_headers_to_llm_api list" + ) # Verify original data fields are preserved assert updated_data["model"] == "text-embedding-ada-002" @@ -4198,9 +4064,7 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry = get_attachment_registry() attachment_registry._attachments = [ PolicyAttachment(policy="global-baseline", scope="*"), # applies to all - PolicyAttachment( - policy="healthcare", teams=["healthcare-team"] - ), # applies to healthcare team + PolicyAttachment(policy="healthcare", teams=["healthcare-team"]), # applies to healthcare team ] attachment_registry._initialized = True @@ -4269,9 +4133,9 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po ) # Verify that 'policies' was removed from the request body - assert ( - "policies" not in data - ), "'policies' should be removed from request body to prevent forwarding to LLM provider" + assert "policies" not in data, ( + "'policies' should be removed from request body to prevent forwarding to LLM provider" + ) # Verify that other fields are preserved assert "model" in data @@ -4316,9 +4180,7 @@ async def test_api_created_global_policy_applies_to_new_key_without_restart(): "runtime-global-policy", Policy(guardrails=PolicyGuardrails(add=["runtime-guardrail"])), ) - attachment_registry.add_attachment( - PolicyAttachment(policy="runtime-global-policy", scope="*") - ) + attachment_registry.add_attachment(PolicyAttachment(policy="runtime-global-policy", scope="*")) await add_guardrails_from_policy_engine( data=data, @@ -4415,9 +4277,7 @@ async def test_bearer_token_not_in_debug_logs(): from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ProxyConfig - secret_token = ( - "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" - ) + secret_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.fakesignature" mock_request = MagicMock(spec=Request) mock_request.headers = { @@ -4463,8 +4323,7 @@ async def test_bearer_token_not_in_debug_logs(): log_output = log_capture.getvalue() assert secret_token not in log_output, ( - f"Bearer token leaked in debug logs. " - f"Found token in log output:\n{log_output[:500]}" + f"Bearer token leaked in debug logs. Found token in log output:\n{log_output[:500]}" ) @@ -4629,9 +4488,7 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, @@ -4642,9 +4499,7 @@ def test_apply_overrides_project_model_specific(setup_test_credentials): } }, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" assert data["api_key"] == "key-hotel-rec-vision" assert data["api_version"] == "2024-06-01" @@ -4657,9 +4512,7 @@ def test_apply_overrides_project_default(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, @@ -4670,9 +4523,7 @@ def test_apply_overrides_project_default(setup_test_credentials): } }, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-rec-app.openai.azure.com/" assert data["api_key"] == "key-hotel-rec" @@ -4684,17 +4535,13 @@ def test_apply_overrides_team_model_specific(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-westus.openai.azure.com/" assert data["api_key"] == "key-hotel-westus" @@ -4706,17 +4553,13 @@ def test_apply_overrides_team_default(setup_test_credentials): api_key="test-key", team_metadata={ "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - }, + "defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}, "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-westus"}}, } }, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" @@ -4729,9 +4572,7 @@ def test_apply_overrides_no_config(setup_test_credentials): team_metadata={}, project_metadata={}, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4747,17 +4588,9 @@ def test_apply_overrides_clientside_credentials_take_precedence( } user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://my-custom-endpoint.openai.azure.com/" assert data["api_key"] == "my-custom-key" @@ -4767,15 +4600,9 @@ def test_apply_overrides_missing_credential_name(setup_test_credentials): data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4": {"azure": {"litellm_credentials": "nonexistent-credential"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -4785,17 +4612,9 @@ def test_apply_overrides_api_version_only_if_present(setup_test_credentials): data = {"model": "gpt-3.5"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" assert "api_version" not in data @@ -4806,15 +4625,9 @@ def test_apply_overrides_no_model_in_data(setup_test_credentials): data = {"messages": [{"role": "user", "content": "hello"}]} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": {"azure": {"litellm_credentials": "some-cred"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "some-cred"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data @@ -4826,9 +4639,7 @@ def test_apply_overrides_none_metadata(setup_test_credentials): team_metadata=None, project_metadata=None, ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data @@ -4837,15 +4648,9 @@ def test_apply_overrides_clientside_api_version_preserved(setup_test_credentials data = {"model": "gpt-4-vision", "api_version": "2025-01-01"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4-vision": {"azure": {"litellm_credentials": "hotel-rec-vision"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) # api_base and api_key should be set from credential assert data["api_base"] == "https://hotel-rec-vision.openai.azure.com/" assert data["api_key"] == "key-hotel-rec-vision" @@ -4858,9 +4663,7 @@ def test_resolve_non_dict_model_config_ignored(): result = _resolve_credential_from_model_config("gpt-4", "not-a-dict", None) assert result is None - result = _resolve_credential_from_model_config( - "gpt-4", None, ["also", "not", "a", "dict"] - ) + result = _resolve_credential_from_model_config("gpt-4", None, ["also", "not", "a", "dict"]) assert result is None # Valid config still works alongside invalid one @@ -4878,9 +4681,7 @@ def test_resolve_pre_alias_model_name_fallback(): "gpt-4": {"azure": {"litellm_credentials": "team-gpt4"}}, } # Post-alias name doesn't match, but pre-alias does (team scope) - result = _resolve_credential_from_model_config( - "azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("azure/gpt-4-0613", None, team_config, pre_alias_model_name="gpt-4") assert result == "team-gpt4" # Same test for project scope @@ -4900,15 +4701,11 @@ def test_resolve_post_alias_name_takes_priority(): "gpt-4o-team-1": {"azure": {"litellm_credentials": "post-alias-cred"}}, } # Team scope - result = _resolve_credential_from_model_config( - "gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("gpt-4o-team-1", None, team_config, pre_alias_model_name="gpt-4") assert result == "post-alias-cred" # Project scope - result = _resolve_credential_from_model_config( - "gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4" - ) + result = _resolve_credential_from_model_config("gpt-4o-team-1", team_config, None, pre_alias_model_name="gpt-4") assert result == "post-alias-cred" @@ -4940,15 +4737,9 @@ def test_apply_overrides_feature_flag_disabled_by_default(): data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}} - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict + team_metadata={"model_config": {"gpt-4": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict) assert "api_base" not in data assert "api_key" not in data @@ -5108,9 +4899,7 @@ async def test_team_guardrail_merges_with_global_policy(): policy_registry = get_policy_registry() policy_registry._policies = { "global-policy": Policy( - guardrails=PolicyGuardrails( - add=["policy-guardrail-1", "policy-guardrail-2"] - ), + guardrails=PolicyGuardrails(add=["policy-guardrail-1", "policy-guardrail-2"]), ), } policy_registry._initialized = True @@ -5131,18 +4920,10 @@ async def test_team_guardrail_merges_with_global_policy(): guardrails = data["metadata"].get("guardrails", []) - assert ( - "team-direct-guardrail" in guardrails - ), f"Team guardrail missing from merged list: {guardrails}" - assert ( - "policy-guardrail-1" in guardrails - ), f"policy-guardrail-1 missing: {guardrails}" - assert ( - "policy-guardrail-2" in guardrails - ), f"policy-guardrail-2 missing: {guardrails}" - assert len(guardrails) == len( - set(guardrails) - ), f"Duplicates in guardrails list: {guardrails}" + assert "team-direct-guardrail" in guardrails, f"Team guardrail missing from merged list: {guardrails}" + assert "policy-guardrail-1" in guardrails, f"policy-guardrail-1 missing: {guardrails}" + assert "policy-guardrail-2" in guardrails, f"policy-guardrail-2 missing: {guardrails}" + assert len(guardrails) == len(set(guardrails)), f"Duplicates in guardrails list: {guardrails}" # Verify get_guardrail_from_metadata returns the merged list even # when litellm_metadata is present (the bug: it returned [] before fix) @@ -5153,9 +4934,9 @@ async def test_team_guardrail_merges_with_global_policy(): dummy = _DummyGuardrail(guardrail_name="team-direct-guardrail") returned = dummy.get_guardrail_from_metadata(data) - assert ( - "team-direct-guardrail" in returned - ), f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + assert "team-direct-guardrail" in returned, ( + f"get_guardrail_from_metadata shadowed by litellm_metadata; got: {returned}" + ) finally: policy_registry._policies = {} @@ -5208,9 +4989,7 @@ def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata(): } result = dummy.get_guardrail_from_metadata(data) - assert result == [ - "my-guardrail" - ], f"Expected guardrails from litellm_metadata fallback, got: {result}" + assert result == ["my-guardrail"], f"Expected guardrails from litellm_metadata fallback, got: {result}" def _build_request_mock_with_headers(headers: dict) -> Request: @@ -5237,9 +5016,7 @@ class TestApplyClientTagPolicyPreAuth: """ def test_merges_header_tags_into_metadata(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme,env:prod"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme,env:prod"}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -5256,9 +5033,7 @@ class TestApplyClientTagPolicyPreAuth: assert data["metadata"]["tags"] == ["tenant:acme", "env:prod"] def test_unions_header_tags_with_existing_metadata_tags(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme,env:prod"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme,env:prod"}) data = { "model": "gpt-3.5-turbo", "metadata": {"tags": ["env:prod", "team:platform"]}, @@ -5283,9 +5058,7 @@ class TestApplyClientTagPolicyPreAuth: # (inside common_checks) enforces per-tag budgets on whatever tags # it sees in request_data, including body tags. The helper only # adds header tags to metadata.tags. - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = { "model": "gpt-3.5-turbo", "tags": ["root-tag"], @@ -5312,9 +5085,7 @@ class TestApplyClientTagPolicyPreAuth: ] def test_uses_litellm_metadata_when_present(self): - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = { "model": "gpt-3.5-turbo", "litellm_metadata": {"foo": "bar"}, @@ -5409,9 +5180,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:paid": return 0.50 return fallback_spend @@ -5446,9 +5215,7 @@ class TestApplyClientTagPolicyPreAuth: from litellm.proxy.auth.auth_checks import _tag_max_budget_check from litellm.proxy.utils import ProxyLogging - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = {"model": "gpt-3.5-turbo"} user_api_key_dict = UserAPIKeyAuth( api_key="hashed-key", @@ -5468,9 +5235,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -5505,9 +5270,7 @@ class TestApplyClientTagPolicyPreAuth: "/v1/messages", ], ) - async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route( - self, route - ): + async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route(self, route): """Regression: on LITELLM_METADATA_ROUTES (bedrock, /v1/messages, ...), common_checks pre-seeds ``litellm_metadata`` and writes key tags there before ``_tag_max_budget_check`` reads from the same key. The auth wrapper @@ -5521,9 +5284,7 @@ class TestApplyClientTagPolicyPreAuth: from litellm.proxy.auth.auth_checks import common_checks from litellm.proxy.utils import ProxyLogging - request_mock = _build_request_mock_with_headers( - {"x-litellm-tags": "tenant:acme"} - ) + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "tenant:acme"}) data = {"model": "us.anthropic.claude-sonnet-4-6"} valid_token = UserAPIKeyAuth( token="test-token", @@ -5548,9 +5309,7 @@ class TestApplyClientTagPolicyPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:tenant:acme": return 0.50 return fallback_spend @@ -5718,9 +5477,7 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:engineering": return 0.50 return fallback_spend @@ -5771,9 +5528,7 @@ class TestApplyKeyTagsPreAuth: litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), ) - async def mock_get_current_spend( - counter_key, fallback_spend, max_budget=None, **kwargs - ): + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): if counter_key == "spend:tag:engineering": return 0.05 return fallback_spend @@ -5854,9 +5609,7 @@ def test_resolve_provider_from_deployment_falls_back_to_pre_alias(): router.get_deployment_by_model_group_name.side_effect = lookup - result = _resolve_provider_from_deployment( - router, "post-alias-name", pre_alias_model_name="pre-alias-name" - ) + result = _resolve_provider_from_deployment(router, "post-alias-name", pre_alias_model_name="pre-alias-name") assert result == "bedrock" @@ -5921,17 +5674,9 @@ def test_apply_overrides_no_router_keeps_legacy_behaviour(setup_test_credentials data = {"model": "gpt-4"} user_api_key_dict = UserAPIKeyAuth( api_key="test-key", - team_metadata={ - "model_config": { - "defaultconfig": { - "azure": {"litellm_credentials": "hotel-azure-eastus"} - } - } - }, - ) - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict, llm_router=None + team_metadata={"model_config": {"defaultconfig": {"azure": {"litellm_credentials": "hotel-azure-eastus"}}}}, ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict, llm_router=None) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" @@ -5957,9 +5702,7 @@ def test_apply_overrides_provider_prefix_in_model_skips_router_lookup( ) router = MagicMock() - _apply_credential_overrides_from_model_config( - data=data, user_api_key_dict=user_api_key_dict, llm_router=router - ) + _apply_credential_overrides_from_model_config(data=data, user_api_key_dict=user_api_key_dict, llm_router=router) assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" router.get_deployment_by_model_group_name.assert_not_called() @@ -6338,9 +6081,7 @@ def test_get_sanitized_user_information_from_key_drops_callback_config(): }, ) - result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict - ) + result = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) auth_metadata = result["user_api_key_auth_metadata"] assert "logging" not in auth_metadata @@ -6380,9 +6121,7 @@ def test_team_alias_targeting_deleted_team_deployment_keeps_requested_model(monk ) with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): - _update_model_if_team_alias_exists( - data=test_data, user_api_key_dict=user_api_key_dict - ) + _update_model_if_team_alias_exists(data=test_data, user_api_key_dict=user_api_key_dict) assert test_data.get("model") == "gpt-4" @@ -6406,9 +6145,7 @@ def test_team_alias_targeting_live_team_deployment_still_rewrites(monkeypatch): ) with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): - _update_model_if_team_alias_exists( - data=test_data, user_api_key_dict=user_api_key_dict - ) + _update_model_if_team_alias_exists(data=test_data, user_api_key_dict=user_api_key_dict) assert test_data.get("model") == "model_name_team-1_live-uuid" @@ -6545,7 +6282,6 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_ assert value not in logged - @pytest.mark.parametrize( "header, expected_redacted", [ @@ -6571,7 +6307,6 @@ def test_redact_credential_headers_classifies_each_header(header, expected_redac assert headers[header] == "secret-value" - @pytest.mark.asyncio async def test_add_litellm_data_to_request_debug_log_does_not_print_credentials(): """The request-header debug line carries values the stdout secret filter does not match.""" @@ -7191,8 +6926,7 @@ AUTHORIZATION_HEADER_CASINGS = ["authorization", "Authorization", "AUTHORIZATION LEAK_TARGET_PROVIDERS = ["bedrock", "bedrock_converse", "vertex_ai"] BEDROCK_ENDPOINT = ( - "https://bedrock-runtime.us-west-2.amazonaws.com" - "/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke" + "https://bedrock-runtime.us-west-2.amazonaws.com/model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/invoke" ) BEDROCK_REGION = "us-west-2" BEDROCK_REQUEST_DATA = {"messages": [{"role": "user", "content": "Say OK"}], "max_tokens": 32} @@ -7250,9 +6984,7 @@ def _signed_headers_component(signature: str, component: str) -> str: @pytest.mark.parametrize("authorization_header_name", AUTHORIZATION_HEADER_CASINGS) @pytest.mark.parametrize("custom_llm_provider", LEAK_TARGET_PROVIDERS) -def test_oauth_credential_is_never_forwarded_to_bedrock_or_vertex( - authorization_header_name, custom_llm_provider -): +def test_oauth_credential_is_never_forwarded_to_bedrock_or_vertex(authorization_header_name, custom_llm_provider): """ A client's Anthropic OAuth credential is meaningless to AWS and Google, and sending it there both breaks the request and hands a third-party cloud a credential it should @@ -7280,9 +7012,7 @@ def test_oauth_credential_entry_is_scoped_to_anthropic_alone(): if not isinstance(scoped_headers, list): scoped_headers = [scoped_headers] - credential_entries = [ - entry for entry in scoped_headers if OAUTH_TOKEN in entry["extra_headers"].values() - ] + credential_entries = [entry for entry in scoped_headers if OAUTH_TOKEN in entry["extra_headers"].values()] assert [entry["custom_llm_provider"] for entry in credential_entries] == ["anthropic"] @@ -7344,9 +7074,7 @@ def test_bedrock_get_request_headers_keeps_the_sigv4_signature(): def test_bedrock_api_key_deployment_keeps_its_own_bearer_token(): forwarded = _headers_forwarded_to(_client_headers(), "bedrock") - signed = _signed_headers_for_bedrock( - {"Content-Type": "application/json", **forwarded}, api_key=BEDROCK_API_KEY - ) + signed = _signed_headers_for_bedrock({"Content-Type": "application/json", **forwarded}, api_key=BEDROCK_API_KEY) assert _authorization_values(signed) == [f"Bearer {BEDROCK_API_KEY}"] @@ -7369,6 +7097,8 @@ def test_vertex_sends_exactly_one_authorization_header(): vertex_request_headers.update(forwarded) assert _authorization_values(vertex_request_headers) == [GOOGLE_ACCESS_TOKEN] + + @pytest.mark.asyncio async def test_newrelic_team_callback_vars_reach_trusted_field(): """A key with a newrelic team callback stamps its vars into the proxy-owned @@ -7737,13 +7467,13 @@ def _request_for(path: str) -> MagicMock: return request -def _spend_log_session_id(data: dict[str, object]) -> str: - """Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id.""" +def _spend_log_session_id(data: dict[str, object], metadata_key: str = "metadata") -> str | None: + """Resolve session_id the way LiteLLM_SpendLogs does, reading the omit decision stamped on the request.""" from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log - metadata = data["metadata"] + metadata = data[metadata_key] assert isinstance(metadata, dict) litellm_params = get_litellm_params( litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None, @@ -7754,7 +7484,12 @@ def _spend_log_session_id(data: dict[str, object]) -> str: logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"), litellm_params=litellm_params, ) - return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id}) + return _get_session_id_for_spend_log( + kwargs={}, + metadata=metadata, + standard_logging_payload={"trace_id": trace_id}, + omit_when_missing=bool(metadata.get(SESSION_ID_OMITTED_METADATA_KEY)), + ) @pytest.mark.asyncio @@ -7780,9 +7515,10 @@ async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ assert isinstance(callback_session_id, str) and len(callback_session_id) == 36 assert _spend_log_session_id(updated) == callback_session_id assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True - assert get_fireworks_session_id( - {"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]} - ) is None + assert ( + get_fireworks_session_id({"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]}) + is None + ) @pytest.mark.asyncio @@ -7800,6 +7536,44 @@ async def test_missing_session_id_unset_keeps_legacy_divergence(): assert _spend_log_session_id(updated) == "per-call-random-trace-id" +@pytest.mark.asyncio +async def test_missing_session_id_omit_leaves_spend_log_session_id_null(): + """Under `omit` a traceparent still becomes the trace id but never a session id, so SpendLogs and + Langfuse agree on having no session.""" + request = _request_for("/v1/chat/completions") + request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert "session_id" not in updated["metadata"] + assert updated["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert updated["metadata"][SESSION_ID_OMITTED_METADATA_KEY] is True + assert _spend_log_session_id(updated) is None + + +@pytest.mark.asyncio +async def test_missing_session_id_omit_keeps_client_supplied_session_id(): + request = _request_for("/v1/chat/completions") + request.headers = {"x-litellm-session-id": "client-session-1"} + + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={"missing_session_id": "omit"}, + ) + + assert updated["metadata"]["session_id"] == "client-session-1" + assert _spend_log_session_id(updated) == "client-session-1" + + @pytest.mark.asyncio async def test_missing_session_id_generate_reuses_traceparent_trace_id(): """A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it.""" @@ -7895,3 +7669,38 @@ async def test_missing_session_id_unknown_value_is_ignored(): ) assert "session_id" not in updated["metadata"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path, sent_in, metadata_key, general_settings", + [ + ("/v1/chat/completions", "metadata", "metadata", {}), + ("/v1/chat/completions", "litellm_metadata", "metadata", {}), + ("/v1/chat/completions", "litellm_metadata", "metadata", {"missing_session_id": "generate"}), + ("/v1/messages", "litellm_metadata", "litellm_metadata", {}), + ("/v1/messages", "metadata", "litellm_metadata", {}), + ("/mcp/tools", "metadata", "metadata", {"missing_session_id": "omit"}), + ("/mcp/tools", "litellm_metadata", "metadata", {"missing_session_id": "omit"}), + ], +) +async def test_client_supplied_omit_marker_never_reaches_the_spend_log( + path: str, sent_in: str, metadata_key: str, general_settings: dict[str, str] +): + """The omit marker is proxy-owned: only the pre-call policy may set it. A caller that sends it in either + metadata bucket, including the one later merged into the route's bucket, must not be able to null out + SpendLogs.session_id on a request the proxy did not omit.""" + updated = await add_litellm_data_to_request( + data={"model": "gpt-4o", "messages": [], sent_in: {SESSION_ID_OMITTED_METADATA_KEY: True}}, + request=_request_for(path), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings=general_settings, + ) + + assert SESSION_ID_OMITTED_METADATA_KEY not in updated[metadata_key] + assert _spend_log_session_id(updated, metadata_key) == ( + updated[metadata_key]["session_id"] + if general_settings.get("missing_session_id") == "generate" + else "per-call-random-trace-id" + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 098b43f6433..5246aaf6904 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25779,9 +25779,9 @@ export interface components { mcp_xff_num_trusted_hops?: number | null; /** * Missing Session Id - * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. + * @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400; 'omit' leaves SpendLogs.session_id null, matching callbacks such as Langfuse that only record a client-established metadata.session_id. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id. */ - missing_session_id?: ("generate" | "reject") | null; + missing_session_id?: ("generate" | "reject" | "omit") | null; /** * Model List Healthy Only * @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called. From 9c7c7a05ac92aebf1694ba087501f4b7c0c1416d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:58:25 -0700 Subject: [PATCH 115/167] test(router): type the deployment affinity JWT test helpers --- .../test_deployment_affinity_check.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 1852d641f0a..b5651062098 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1,11 +1,12 @@ import asyncio +import itertools +import json +from collections.abc import Sequence +from typing import Final from unittest.mock import AsyncMock, patch import pytest - -import json - import litellm from litellm.caching.dual_cache import DualCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( @@ -102,11 +103,10 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): # Deterministic routing: first selection uses seq[0], second selection attempts seq[1] # unless the list has been filtered to length=1 by deployment affinity. - choice_calls = {"count": 0} + choice_calls: Final = itertools.count(1) - def deterministic_choice(seq): - choice_calls["count"] += 1 - if choice_calls["count"] == 1: + def deterministic_choice(seq: Sequence[dict[str, object]]) -> dict[str, object]: + if next(choice_calls) == 1: return seq[0] return seq[1] if len(seq) > 1 else seq[0] @@ -1000,7 +1000,7 @@ async def test_model_group_affinity_config_overrides_global(): assert len(filtered) == 2 -def _jwt_metadata(user_id: str) -> dict: +def _jwt_metadata(user_id: str) -> dict[str, str | None]: return {"user_api_key_hash": None, "user_api_key_user_id": user_id} @@ -1090,7 +1090,7 @@ async def test_proxy_jwt_auth_metadata_pins_per_user(): enable_responses_api_affinity=False, ) - def proxy_request(user_id: str) -> dict: + def proxy_request(user_id: str) -> dict[str, object]: return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data={"model": model_group, "messages": [{"role": "user", "content": "hi"}], "metadata": {}}, user_api_key_dict=UserAPIKeyAuth(api_key=None, user_id=user_id), @@ -1098,12 +1098,14 @@ async def test_proxy_jwt_auth_metadata_pins_per_user(): ) alice_request = proxy_request("jwt-user-alice") - assert alice_request["metadata"]["user_api_key_hash"] is None + alice_metadata = alice_request["metadata"] + assert isinstance(alice_metadata, dict) + assert alice_metadata["user_api_key_hash"] is None await callback.async_pre_call_deployment_hook( kwargs={ **alice_request, - "metadata": {**alice_request["metadata"], "deployment_model_name": model_group}, + "metadata": {**alice_metadata, "deployment_model_name": model_group}, "model_info": {"id": "openai-deployment-b"}, }, call_type=None, From 92122086ecc9271640d284fe4e1db32bdfc205ff Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 19:12:11 +0000 Subject: [PATCH 116/167] fix: stop a cleared Organization field from failing key creation (#39316) * fix: stop a cleared Organization field from failing key creation Clearing the Organization combobox in the Create Key modal left organization_id set to an empty string, so /key/generate looked up an organization named "" and failed with "Organization doesn't exist in db. Organization=". OrganizationDropdown now emits null on clear, and GenerateKeyRequest normalizes an empty organization_id or project_id to None the same way it already does for team_id. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: drop customer-specific docstring from key request normalization test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin --- litellm/proxy/_types.py | 2 +- .../test_key_management_endpoints.py | 12 ++++++++++++ .../common_components/OrganizationDropdown.test.tsx | 11 +++++++++++ .../common_components/OrganizationDropdown.tsx | 4 ++-- .../src/components/organisms/create_key_button.tsx | 6 +++--- .../src/components/templates/key_edit_view.tsx | 6 +++--- 6 files changed, 32 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0aea72be1e2..98464d3a127 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,7 +1216,7 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None - @field_validator("team_id", "organization_id", mode="before") + @field_validator("team_id", "organization_id", "project_id", mode="before") @classmethod def treat_cleared_id_as_unset(cls, v: object) -> object: if v == "": diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 68704f476b5..7954a4693cc 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17505,6 +17505,18 @@ def test_generate_key_request_blank_team_id_is_personal(): assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" +def test_generate_key_request_blank_organization_and_project_id_are_unset(): + from litellm.proxy._types import RegenerateKeyRequest + + cleared = GenerateKeyRequest(organization_id="", project_id="") + assert cleared.organization_id is None + assert cleared.project_id is None + assert "organization_id" not in cleared.model_dump(exclude_none=True) + assert RegenerateKeyRequest(organization_id="").organization_id is None + assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").organization_id == "org-1" + assert GenerateKeyRequest(organization_id="org-1", project_id="proj-1").project_id == "proj-1" + + def test_key_request_blank_organization_id_is_unset(): from litellm.proxy._types import RegenerateKeyRequest, UpdateKeyRequest diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx index e0b2b897b36..524ec6a715e 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.test.tsx @@ -58,6 +58,17 @@ describe("OrganizationDropdown", () => { expect(onChange.mock.calls[0][0]).toBe("org-1"); }); + it("emits null, never the empty string, when the selection is cleared", async () => { + const onChange = vi.fn(); + const user = userEvent.setup(); + render(); + + await user.click(screen.getByRole("button", { name: "Clear" })); + + expect(onChange).toHaveBeenCalledTimes(1); + expect(onChange).toHaveBeenCalledWith(null); + }); + it("should filter options by organization id", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx index 8da35ecd02e..663028b2d92 100644 --- a/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/OrganizationDropdown.tsx @@ -5,7 +5,7 @@ import { Organization } from "../networking"; interface OrganizationDropdownProps { organizations?: Organization[] | null; value?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; disabled?: boolean; loading?: boolean; style?: React.CSSProperties; @@ -32,7 +32,7 @@ const OrganizationDropdown: React.FC = ({ sublabel: org.organization_id, }))} value={value} - onValueChange={(organizationId) => onChange?.(organizationId)} + onValueChange={(organizationId) => onChange?.(organizationId || null)} placeholder={placeholder} emptyText={loading ? "Loading organizations…" : "No organizations found"} disabled={disabled} diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 5749541dcea..ee3a88acba0 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -587,9 +587,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp } }; - const changeOrganization = (write: FieldWrite) => (orgId: string) => { - write(orgId || undefined); - setSelectedOrganizationId(orgId || null); + const changeOrganization = (write: FieldWrite) => (orgId: string | null) => { + write(orgId ?? undefined); + setSelectedOrganizationId(orgId); // Clear team and project when org changes setSelectedCreateKeyTeam(null); setSelectedProjectId(null); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 3e772fd0e9b..1330be2788b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -303,9 +303,9 @@ export function KeyEditView({ } }; - const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | undefined) => { - setField(orgId || null); - setSelectedOrganizationId(orgId || null); + const handleOrganizationChange = (setField: (value: string | null) => void, orgId: string | null) => { + setField(orgId); + setSelectedOrganizationId(orgId); form.setValue("team_id", undefined); }; From 72273ea88b07d158097fa7fa736be2f8d52a4ed1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 12:13:12 -0700 Subject: [PATCH 117/167] refactor(proxy): compose the /user/list search predicate without mutating the where dict Greptile flagged the new search branch in get_users for extending the endpoint's indexed-assignment pattern. The search predicate now comes from a small pure helper and is merged in the one-shot comprehension that already strips unset Query params, so the where clause Prisma receives is unchanged Claude-Session: https://claude.ai/code/session_018yW93iDaEMhoQUXcYjus7D --- .../internal_user_endpoints.py | 36 ++++++++++++------- 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 73d6761e17f..ad6fe0f67e5 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,6 +17,7 @@ import json import traceback from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime, timezone +from types import MappingProxyType from typing import Any, Final, Literal, Protocol, cast, overload import fastapi @@ -2069,6 +2070,22 @@ async def _authorize_user_list_request( return ",".join(allowed_org_ids) +_NO_SEARCH_WHERE: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _user_search_where(search: str | None) -> Mapping[str, object]: + """Prisma predicate for `/user/list?search=`: user_id or user_email contains it, case-insensitive.""" + if not search: + return _NO_SEARCH_WHERE + search_where: Final[UserSearchWhere] = { + "OR": ( + {"user_id": {"contains": search, "mode": "insensitive"}}, + {"user_email": {"contains": search, "mode": "insensitive"}}, + ) + } + return search_where + + @router.get( "/user/list", tags=["Internal User management"], @@ -2175,15 +2192,6 @@ async def get_users( "mode": "insensitive", # Case-insensitive search } - if search: - search_where: Final[UserSearchWhere] = { - "OR": ( - {"user_id": {"contains": search, "mode": "insensitive"}}, - {"user_email": {"contains": search, "mode": "insensitive"}}, - ) - } - where_conditions["OR"] = search_where["OR"] - if team is not None and isinstance(team, str): where_conditions["teams"] = { "has": team # Array contains for string arrays in Prisma @@ -2201,7 +2209,11 @@ async def get_users( where_conditions["organization_memberships"] = {"some": {"organization_id": {"in": org_id_list}}} ## Filter any none fastapi.Query params - e.g. where_conditions: {'user_email': {'contains': Query(None), 'mode': 'insensitive'}, 'teams': {'has': Query(None)}} - where_conditions = {k: v for k, v in where_conditions.items() if v is not None} + where: Final[Mapping[str, object]] = { + key: value + for key, value in (*where_conditions.items(), *_user_search_where(search).items()) + if value is not None + } # Build order_by conditions @@ -2210,14 +2222,14 @@ async def get_users( ) users: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await UserRepository(prisma_client).table.find_many( - where=where_conditions, + where=where, skip=skip, take=page_size, order=(order_by if order_by else {"created_at": "desc"}), # Default to created_at desc if no sort specified ) # Get total count of user rows - total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where) # Get key count for each user user_key_counts: Final = await get_user_key_counts(prisma_client, [user.user_id for user in users]) From fb4b1e728a18a56f8ec93d72ade3c237c6bbddb6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 3 Sep 2026 12:16:08 -0700 Subject: [PATCH 118/167] test(team-race): wait on pg_locks instead of a fixed sleep The three race tests claimed to pin the interleaving deterministically, but lock_acquired.set() ran as the first statement of the task coroutine, before the awaited endpoint call, so it only signalled that the task had started. The real synchronisation was `await asyncio.sleep(0.2)` followed by `assert not task.done()`, which is a timing assumption on a box running four xdist workers against one Postgres. Take the blocking connection's advisory lock key straight out of pg_locks, then poll from the unblocked watcher connection until a non-granted lock on that same key appears. That is the condition the sleep was standing in for, and it holds however slow the machine is. If the endpoint returns without ever queueing, the helper now fails with the endpoint's own exception chained on instead of a bare assert. The two tests reported failing in CI used sleep(0.2); the third used sleep(0.3) and was not reported, which is consistent with the margin being the cause. --- .../test_team_delete_member_add_race.py | 74 +++++++++++++------ 1 file changed, 53 insertions(+), 21 deletions(-) diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py index 7577570be48..9af742b0dfe 100644 --- a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py +++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py @@ -9,14 +9,16 @@ be forced by a sequential script: it needs one request to be genuinely mid-fligh other commits. A mocked prisma cannot arbitrate that either, since the property under test is whether Postgres's own advisory lock actually serializes the two requests. -These tests pin the interleaving the same way test_access_group_team_sync.py does: a second -real connection holds the team's advisory lock in its own transaction, so the function under -test is provably blocked on it rather than hoping a sleep lands in the right gap. +These tests pin the interleaving without a timing assumption: a second real connection holds +the team's advisory lock in its own transaction, and the test then waits for Postgres itself +to report the endpoint queued behind that exact lock. A sleep can only guess whether the +endpoint has reached the lock yet; pg_locks answers it. """ import asyncio import json import os +import time import uuid from contextlib import asynccontextmanager from datetime import timedelta @@ -39,6 +41,48 @@ _DELETE_SEEDED = 'DELETE FROM "LiteLLM_TeamMembership" WHERE team_id = $1' _DELETE_USER = 'DELETE FROM "LiteLLM_UserTable" WHERE user_id = $1' _DELETE_TEAM = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = $1' _LOCK_SQL = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked" +_HELD_LOCK_KEY_SQL = ( + "SELECT classid::bigint AS classid, objid::bigint AS objid FROM pg_locks " + "WHERE locktype = 'advisory' AND granted AND pid = pg_backend_pid()" +) +_LOCK_WAITER_SQL = ( + "SELECT count(*)::int AS waiters FROM pg_locks " + "WHERE locktype = 'advisory' AND NOT granted " + "AND classid::bigint = $1 AND objid::bigint = $2" +) +_LOCK_WAIT_TIMEOUT_SECONDS = 20.0 +_LOCK_POLL_SECONDS = 0.01 + + +async def _hold_team_lock(held, team_id: str) -> tuple[int, int]: + """Take the team's advisory lock and return its pg_locks key. + + Reading the key back off our own backend avoids re-deriving hashtext()'s signed + 32-bit split here, and pins the watcher to this lock rather than to any advisory + lock another xdist worker happens to hold on the same database.""" + await held.query_raw(_LOCK_SQL, team_id) + rows = await held.query_raw(_HELD_LOCK_KEY_SQL) + assert len(rows) == 1, f"expected exactly one advisory lock on the blocking connection, got {rows}" + return rows[0]["classid"], rows[0]["objid"] + + +async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what: str) -> None: + """Block until Postgres reports `task` queued behind the held lock. + + This is the assertion that the endpoint serializes on the team's advisory lock, and it + is what a fixed sleep was standing in for: the endpoint is only provably waiting once a + non-granted advisory lock on the same key exists. `watcher` must be a connection that is + not itself blocked, so it can observe the queue.""" + classid, objid = lock_key + deadline = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS + while time.monotonic() < deadline: + if task.done(): + raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() + rows = await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid) + if rows[0]["waiters"]: + return + await asyncio.sleep(_LOCK_POLL_SECONDS) + raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s") def _race_ids() -> tuple[str, str]: @@ -110,10 +154,8 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def add_member(): - lock_acquired.set() await _add_team_members_to_team( data=TeamMemberAddRequest( team_id=team_id, @@ -128,11 +170,9 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(add_member()) - await lock_acquired.wait() - await asyncio.sleep(0.2) - assert not task.done(), "member_add did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "member_add") # the delete wins the race: strip the team row while the lock is held await held.execute_raw(_DELETE_TEAM, team_id) @@ -185,10 +225,8 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def run_delete(): - lock_acquired.set() return await team_member_delete( data=TeamMemberDeleteRequest(team_id=team_id, user_id=user_id), user_api_key_dict=_admin_auth(), @@ -196,11 +234,9 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await lock_acquired.wait() - await asyncio.sleep(0.2) - assert not task.done(), "member_delete did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "member_delete") # member_add wins the race: it adds `other_user` while holding the lock await held.litellm_teamtable.update( @@ -265,10 +301,8 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): blocker = Prisma() await blocker.connect() - lock_acquired = asyncio.Event() async def run_delete(): - lock_acquired.set() return await delete_team( data=DeleteTeamRequest(team_ids=[team_id]), http_request=MagicMock(), @@ -278,11 +312,9 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): try: async with blocker.tx(timeout=timedelta(seconds=30)) as held: - await held.query_raw(_LOCK_SQL, team_id) + lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await lock_acquired.wait() - await asyncio.sleep(0.3) - assert not task.done(), "delete_team did not wait on the team's advisory lock" + await _await_lock_contention(db, lock_key, task, "delete_team") # member_add wins the race: write the reference while holding the lock await held.litellm_usertable.upsert( From 8ac704a855cc2a27887138a5a2dd121f4a81b1ae Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 3 Sep 2026 12:16:16 -0700 Subject: [PATCH 119/167] test(responses): cancel the streaming response while it is still in flight test_cancel_streaming_response drained the whole stream and only then called cancel, by which point the response had finished and the cancel was expected to fail. The test passed only because the surrounding except matched the literal string "Cannot cancel a completed response", which is upstream OpenAI's wording, not ours: it appears nowhere in this repo. Any change to that text, a different status code, or the background job still running flipped the result. The test also never exercised cancellation, and its only positive assertion was hasattr(cancel_response, "id"). Break out of the stream at the first chunk carrying a response id and cancel there, with a prompt long enough that the response cannot have completed in the meantime. Both proxy cancel paths, the polling handler and the provider passthrough, settle on status "cancelled", so assert that and the returned id rather than a provider error string. --- .../test_e2e_openai_responses_api.py | 54 +++++++++---------- 1 file changed, 24 insertions(+), 30 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index abae26e02cd..b24ac0bdb96 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -154,42 +154,36 @@ def test_cancel_response(): def test_cancel_streaming_response(): - try: - client = get_test_client() - from litellm.types.llms.openai import ResponsesAPIResponse + """Cancel a background streaming response while it is still generating. - stream = client.responses.create( - model="gpt-5.5", - input="just respond with the word 'ping'", - stream=True, - background=True, - ) + The prompt is deliberately long-running and the stream is abandoned at the first + chunk carrying a response id, so the response is provably still in flight when the + cancel lands. Draining the stream first would finish the response, making the cancel + fail and leaving nothing but the provider's error wording to assert on. + """ + client = get_test_client() - collected_chunks = [] + with client.responses.create( + model="gpt-5.5", + input="write a 2000 word essay on the history of the printing press", + stream=True, + background=True, + ) as stream: + chunk_count = 0 response_id = None for chunk in stream: - print("stream chunk=", chunk) - collected_chunks.append(chunk) - # Extract response ID from the first chunk that has it - if ( - response_id is None - and hasattr(chunk, "response") - and hasattr(chunk.response, "id") - ): - response_id = chunk.response.id + chunk_count += 1 + response_id = getattr(getattr(chunk, "response", None), "id", None) + if response_id is not None: + break - assert len(collected_chunks) > 0 + assert chunk_count > 0, "stream produced no chunks" + assert response_id is not None, "no streamed chunk carried a response id to cancel" - # cancel the response if we got a response ID - if response_id: - cancel_response = client.responses.cancel(response_id) - print("CANCEL streaming response=", cancel_response) - assert hasattr(cancel_response, "id") - except Exception as e: - if "Cannot cancel a completed response" in str(e): - pass - else: - raise e + cancel_response = client.responses.cancel(response_id) + print("CANCEL streaming response=", cancel_response) + assert cancel_response.id == response_id + assert cancel_response.status == "cancelled" def test_cancel_invalid_response_id(): From ae62311dfd1c3ffe77a6310a566aff9dfff0af9a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 12:17:50 -0700 Subject: [PATCH 120/167] chore(lint): ratchet the LIT010 ceiling down by the rebind this branch removed Claude-Session: https://claude.ai/code/session_018yW93iDaEMhoQUXcYjus7D --- type-discipline-budget.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index cbcb5dca443..4a65be15b14 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16477 + "limit": 16476 }, "LIT011": { "limit": 5519 From 51d821ae45aef7fa95e222ebcb55486210d08543 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:21:54 -0700 Subject: [PATCH 121/167] fix(cost): bill bedrock_mantle web search at $12 per 1k queries using Bedrock's reported count --- .../llm_cost_calc/tool_call_cost_tracking.py | 34 +++-- ...odel_prices_and_context_window_backup.json | 25 ++++ litellm/types/llms/openai.py | 13 ++ model_prices_and_context_window.json | 25 ++++ .../test_tool_call_cost_tracking.py | 118 ++++++++++++++++++ 5 files changed, 207 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 5504756ceb8..bf99035a6b1 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -5,6 +5,8 @@ Helper utilities for tracking the cost of built-in tools. from collections.abc import Mapping from typing import Final, Literal +from pydantic import ValidationError + import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -13,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, + ResponsesToolUsage, WebSearchOptions, ) from litellm.types.utils import ( @@ -32,6 +35,17 @@ def _output_item_type(output_item: object) -> str | None: return item_type if isinstance(item_type, str) else None +def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: + tool_usage: Final = getattr(response_object, "tool_usage", None) + if tool_usage is None: + return None + try: + web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search + except ValidationError: + return None + return None if web_search is None else web_search.num_requests + + def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: details: Final = getattr(usage, "server_side_tool_usage_details", None) if not isinstance(details, Mapping): @@ -182,15 +196,19 @@ class StandardBuiltInToolCostTracking: Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by get_cost_for_web_search_request and never reach here. This path prices per call, so it must count - the web_search_call items. Chat-completions responses only expose url_citation annotations with no - count, so they floor to a single billable search. + the web_search_call items, unless the response reports the billable count itself + (Bedrock's tool_usage.web_search.num_requests, which excludes open_page fetches). Chat-completions + responses only expose url_citation annotations with no count, so they floor to a single billable search. """ - if isinstance(response_object, ResponsesAPIResponse): - count = sum( - 1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call" - ) - return max(count, 1) - return 1 + if not isinstance(response_object, ResponsesAPIResponse): + return 1 + reported: Final = _reported_web_search_requests(response_object) + if reported is not None: + return reported + count: Final = sum( + 1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call" + ) + return max(count, 1) @staticmethod def _handle_file_search_cost( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9b358a1cedd..ffac34ea37e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -52911,6 +52911,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -52945,6 +52950,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53007,6 +53017,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53195,6 +53210,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53226,6 +53246,11 @@ "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "output_cost_per_token": 1.65e-05, "output_cost_per_token_above_272k_tokens": 2.475e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 32d88da0085..b33ff954c35 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -66,6 +66,7 @@ from pydantic import ( ConfigDict, Discriminator, Field, + NonNegativeInt, PrivateAttr, SerializerFunctionWrapHandler, field_serializer, @@ -1321,6 +1322,18 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject): model_config = {"extra": "allow"} +class WebSearchToolUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + num_requests: NonNegativeInt + + +class ResponsesToolUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + web_search: WebSearchToolUsage | None = None + + ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] """ The status of the response generation. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9b358a1cedd..ffac34ea37e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -52911,6 +52911,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -52945,6 +52950,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53007,6 +53017,11 @@ "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53195,6 +53210,11 @@ "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, @@ -53226,6 +53246,11 @@ "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, "output_cost_per_token": 1.65e-05, "output_cost_per_token_above_272k_tokens": 2.475e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.012, + "search_context_size_low": 0.012, + "search_context_size_medium": 0.012 + }, "litellm_provider": "bedrock_mantle", "max_input_tokens": 1050000, "max_output_tokens": 128000, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index fd795ffcc96..0d75301a57a 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -928,3 +928,121 @@ def test_web_search_gate_reads_server_side_tool_usage_details_without_citations( standard_built_in_tools_params=None, ) assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL + + +_BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + "bedrock_mantle/openai.gpt-5.5", + "bedrock_mantle/openai.gpt-5.4", +) + +_BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None): + from litellm.types.llms.openai import ResponsesAPIResponse + + payload = { + "id": "resp_1", + "created_at": 1756900000, + "model": model.split("/", 1)[-1], + "object": "response", + "status": "completed", + "output": [ + {"type": "web_search_call", "id": f"ws_{i}", "status": "completed", "action": action} + for i, action in enumerate(actions) + ], + } + return ResponsesAPIResponse.model_validate( + payload if tool_usage is None else {**payload, "tool_usage": tool_usage} + ) + + +def _bedrock_mantle_web_search_cost(model, response): + from litellm.types.utils import Usage + + return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + custom_llm_provider="bedrock_mantle", + standard_built_in_tools_params=None, + ) + + +@pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS) +def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model): + """ + Regression for LIT-6870: the bedrock_mantle GPT ids forward the web_search tool but carried no + search_context_cost_per_query, so every Bedrock web search (billed at $12 per 1,000 queries) + was costed at $0. Two reported queries must bill 2 x $0.012, whichever model prefix shape the + cost path resolves the deployment under. + """ + pricing = litellm.get_model_info(model)["search_context_cost_per_query"] + assert pricing == { + "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + "search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE, + } + + response = _bedrock_mantle_responses_with_web_search( + model, + actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], + tool_usage={"web_search": {"num_requests": 2}}, + ) + for cost_model in (model, model.split("/", 1)[1]): + cost = _bedrock_mantle_web_search_cost(cost_model, response) + assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}" + ) + + +@pytest.mark.parametrize("num_requests", [1, 0]) +def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests): + """ + Regression for LIT-6870: Bedrock bills one query per search and reports the billable count as + tool_usage.web_search.num_requests, while its open_page fetches share the web_search_call item + type. A search plus an open_page must bill the reported count, never the two items. + """ + model = "bedrock_mantle/openai.gpt-5.6-sol" + response = _bedrock_mantle_responses_with_web_search( + model, + actions=[ + {"type": "search", "query": "litellm"}, + {"type": "open_page", "url": "https://docs.litellm.ai/"}, + ], + tool_usage={"web_search": {"num_requests": num_requests}}, + ) + + cost = _bedrock_mantle_web_search_cost(model, response) + + assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"{num_requests} reported web search requests must bill {num_requests} x " + f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" + ) + + +@pytest.mark.parametrize( + "tool_usage", + [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}], +) +def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage): + """ + Without a usable reported count (no tool_usage, no web_search block, or a malformed one) the + per-call path must keep counting web_search_call items instead of raising or billing zero. + """ + model = "bedrock_mantle/openai.gpt-5.6-sol" + response = _bedrock_mantle_responses_with_web_search( + model, + actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], + tool_usage=tool_usage, + ) + + cost = _bedrock_mantle_web_search_cost(model, response) + + assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( + f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x " + f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" + ) From 36143b53f900d8962b1cf72749f85f20e870e961 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:34:27 -0700 Subject: [PATCH 122/167] test(responses): bound the background stream cancel e2e so an upstream stall skips fast test_cancel_streaming_response drained the whole background stream before cancelling, so an OpenAI keepalive stall held the e2e_openai_endpoints job for 301s and failed it on a generic APIError, and on a healthy day it cancelled an already completed response and swallowed the 400 without verifying a cancel. Cancel at the first event carrying a response id, bound admission to 90s, skip naming the stall when only keepalives arrived, and assert status == cancelled --- .../test_e2e_openai_responses_api.py | 70 +++++++++++-------- 1 file changed, 39 insertions(+), 31 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index abae26e02cd..604f1c84d20 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,6 +1,10 @@ +import time + import httpx -from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError import pytest +from openai import APIStatusError, BadRequestError, NotFoundError, OpenAI + +BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS = 90 def generate_key(): @@ -154,42 +158,46 @@ def test_cancel_response(): def test_cancel_streaming_response(): - try: - client = get_test_client() - from litellm.types.llms.openai import ResponsesAPIResponse + client = get_test_client() + started = time.monotonic() + stream = client.responses.create( + model="gpt-5.5", + input="just respond with the word 'ping'", + stream=True, + background=True, + timeout=BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS, + ) - stream = client.responses.create( - model="gpt-5.5", - input="just respond with the word 'ping'", - stream=True, - background=True, - ) - - collected_chunks = [] - response_id = None + keepalive_events = 0 + response_id = None + with stream: for chunk in stream: print("stream chunk=", chunk) - collected_chunks.append(chunk) - # Extract response ID from the first chunk that has it - if ( - response_id is None - and hasattr(chunk, "response") - and hasattr(chunk.response, "id") - ): + if chunk.type == "keepalive": + keepalive_events += 1 + elif getattr(chunk, "response", None) is not None: response_id = chunk.response.id + break + if time.monotonic() - started > BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS: + break - assert len(collected_chunks) > 0 + elapsed = time.monotonic() - started + if response_id is None and keepalive_events: + pytest.skip( + f"OpenAI held the background stream in keepalive for {elapsed:.0f}s " + f"({keepalive_events} keepalive events) without creating the response" + ) + assert response_id is not None, f"no response event within {elapsed:.0f}s of streaming a background response" - # cancel the response if we got a response ID - if response_id: - cancel_response = client.responses.cancel(response_id) - print("CANCEL streaming response=", cancel_response) - assert hasattr(cancel_response, "id") - except Exception as e: - if "Cannot cancel a completed response" in str(e): - pass - else: - raise e + try: + cancel_response = client.responses.cancel(response_id) + except BadRequestError as e: + if "Cannot cancel a completed response" not in str(e): + raise + print("response completed before cancel=", e) + return + print("CANCEL streaming response=", cancel_response) + assert cancel_response.status == "cancelled" def test_cancel_invalid_response_id(): From 339da4183d998b1ff2db54bc0468fa994a620095 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:34:54 -0700 Subject: [PATCH 123/167] test(cost): type the web search cost helpers and cover OpenAI-shaped tool_usage --- .../test_tool_call_cost_tracking.py | 59 ++++++++++--------- 1 file changed, 32 insertions(+), 27 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 0d75301a57a..6cc3dcceebc 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,4 +1,5 @@ import os +from collections.abc import Mapping, Sequence import pytest @@ -6,7 +7,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) -from litellm.types.llms.openai import FileSearchTool, WebSearchOptions +from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams @@ -941,9 +942,9 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 -def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None): - from litellm.types.llms.openai import ResponsesAPIResponse - +def _responses_with_web_search( + model: str, actions: Sequence[Mapping[str, str]], tool_usage: Mapping[str, object] | None = None +) -> ResponsesAPIResponse: payload = { "id": "resp_1", "created_at": 1756900000, @@ -960,26 +961,21 @@ def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None): ) -def _bedrock_mantle_web_search_cost(model, response): +def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_provider: str) -> float: from litellm.types.utils import Usage return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( model=model, response_object=response, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - custom_llm_provider="bedrock_mantle", + custom_llm_provider=custom_llm_provider, standard_built_in_tools_params=None, ) @pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS) def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model): - """ - Regression for LIT-6870: the bedrock_mantle GPT ids forward the web_search tool but carried no - search_context_cost_per_query, so every Bedrock web search (billed at $12 per 1,000 queries) - was costed at $0. Two reported queries must bill 2 x $0.012, whichever model prefix shape the - cost path resolves the deployment under. - """ + """Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike.""" pricing = litellm.get_model_info(model)["search_context_cost_per_query"] assert pricing == { "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE, @@ -987,13 +983,13 @@ def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model) "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE, } - response = _bedrock_mantle_responses_with_web_search( + response = _responses_with_web_search( model, actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], tool_usage={"web_search": {"num_requests": 2}}, ) for cost_model in (model, model.split("/", 1)[1]): - cost = _bedrock_mantle_web_search_cost(cost_model, response) + cost = _web_search_cost(cost_model, response, "bedrock_mantle") assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}" ) @@ -1001,13 +997,9 @@ def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model) @pytest.mark.parametrize("num_requests", [1, 0]) def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests): - """ - Regression for LIT-6870: Bedrock bills one query per search and reports the billable count as - tool_usage.web_search.num_requests, while its open_page fetches share the web_search_call item - type. A search plus an open_page must bill the reported count, never the two items. - """ + """A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items.""" model = "bedrock_mantle/openai.gpt-5.6-sol" - response = _bedrock_mantle_responses_with_web_search( + response = _responses_with_web_search( model, actions=[ {"type": "search", "query": "litellm"}, @@ -1016,7 +1008,7 @@ def test_web_search_call_count_prefers_provider_reported_num_requests(local_mode tool_usage={"web_search": {"num_requests": num_requests}}, ) - cost = _bedrock_mantle_web_search_cost(model, response) + cost = _web_search_cost(model, response, "bedrock_mantle") assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( f"{num_requests} reported web search requests must bill {num_requests} x " @@ -1029,20 +1021,33 @@ def test_web_search_call_count_prefers_provider_reported_num_requests(local_mode [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}], ) def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage): - """ - Without a usable reported count (no tool_usage, no web_search block, or a malformed one) the - per-call path must keep counting web_search_call items instead of raising or billing zero. - """ + """Without a usable reported count the per-call path keeps counting web_search_call items.""" model = "bedrock_mantle/openai.gpt-5.6-sol" - response = _bedrock_mantle_responses_with_web_search( + response = _responses_with_web_search( model, actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], tool_usage=tool_usage, ) - cost = _bedrock_mantle_web_search_cost(model, response) + cost = _web_search_cost(model, response, "bedrock_mantle") assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x " f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" ) + + +def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map): + """OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count.""" + response = _responses_with_web_search( + "gpt-5.6", + actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}], + tool_usage={ + "image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + "web_search": {"num_requests": 1}, + }, + ) + + cost = _web_search_cost("gpt-5.6", response, "openai") + + assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}" From 1e75668a259ee9b6c7d77abd17ef51df70250b1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:37:17 -0700 Subject: [PATCH 124/167] fix(openai): default stream usage on PrivateLink and regional api.openai.com hosts --- litellm/llms/openai/common_utils.py | 9 ++++ litellm/llms/openai/openai.py | 8 ++- tests/test_litellm/llms/openai/test_openai.py | 52 +++++++++++++++++++ .../llms/openai/test_openai_common_utils.py | 21 +++++++- 4 files changed, 84 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/llms/openai/test_openai.py diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index bcd4ea43243..2db6d78a218 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -11,6 +11,7 @@ import time import uuid from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Optional +from urllib.parse import urlsplit import httpx import openai @@ -43,6 +44,14 @@ _OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(OpenAI) _AZURE_OPENAI_INIT_PARAMS: Final[tuple[str, ...]] = _get_client_init_params(AzureOpenAI) +_OPENAI_API_HOST: Final[str] = "api.openai.com" + + +def is_openai_backed_api_base(api_base: str) -> bool: + hostname: Final = urlsplit(api_base).hostname + return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}")) + + class OpenAIError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 1cfc6e06ee9..edc8d64d9c2 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -2,7 +2,6 @@ import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast -from urllib.parse import urlparse import httpx @@ -55,6 +54,7 @@ from .common_utils import ( OpenAIError, build_output_token_limit_response, drop_params_from_unprocessable_entity_error, + is_openai_backed_api_base, is_output_token_limit_error, ) from .workload_identity import resolve_openai_workload_identity_config @@ -1190,10 +1190,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ if stream_options is not None: return {"stream_options": stream_options} - else: - # by default litellm will include usage for openai endpoints - if api_base is None or urlparse(api_base).hostname == "api.openai.com": - return {"stream_options": {"include_usage": True}} + if api_base is None or is_openai_backed_api_base(api_base): + return {"stream_options": {"include_usage": True}} return {} # Embedding diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/test_litellm/llms/openai/test_openai.py new file mode 100644 index 00000000000..136b837f191 --- /dev/null +++ b/tests/test_litellm/llms/openai/test_openai.py @@ -0,0 +1,52 @@ +import pytest + +from litellm.llms.openai.openai import OpenAIChatCompletion + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.openai.com/v1", + "https://api.openai.com:443/v1", + "https://southcentralus.privatelink.api.openai.com/v1", + "https://eu.api.openai.com/v1", + "https://us.api.openai.com/v1", + "HTTPS://API.OPENAI.COM/v1/", + ], +) +def test_get_stream_options_defaults_include_usage_on_every_openai_backed_host(api_base): + """ + PrivateLink and regional hostnames reach the real OpenAI backend, so a stream with no caller + stream_options must ask for the usage chunk exactly as the default base does. Regression guard + for LIT-6875: spend for those deployments fell back to local token counting. + """ + assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == { + "stream_options": {"include_usage": True} + } + + +@pytest.mark.parametrize( + "api_base", + [ + "https://my-gateway.example/v1", + "https://api.openai.com.evil.example/v1", + "https://notapi.openai.com/v1", + "https://gateway.example/v1?upstream=api.openai.com", + "https://openai.internal.example/api.openai.com/v1", + ], +) +def test_get_stream_options_leaves_foreign_hosts_without_a_usage_default(api_base): + """Only the host decides: an OpenAI-compatible backend elsewhere may not support stream_options at all.""" + assert OpenAIChatCompletion().get_stream_options(stream_options=None, api_base=api_base) == {} + + +@pytest.mark.parametrize( + "api_base", + ["https://southcentralus.privatelink.api.openai.com/v1", "https://my-gateway.example/v1"], +) +def test_get_stream_options_passes_caller_stream_options_through_on_any_host(api_base): + caller_options = {"include_usage": False} + assert OpenAIChatCompletion().get_stream_options(stream_options=caller_options, api_base=api_base) == { + "stream_options": caller_options + } diff --git a/tests/test_litellm/llms/openai/test_openai_common_utils.py b/tests/test_litellm/llms/openai/test_openai_common_utils.py index 3ae29e411e8..d3c21c5bd5a 100644 --- a/tests/test_litellm/llms/openai/test_openai_common_utils.py +++ b/tests/test_litellm/llms/openai/test_openai_common_utils.py @@ -7,7 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.token_counter import token_counter -from litellm.llms.openai.common_utils import BaseOpenAILLM +from litellm.llms.openai.common_utils import BaseOpenAILLM, is_openai_backed_api_base # Test parameters for different API functions API_FUNCTION_PARAMS = [ @@ -392,3 +392,22 @@ async def test_async_genuine_bad_request_still_raises(provider, stream): with pytest.raises(litellm.BadRequestError): await _call_and_drain() + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ("https://api.openai.com/v1", True), + ("https://api.openai.com:443/v1/", True), + ("https://southcentralus.privatelink.api.openai.com/v1", True), + ("https://eu.api.openai.com/v1", True), + ("HTTPS://API.OPENAI.COM/v1", True), + ("https://my-gateway.example/v1", False), + ("https://api.openai.com.evil.example/v1", False), + ("https://notapi.openai.com/v1", False), + ("https://gateway.example/v1?upstream=api.openai.com", False), + ("not a url", False), + ], +) +def test_is_openai_backed_api_base_decides_by_hostname_only(api_base, expected): + assert is_openai_backed_api_base(api_base) is expected From b38516da88ba4eb7ba3fa584e08133881f4d1108 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:37:31 -0700 Subject: [PATCH 125/167] test(responses): make the background stream cancel deterministic A five-token response can complete before the cancel lands, which put the test back on the "Cannot cancel a completed response" path it used to swallow. Ask for a long generation so the cancel always beats completion, and assert the cancelled status unconditionally --- .../test_e2e_openai_responses_api.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 604f1c84d20..0dbecbe2801 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -162,7 +162,7 @@ def test_cancel_streaming_response(): started = time.monotonic() stream = client.responses.create( model="gpt-5.5", - input="just respond with the word 'ping'", + input="count from 1 to 500, one number per line", stream=True, background=True, timeout=BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS, @@ -189,13 +189,7 @@ def test_cancel_streaming_response(): ) assert response_id is not None, f"no response event within {elapsed:.0f}s of streaming a background response" - try: - cancel_response = client.responses.cancel(response_id) - except BadRequestError as e: - if "Cannot cancel a completed response" not in str(e): - raise - print("response completed before cancel=", e) - return + cancel_response = client.responses.cancel(response_id) print("CANCEL streaming response=", cancel_response) assert cancel_response.status == "cancelled" From 425e3069b93ec005e3186f5d32fa21ef70effe17 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 3 Sep 2026 12:40:22 -0700 Subject: [PATCH 126/167] fix(proxy): expose configured model mode --- litellm/proxy/utils.py | 3 ++ litellm/router.py | 11 ++++++ litellm/types/proxy/model_listing.py | 4 +- tests/test_litellm/proxy/test_proxy_utils.py | 41 ++++++++++++++++++++ 4 files changed, 57 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cab2bd6d9db..24a65457928 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7531,6 +7531,9 @@ def create_model_info_response( max_input_tokens = configured_input if configured_output is not None: max_output_tokens = configured_output + configured_mode: Final = llm_router.get_configured_mode(model_id) + if isinstance(configured_mode, str): + base["mode"] = configured_mode if max_input_tokens is not None: base["max_input_tokens"] = max_input_tokens diff --git a/litellm/router.py b/litellm/router.py index 2b8b342d253..cfb080a24a0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9981,6 +9981,17 @@ class Router: coerce_token_limit(model_info.get("max_output_tokens")), ) + def get_configured_mode(self, model_name: str) -> "str | None": + """Return the mode explicitly configured for a concrete deployment.""" + deployment: Final = self.get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + + mode: Final = deployment.model_info.get("mode") + if isinstance(mode, str) and mode.strip(): + return mode + return None + def get_configured_display_name(self, model_name: str) -> "str | None": """ Return the display_name explicitly configured in a concrete deployment's diff --git a/litellm/types/proxy/model_listing.py b/litellm/types/proxy/model_listing.py index b59c0f2cf19..24cfa85eee4 100644 --- a/litellm/types/proxy/model_listing.py +++ b/litellm/types/proxy/model_listing.py @@ -11,8 +11,8 @@ class ModelInfoMetadata(TypedDict): class ModelInfoResponse(TypedDict): """OpenAI-compatible model object. `mode`, `max_input_tokens`, and - `max_output_tokens` are attached when the cost map knows them; `metadata` - is present only when the endpoint is called with include_metadata=true. + `max_output_tokens` are attached when the cost map or deployment config + knows them; `metadata` is present only with include_metadata=true. """ id: str diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index dcaad968663..f16c6c937d0 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -942,6 +942,47 @@ def test_create_model_info_response_uses_deployment_limits_when_not_in_cost_map( assert response["max_output_tokens"] == 8000 +def test_create_model_info_response_uses_deployment_mode_for_auto_router(): + router = litellm.Router( + model_list=[ + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + }, + { + "model_name": "claude-auto", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": "claude-sonnet", + "MEDIUM": "claude-sonnet", + "COMPLEX": "claude-sonnet", + } + }, + "complexity_router_default_model": "claude-sonnet", + }, + "model_info": { + "mode": "chat", + "max_input_tokens": 1_000_000, + "max_output_tokens": 128_000, + }, + }, + ] + ) + + response = create_model_info_response( + model_id="claude-auto", + provider="openai", + llm_router=router, + get_model_info=_raise_unmapped, + ) + + assert response["mode"] == "chat" + assert response["max_input_tokens"] == 1_000_000 + assert response["max_output_tokens"] == 128_000 + + def test_create_model_info_response_deployment_limits_override_cost_map(): router = MagicMock() router.get_configured_token_limits.return_value = (200000, None) From 897fba08c8ddb6dba99288b040ae5c7cad8a5757 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 11:08:03 -0700 Subject: [PATCH 127/167] feat(models): add gpt-6-astra pricing and metadata Adds the OpenAI gpt-6-astra entry to both price files with standard, flex, priority (fast mode), batch, and above-272K long-context rates, and regression tests covering each tier and the batch rates. --- ...odel_prices_and_context_window_backup.json | 68 +++++++++++++++++++ model_prices_and_context_window.json | 68 +++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 48 +++++++++++++ tests/test_litellm/test_cost_calculator.py | 14 ++++ ...penai_service_tier_long_context_pricing.py | 7 ++ 5 files changed, 205 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9b358a1cedd..04a67200ced 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29277,6 +29277,74 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + "cache_creation_input_token_cost_flex": 6.25e-06, + "cache_creation_input_token_cost_priority": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_read_input_token_cost_flex": 5e-07, + "cache_read_input_token_cost_priority": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "input_cost_per_token_batches": 5e-06, + "input_cost_per_token_flex": 5e-06, + "input_cost_per_token_priority": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "output_cost_per_token_batches": 2.5e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 0.0001, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9b358a1cedd..04a67200ced 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29277,6 +29277,74 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": true }, + "gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + "cache_creation_input_token_cost_flex": 6.25e-06, + "cache_creation_input_token_cost_priority": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "cache_read_input_token_cost_above_272k_tokens_flex": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_read_input_token_cost_flex": 5e-07, + "cache_read_input_token_cost_priority": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_cost_per_token_above_272k_tokens_flex": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "input_cost_per_token_batches": 5e-06, + "input_cost_per_token_flex": 5e-06, + "input_cost_per_token_priority": 2e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "output_cost_per_token_above_272k_tokens_flex": 3.75e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "output_cost_per_token_batches": 2.5e-05, + "output_cost_per_token_flex": 2.5e-05, + "output_cost_per_token_priority": 0.0001, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_272k_tokens": 1e-05, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index b7f0ca1efe1..0b6832d4bef 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -1662,6 +1662,54 @@ def test_generic_cost_per_token_gpt56_cyber( assert completion_cost == pytest.approx(completion_tokens * output_rate) +@pytest.mark.parametrize( + "service_tier,tier_multiplier", + [(None, 1.0), ("flex", 0.5), ("priority", 2.0), ("fast", 2.0)], +) +@pytest.mark.parametrize( + "prompt_tokens,input_side_multiplier,output_multiplier", + [(100000, 1.0, 1.0), (300000, 2.0, 1.5)], +) +def test_generic_cost_per_token_gpt_6_astra_price_sheet( + _local_model_cost_map, + service_tier, + tier_multiplier, + prompt_tokens, + input_side_multiplier, + output_multiplier, +): + """gpt-6-astra launch price sheet: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens. + + Above 272K prompt tokens the input-side rates double and the output rate is 1.5x on the whole + request. Flex is half the applicable rate and fast mode, billed as priority, is double it. + """ + cached_tokens = 50000 + cache_write_tokens = 40000 + text_tokens = prompt_tokens - cached_tokens - cache_write_tokens + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=cached_tokens, cache_write_tokens=cache_write_tokens + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="gpt-6-astra", + usage=usage, + custom_llm_provider="openai", + service_tier=service_tier, + ) + + input_side = tier_multiplier * input_side_multiplier + assert prompt_cost == pytest.approx( + input_side * (text_tokens * 1e-5 + cached_tokens * 1e-6 + cache_write_tokens * 1.25e-5) + ) + assert completion_cost == pytest.approx(tier_multiplier * output_multiplier * completion_tokens * 5e-5) + + @pytest.mark.parametrize( "model,input_cost,output_cost,cache_read_cost", [ diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7c2174018e8..6dc3b2790c9 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4473,3 +4473,17 @@ def test_explicit_pricing_precedes_private_provider_response_model( ) assert selected == expected + + +def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_model_cost_map): + """gpt-6-astra batch pricing is 50% off the standard $10 input and $50 output rates per 1M tokens.""" + from litellm.cost_calculator import batch_cost_calculator + + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, model="gpt-6-astra", custom_llm_provider="openai" + ) + + assert prompt_cost == pytest.approx(1000 * 5e-6) + assert completion_cost == pytest.approx(500 * 2.5e-5) diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py index c0860a5b55f..70e9c2720b8 100644 --- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -52,6 +52,12 @@ PRIORITY_LONG_CONTEXT = { "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, }, + "gpt-6-astra": { + "input_cost_per_token_above_272k_tokens_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.00015, + "cache_read_input_token_cost_above_272k_tokens_priority": 4e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5e-05, + }, } EXPECTED = {**FLEX_LONG_CONTEXT, **PRIORITY_LONG_CONTEXT} @@ -114,6 +120,7 @@ TIERED_COST_CASES = [ ("gpt-5.6-sol", "priority", 1.6e-05, 6e-05), ("gpt-5.6-terra", "priority", 8e-06, 3.6e-05), ("gpt-5.6-luna", "priority", 8e-07, 3.6e-06), + ("gpt-6-astra", "priority", 4e-05, 0.00015), ] From 4991d0bf3e58fc1022d97ce8baf1a517c656d5a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:47:13 -0700 Subject: [PATCH 128/167] fix(models): match gpt-6-astra reasoning effort levels to OpenAI docs OpenAI documents low, medium, high, xhigh, and max for gpt-6-astra, with no none level, so the entry stops advertising none and starts advertising max. --- .../model_prices_and_context_window_backup.json | 3 ++- model_prices_and_context_window.json | 3 ++- .../test_reasoning_effort_capability.py | 17 +++++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 04a67200ced..385cf22d06c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29330,9 +29330,10 @@ ], "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_cache_breakpoint": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 04a67200ced..385cf22d06c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29330,9 +29330,10 @@ ], "supports_computer_use": true, "supports_function_calling": true, + "supports_max_reasoning_effort": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, - "supports_none_reasoning_effort": true, + "supports_none_reasoning_effort": false, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_cache_breakpoint": true, diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index a0dbf3b6637..7b2e45ab3ed 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -371,3 +371,20 @@ class TestKimiK3AdvertisesItsDocumentedLevels: "low", "high", ) + + +class TestGpt6AstraAdvertisesItsDocumentedLevels: + def test_the_entry_advertises_low_through_max_without_none(self, local_model_cost_map): + """OpenAI documents low, medium, high, xhigh and max for gpt-6-astra. Unlike gpt-5.6-sol it + does not take none, so a group must not offer none and must offer max.""" + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="gpt-6-astra", custom_llm_provider="openai")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "low", + "medium", + "high", + "xhigh", + "max", + ) From 0904a9223bb4e01c2ad9d6ceb56529fe4637b4b9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:56:56 -0700 Subject: [PATCH 129/167] test(responses): collect the admitted stream events without local mutation --- .../test_e2e_openai_responses_api.py | 46 +++++++++++-------- 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index 0dbecbe2801..7e338dafb86 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -1,10 +1,13 @@ import time +from collections.abc import Iterator +from typing import Final import httpx import pytest -from openai import APIStatusError, BadRequestError, NotFoundError, OpenAI +from openai import APIStatusError, BadRequestError, NotFoundError, OpenAI, Stream +from openai.types.responses import ResponseStreamEvent -BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS = 90 +BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS: Final = 90 def generate_key(): @@ -157,10 +160,25 @@ def test_cancel_response(): raise e +def admitted_response_id(chunk: ResponseStreamEvent) -> str | None: + response: Final = getattr(chunk, "response", None) + return None if response is None else response.id + + +def events_until_admission(stream: Stream[ResponseStreamEvent], started: float) -> Iterator[ResponseStreamEvent]: + for chunk in stream: + print("stream chunk=", chunk) + yield chunk + if admitted_response_id(chunk) is not None: + return + if time.monotonic() - started > BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS: + return + + def test_cancel_streaming_response(): - client = get_test_client() - started = time.monotonic() - stream = client.responses.create( + client: Final = get_test_client() + started: Final = time.monotonic() + stream: Final = client.responses.create( model="gpt-5.5", input="count from 1 to 500, one number per line", stream=True, @@ -168,20 +186,12 @@ def test_cancel_streaming_response(): timeout=BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS, ) - keepalive_events = 0 - response_id = None with stream: - for chunk in stream: - print("stream chunk=", chunk) - if chunk.type == "keepalive": - keepalive_events += 1 - elif getattr(chunk, "response", None) is not None: - response_id = chunk.response.id - break - if time.monotonic() - started > BACKGROUND_STREAM_ADMISSION_DEADLINE_SECONDS: - break + events: Final = tuple(events_until_admission(stream, started)) - elapsed = time.monotonic() - started + elapsed: Final = time.monotonic() - started + keepalive_events: Final = sum(1 for chunk in events if chunk.type == "keepalive") + response_id: Final = next((rid for rid in map(admitted_response_id, events) if rid is not None), None) if response_id is None and keepalive_events: pytest.skip( f"OpenAI held the background stream in keepalive for {elapsed:.0f}s " @@ -189,7 +199,7 @@ def test_cancel_streaming_response(): ) assert response_id is not None, f"no response event within {elapsed:.0f}s of streaming a background response" - cancel_response = client.responses.cancel(response_id) + cancel_response: Final = client.responses.cancel(response_id) print("CANCEL streaming response=", cancel_response) assert cancel_response.status == "cancelled" From 19c819a69eb3d89cb6b113a15ea703caab92557e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:04:06 -0700 Subject: [PATCH 130/167] fix(vertex): add the API version to versionless project routes on the Vertex passthrough --- litellm/llms/vertex_ai/common_utils.py | 23 +++++++--- .../vertex_ai/test_vertex_ai_common_utils.py | 42 +++++++++++++++++++ 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index a36c920dda0..970759479fe 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -998,6 +998,16 @@ def replace_project_and_location_in_route(requested_route: str, vertex_project: return modified_route +def _api_version_for_route(requested_route: str) -> Literal["v1", "v1beta1"]: + return "v1beta1" if "cachedContent" in requested_route else "v1" + + +def _with_api_version(requested_route: str) -> str: + if not requested_route.startswith("/projects/"): + return requested_route + return f"/{_api_version_for_route(requested_route)}{requested_route}" + + def construct_target_url( base_url: str, requested_route: str, @@ -1017,18 +1027,19 @@ def construct_target_url( new_base_url: Final = httpx.URL(base_url) if "locations" in requested_route: # contains the target project id + location - if vertex_project and vertex_location: - requested_route = replace_project_and_location_in_route(requested_route, vertex_project, vertex_location) - return new_base_url.copy_with(path=requested_route) + targeted_route: Final = ( + replace_project_and_location_in_route(requested_route, vertex_project, vertex_location) + if vertex_project and vertex_location + else requested_route + ) + return new_base_url.copy_with(path=_with_api_version(targeted_route)) """ - Add endpoint version (e.g. v1beta for cachedContent, v1 for rest) - Add default project id - Add default location """ - vertex_version: Literal["v1", "v1beta1"] = "v1" - if "cachedContent" in requested_route: - vertex_version = "v1beta1" + vertex_version: Literal["v1", "v1beta1"] = _api_version_for_route(requested_route) # Check if the requested route starts with a version # e.g. /v1beta1/publishers/google/models/gemini-3-pro-preview:streamGenerateContent diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index d1d751989ea..624646ab328 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -964,6 +964,48 @@ def test_construct_target_url_with_version_prefix(): assert str(target_url) == expected_url +@pytest.mark.parametrize( + ("requested_route", "expected_url"), + [ + ( + "/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + ), + ( + "/projects/test-project/locations/global/publishers/anthropic/models/count-tokens:rawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/count-tokens:rawPredict", + ), + ( + "/projects/other-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-6:rawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:rawPredict", + ), + ( + "/projects/test-project/locations/global/cachedContents", + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents", + ), + ( + "/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", + ), + ( + "/v1beta1/projects/test-project/locations/global/cachedContents", + "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents", + ), + ], +) +def test_construct_target_url_versionless_project_route_gets_api_version(requested_route, expected_url): + from litellm.llms.vertex_ai.common_utils import construct_target_url + + target_url = construct_target_url( + base_url="https://aiplatform.googleapis.com", + requested_route=requested_route, + vertex_project="test-project", + vertex_location="global", + ) + + assert str(target_url) == expected_url + + def test_fix_enum_types(): """ Test _fix_enum_types function removes enum fields when type is not string. From f40f14ae39b4452cb55f117ab3b06949bb3135a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:09:46 -0700 Subject: [PATCH 131/167] fix(tests): fold the local price map into the provider model sets CI unit shards load the price map from main at import, so a model that only exists on the branch never reaches open_ai_chat_completion_models and cost_per_token cannot infer its provider. Refresh the sets after swapping in the local map so the tier pricing cases resolve gpt-6-astra before merge --- .../test_openai_service_tier_long_context_pricing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py index 70e9c2720b8..bdb2dc26813 100644 --- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -69,6 +69,7 @@ NO_PUBLISHED_PRIORITY_LONG_CONTEXT = ("gpt-5.4", "gpt-5.5") def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.add_known_models() @lru_cache(maxsize=2) From 60b725cfd85ff6e85872bffc62b340c5e97af28a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:16:57 -0700 Subject: [PATCH 132/167] test(vertex): type the parametrized versionless route test parameters --- .../test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 624646ab328..dddc95bf54a 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -993,7 +993,7 @@ def test_construct_target_url_with_version_prefix(): ), ], ) -def test_construct_target_url_versionless_project_route_gets_api_version(requested_route, expected_url): +def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None: from litellm.llms.vertex_ai.common_utils import construct_target_url target_url = construct_target_url( From 1f20b381151d1c01d38edb56743dee4b771fa25a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:17:48 -0700 Subject: [PATCH 133/167] fix(vector_stores): only list vector stores the caller was granted (#39612) * fix(vector_stores): only list vector stores the caller was granted /vector_store/list returned every managed vector store with no team_id to any key, and let a dashboard session see stores created from the dashboard because every session shares the litellm-dashboard team id. Non-admin listings now show a store only when the key or one of the caller's real teams is allowlisted for it via object_permission.vector_stores, or the team owns it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vector_stores): keep a dashboard session key's own grants when the user has no teams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints.py | 14 ++- litellm/proxy/vector_store_endpoints/utils.py | 78 +++++++++++- .../test_vector_store_access_control.py | 116 +++++++++++++++++- 3 files changed, 192 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index c928398a87f..fe4732c6492 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -30,7 +30,10 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user -from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.proxy.vector_store_endpoints.utils import ( + can_user_access_vector_store, + filter_listable_vector_stores, +) from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.types.vector_stores import ( @@ -390,11 +393,10 @@ async def list_vector_stores( # Filter vector stores based on access control accessible_vector_stores: Final = [] - for vs in vector_store_map.values(): - if await _check_vector_store_access(vs, user_api_key_dict): - redacted = LiteLLM_ManagedVectorStore(**vs) - redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) - accessible_vector_stores.append(redacted) + for vs in await filter_listable_vector_stores(vector_store_map.values(), user_api_key_dict): + redacted = LiteLLM_ManagedVectorStore(**vs) + redacted["litellm_params"] = _redact_sensitive_litellm_params(vs.get("litellm_params")) + accessible_vector_stores.append(redacted) total_count: Final = len(accessible_vector_stores) total_pages: Final = (total_count + page_size - 1) // page_size diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 93f1510bf22..6e94a5a88ac 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,11 +1,17 @@ import json import re +from collections.abc import Iterable +from types import MappingProxyType from typing import Any, Final, Literal from fastapi import HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger +from litellm.proxy._experimental.mcp_server.ui_session_utils import ( + is_ui_session_credential, + resolve_ui_session_team_ids, +) from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, LitellmUserRoles, @@ -160,10 +166,16 @@ async def can_user_access_vector_store( if _is_proxy_admin(user_api_key_dict): return True - vector_store_team_id: Final = vector_store.get("team_id") - if vector_store_team_id is None: + if vector_store.get("team_id") is None: return True + return await _is_vector_store_granted(vector_store, user_api_key_dict) + + +async def _is_vector_store_granted( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth, +) -> bool: vector_store_id: Final = vector_store.get("vector_store_id") or "" key_object_permission = user_api_key_dict.object_permission @@ -178,12 +190,70 @@ async def can_user_access_vector_store( if _object_permission_allows_vector_store(team_object_permission, vector_store_id): return True - if user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store_team_id: - return True + return user_api_key_dict.team_id is not None and user_api_key_dict.team_id == vector_store.get("team_id") + +async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + team: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return user_api_key_dict.model_copy( + update=MappingProxyType( + { + "team_id": team_id, + "team_object_permission": team.object_permission, + "team_object_permission_id": team.object_permission_id, + } + ) + ) + + +async def _vector_store_listing_auth_contexts( + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[UserAPIKeyAuth, ...]: + if not is_ui_session_credential(user_api_key_dict): + return (user_api_key_dict,) + session_key_context: Final = user_api_key_dict.model_copy( + update=MappingProxyType({"team_id": None, "team_object_permission": None, "team_object_permission_id": None}) + ) + team_ids: Final = await resolve_ui_session_team_ids(user_api_key_dict) + team_contexts: Final = tuple([await _team_auth_context(team_id, user_api_key_dict) for team_id in team_ids]) + return (session_key_context, *team_contexts) + + +async def _is_vector_store_granted_to_any( + vector_store: LiteLLM_ManagedVectorStore, + auth_contexts: tuple[UserAPIKeyAuth, ...], +) -> bool: + for auth_context in auth_contexts: + if await _is_vector_store_granted(vector_store, auth_context): + return True return False +async def filter_listable_vector_stores( + vector_stores: Iterable[LiteLLM_ManagedVectorStore], + user_api_key_dict: UserAPIKeyAuth, +) -> tuple[LiteLLM_ManagedVectorStore, ...]: + """Non-admins only see stores their key, one of their teams' object_permission, or team ownership grants.""" + if _is_proxy_admin(user_api_key_dict): + return tuple(vector_stores) + + auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict) + return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)]) + + async def get_litellm_managed_vector_store( vector_store_id: str, ) -> LiteLLM_ManagedVectorStore | None: diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py index 7d72121456a..93049b21460 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_access_control.py @@ -120,9 +120,7 @@ async def test_delete_vector_store_checks_access(): "team_id": "team_456", } ) - mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=mock_vector_store - ) + mock_prisma.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=mock_vector_store) # User from different team should get 403 user_api_key_dict = UserAPIKeyAuth(team_id="team_789") @@ -134,9 +132,115 @@ async def test_delete_vector_store_checks_access(): ): with patch("litellm.vector_store_registry", None): with pytest.raises(HTTPException) as exc_info: - await delete_vector_store( - data=request, user_api_key_dict=user_api_key_dict - ) + await delete_vector_store(data=request, user_api_key_dict=user_api_key_dict) assert exc_info.value.status_code == 403 assert "Access denied" in exc_info.value.detail + + +_UNSCOPED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_unscoped", + "custom_llm_provider": "openai", + "team_id": None, +} +_TEAM_A_OWNED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_team_a", + "custom_llm_provider": "openai", + "team_id": "team_a", +} +_UI_CREATED: LiteLLM_ManagedVectorStore = { + "vector_store_id": "vs_ui_created", + "custom_llm_provider": "openai", + "team_id": "litellm-dashboard", +} + + +async def _listed_ids(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + list_vector_stores, + ) + + with patch( # test-quality-ok: the list route reads rows through this module-level DB helper, no injection seam + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[_UNSCOPED, _TEAM_A_OWNED, _UI_CREATED]), + ): + response = await list_vector_stores(user_api_key_dict=user_api_key_dict) + return sorted(vs["vector_store_id"] for vs in response["data"]) + + +@pytest.mark.asyncio +async def test_list_vector_stores_hides_ungranted_stores_from_non_admin_keys(): + """A store with no team_id and no allowlist entry is not listed for a key it was never granted to; + only team ownership or an explicit object_permission grant makes a store visible.""" + assert await _listed_ids(UserAPIKeyAuth()) == [] + assert await _listed_ids(UserAPIKeyAuth(team_id="team_a")) == ["vs_team_a"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-1", vector_stores=["vs_unscoped"]), + ) + ) == ["vs_unscoped"] + assert await _listed_ids( + UserAPIKeyAuth( + team_id="team_b", + team_object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="op-2", vector_stores=["vs_unscoped"] + ), + ) + ) == ["vs_unscoped"] + assert await _listed_ids(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)) == [ + "vs_team_a", + "vs_ui_created", + "vs_unscoped", + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_team_ids", "session_key_grants", "expected"), + [ + ([], None, []), + ([], ["vs_unscoped"], ["vs_unscoped"]), + (["team_a"], None, ["vs_team_a"]), + (["team_a", "team_granted"], None, ["vs_team_a", "vs_unscoped"]), + ], +) +async def test_list_vector_stores_dashboard_session_resolves_real_teams( + user_team_ids: list[str], session_key_grants: list[str] | None, expected: list[str] +): + """A dashboard session lists through the user's real teams plus the session key's own grants: stores created + from the dashboard (team_id litellm-dashboard) are not visible just because every session shares that team id, + while stores owned by or granted to one of the user's teams, or granted to the session key itself, are.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + + alice = UserAPIKeyAuth( + team_id="litellm-dashboard", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + object_permission=( + LiteLLM_ObjectPermissionTable(object_permission_id="op-4", vector_stores=session_key_grants) + if session_key_grants is not None + else None + ), + ) + teams = { + "team_a": LiteLLM_TeamTableCachedObj(team_id="team_a"), + "team_granted": LiteLLM_TeamTableCachedObj( + team_id="team_granted", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-3", vector_stores=["vs_unscoped"]), + ), + } + + async def fake_get_team_object(team_id: str, **_kwargs: object) -> LiteLLM_TeamTableCachedObj: + return teams[team_id] + + with ( + patch( # test-quality-ok: team rows come from the module-level prisma client, no injection seam + "litellm.proxy.auth.auth_checks.get_team_object", new=fake_get_team_object + ), + patch( # test-quality-ok: the user row comes from the module-level prisma client, no injection seam + "litellm.proxy.vector_store_endpoints.utils.resolve_ui_session_team_ids", + new=AsyncMock(return_value=user_team_ids), + ), + ): + assert await _listed_ids(alice) == expected From ab515dbc90ef6aad53f4d064ee61348db37a3410 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:20:13 -0700 Subject: [PATCH 134/167] fix: treat gpt-6 names as the gpt-5 request family in OpenAI and Azure configs --- .../llms/azure/chat/gpt_5_transformation.py | 22 ++-------------- litellm/llms/azure/chat/gpt_transformation.py | 3 ++- .../llms/openai/chat/gpt_5_transformation.py | 26 ++++++++----------- .../llms/openai/responses/transformation.py | 3 ++- .../chat/test_azure_gpt5_transformation.py | 12 +++++++++ .../test_openai_responses_transformation.py | 2 ++ .../llms/openai/test_gpt5_transformation.py | 14 ++++++++++ .../llms/openai/test_is_model_gpt_5_model.py | 4 +++ tests/test_litellm/test_main.py | 15 +++++++++++ 9 files changed, 64 insertions(+), 37 deletions(-) diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 6fdd277a04f..3189f5b57ac 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -7,6 +7,7 @@ from litellm.exceptions import UnsupportedParamsError from litellm.llms.openai.chat.gpt_5_transformation import ( OpenAIGPT5Config, _get_effort_level, + is_gpt_reasoning_series_name, ) from litellm.types.llms.openai import AllMessageValues @@ -35,26 +36,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - """Check if the Azure model string refers to a gpt-5 variant. - - Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix - used for manual routing. - """ - # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07, - # …) are regular chat models: they support temperature and tool_choice but NOT - # reasoning_effort. They must NOT be routed through the GPT-5 reasoning path. - # - # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning - # models and must stay on the GPT-5 path. The distinguishing feature is that - # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5" - # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version - # number (i.e. "gpt-5.-chat"). - # - # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather - # than a substring check) makes this boundary explicit and avoids any ambiguity - # if future model names coincidentally contain "gpt-5-chat" as an interior run. - _normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "azure/" - return ("gpt-5" in model and not _normalized.startswith("gpt-5-chat")) or "gpt5_series" in model + return is_gpt_reasoning_series_name(model) or "gpt5_series" in model def get_supported_openai_params(self, model: str) -> list[str]: """Get supported parameters for Azure OpenAI GPT-5 models. diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 0ac0662205a..880a51eb584 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, ) from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.openai.chat.gpt_5_transformation import GPT_REASONING_SERIES_MARKERS from litellm.types.llms.azure import ( API_VERSION_MONTH_SUPPORTED_RESPONSE_FORMAT, API_VERSION_YEAR_SUPPORTED_RESPONSE_FORMAT, @@ -139,7 +140,7 @@ class AzureOpenAIConfig(BaseConfig): name family needs the rename, including the ``gpt-5-chat*`` models that are excluded from the reasoning path by https://github.com/BerriAI/litellm/issues/13781. """ - return "gpt-5" in model or "gpt5_series" in model + return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) or "gpt5_series" in model def _is_response_format_supported_model(self, model: str) -> bool: """ diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index 0223be300b0..b02f953425d 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -61,6 +61,14 @@ def _get_effort_level(value: str | dict | None) -> str | None: return None +GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6") + + +def is_gpt_reasoning_series_name(model: str) -> bool: + normalized: Final = model.split("/")[-1] + return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat") + + class OpenAIGPT5Config(OpenAIGPTConfig): """Configuration for gpt-5 models including GPT-5-Codex variants. @@ -73,21 +81,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig): @classmethod def is_model_gpt_5_model(cls, model: str) -> bool: - # The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07, - # …) are regular chat models: they support temperature and tool_choice but NOT - # reasoning_effort. They must NOT be routed through the GPT-5 reasoning path. - # - # Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning - # models and must stay on the GPT-5 path. The distinguishing feature is that - # the gpt-5-chat family has a literal "-chat" immediately after "gpt-5" - # (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version - # number (i.e. "gpt-5.-chat"). - # - # Using a startswith("gpt-5-chat") prefix check on the normalized name (rather - # than a substring check) makes this boundary explicit and avoids any ambiguity - # if future model names coincidentally contain "gpt-5-chat" as an interior run. - _normalized: Final = model.split("/")[-1] # strip provider prefix, e.g. "openai/" - return "gpt-5" in model and not _normalized.startswith("gpt-5-chat") + return is_gpt_reasoning_series_name(model) @classmethod def is_model_gpt_5_search_model(cls, model: str) -> bool: @@ -122,6 +116,8 @@ class OpenAIGPT5Config(OpenAIGPTConfig): def is_model_gpt_5_4_plus_model(cls, model: str) -> bool: """Check if the model is gpt-5.4 or newer (5.4, 5.5, 5.6, etc., including pro).""" model_name: Final = model.split("/")[-1] + if model_name.startswith("gpt-6"): + return True if not model_name.startswith("gpt-5."): return False try: diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 01313e95878..b97521b90c2 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import * @@ -88,7 +89,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): parts: Final = model.split("/") if len(parts) > 1 and parts[0] not in ("openai",): return False - return "gpt-5" in model and "gpt-5-chat" not in model + return is_gpt_reasoning_series_name(model) @staticmethod def _supports_reasoning_effort_none(model: str) -> bool: diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index e06cae97283..bd0f16a695b 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -336,3 +336,15 @@ class TestAzureResolvesTheDeclaredDefaultEffort: drop_params=True, ) assert ("temperature" in mapped) is temperature_survives + + +def test_azure_gpt_6_astra_takes_the_reasoning_series_request_shape(): + params = litellm.get_optional_params( + model="gpt-6-astra", + custom_llm_provider="azure", + max_tokens=100, + reasoning_effort="max", + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + assert params["reasoning_effort"] == "max" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index b0ffd1845fe..4ac072d0ca6 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -1718,6 +1718,8 @@ class TestResponsesSurfaceSharesTheEffortRule: ("gpt-5.6-sol", None, False), ("gpt-5.6-terra", "none", True), ("gpt-5.6-terra", "medium", False), + ("gpt-6-astra", None, False), + ("gpt-6-astra", "low", False), ], ) def test_temperature_follows_the_resolved_effort( diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index 9c5bd34d59a..c86ce4df2ac 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1505,3 +1505,17 @@ class TestACatalogueOlderThanTheCodeDoesNotStripTemperature: drop_params=True, ) assert "temperature" not in mapped + + +def test_gpt_6_astra_takes_the_reasoning_series_request_shape(): + params = litellm.get_optional_params( + model="gpt-6-astra", + custom_llm_provider="openai", + max_tokens=100, + reasoning_effort="max", + verbosity="low", + ) + assert params["max_completion_tokens"] == 100 + assert "max_tokens" not in params + assert params["reasoning_effort"] == "max" + assert params["verbosity"] == "low" diff --git a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py index 1095819c98c..107a1afb2c6 100644 --- a/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py +++ b/tests/test_litellm/llms/openai/test_is_model_gpt_5_model.py @@ -41,6 +41,8 @@ from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config # Models that MUST be classified as GPT-5 (routed through GPT-5 reasoning path) GPT5_MODELS = [ + "gpt-6-astra", + "openai/gpt-6-astra", "gpt-5", "gpt-5.1", "gpt-5.2", @@ -120,6 +122,8 @@ class TestOpenAIGPT5ConfigIsModelGpt5Model: # /v1/responses bridge (when reasoning_effort is set and tools are passed) on # is_model_gpt_5_4_plus_model, so the gpt-5.6 family must land on the True side. GPT5_4_PLUS_MODELS = [ + "gpt-6-astra", + "openai/gpt-6-astra", "gpt-5.4", "gpt-5.5", "gpt-5.5-pro", diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 3c8bf142835..c9acbe2d884 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -784,6 +784,21 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo assert model_info.get("mode") == "responses" +def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): + from litellm.main import responses_api_bridge_check + + with patch("litellm.main._get_model_info_helper") as mock_get_model_info: + mock_get_model_info.return_value = {"max_tokens": 128000} + model_info, model = responses_api_bridge_check( + model="gpt-6-astra", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + ) + + assert model == "gpt-6-astra" + assert model_info.get("mode") == "responses" + + def test_responses_api_bridge_check_gpt_5_5_tools_plus_reasoning_routes_to_responses(): """gpt-5.5+ with both tools and reasoning_effort should route to Responses API.""" from litellm.main import responses_api_bridge_check From b86a0b5562a1a5fbedcd1a5c2b6a02e7f89e64ee Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 16:26:51 -0400 Subject: [PATCH 135/167] test(router): cover get_configured_mode so router_code_coverage passes 425e3069b9 added Router.get_configured_mode but only exercised it through create_model_info_response, which the router coverage gate does not count. The code-quality workflow has been failing on staging and on every open PR since. --- tests/test_litellm/test_router.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3b4c80b6b4f..f91920a636b 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12553,3 +12553,33 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo bucket = captured.get("litellm_metadata") or captured["metadata"] assert captured["model_info"]["id"] == "provisional-dep" assert bucket["litellm_gateway_injected_cache"] == "" + + +def test_get_configured_mode_reads_deployment_model_info(): + router = Router( + model_list=[ + { + "model_name": "my-tts", + "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"}, + "model_info": {"mode": "audio_speech"}, + } + ] + ) + + assert router.get_configured_mode("my-tts") == "audio_speech" + + +@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": " "}, {"mode": 123}]) +def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info): + router = Router( + model_list=[ + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + "model_info": model_info, + } + ] + ) + + assert router.get_configured_mode("plain-model") is None + assert router.get_configured_mode("unknown-model") is None From df73c623b231b68d690f349a7bb70a05b4c82333 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 3 Sep 2026 13:39:58 -0700 Subject: [PATCH 136/167] feat(router): limit heuristic_v2 auto-routers to one without the auto_router license feature (#39468) Without the auto_router feature in the signed enterprise license a proxy may hold one complexity router with classifier_type heuristic_v2 across config.yaml and the DB; with it the limit is lifted. The ceiling is derived once from LicenseCheck and handed to the Router, which refuses the extra router at registration. config.yaml over the limit refuses to start, and /model/new, /model/update and PATCH /model/{id}/update refuse the write with a 403 before touching the DB. Expiry follows the existing max_users/max_teams pattern: judged when the license is verified, not on every call, and a verify that rejects the license (expired or unreadable) leaves no signed payload behind. The rollback after a failed upsert re-admits state that was already serving, so it is exempt from the ceiling: an edit that fails, including one refused by a ceiling that has since tightened, leaves the router serving its previous configuration. A write that leaves a row on heuristic_v2 under a limited license runs in one transaction that takes a Postgres advisory lock before counting the DB rows plus this proxy's config.yaml routers, so concurrent writes on any pod cannot both claim the sole slot and no surplus row is ever persisted. Only the row insert runs under that lock: the team model bookkeeping, which needs a second pool connection, runs after the transaction has committed. PATCH /model/{id}/update follows the same order as create: the row is written through the slot first and the team's model list is updated only afterwards, so a refused write leaves the team as it was. The slot transaction bypasses the repository's publish-on-write, so it publishes the config change once after commit, as delete_team_models does. --- litellm/constants.py | 1 + litellm/proxy/auth/litellm_license.py | 23 +- .../model_management_endpoints.py | 198 +++++--- litellm/proxy/proxy_server.py | 20 +- litellm/router.py | 46 +- .../router_utils/auto_router_model_naming.py | 34 +- litellm/types/router.py | 12 + .../proxy/auth/test_litellm_license.py | 69 +++ .../test_model_management_endpoints.py | 421 ++++++++++++++++-- .../test_ptu_model_settings.py | 6 + .../proxy/proxy_server/test_proxy_config.py | 115 +++++ .../router_strategy/test_complexity_router.py | 160 +++++++ .../test_auto_router_model_naming.py | 57 +++ 13 files changed, 1067 insertions(+), 95 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1063b6ddeeb..f5acadc32ab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", + "heuristic_v2_router_limit", } ) DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) diff --git a/litellm/proxy/auth/litellm_license.py b/litellm/proxy/auth/litellm_license.py index 677f1a0fdda..55bb1e3925a 100644 --- a/litellm/proxy/auth/litellm_license.py +++ b/litellm/proxy/auth/litellm_license.py @@ -16,6 +16,10 @@ if TYPE_CHECKING: from litellm.proxy._types import EnterpriseLicenseData +AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router" +HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit." + + class LicenseCheck: """ - Check if license in env @@ -149,6 +153,19 @@ class LicenseCheck: return False return team_count > _max_teams_in_license + def heuristic_v2_router_limit(self) -> int | None: + """ + How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the + signed license lists the auto_router feature, otherwise one. A license verified through + the API carries no feature list, so it does not lift the limit either. + """ + if self.airgapped_license_data is None: + return 1 + allowed_features: Final = self.airgapped_license_data.get("allowed_features") + if isinstance(allowed_features, list) and AUTO_ROUTER_LICENSE_FEATURE in allowed_features: + return None + return 1 + def verify_license_without_api_request(self, public_key, license_key): try: from cryptography.hazmat.primitives import hashes @@ -179,19 +196,21 @@ class LicenseCheck: # Decode and parse the data license_data: Final = json.loads(message.decode()) - self.airgapped_license_data = EnterpriseLicenseData(**license_data) - # debug information provided in license data verbose_proxy_logger.debug("License data: %s", license_data) # Check expiration date expiration_date: Final = datetime.strptime(license_data["expiration_date"], "%Y-%m-%d") if expiration_date < datetime.now(): + self.airgapped_license_data = None return False, "License has expired" + self.airgapped_license_data = EnterpriseLicenseData(**license_data) + return True except Exception as e: + self.airgapped_license_data = None verbose_proxy_logger.debug( "litellm.proxy.auth.litellm_license.py::verify_license_without_api_request - Unable to verify License locally. - %s", e, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 613e726f89d..82ee33cbc39 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,10 +13,11 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager from json import JSONDecodeError from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeVar, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator @@ -49,6 +50,7 @@ from litellm.proxy._types import ( TeamModelDeleteRequest, UserAPIKeyAuth, ) +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, @@ -96,6 +98,9 @@ from litellm.router_strategy.complexity_router import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, @@ -153,6 +158,8 @@ class _ProxyModelTable(Protocol): def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + def create(self, *, data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + def update( self, *, where: Mapping[str, object], data: Mapping[str, object] ) -> Awaitable[_ProxyModelRow | None]: ... @@ -166,6 +173,9 @@ class _TxModelTables(Protocol): litellm_proxymodeltable: _ProxyModelTable +_RowT = TypeVar("_RowT") + + class _ExistingModelRow(Protocol): @property def litellm_params(self) -> Mapping[str, object]: ... @@ -269,6 +279,66 @@ def _raise_on_strategy_router_write_violation( ) +HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301 +_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_HEURISTIC_V2_DB_ROWS_SQL: Final = """ +SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable" +WHERE model_id <> $1 + AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END) + -> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2' +""" + + +def _effective_complexity_router_config( + incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None +) -> object: + """The complexity config a write leaves on the row: the incoming one when the write carries it, else the stored one.""" + incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config + if incoming is not None or existing_params is None: + return incoming + return existing_params.complexity_router_config + + +@asynccontextmanager +async def _heuristic_v2_slot( + prisma_client: PrismaClient, *, effective_config: object, model_id: str | None +) -> AsyncGenerator[_ProxyModelTable, None]: + """Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled. + + A write that leaves the row on classifier_type heuristic_v2 under a limited license runs + inside one transaction that takes an advisory lock in its own statement before counting + (a statement's snapshot predates anything it locks), so pods cannot both pass the count: + the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged + against the license limit and the write is refused with a 403 before it happens. The row + being edited keeps its own slot through ``model_id``. Every other write, and every write on + an unlimited license, goes through the repository table with no lock. Only the row write + itself may run inside: anything that needs a second connection (the team model bookkeeping) + must wait until the transaction has committed and the lock is released. The transaction + writes bypass the repository's publish-on-write, so the config change is published once + after commit, the way delete_team_models does. + """ + from litellm.proxy.proxy_server import _license_check, llm_router + + limit: Final = _license_check.heuristic_v2_router_limit() + if limit is None or not uses_heuristic_v2_classifier(effective_config): + yield _proxy_model_table(prisma_client) + return + async with prisma_client.db.tx() as tx_ctx: + tables: Final[_TxModelTables] = tx_ctx + await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY) + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "") + db_held: Final = rows[0].get("held") if rows else 0 + config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows) + violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit) + if violation is not None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}" + ) + yield tables.litellm_proxymodeltable + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + + ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add" _REQUIRED_RATE_LIMIT_FIELDS: Final = ("rpm", "tpm") @@ -720,22 +790,29 @@ async def patch_model( ) requested_model_name: Final = patch_data.model_name + stored_model_name: str | None = None + + async def write_row(update_data: PrismaCompatibleUpdateDBModel) -> _ProxyModelRow | None: + nonlocal stored_model_name + stored_model_name = update_data.get("model_name") + update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name + update_data["updated_at"] = cast(str, get_utc_datetime()) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + patch_data.litellm_params, db_model.litellm_params + ), + model_id=model_id, + ) as table: + return await table.update(where={"model_id": model_id}, data=update_data) + # Handle team model updates with proper alias management - update_data: Final = await _update_team_model_in_db( + updated_model: Final = await _update_team_model_in_db( db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, - ) - - # Add metadata about update - update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name - update_data["updated_at"] = cast(str, get_utc_datetime()) - - # Perform partial update - updated_model: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": model_id}, - data=update_data, + write_row=write_row, ) if updated_model is None: @@ -746,7 +823,6 @@ async def patch_model( param=None, ) - stored_model_name: Final = update_data.get("model_name") if ( stored_model_name is not None and stored_model_name == requested_model_name @@ -980,7 +1056,8 @@ async def _add_model_to_db( prisma_client: PrismaClient, new_encryption_key: str | None = None, should_create_model_in_db: bool = True, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": # encrypt litellm params # _litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True) _original_litellm_model_name: Final = model_params.litellm_params.model @@ -998,18 +1075,20 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id _create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above - if should_create_model_in_db: - model_response = await ModelRepository(prisma_client).table.create(data=_create_data) - else: - model_response = LiteLLM_ProxyModelTable(**_data) - return model_response + if not should_create_model_in_db: + return LiteLLM_ProxyModelTable(**_data) + if slot is None: + return await _proxy_model_table(prisma_client).create(data=_create_data) + async with slot as table: + return await table.create(data=_create_data) async def _add_team_model_to_db( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None": + slot: AbstractAsyncContextManager[_ProxyModelTable] | None = None, +) -> "_ProxyModelRow | LiteLLM_ProxyModelTable": """ If 'team_id' is provided, @@ -1040,6 +1119,7 @@ async def _add_team_model_to_db( model_params=model_params, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, + slot=slot, ) if original_model_name: @@ -1060,7 +1140,8 @@ async def _update_team_model_in_db( patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -) -> PrismaCompatibleUpdateDBModel: + write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], +) -> _RowT: """ Handle team model updates with proper alias management. @@ -1068,6 +1149,9 @@ async def _update_team_model_in_db( - Creates unique internal model_name and team alias - Adds model to team object - Preserves team_public_model_name for external reference + + The row is written through ``write_row`` before the team's model list is touched, so a + refused or failed write leaves the team as it was (the create path orders itself the same way). """ # Validate team_id if present in patch_data from litellm.proxy.proxy_server import premium_user @@ -1079,9 +1163,7 @@ async def _update_team_model_in_db( premium_user=premium_user, ) - # Validated before any write, beside the premium check the create path already runs - # here. The team ACL is updated below and autocommits, so a validator that raises - # further down would leave the team mutated and the deployment row never written. + # Validated before the row write, beside the premium check the create path already runs here. # # The merged view is what gets stored, so that is what has to satisfy the invariants. # Validating the patch alone rejected a partial edit of an already valid deployment: @@ -1101,7 +1183,7 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return update_db_model(db_model=db_model, updated_patch=patch_data) + return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -1120,11 +1202,14 @@ async def _update_team_model_in_db( db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None is_new_team_assignment: Final = db_team_id != patch_team_id + # Team rows keep their internal UUID-based model_name; the public name lives in model_info + patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None + row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + if is_new_team_assignment: await _setup_new_team_model_assignment( team_id=patch_team_id, public_model_name=public_model_name, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, ) else: @@ -1132,12 +1217,11 @@ async def _update_team_model_in_db( team_id=patch_team_id, public_model_name=public_model_name, db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) - return update_db_model(db_model=db_model, updated_patch=patch_data) + return row def _get_public_model_name( @@ -1189,13 +1273,9 @@ def _get_public_model_name( async def _setup_new_team_model_assignment( team_id: str, public_model_name: str, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, ) -> None: - """Set up a new team model with unique name and team membership.""" - unique_model_name: Final = f"model_name_{team_id}_{uuid.uuid4()}" - patch_data.model_name = unique_model_name - + """Register a newly team-assigned model's public name on the team.""" await team_model_add( data=TeamModelAddRequest( team_id=team_id, @@ -1385,7 +1465,6 @@ async def _update_existing_team_model_assignment( team_id: str, public_model_name: str, db_model: Deployment, - patch_data: updateDeployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient | None, ) -> None: @@ -1409,9 +1488,6 @@ async def _update_existing_team_model_assignment( old_public_name: Final = db_model.model_info.team_public_model_name if db_model.model_info else None if old_public_name and public_model_name != old_public_name: - # Clear user-supplied public name from patch before any early return so the - # caller does not overwrite the internal UUID-based model_name in the DB. - patch_data.model_name = None if prisma_client is None: verbose_proxy_logger.warning( "prisma_client not initialized; skipping public name update entirely to avoid orphaned entries" @@ -1459,10 +1535,6 @@ async def _update_existing_team_model_assignment( # else: old_public_name == public_model_name (no rename needed) # No team_model_add/delete calls required; public name is already registered - # Always clear patch_data.model_name to prevent caller from overwriting - # the internal UUID-based model_name in the DB with the user-supplied public name - patch_data.model_name = None - class ModelManagementAuthChecks: """ @@ -1878,18 +1950,19 @@ async def add_new_model( reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None) try: _original_litellm_model_name: Final = model_params.model_name - if model_params.model_info.team_id is None: - model_response = await _add_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - else: - model_response = await _add_team_model_to_db( - model_params=priced_model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) + add_model: Final = ( + _add_model_to_db if model_params.model_info.team_id is None else _add_team_model_to_db + ) + model_response = await add_model( + model_params=priced_model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + slot=_heuristic_v2_slot( + prisma_client, + effective_config=priced_model_params.litellm_params.complexity_router_config, + model_id=priced_model_params.model_info.id, + ), + ) reload_outcome = await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) @@ -1903,6 +1976,8 @@ async def add_new_model( passed_model_info=priced_model_params.model_info, ) except Exception as e: + if isinstance(e, HTTPException): + raise verbose_proxy_logger.exception("Exception in add_new_model: %s", e) else: @@ -2070,10 +2145,17 @@ async def update_model( "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, **({} if renamed_to is None else {"model_name": renamed_to}), } - model_response: Final = await _proxy_model_table(prisma_client).update( - where={"model_id": _model_id}, - data=_data, - ) + async with _heuristic_v2_slot( + prisma_client, + effective_config=_effective_complexity_router_config( + model_params.litellm_params, deployment.litellm_params + ), + model_id=_model_id, + ) as table: + model_response: Final = await table.update( + where={"model_id": _model_id}, + data=_data, + ) if renamed_to is not None: await sync_access_groups_for_renamed_model( prisma_client=prisma_client, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..96b425f2a24 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -120,6 +120,8 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( STRATEGY_ROUTER_PARAM_FIELDS, carries_complexity_router_settings, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, validate_complexity_router_config_placement, ) from litellm.types.utils import ( @@ -301,7 +303,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler -from litellm.proxy.auth.litellm_license import LicenseCheck +from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck from litellm.proxy.auth.model_checks import ( expand_wildcard_deployments_for_model_info, get_all_fallbacks, @@ -4316,6 +4318,19 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object]) raise ValueError(f"model {model.get('model_name', '')!r}: {violation}") +def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None: + """ + Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows. + + Checked here rather than left to router registration for the same reason as the two + validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so + the router's own refusal would turn the extra router into a silently missing model. + """ + violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit) + if violation is not None: + raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}") + + def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place """ Stamps `model_info.id` from the raw litellm_params before plugin resolution swaps @@ -5721,6 +5736,7 @@ class ProxyConfig: model_list: Final = config.get("model_list", None) if model_list: router_params["model_list"] = model_list + validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit()) print( # noqa: T201 "\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m" ) @@ -5810,6 +5826,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) if redis_usage_cache is not None and router.cache.redis_cache is None: @@ -6270,6 +6287,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, + heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) else: diff --git a/litellm/router.py b/litellm/router.py index cfb080a24a0..dc750941559 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,7 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -117,6 +117,9 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + uses_heuristic_v2_classifier, ) from litellm.router_utils.batch_utils import ( _get_router_metadata_variable_name, @@ -211,6 +214,7 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, GuardrailTypedDict, + HeuristicV2RouterLimit, LiteLLM_Params, MockRouterTestingParams, ModelGroupInfo, @@ -683,6 +687,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, + heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -759,6 +764,7 @@ class Router: self.set_verbose = set_verbose self.ignore_invalid_deployments = ignore_invalid_deployments + self.heuristic_v2_router_limit = heuristic_v2_router_limit self.fallback_access_check: Final = fallback_access_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks @@ -8796,6 +8802,30 @@ class Router: """ return classify_strategy_router_model(litellm_params.model) == "complexity" + def config_deployments(self) -> Iterator[Mapping[str, object]]: + """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset).""" + for deployment in self.model_list: + if not isinstance(deployment, Mapping): + continue + model_info = deployment.get("model_info") + if not (isinstance(model_info, Mapping) and model_info.get("db_model")): + yield deployment + + def heuristic_v2_router_limit_violation(self) -> str | None: + """ + Why one more heuristic_v2 router cannot join this router, or None when it can. + + Judged against every deployment currently on the model_list; an upsert pops the row being + edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is + resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which + is the SDK default, and the proxy injects a resolver backed by its license. + """ + limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None + others: Final = count_heuristic_v2_routers( + deployment for deployment in self.model_list if isinstance(deployment, Mapping) + ) + return heuristic_v2_limit_violation(held=others + 1, limit=limit) + def init_complexity_router_deployment(self, deployment: Deployment): """ Initialize the complexity-router deployment. @@ -8813,6 +8843,10 @@ class Router: ) complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config + if uses_heuristic_v2_classifier(complexity_router_config): + limit_violation: Final = self.heuristic_v2_router_limit_violation() + if limit_violation is not None: + raise ValueError(limit_violation) default_model: str | None = deployment.litellm_params.complexity_router_default_model @@ -9636,8 +9670,16 @@ class Router: raise e def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: + """Put a deployment back the way it was before a failed upsert popped it. + + A rollback re-admits state that was already serving, so it does not go through the + heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first + registered, judging the rollback would drop a serving router over an unrelated failed edit. + """ if previous_deployment is None or self.has_model_id(model_id): return + limit_resolver: Final = self.heuristic_v2_router_limit + self.heuristic_v2_router_limit = None try: self.add_deployment(deployment=previous_deployment) verbose_router_logger.info( @@ -9652,6 +9694,8 @@ class Router: model_id, restore_error, ) + finally: + self.heuristic_v2_router_limit = limit_resolver @staticmethod def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index a8aa543d735..2efbfb5782e 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under ``ignore_invalid_deployments``. """ -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -163,6 +163,38 @@ def strategy_router_dependencies( ) +def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool: + """Whether this complexity config classifies with the bundled heuristic_v2 model.""" + return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2" + + +def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool: + """Whether this deployment is a complexity router that classifies with heuristic_v2.""" + return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and ( + uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config")) + ) + + +def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int: + """How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers.""" + return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params")))) + + +def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None: + """Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits. + + ``limit`` None means unlimited. The message is shared by every enforcement point (config + load, model writes, router registration) and stays SDK-neutral: it names the cap and what + the caller can change; the proxy appends how its license lifts the cap. + """ + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make " + f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router." + ) + + def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None: """Reject a complexity config the router would refuse to build a deployment from. diff --git a/litellm/types/router.py b/litellm/types/router.py index 4f4df1a8d2e..7ebd50f1328 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -885,6 +885,18 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... +class HeuristicV2RouterLimit(Protocol): + """ + Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited. + + The Router calls it on every registration and limit query instead of caching the answer, so the + proxy can keep the limit on its license object (re-verified on config load) rather than hand + over a snapshot. + """ + + def __call__(self) -> int | None: ... + + class LiteLLM_RouterFileObject(TypedDict, total=False): """ Tracking the litellm params hash, used for mapping the file id to the right model diff --git a/tests/test_litellm/proxy/auth/test_litellm_license.py b/tests/test_litellm/proxy/auth/test_litellm_license.py index 8da365cb587..1db53638070 100644 --- a/tests/test_litellm/proxy/auth/test_litellm_license.py +++ b/tests/test_litellm/proxy/auth/test_litellm_license.py @@ -2,6 +2,8 @@ import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey + from litellm.proxy.auth.litellm_license import LicenseCheck @@ -30,3 +32,70 @@ def test_is_over_limit(): assert license_check.is_over_limit(101) is False assert license_check.is_over_limit(100) is False assert license_check.is_over_limit(99) is False + + +def test_heuristic_v2_router_limit() -> None: + """Only the signed license's auto_router feature lifts the one-router limit; an API-verified + license (no airgapped data) and an airgapped license without the feature keep it.""" + license_check = LicenseCheck() + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]} + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = { + "expiration_date": "2999-01-01", + "allowed_features": ["sso", "auto_router", "audit_logs"], + } + assert license_check.heuristic_v2_router_limit() is None + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = {"expiration_date": "2999-01-01"} + assert license_check.heuristic_v2_router_limit() == 1 + + license_check.airgapped_license_data = None + assert license_check.heuristic_v2_router_limit() == 1 + + +def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]: + import base64 + + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import padding, rsa + + private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + message = json.dumps( + {"expiration_date": expiration_date, "user_id": "u", "allowed_features": ["auto_router"]} + ).encode() + signature = private_key.sign( + message, + padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH), + hashes.SHA256(), + ) + return private_key.public_key(), base64.b64encode(message + b"." + signature).decode() + + +def test_expired_or_unreadable_license_grants_no_features() -> None: + """The verifier stores the signed payload only after the expiry check passes and clears it when a + later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit.""" + license_check = LicenseCheck() + public_key, valid_key = _signed_license("2999-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.heuristic_v2_router_limit() is None + + _, expired_key = _signed_license("2000-01-01") + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True + assert license_check.airgapped_license_data is None + assert license_check.heuristic_v2_router_limit() == 1 + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True + assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True + assert license_check.airgapped_license_data is None + + +def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None: + license_check = LicenseCheck() + public_key, license_key = _signed_license("2999-01-01") + + assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True + assert license_check.heuristic_v2_router_limit() is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 5fa59a85c9d..c69f8f20a13 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -28,9 +28,18 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_team_models, ) from litellm.proxy.utils import PrismaClient +from litellm.router import Router from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +async def _passthrough_row(update_data): + return update_data + + +async def _write_empty_row(**kwargs): + return await kwargs["write_row"]({}) + + class MockPrismaClient: def __init__( self, @@ -1191,7 +1200,7 @@ class TestTeamModelSiblingRouting: team_id = "team_no_alias" public_name = "gpt-4.1-mini" - async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client): + async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client, slot=None): return MagicMock(model_id=str(uuid.uuid4())) mock_team_model_add = AsyncMock() @@ -1372,7 +1381,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert result.get("model_name", "").startswith("model_name_test_team_123_") @@ -1435,7 +1445,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1481,7 +1490,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=None, ) @@ -1490,39 +1498,72 @@ class TestTeamModelUpdate: mock_delete.assert_not_called() @pytest.mark.asyncio - async def test_rename_with_prisma_none_clears_patch_model_name(self): - """Rename path must clear patch_data.model_name even when prisma is unavailable (P1).""" + async def test_a_refused_row_write_leaves_the_team_untouched(self): + """The team's model list autocommits, so it is written only after the row write succeeded: a + refused write (the heuristic_v2 slot 403, a DB error) must not leave the team listing a name + whose row never changed.""" + from fastapi import HTTPException + from litellm.proxy.management_endpoints.model_management_endpoints import ( - _update_existing_team_model_assignment, + _update_team_model_in_db, ) from litellm.types.router import ModelInfo db_model = Deployment( - model_name="model_name_team_123_uuid1", + model_name="gpt-4o", litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), - model_info=ModelInfo( - team_id="team_123", team_public_model_name="old-public-name" + model_info=ModelInfo(), + ) + user_api_key_dict = UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN) + events: list[str] = [] + written: dict[str, object] = {} + + def patch_data() -> updateDeployment: + return updateDeployment(model_name="team-public", model_info=ModelInfo(team_id="team_123")) + + async def refuse_row(update_data): + events.append("row") + raise HTTPException(status_code=403, detail="slot held") + + async def accept_row(update_data): + events.append("row") + written.update(update_data) + return update_data + + async def team_add(**_): + events.append("team_model_add") + + with ( + patch( # test-quality-ok: the team auth check needs a live DB; the write order is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.allow_team_model_action", + AsyncMock(return_value=True), ), - ) - patch_data = updateDeployment( - model_name="new-public-name", - model_info=ModelInfo(team_id="team_123"), - ) - user_api_key_dict = UserAPIKeyAuth( - user_id="test_user", - user_role=LitellmUserRoles.PROXY_ADMIN, - ) + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: team models are premium-gated through a proxy global with no injection seam + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_add, + ), + ): + with pytest.raises(HTTPException): + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=refuse_row, + ) + assert events == ["row"] - await _update_existing_team_model_assignment( - team_id="team_123", - public_model_name="new-public-name", - db_model=db_model, - patch_data=patch_data, - user_api_key_dict=user_api_key_dict, - prisma_client=None, - ) - - assert patch_data.model_name is None + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data(), + user_api_key_dict=user_api_key_dict, + prisma_client=MockPrismaClient(team_exists=True), # type: ignore + write_row=accept_row, + ) + assert events == ["row", "row", "team_model_add"] + assert str(written["model_name"]).startswith("model_name_team_123_") + assert "team-public" in str(written["model_info"]) @pytest.mark.asyncio async def test_rename_handles_legacy_string_model_info(self): @@ -1574,7 +1615,6 @@ class TestTeamModelUpdate: team_id="team_123", public_model_name="new-public-name", db_model=db_model, - patch_data=patch_data, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, # type: ignore ) @@ -1614,7 +1654,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) assert "403" in str(exc_info.value) @@ -1900,7 +1941,8 @@ class TestTeamModelUpdate: db_model=db_model, patch_data=patch_data, user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, # type: ignore + prisma_client=prisma_client, # type: ignore, + write_row=_passthrough_row, ) # team ACL must not be touched on a no-op edit @@ -4311,6 +4353,321 @@ class TestStrategyRouterWriteValidation: is None ) + @staticmethod + def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router: + return Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}}, + { + "model_name": "held-v2", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": "held-id"}, + }, + ], + heuristic_v2_router_limit=lambda: limit, + ) + + class _FakeTx: + """Stands in for a prisma transaction: records the raw statements and exposes the model table.""" + + def __init__(self, db_held: int) -> None: + self.db_held = db_held + self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] + self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) + + async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: + self.raw_calls.append((sql, args)) + return [{"held": self.db_held}] if "count(*)" in sql else [] + + async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + class _FakeDb: + """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" + + def __init__(self, db_held: int, existing_row: object = None) -> None: + self.db = self + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held) + self.litellm_proxymodeltable = MagicMock( + create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) + ) + + def tx(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self.tx_obj + + _V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}} + _V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}} + + @pytest.mark.parametrize( + "incoming,existing,expected", + [ + (_V2, None, _V2), + (_V2, _V1, _V2), + (None, _V1, _V1), + (None, None, None), + ("no-config", _V2, _V2), + ], + ) + def test_effective_complexity_router_config( + self, incoming: object, existing: object, expected: object + ) -> None: + """A write is judged on the config it leaves on the row: the incoming one when it carries one, else the stored one.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _effective_complexity_router_config, + ) + from litellm.types.router import updateLiteLLMParams + + incoming_params = None if incoming is None else updateLiteLLMParams( + complexity_router_config=None if incoming == "no-config" else incoming + ) + existing_params = None if existing is None else updateLiteLLMParams(complexity_router_config=existing) + assert _effective_complexity_router_config(incoming_params, existing_params) == expected + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,effective_config,db_held,config_holds_one,model_id,expected", + [ + (1, _V2, 1, False, None, "refused"), + (1, _V2, 0, True, None, "refused"), + (1, _V2, 0, False, None, "reserved"), + (1, _V2, 0, False, "held-id", "reserved"), + (2, _V2, 1, False, None, "reserved"), + (1, _V1, 5, True, None, "plain"), + (1, None, 5, True, None, "plain"), + (None, _V2, 5, True, None, "plain"), + ], + ) + async def test_heuristic_v2_slot_matrix( + self, + limit: int | None, + effective_config: object, + db_held: int, + config_holds_one: bool, + model_id: str | None, + expected: str, + ) -> None: + """The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows + (other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL + parameter, and every other write runs on the plain client with no lock.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + HEURISTIC_V2_SLOT_LOCK_KEY, + _heuristic_v2_slot, + ) + + fake = self._FakeDb(db_held) + live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None + with ( + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here + "litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change", + new=AsyncMock(), + ) as published, + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id): + pass + assert exc_info.value.status_code == 403 + assert "At most 1 auto-router" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables: + handle = tables + if expected == "plain": + await handle.create(data={}) + fake.litellm_proxymodeltable.create.assert_awaited_once_with(data={}) + assert fake.tx_obj.raw_calls == [] + return + assert handle is fake.tx_obj.litellm_proxymodeltable + published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable") + (lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls + assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql + assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,) + assert count_params == (model_id or "",) + + @pytest.mark.asyncio + async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: + """team_model_add needs a second pool connection, so it must run only after the slot transaction + (and its advisory lock) has closed; a pool-sized burst of team creates would otherwise stall on the + lock holder waiting for a connection the waiters are occupying.""" + from contextlib import asynccontextmanager + + from litellm.proxy.management_endpoints.model_management_endpoints import _add_team_model_to_db + from litellm.types.router import ModelInfo + + events: list[str] = [] + created = MagicMock(model_id="row-1") + + @asynccontextmanager + async def slot(): + events.append("slot-enter") + yield MagicMock(create=AsyncMock(return_value=created)) + events.append("slot-exit") + + async def team_model_add(**_: object) -> None: + events.append("team_model_add") + + deployment = Deployment( + model_name="public-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + model_info=ModelInfo(id="row-1", team_id="team-1"), + ) + with ( + patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add", + side_effect=team_model_add, + ), + ): + result = await _add_team_model_to_db( + model_params=deployment, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + prisma_client=MagicMock(), + slot=slot(), + ) + + assert result is created + assert events == ["slot-enter", "slot-exit", "team_model_add"] + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_heuristic_v2_router_before_the_db_write(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + ) + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="second-v2", + litellm_params=LiteLLM_Params(model="auto_router/complexity_router", complexity_router_config=self._V2), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "At most 1 auto-router" in str(exc_info.value.message) + fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() + fake.litellm_proxymodeltable.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + """patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + patch_model, + ) + from litellm.types.router import updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + fake = self._FakeDb(db_held=1) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: the write must be refused before this DB step runs + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=self._db_complexity_router(model_id)), + ), + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: the helper's team bookkeeping needs a live DB; the row writer it is handed is what is under test + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=_write_empty_row), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment(litellm_params=updateLiteLLMParams(complexity_router_config=self._V2)), + user_api_key_dict=admin, + ) + assert exc_info.value.status_code == 403 + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_update_model_refuses_switching_another_router_to_heuristic_v2(self) -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import ( + update_model, + ) + from litellm.types.router import ModelInfo, updateLiteLLMParams + + model_id = "other-id" + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "model_name": "my-auto-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + "model_info": {"id": model_id}, + } + existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"] + fake = self._FakeDb(db_held=1, existing_row=existing_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + ): + with pytest.raises(ProxyException) as exc_info: + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(complexity_router_config=self._V2), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + fake.tx_obj.litellm_proxymodeltable.update.assert_not_awaited() + fake.litellm_proxymodeltable.update.assert_not_awaited() + @pytest.mark.asyncio async def test_update_model_rejects_prefix_strip(self): from litellm.proxy._types import ProxyException diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index a1c38d26b9d..d35b77f732c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -45,6 +45,10 @@ from litellm.types.router import ( from litellm.types.utils import Usage +async def _passthrough_row(update_data): + return update_data + + def test_model_info_accepts_valid_ptu_fields(): info = ModelInfo( id="x", @@ -385,6 +389,7 @@ class TestTeamModelUpdateValidatesBeforeWriting: patch_data=patch_data, user_api_key_dict=MagicMock(), prisma_client=MagicMock(), + write_row=_passthrough_row, ) return result, touched @@ -914,6 +919,7 @@ class TestPtuDeploymentsAreNotBilledPerToken: patch_data=patch, user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN), prisma_client=MagicMock(), + write_row=_passthrough_row, ) assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 1ab18639fff..dcfad8f6815 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -28,6 +28,7 @@ from litellm.proxy.proxy_server import ( resolve_routing_plugins, validate_deployment_complexity_router_placement, validate_deployment_max_agentic_loops, + validate_heuristic_v2_router_limit, ) from .conftest import normalize @@ -193,6 +194,120 @@ def test_validate_deployment_complexity_router_placement_leaves_valid_deployment assert model["litellm_params"] == litellm_params +def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}}, + }, + } + + +def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None: + """Same reason as the two validators above: the proxy router swallows registration errors, so + an over-limit config.yaml must fail here instead of booting with a silently missing router.""" + with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info: + validate_heuristic_v2_router_limit( + [_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1 + ) + assert "'auto_router' feature lifts the limit" in str(exc_info.value) + + +@pytest.mark.parametrize( + "model_list,limit", + [ + ([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None), + ([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1), + ([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1), + ], +) +def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone( + model_list: list[dict[str, object]], limit: int | None +) -> None: + assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None + + +_TWO_HEURISTIC_V2_ROUTERS_YAML = ( + "model_list:\n" + " - model_name: gpt-4o-mini\n" + " litellm_params:\n" + " model: openai/gpt-4o-mini\n" + " api_key: k\n" + " - model_name: v2-a\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + " - model_name: v2-b\n" + " litellm_params:\n" + " model: auto_router/complexity_router\n" + " complexity_router_config:\n" + " classifier_type: heuristic_v2\n" + " tiers: {SIMPLE: gpt-4o-mini}\n" + "router_settings:\n" + " heuristic_v2_router_limit: 99\n" +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("license_limit", [1, None]) +async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only( + tmp_path, monkeypatch, license_limit: int | None +) -> None: + """`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator + cannot grant the entitlement by editing the config, and a licensed proxy boots both routers.""" + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr( + "litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit + ) + + if license_limit is None: + router, _model_list, _general_settings = await ProxyConfig().load_config( + router=None, config_file_path=str(f) + ) + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + return + + with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license( + tmp_path, monkeypatch +) -> None: + """config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB + is refused at registration because the router was built with the license's ceiling.""" + from litellm.types.router import Deployment + + f = tmp_path / "c.yaml" + f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML.replace(" - model_name: v2-b\n", " - model_name: v1-b\n", 1).replace( + "classifier_type: heuristic_v2\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + "classifier_type: heuristic\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings", + )) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1) + + router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert router.heuristic_v2_router_limit is not None + assert router.heuristic_v2_router_limit() == 1 + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"}) + assert router.upsert_deployment(db_row) is None + assert sorted(router.complexity_routers) == ["v1-b", "v2-a"] + + def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key(): model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}} diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index aa1b51afe10..da3791da39a 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -14,6 +14,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY @@ -1085,6 +1086,165 @@ class TestRouterComplexityDeploymentMethods: router.init_complexity_router_deployment(deployment) assert "auto_router/complexity_router/test-router" in router.complexity_routers + @staticmethod + def _router_row(model_name: str, model_id: str, classifier_type: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": classifier_type, + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + }, + }, + "model_info": {"id": model_id}, + } + + _POOL: dict[str, object] = { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}, + } + + def test_heuristic_v2_ceiling_keeps_the_first_router_and_drops_the_rest(self) -> None: + """The proxy runs with ignore_invalid_deployments, so the second heuristic_v2 router is dropped + at registration while a heuristic (v1) sibling and the first v2 router stay routable.""" + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + self._router_row("v1-c", "id-c", "heuristic"), + ], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert sorted(router.complexity_routers) == ["v1-c", "v2-a"] + assert router.get_deployment(model_id="id-b") is None + + def test_heuristic_v2_ceiling_raises_without_ignore_invalid_deployments(self) -> None: + with pytest.raises(ValueError, match="At most 1 auto-router"): + Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: 1, + ) + + def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None: + """The Router never caches the limit: when the resolver's answer moves (the proxy re-verified + its license), the next registration and the next limit query see the new value.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + limits["value"] = 1 + assert router.heuristic_v2_router_limit_violation() is not None + assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + + def test_heuristic_v2_ceiling_tightening_refuses_the_edit_and_keeps_the_live_router(self) -> None: + """Two heuristic_v2 routers registered under an unlimited ceiling, then the ceiling drops to one: + an edit to either must be refused before its live row is popped, or the failed re-add and + the failed restore would drop a serving router while the write reports success.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + assert router.upsert_deployment(Deployment(**self._router_row("v2-a-renamed", "id-a", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.get_deployment(model_id="id-a") is not None + + assert router.upsert_deployment(Deployment(**self._router_row("v1-a", "id-a", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-a", "v2-b"] + + def test_config_deployments_excludes_db_rows(self) -> None: + """The proxy counts config.yaml routers from here and DB rows from the database, so a DB-loaded + row (``model_info.db_model``) must not show up twice.""" + router = Router(model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")]) + db_row = self._router_row("v2-db", "id-db", "heuristic_v2") + db_row["model_info"] = {"id": "id-db", "db_model": True} + assert router.upsert_deployment(Deployment(**db_row)) is not None + + assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"] + assert count_heuristic_v2_routers(router.config_deployments()) == 1 + + def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None: + """A rollback after a failed upsert re-admits state that was already serving, so it must not be + judged by a ceiling that tightened since: converting one of two live heuristic_v2 routers to a + config whose registration fails must leave it serving its previous v2 configuration.""" + limits = {"value": None} + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ], + heuristic_v2_router_limit=lambda: limits["value"], + ignore_invalid_deployments=True, + ) + limits["value"] = 1 + + broken = self._router_row("v1-a", "id-a", "heuristic") + broken["litellm_params"]["complexity_router_config"]["tiers"] = {} + assert router.upsert_deployment(Deployment(**broken)) is None + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + live = router.get_deployment(model_id="id-a") + assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2" + assert router.heuristic_v2_router_limit_violation() is not None + + def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None: + router = Router( + model_list=[ + self._POOL, + self._router_row("v2-a", "id-a", "heuristic_v2"), + self._router_row("v2-b", "id-b", "heuristic_v2"), + ] + ) + + assert sorted(router.complexity_routers) == ["v2-a", "v2-b"] + assert router.heuristic_v2_router_limit_violation() is None + + def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None: + """A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot + while a different deployment switching to heuristic_v2 is refused.""" + router = Router( + model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")], + heuristic_v2_router_limit=lambda: 1, + ignore_invalid_deployments=True, + ) + + assert router.heuristic_v2_router_limit_violation() is not None + + edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2") + assert router.upsert_deployment(Deployment(**edited)) is not None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + + assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None + assert sorted(router.complexity_routers) == ["v2-a-renamed"] + assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None + assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"] + def test_hybrid_initialization_waits_for_later_pool_deployments(self): router = Router( model_list=[ diff --git a/tests/test_litellm/router_utils/test_auto_router_model_naming.py b/tests/test_litellm/router_utils/test_auto_router_model_naming.py index 0007f09896a..238d0546518 100644 --- a/tests/test_litellm/router_utils/test_auto_router_model_naming.py +++ b/tests/test_litellm/router_utils/test_auto_router_model_naming.py @@ -1,8 +1,13 @@ +from collections.abc import Mapping + import pytest from litellm.router_utils.auto_router_model_naming import ( carries_complexity_router_settings, classify_strategy_router_model, + count_heuristic_v2_routers, + heuristic_v2_limit_violation, + is_heuristic_v2_router, strategy_router_dependencies, validate_complexity_router_config_placement, validate_complexity_router_config_write, @@ -369,3 +374,55 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie flat param on an s3_vectors vector store, so an unscoped gate would reject a valid deployment. Either complexity field names one on its own, which is what the load itself requires.""" assert carries_complexity_router_settings(model, present_fields) is scoped + + +@pytest.mark.parametrize( + "litellm_params,expected", + [ + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False), + ({"model": "auto_router/complexity_router"}, False), + ({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False), + ({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False), + ({}, False), + ], +) +def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None: + """Only a complexity router whose config selects heuristic_v2 counts toward the license limit.""" + assert is_heuristic_v2_router(litellm_params) is expected + + +def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None: + v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}} + rows: list[Mapping[str, object]] = [ + {"model_name": "a", "litellm_params": v2}, + {"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "c", "litellm_params": v2}, + {"model_name": "d"}, + {"model_name": "e", "litellm_params": "not a mapping"}, + ] + assert count_heuristic_v2_routers(rows) == 2 + assert count_heuristic_v2_routers(()) == 0 + + +@pytest.mark.parametrize( + "held,limit,violates", + [ + (1, 1, False), + (2, 1, True), + (0, 1, False), + (5, None, False), + (3, 3, False), + (4, 3, True), + ], +) +def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None: + violation = heuristic_v2_limit_violation(held=held, limit=limit) + assert (violation is not None) is violates + if violation is not None: + assert f"At most {limit} auto-router" in violation + assert f"would make {held}" in violation + assert "license" not in violation From 4e0907fb2dd9f0656f2881cecfabff336751feaf Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 16:40:58 -0400 Subject: [PATCH 137/167] test(router): use an unmapped model so get_configured_mode tests do not write into the global cost map --- tests/test_litellm/test_router.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f91920a636b..26c0209bc05 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12560,7 +12560,7 @@ def test_get_configured_mode_reads_deployment_model_info(): model_list=[ { "model_name": "my-tts", - "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"}, + "litellm_params": {"model": "openai/some-unmapped-mode-model"}, "model_info": {"mode": "audio_speech"}, } ] @@ -12575,7 +12575,7 @@ def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info) model_list=[ { "model_name": "plain-model", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + "litellm_params": {"model": "openai/some-unmapped-mode-model"}, "model_info": model_info, } ] From a264c62b04ea357ce68bbd686bdf98aa81294bde Mon Sep 17 00:00:00 2001 From: moe-berri Date: Thu, 3 Sep 2026 13:43:30 -0700 Subject: [PATCH 138/167] test(router): cover configured mode lookup --- tests/test_litellm/test_router.py | 65 +++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3b4c80b6b4f..477eb7305ca 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7358,6 +7358,71 @@ def test_get_configured_token_limits_coerces_numeric_strings(): assert router.get_configured_token_limits("quoted-limits-model") == (32000, 8000) +def test_get_configured_mode_reads_deployment_model_info(): + router = litellm.Router( + model_list=[ + { + "model_name": "chat-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"mode": "chat"}, + } + ] + ) + + assert router.get_configured_mode("chat-model") == "chat" + + +def test_get_configured_mode_returns_none_for_unset_or_unknown(): + router = litellm.Router( + model_list=[ + { + "model_name": "no-mode-model", + "litellm_params": {"model": "openai/some-unmapped-model"}, + } + ] + ) + + assert router.get_configured_mode("no-mode-model") is None + assert router.get_configured_mode("not-a-real-model") is None + + +def test_get_configured_mode_skips_wildcard_pattern_matching(): + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock/*", + "litellm_params": {"model": "bedrock/*"}, + "model_info": {"mode": "chat"}, + } + ] + ) + + with patch.object( + router.pattern_router, "route", side_effect=AssertionError("pattern route called") + ): + assert ( + router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") + is None + ) + + +def test_get_configured_mode_treats_malformed_values_as_absent(): + malformed = ["", " ", 12345, ["chat"], {"mode": "chat"}, True] + router = litellm.Router( + model_list=[ + { + "model_name": f"bad-mode-{i}", + "litellm_params": {"model": "openai/some-unmapped-model"}, + "model_info": {"mode": bad}, + } + for i, bad in enumerate(malformed) + ] + ) + + for i in range(len(malformed)): + assert router.get_configured_mode(f"bad-mode-{i}") is None + + def test_get_configured_display_name_reads_deployment_model_info(): router = litellm.Router( model_list=[ From e29796882668a7359a05acab53fb2810f2b12009 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:45:06 -0700 Subject: [PATCH 139/167] feat(cost): honor off_peak_pricing reasoning and cache-creation rates The block accepts output_cost_per_reasoning_token and cache_creation_input_token_cost. The generic cost path and the DashScope calculator swap them in while a window is open, and unset keys keep the standard rate. One shared TokenRates value replaces the DashScope-local copy, and apply_off_peak_pricing takes and returns it. --- .../litellm_core_utils/llm_cost_calc/utils.py | 155 +++++++--- litellm/llms/dashscope/cost_calculator.py | 25 +- litellm/types/utils.py | 2 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 267 ++++++++++++++++++ .../test_dashscope_cost_calculator.py | 89 ++++++ 5 files changed, 470 insertions(+), 68 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 21587af73aa..e03c2c93c26 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -415,40 +415,69 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = return False -def _coerce_off_peak_rate(value: object, default: float) -> float: +@dataclass(frozen=True, slots=True) +class TokenRates: + """The per-token rates one request bills at. reasoning_rate is None when reasoning bills at + output_rate: the model has no dedicated reasoning rate, or the caller resolves reasoning on + its own. + """ + + input_rate: float + output_rate: float + cache_read_rate: float + cache_creation_rate: float + reasoning_rate: float | None + + @property + def billed_reasoning_rate(self) -> float: + return self.output_rate if self.reasoning_rate is None else self.reasoning_rate + + +def _parse_off_peak_rate(value: object) -> float | None: if isinstance(value, bool): - return default + return None if isinstance(value, (int, float)): return float(value) if isinstance(value, str): try: return float(value) except ValueError: - return default - return default + return None + return None -def apply_off_peak_pricing( - model_info: ModelInfo, - current_time: datetime | None, - prompt_base_cost: float, - completion_base_cost: float, - cache_read_cost: float, -) -> tuple[float, float, float]: +def _off_peak_rate(off_peak: Mapping[str, object], key: str, standard_rate: float) -> float: + parsed: Final = _parse_off_peak_rate(off_peak.get(key)) + return standard_rate if parsed is None else parsed + + +def _open_off_peak_block(model_info: ModelInfo, current_time: datetime | None) -> Mapping[str, object] | None: + off_peak: Final = model_info.get("off_peak_pricing") + if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time): + return None + return off_peak + + +def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: """Swap in off-peak per-token rates when the current UTC time is inside one of the model's off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in windows. An off-peak rate replaces the rate that would otherwise apply rather than discounting it, so a model that also has tiered or above-threshold pricing bills the flat off-peak rate for the whole request while the window is open. Any rate left unset in - off_peak_pricing falls back to the standard rate. + off_peak_pricing falls back to the standard rate, so a block without + output_cost_per_reasoning_token keeps the model's own reasoning rate, or its off-peak output + rate when reasoning has no dedicated rate at all. """ - off_peak: Final = model_info.get("off_peak_pricing") - if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time): - return prompt_base_cost, completion_base_cost, cache_read_cost - return ( - _coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost), - _coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost), - _coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost), + off_peak: Final = _open_off_peak_block(model_info, current_time) + if off_peak is None: + return rates + off_peak_reasoning_rate: Final = _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token")) + return TokenRates( + input_rate=_off_peak_rate(off_peak, "input_cost_per_token", rates.input_rate), + output_rate=_off_peak_rate(off_peak, "output_cost_per_token", rates.output_rate), + cache_read_rate=_off_peak_rate(off_peak, "cache_read_input_token_cost", rates.cache_read_rate), + cache_creation_rate=_off_peak_rate(off_peak, "cache_creation_input_token_cost", rates.cache_creation_rate), + reasoning_rate=rates.reasoning_rate if off_peak_reasoning_rate is None else off_peak_reasoning_rate, ) @@ -458,14 +487,28 @@ def _apply_off_peak_to_base_costs( base_costs: tuple[float, float, float, float, float], ) -> tuple[float, float, float, float, float]: """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path - produced them. Cache-creation rates are passed through untouched, since off_peak_pricing - has no field for them. + produced them. The one-hour cache-creation rate passes through untouched, since + off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate. """ prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs - off_peak_prompt, off_peak_completion, off_peak_cache_read = apply_off_peak_pricing( - model_info, current_time, prompt, completion, cache_read + rates: Final = apply_off_peak_pricing( + model_info, + current_time, + TokenRates( + input_rate=prompt, + output_rate=completion, + cache_read_rate=cache_read, + cache_creation_rate=cache_creation, + reasoning_rate=None, + ), + ) + return ( + rates.input_rate, + rates.output_rate, + rates.cache_creation_rate, + cache_creation_above_1hr, + rates.cache_read_rate, ) - return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read) def _get_token_base_cost( @@ -1029,6 +1072,29 @@ def _resolve_reasoning_token_cost( return standard_reasoning_cost if standard_reasoning_cost is not None else completion_base_cost +def _resolve_billed_reasoning_rate( + model_info: ModelInfo, + usage: Usage, + service_tier: str | None, + completion_base_cost: float, + current_time: datetime | None, +) -> float: + off_peak: Final = _open_off_peak_block(model_info, current_time) + off_peak_reasoning_rate: Final = ( + None if off_peak is None else _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token")) + ) + if off_peak_reasoning_rate is not None: + return off_peak_reasoning_rate + tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) + if tiered_reasoning_rate is not None: + return tiered_reasoning_rate + return _resolve_reasoning_token_cost( + model_info=model_info, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + ) + + def generic_cost_per_token( model: str, usage: Usage, @@ -1037,6 +1103,7 @@ def generic_cost_per_token( data_residency: str | None = None, model_info: ModelInfo | None = None, vertex_location: str | None = None, + current_time: datetime | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -1051,6 +1118,7 @@ def generic_cost_per_token( - vertex_location: optional Vertex AI location the request was served from (e.g. "us-east5", "global"), used to apply the per-model regional-endpoint uplift multiplier when non-global. + - current_time: the moment the request is billed at, for off_peak_pricing; defaults to now, UTC Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd @@ -1117,6 +1185,7 @@ def generic_cost_per_token( usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0 ) + billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) ( prompt_base_cost, completion_base_cost, @@ -1127,6 +1196,7 @@ def generic_cost_per_token( model_info=model_info, usage=usage, service_tier=service_tier, + current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) @@ -1185,17 +1255,13 @@ def generic_cost_per_token( ## REASONING COST if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0: - tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) - _output_cost_per_reasoning_token = ( - tiered_reasoning_rate - if tiered_reasoning_rate is not None - else _resolve_reasoning_token_cost( - model_info=model_info, - service_tier=service_tier, - completion_base_cost=completion_base_cost, - ) + completion_cost += float(reasoning_tokens) * _resolve_billed_reasoning_rate( + model_info=model_info, + usage=usage, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + current_time=billing_time, ) - completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token ## IMAGE COST if not is_text_tokens_total and image_tokens and image_tokens > 0: @@ -1247,6 +1313,7 @@ def get_token_type_cost_breakdown( service_tier: str | None = None, data_residency: str | None = None, vertex_location: str | None = None, + current_time: datetime | None = None, ) -> TokenTypeCostBreakdown: """ Provider-agnostic cost of reasoning and cache tokens, derived from the usage @@ -1265,6 +1332,7 @@ def get_token_type_cost_breakdown( except Exception: return TokenTypeCostBreakdown(0.0, 0.0, 0.0) + billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc) ( _prompt_base_cost, completion_base_cost, @@ -1275,6 +1343,7 @@ def get_token_type_cost_breakdown( model_info=model_info, usage=usage, service_tier=service_tier, + current_time=billing_time, threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) @@ -1284,18 +1353,12 @@ def get_token_type_cost_breakdown( if not reasoning_tokens: reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0)) - # Reasoning is billed at the selected tier's reasoning rate for tiered models, - # else at the service-tier-aware per-reasoning-token rate - this mirrors how the - # total completion cost is computed, so the breakdown can never diverge from it. - tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage) - reasoning_rate: Final = ( - tiered_reasoning_rate - if tiered_reasoning_rate is not None - else _resolve_reasoning_token_cost( - model_info=model_info, - service_tier=service_tier, - completion_base_cost=completion_base_cost, - ) + reasoning_rate: Final = _resolve_billed_reasoning_rate( + model_info=model_info, + usage=usage, + service_tier=service_tier, + completion_base_cost=completion_base_cost, + current_time=billing_time, ) reasoning_cost = float(reasoning_tokens) * reasoning_rate diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index d8eb1f9f8d7..17f70ec5db7 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -7,12 +7,13 @@ cached, cache-creation, output, reasoning) is billed at that one tier's rate. See https://help.aliyun.com/zh/model-studio/billing-for-model-studio """ -from dataclasses import dataclass, replace +from dataclasses import dataclass from datetime import datetime from typing import Final from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.litellm_core_utils.llm_cost_calc.utils import ( + TokenRates, apply_off_peak_pricing, parse_completion_tokens_details, parse_prompt_tokens_details, @@ -34,19 +35,6 @@ class TokenBreakdown: return self.text_tokens + self.cached_tokens + self.cache_creation_tokens -@dataclass(frozen=True, slots=True) -class TokenRates: - input_rate: float - cache_read_rate: float - cache_creation_rate: float - output_rate: float - reasoning_rate: float | None - - @property - def billed_reasoning_rate(self) -> float: - return self.output_rate if self.reasoning_rate is None else self.reasoning_rate - - def _extract_token_breakdown(usage: Usage) -> TokenBreakdown: prompt_details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = prompt_details["cache_hit_tokens"] @@ -105,13 +93,6 @@ def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates: ) -def _off_peak_rates(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates: - input_rate, output_rate, cache_read_rate = apply_off_peak_pricing( - model_info, current_time, rates.input_rate, rates.output_rate, rates.cache_read_rate - ) - return replace(rates, input_rate=input_rate, output_rate=output_rate, cache_read_rate=cache_read_rate) - - def _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]: prompt_cost: Final = ( (breakdown.text_tokens * rates.input_rate) @@ -155,6 +136,6 @@ def cost_per_token( else None ) standard_rates: Final = _flat_rates(model_info) if tier is None else _tier_rates(model_info, tier) - rates: Final = _off_peak_rates(model_info, current_time, standard_rates) + rates: Final = apply_off_peak_pricing(model_info, current_time, standard_rates) return _bill(breakdown, rates) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 569fce4f7b8..a8e1f1b7ee8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -222,7 +222,9 @@ class OffPeakPricing(TypedDict, total=False): weekday_timezone: ReadOnly[str] input_cost_per_token: ReadOnly[float] output_cost_per_token: ReadOnly[float] + output_cost_per_reasoning_token: ReadOnly[float] cache_read_input_token_cost: ReadOnly[float] + cache_creation_input_token_cost: ReadOnly[float] class ModelInfoBase(ProviderSpecificModelInfo, total=False): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index b7f0ca1efe1..7bc02145841 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -29,11 +29,13 @@ from litellm.types.utils import ( from litellm.litellm_core_utils.llm_cost_calc.utils import ( CostCalculatorUtils, PromptTokensDetailsResult, + TokenRates, TokenTypeCostBreakdown, _calculate_input_cost, _get_token_base_cost, _is_off_peak, _is_within_off_peak_window, + apply_off_peak_pricing, calculate_cache_writing_cost, generic_cost_per_token, get_token_type_cost_breakdown, @@ -782,6 +784,271 @@ def test_get_token_base_cost_off_peak_wins_over_tiered_pricing(): assert outside[:2] == (3e-6, 6e-6) +def _register_off_peak_reasoning_model( + model_name: str, off_peak_pricing: dict, reasoning_rate: float | None = 4e-6, **service_tier_rates: float +) -> None: + reasoning_entry = {} if reasoning_rate is None else {"output_cost_per_reasoning_token": reasoning_rate} + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 1.25e-6, + "off_peak_pricing": off_peak_pricing, + **reasoning_entry, + **service_tier_rates, + } + } + ) + + +def _off_peak_reasoning_usage() -> Usage: + return Usage( + prompt_tokens=100, + completion_tokens=80, + total_tokens=180, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=30, text_tokens=50), + ) + + +def test_generic_cost_per_token_off_peak_reasoning_rate(): + """Regression (LIT-6887): the block's output_cost_per_reasoning_token used to be ignored, so + reasoning tokens billed at the model's standard reasoning rate all through the window.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-reasoning" + _register_off_peak_reasoning_model( + model_name, + {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6, "output_cost_per_reasoning_token": 5e-7}, + ) + + _, inside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7) + + _, outside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside == pytest.approx(50 * 2e-6 + 30 * 4e-6) + + +def test_generic_cost_per_token_off_peak_block_without_reasoning_rate(): + """A block that leaves output_cost_per_reasoning_token unset keeps the model's own reasoning + rate, and a model with no reasoning rate at all follows the off-peak output rate.""" + from datetime import datetime, timezone + + inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + block = {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6} + + _register_off_peak_reasoning_model("litellm-test-off-peak-model-reasoning-rate", block) + _, with_model_rate = generic_cost_per_token( + model="litellm-test-off-peak-model-reasoning-rate", + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=inside_window, + ) + assert with_model_rate == pytest.approx(50 * 1e-6 + 30 * 4e-6) + + _register_off_peak_reasoning_model("litellm-test-off-peak-no-reasoning-rate", block, reasoning_rate=None) + _, without_model_rate = generic_cost_per_token( + model="litellm-test-off-peak-no-reasoning-rate", + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=inside_window, + ) + assert without_model_rate == pytest.approx(80 * 1e-6) + + +def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier(): + """Tiered models resolve reasoning on their own path, so the block has to win there too.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-tiered-reasoning" + litellm.register_model( + { + model_name: { + "litellm_provider": "openai", + "mode": "chat", + "tiered_pricing": [ + { + "range": [0, 128000], + "input_cost_per_token": 3e-6, + "output_cost_per_token": 6e-6, + "output_cost_per_reasoning_token": 8e-6, + }, + ], + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "output_cost_per_token": 1e-6, + "output_cost_per_reasoning_token": 5e-7, + }, + } + } + ) + + _, inside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7) + + _, outside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside == pytest.approx(50 * 6e-6 + 30 * 8e-6) + + +def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_tier(): + """A priority request bills its service-tier reasoning rate outside the window and the block's + rate inside it.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-reasoning-service-tier" + _register_off_peak_reasoning_model( + model_name, + {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6, "output_cost_per_reasoning_token": 5e-7}, + output_cost_per_token_priority=3e-6, + output_cost_per_reasoning_token_priority=6e-6, + ) + + _, inside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + service_tier="priority", + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7) + + _, outside = generic_cost_per_token( + model=model_name, + usage=_off_peak_reasoning_usage(), + custom_llm_provider="openai", + service_tier="priority", + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside == pytest.approx(50 * 3e-6 + 30 * 6e-6) + + +def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings(): + """A YAML true never turns into a rate of 1.0, and a quoted number still counts.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-odd-values" + _register_off_peak_reasoning_model( + model_name, + { + "hours_utc": "16:30-00:30", + "cache_creation_input_token_cost": True, + "output_cost_per_reasoning_token": "5e-7", + }, + ) + standard = TokenRates( + input_rate=1e-6, output_rate=2e-6, cache_read_rate=1e-7, cache_creation_rate=1.25e-6, reasoning_rate=4e-6 + ) + + rates = apply_off_peak_pricing( + litellm.get_model_info(model_name, custom_llm_provider="openai"), + datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + standard, + ) + assert rates.cache_creation_rate == 1.25e-6 + assert rates.reasoning_rate == 5e-7 + + +def test_get_token_base_cost_off_peak_cache_creation_rate(): + """Regression (LIT-6887): the block's cache_creation_input_token_cost used to be ignored. It + replaces the five-minute cache-creation rate inside the window; the one-hour rate, and a + block without the key, keep the standard rate.""" + from datetime import datetime, timezone + from typing import cast + + from litellm.types.utils import ModelInfo + + model_info = cast( + ModelInfo, + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "cache_creation_input_token_cost": 1.25e-6, + "cache_creation_input_token_cost_above_1hr": 2e-6, + "off_peak_pricing": {"hours_utc": "16:30-00:30", "cache_creation_input_token_cost": 5e-7}, + }, + ) + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) + + inside = _get_token_base_cost(model_info, usage, current_time=inside_window) + assert inside[2] == 5e-7 + assert inside[3] == 2e-6 + + outside = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc)) + assert outside[2] == 1.25e-6 + + without_key = cast( + ModelInfo, + {**model_info, "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}}, + ) + assert _get_token_base_cost(without_key, usage, current_time=inside_window)[2] == 1.25e-6 + + +def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates(): + """The per-token-type breakdown feeds the spend logs, so it has to bill the new keys the same + way the total does.""" + from datetime import datetime, timezone + + model_name = "litellm-test-off-peak-breakdown" + _register_off_peak_reasoning_model( + model_name, + { + "hours_utc": "16:30-00:30", + "output_cost_per_reasoning_token": 5e-7, + "cache_creation_input_token_cost": 5e-7, + }, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=80, + total_tokens=1080, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=30, text_tokens=50), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, cache_creation_tokens=400, text_tokens=500), + ) + + inside = get_token_type_cost_breakdown( + model=model_name, + custom_llm_provider="openai", + usage=usage, + current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc), + ) + assert inside.reasoning_cost == pytest.approx(30 * 5e-7) + assert inside.cache_creation_cost == pytest.approx(400 * 5e-7) + assert inside.cache_read_cost == pytest.approx(100 * 1e-7) + + outside = get_token_type_cost_breakdown( + model=model_name, + custom_llm_provider="openai", + usage=usage, + current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), + ) + assert outside.reasoning_cost == pytest.approx(30 * 4e-6) + assert outside.cache_creation_cost == pytest.approx(400 * 1.25e-6) + + def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map): """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output.""" model = "gpt-5.4" diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index b6281834f24..f0949ce041a 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -649,6 +649,95 @@ class TestDashscopeCostCalculator: assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) + def test_dashscope_off_peak_reasoning_rate_replaces_the_dedicated_reasoning_rate(self): + """Regression (LIT-6887): a block carrying output_cost_per_reasoning_token bills reasoning + tokens at it inside the window, over the model's own reasoning rate, which returns outside.""" + self._register_off_peak_flat_model( + "dashscope/qwen-reasoning-rate-off-peak-test", + { + "hours_utc": self.OFF_PEAK_WINDOW, + "output_cost_per_token": 2.4e-06, + "output_cost_per_reasoning_token": 4.5e-06, + }, + ) + litellm.model_cost["dashscope/qwen-reasoning-rate-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-06 + usage = Usage( + prompt_tokens=100, + completion_tokens=200, + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=50), + ) + + _, completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-rate-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + assert math.isclose(completion_cost, (150 * 2.4e-06) + (50 * 4.5e-06), rel_tol=1e-10) + + _, peak_completion_cost = dashscope_cost_per_token( + model="qwen-reasoning-rate-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + assert math.isclose(peak_completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_cache_creation_rate_replaces_the_standard_rate(self): + """Regression (LIT-6887): a block carrying cache_creation_input_token_cost bills cache-creation + tokens at it inside the window, while the cache-read rate it leaves unset stays standard.""" + self._register_off_peak_flat_model( + "dashscope/qwen-cache-creation-off-peak-test", + {"hours_utc": self.OFF_PEAK_WINDOW, "cache_creation_input_token_cost": 1.5e-06}, + ) + usage = Usage( + prompt_tokens=1000, + completion_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100), + ) + + prompt_cost, _ = dashscope_cost_per_token( + model="qwen-cache-creation-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + assert math.isclose(prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 1.5e-06), rel_tol=1e-10) + + peak_prompt_cost, _ = dashscope_cost_per_token( + model="qwen-cache-creation-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10) + + def test_dashscope_off_peak_reasoning_and_cache_creation_rates_override_the_selected_tier(self): + """The new keys override the selected tier the way the input and output rates already do.""" + self._register_tiered_model( + "dashscope/qwen-tiered-reasoning-off-peak-test", + [ + { + "range": [0, 1000], + "input_cost_per_token": 4e-07, + "cache_creation_input_token_cost": 3e-07, + "output_cost_per_token": 1.6e-06, + "output_cost_per_reasoning_token": 3.2e-06, + }, + ], + ) + litellm.model_cost["dashscope/qwen-tiered-reasoning-off-peak-test"]["off_peak_pricing"] = { + "hours_utc": self.OFF_PEAK_WINDOW, + "cache_creation_input_token_cost": 1e-07, + "output_cost_per_reasoning_token": 8e-07, + } + usage = Usage( + prompt_tokens=500, + completion_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper(cache_creation_tokens=200), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=40), + ) + + prompt_cost, completion_cost = dashscope_cost_per_token( + model="qwen-tiered-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW + ) + assert math.isclose(prompt_cost, (300 * 4e-07) + (200 * 1e-07), rel_tol=1e-10) + assert math.isclose(completion_cost, (60 * 1.6e-06) + (40 * 8e-07), rel_tol=1e-10) + + peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token( + model="qwen-tiered-reasoning-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW + ) + assert math.isclose(peak_prompt_cost, (300 * 4e-07) + (200 * 3e-07), rel_tol=1e-10) + assert math.isclose(peak_completion_cost, (60 * 1.6e-06) + (40 * 3.2e-06), rel_tol=1e-10) + def test_dashscope_off_peak_defaults_to_the_current_time(self): """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the default current time.""" From 108f55894652b1a995e86a928d8ceb6a66c4c673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:46:33 -0700 Subject: [PATCH 140/167] test: drop the internal patch from the gpt-6-astra bridge test --- tests/test_litellm/test_main.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index c9acbe2d884..2df1c2f4ca5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -787,13 +787,11 @@ def test_responses_api_bridge_check_gpt_5_4_tools_plus_reasoning_routes_to_respo def test_responses_api_bridge_check_gpt_6_astra_tools_with_default_reasoning_routes_to_responses(): from litellm.main import responses_api_bridge_check - with patch("litellm.main._get_model_info_helper") as mock_get_model_info: - mock_get_model_info.return_value = {"max_tokens": 128000} - model_info, model = responses_api_bridge_check( - model="gpt-6-astra", - custom_llm_provider="openai", - tools=[{"type": "function", "function": {"name": "get_capital"}}], - ) + model_info, model = responses_api_bridge_check( + model="gpt-6-astra", + custom_llm_provider="openai", + tools=[{"type": "function", "function": {"name": "get_capital"}}], + ) assert model == "gpt-6-astra" assert model_info.get("mode") == "responses" From 52e24aebbabdd4d889dda96f3ebeab0e3bd8c3b1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 3 Sep 2026 13:49:43 -0700 Subject: [PATCH 141/167] refactor(tests): assign the streamed id and lock poll once instead of rebinding The cancel test accumulated chunk_count and reassigned response_id on every iteration, and the lock watcher rebound its query result on every poll. Both are the mutable-local pattern the repo avoids. The stream now drains through a generator that stops at the first chunk carrying a response id, so the caller binds streamed_ids once and reads the id off the tail. Empty stream, no-id stream and first-chunk-id all behave exactly as the loop did. The watcher inlines its poll result. --- .../test_e2e_openai_responses_api.py | 20 +++++++++++-------- .../test_team_delete_member_add_race.py | 3 +-- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index b24ac0bdb96..755f17c394e 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -153,6 +153,15 @@ def test_cancel_response(): raise e +def _response_ids_until_first(stream): + """Yield each streamed chunk's response id, stopping at the first chunk that carries one.""" + for chunk in stream: + response_id = getattr(getattr(chunk, "response", None), "id", None) + yield response_id + if response_id is not None: + return + + def test_cancel_streaming_response(): """Cancel a background streaming response while it is still generating. @@ -169,15 +178,10 @@ def test_cancel_streaming_response(): stream=True, background=True, ) as stream: - chunk_count = 0 - response_id = None - for chunk in stream: - chunk_count += 1 - response_id = getattr(getattr(chunk, "response", None), "id", None) - if response_id is not None: - break + streamed_ids = tuple(_response_ids_until_first(stream)) - assert chunk_count > 0, "stream produced no chunks" + assert streamed_ids, "stream produced no chunks" + response_id = streamed_ids[-1] assert response_id is not None, "no streamed chunk carried a response id to cancel" cancel_response = client.responses.cancel(response_id) diff --git a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py index 9af742b0dfe..d3ffcf2445e 100644 --- a/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py +++ b/tests/proxy_admin_ui_tests/test_team_delete_member_add_race.py @@ -78,8 +78,7 @@ async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what: while time.monotonic() < deadline: if task.done(): raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() - rows = await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid) - if rows[0]["waiters"]: + if (await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid))[0]["waiters"]: return await asyncio.sleep(_LOCK_POLL_SECONDS) raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s") From aff6b7e21296f4b6b97883ce8302b640992ec0b1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:51:04 -0700 Subject: [PATCH 142/167] fix(ui): clear agents when updating team permissions (#39600) Always serialize object_permission.agents and agent_access_groups in the team update payload so removing the last agent in the dashboard sends an explicit empty array instead of omitting the key, which the backend merge treats as no change Resolves LIT-6861 Co-authored-by: yassin --- .../src/components/team/TeamInfo.test.tsx | 45 +++++++++++++++++++ .../src/components/team/TeamInfo.tsx | 8 +--- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a9c1077e96b..aedc04283f5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1888,6 +1888,8 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { mcp_access_groups: [], mcp_tool_permissions: {}, mcp_toolsets: [], + agents: [], + agent_access_groups: [], vector_stores: ["vs-1"], }; @@ -1908,6 +1910,49 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { }); }); + const openEditorWithAgents = async (user: ReturnType) => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + models: ["gpt-4"], + object_permission: { agents: ["agent-1"], agent_access_groups: ["group-a"] }, + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0)); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + }; + + it("resends the stored agents and agent_access_groups when the selector is left untouched", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithAgents(user); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.agents).toStrictEqual(["agent-1"]); + expect(objectPermission.agent_access_groups).toStrictEqual(["group-a"]); + }); + + it("sends empty agents and agent_access_groups arrays after the last agent chip is removed", async () => { + const user = userEvent.setup({ delay: null }); + await openEditorWithAgents(user); + + await user.click(within(screen.getByLabelText("agent-1")).getByRole("button")); + await user.click(within(screen.getByLabelText("group:group-a")).getByRole("button")); + expect(screen.queryByLabelText("agent-1")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("group:group-a")).not.toBeInTheDocument(); + + const payload = await save(user); + + const objectPermission = wireBody(payload).object_permission as Record; + expect(objectPermission.agents).toStrictEqual([]); + expect(objectPermission.agent_access_groups).toStrictEqual([]); + }); + it("resends every stored value once both sections are opened", async () => { const user = userEvent.setup({ delay: null }); await openEditor(user); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 3f6d6a96972..c2b8cd3cc56 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -863,12 +863,8 @@ const TeamInfoView: React.FC = ({ agents: [], accessGroups: [], }; - if (agents && agents.length > 0) { - updateData.object_permission.agents = agents; - } - if (agentAccessGroups && agentAccessGroups.length > 0) { - updateData.object_permission.agent_access_groups = agentAccessGroups; - } + updateData.object_permission.agents = agents; + updateData.object_permission.agent_access_groups = agentAccessGroups; delete values.agents_and_groups; // Handle vector stores permissions From eb6c24a2a036ee28aa557e20673d4cf603c5487f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:53:30 -0700 Subject: [PATCH 143/167] fix(auto_router): bill the routing embedding to the caller's key and team (#39532) * fix(auto_router): bill the routing embedding to the caller's key and team Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(auto_router): validate the forwarded caller metadata with a pydantic model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../internal_call_metadata.py | 15 +++ .../auto_router/auto_router.py | 59 +++++++-- .../router_strategy/test_auto_router.py | 120 ++++++++++++------ 3 files changed, 146 insertions(+), 48 deletions(-) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 4d043701f40..87f007ca1d5 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -18,9 +18,11 @@ caller's identity metadata, minus two things that must never be forwarded as-is: from __future__ import annotations from collections.abc import Mapping +from types import MappingProxyType from typing import Final from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, NON_INFERENCE_CALL_TYPES +from litellm.litellm_core_utils.initialize_dynamic_callback_params import initialize_standard_callback_dynamic_params from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, InternalCallOrigin BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) @@ -142,6 +144,19 @@ def forwarded_internal_call_metadata( } +def parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, str]: + kwargs: Final = request_kwargs or MappingProxyType({}) + return MappingProxyType( + {k: v for k in ("litellm_session_id", "litellm_trace_id") if isinstance(v := kwargs.get(k), str)} + ) + + +def effective_turn_off_message_logging(request_kwargs: Mapping[str, object] | None) -> bool | None: + return initialize_standard_callback_dynamic_params(dict(request_kwargs) if request_kwargs else None).get( + "turn_off_message_logging" + ) + + def sanitized_forwardable_call_metadata( parent_metadata: Mapping[str, object], call_origin: InternalCallOrigin, diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index c77745a498d..6b443026f61 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -2,23 +2,41 @@ Auto-Routing Strategy that works with a Semantic Router Config """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final, Optional +from pydantic import BaseModel, ConfigDict + from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import ( + effective_turn_off_message_logging, + forwarded_internal_call_metadata, + parent_session_kwargs, +) +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN if TYPE_CHECKING: from semantic_router.routers import SemanticRouter from semantic_router.routers.base import Route from litellm.router import Router + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder from litellm.types.router import PreRoutingHookResponse else: Router = Any PreRoutingHookResponse = Any Route = Any SemanticRouter = Any + LiteLLMRouterEncoder = Any + + +class _CallerMetadata(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + metadata: Mapping[str, object] | None = None + litellm_metadata: Mapping[str, object] | None = None class AutoRouter(CustomLogger): @@ -50,6 +68,8 @@ class AutoRouter(CustomLogger): """ from semantic_router.routers import SemanticRouter + from litellm.router_strategy.auto_router.litellm_encoder import LiteLLMRouterEncoder + self.auto_router_config_path: str | None = auto_router_config_path self.auto_router_config: str | None = auto_router_config self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE @@ -59,6 +79,11 @@ class AutoRouter(CustomLogger): self.embedding_model: str = embedding_model self.max_input_chars: int = max_input_chars self.litellm_router_instance: Router = litellm_router_instance + self.encoder: LiteLLMRouterEncoder = LiteLLMRouterEncoder( + litellm_router_instance=litellm_router_instance, + model_name=embedding_model, + max_input_chars=max_input_chars, + ) def _load_semantic_routing_routes(self) -> list[Route]: from semantic_router.routers import SemanticRouter @@ -129,9 +154,6 @@ class AutoRouter(CustomLogger): from semantic_router.routers import SemanticRouter from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages - from litellm.router_strategy.auto_router.litellm_encoder import ( - LiteLLMRouterEncoder, - ) from litellm.types.router import PreRoutingHookResponse resolved_messages: Final = ( @@ -149,34 +171,47 @@ class AutoRouter(CustomLogger): ####################### routelayer = SemanticRouter( routes=self.loaded_routes, - encoder=LiteLLMRouterEncoder( - litellm_router_instance=self.litellm_router_instance, - model_name=self.embedding_model, - max_input_chars=self.max_input_chars, - ), + encoder=self.encoder, auto_sync=self.auto_sync_value, ) self.routelayer = routelayer message_content: Final = self._extract_text_from_messages(resolved_messages) - route_name: Final = self._matched_route_name(routelayer, message_content) + route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs) return PreRoutingHookResponse( model=route_name or self.default_model, messages=messages, ) - def _matched_route_name(self, routelayer: "SemanticRouter", text: str) -> str | None: + async def _matched_route_name( + self, routelayer: "SemanticRouter", text: str, request_kwargs: Mapping[str, object] + ) -> str | None: """Name of the route `text` matches, or None when nothing matched or the match failed. - The route layer embeds `text` to compare it against the routes, and that embedding call can + `text` is embedded here rather than by `routelayer(text=...)` so the caller's metadata reaches + `aembedding()` and the embedding's spend lands on the key/team that sent the request; + SemanticRouter has no way to pass kwargs through to its encoder. That embedding call can fail (context limit, timeout, provider error). Choosing a model is a routing decision, so a failure here falls back to the default model rather than failing the user's request. """ from semantic_router.schema import RouteChoice try: - route_choice: Final = routelayer(text=text) + caller: Final = _CallerMetadata.model_validate(request_kwargs) + query_vector: Final = ( + await self.encoder.aencode_queries( + [text], + metadata=forwarded_internal_call_metadata(caller.metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN), + litellm_metadata=forwarded_internal_call_metadata( + caller.litellm_metadata, AUTOROUTER_CLASSIFIER_CALL_ORIGIN + ), + proxy_server_request={"body": {"model": self.embedding_model, "input": [text]}}, + turn_off_message_logging=effective_turn_off_message_logging(request_kwargs), + **parent_session_kwargs(request_kwargs), + ) + )[0] + route_choice: Final = await routelayer.acall(vector=query_vector) except Exception as e: # noqa: BLE001 -- the embedding call behind the route layer can fail many ways (context limit, timeout, provider/network error); none of them may fail the request verbose_router_logger.warning( "AutoRouter: semantic routing failed (%s), falling back to default model %s", e, self.default_model diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index 36199b45847..123ada83ca4 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -330,36 +330,46 @@ ROUTER_CONFIG: Final = json.dumps( ) -class FailingRouteLayer: - """Route layer whose embedding call fails, as it does when the prompt exceeds the encoder's window.""" - - def __call__(self, text: str) -> Any: - raise ValueError( - "Internal_litellm_router API call failed. Error: litellm.InternalServerError: " - "input is too large to process. increase the physical batch size" - ) - - class FixedRouteLayer: - """Route layer that returns whatever the test tells it to, recording the text it was asked about.""" + """Route layer that returns whatever the test tells it to for the query vector it is handed.""" def __init__(self, route_choice: Any) -> None: self.route_choice = route_choice - self.seen_text: str | None = None - def __call__(self, text: str) -> Any: - self.seen_text = text + async def acall(self, vector: Any) -> Any: return self.route_choice +def _embedding_response(input: List[str]) -> Any: + import litellm + + return litellm.EmbeddingResponse( + data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))] + ) + + class StubEmbeddingRouter: - """Stands in for the LiteLLM Router when the route index has to be built for real.""" + """Stands in for the LiteLLM Router, recording the text and kwargs each query embedding was made with.""" + + def __init__(self) -> None: + self.seen_text: str | None = None + self.aembedding_kwargs: Dict[str, Any] | None = None def embedding(self, input: List[str], model: str, **kwargs: Any) -> Any: - import litellm + return _embedding_response(input) - return litellm.EmbeddingResponse( - data=[{"embedding": [0.1, 0.2], "index": i, "object": "embedding"} for i in range(len(input))] + async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any: + self.seen_text = input[0] + self.aembedding_kwargs = kwargs + return _embedding_response(input) + + +class FailingEmbeddingRouter(StubEmbeddingRouter): + """Router whose query embedding fails, as it does when the prompt exceeds the encoder's window.""" + + async def aembedding(self, input: List[str], model: str, **kwargs: Any) -> Any: + raise ValueError( + "litellm.InternalServerError: input is too large to process. increase the physical batch size" ) @@ -369,7 +379,7 @@ def _auto_router(routelayer: Any, litellm_router_instance: Any = None, **kwargs: auto_router_config=ROUTER_CONFIG, default_model="fallback-model", embedding_model="text-embedding-3-small", - litellm_router_instance=litellm_router_instance or MagicMock(), + litellm_router_instance=litellm_router_instance or StubEmbeddingRouter(), **kwargs, ) auto_router.routelayer = routelayer @@ -381,7 +391,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: @pytest.mark.asyncio async def test_should_fall_back_to_default_model_when_the_embedding_call_fails(self): - auto_router: Final = _auto_router(FailingRouteLayer()) + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=FailingEmbeddingRouter()) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -440,8 +450,8 @@ class TestAutoRouterAlwaysResolvesARoutableModel: async def test_should_still_route_to_the_matched_route_when_one_matches(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -451,7 +461,7 @@ class TestAutoRouterAlwaysResolvesARoutableModel: assert result is not None assert result.model == "code-model" - assert layer.seen_text == "fix this stack trace" + assert router.seen_text == "fix this stack trace" class TestAutoRouterEmbeddingInputCap: @@ -483,8 +493,8 @@ class TestAutoRouterRoutesResponsesApiInput: async def test_should_route_a_string_input_when_messages_is_none(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -498,14 +508,14 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "code-model" assert result.messages is None - assert layer.seen_text == "fix this stack trace" + assert router.seen_text == "fix this stack trace" @pytest.mark.asyncio async def test_should_route_a_list_input_with_instructions_when_messages_is_none(self): from semantic_router.schema import RouteChoice - layer: Final = FixedRouteLayer(RouteChoice(name="code-model")) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(RouteChoice(name="code-model")), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -525,13 +535,13 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "code-model" - assert layer.seen_text is not None - assert "fix this stack trace" in layer.seen_text + assert router.seen_text is not None + assert "fix this stack trace" in router.seen_text @pytest.mark.asyncio async def test_should_skip_routing_when_neither_messages_nor_input_is_present(self): - layer: Final = FixedRouteLayer(None) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -540,12 +550,12 @@ class TestAutoRouterRoutesResponsesApiInput: ) assert result is None - assert layer.seen_text is None + assert router.seen_text is None @pytest.mark.asyncio async def test_should_keep_routing_an_empty_messages_list_to_the_default_model(self): - layer: Final = FixedRouteLayer(None) - auto_router: Final = _auto_router(layer) + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(FixedRouteLayer(None), litellm_router_instance=router) result: Final = await auto_router.async_pre_routing_hook( model="my-auto-router", @@ -555,4 +565,42 @@ class TestAutoRouterRoutesResponsesApiInput: assert result is not None assert result.model == "fallback-model" - assert layer.seen_text == "" + assert router.seen_text == "" + + +class TestAutoRouterAttributesItsEmbeddingSpend: + """The query embedding is billed to the key that sent the request, like any other call it made.""" + + @pytest.mark.asyncio + async def test_should_forward_the_callers_identity_to_the_query_embedding_minus_its_budget_reservation(self): + from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY + + router: Final = StubEmbeddingRouter() + auto_router: Final = _auto_router(None, litellm_router_instance=router) + request_kwargs: Final = { + "metadata": { + "user_api_key": "hashed-key", + "user_api_key_team_id": "team-1", + "user_api_key_budget_reservation": {"reservation_id": "r-1"}, + }, + "litellm_session_id": "session-1", + } + + result: Final = await auto_router.async_pre_routing_hook( + model="my-auto-router", + request_kwargs=request_kwargs, + messages=[{"role": "user", "content": "fix this stack trace"}], + ) + + assert result is not None + assert router.seen_text == "fix this stack trace" + assert router.aembedding_kwargs is not None + forwarded: Final = router.aembedding_kwargs["metadata"] + assert forwarded["user_api_key"] == "hashed-key" + assert forwarded["user_api_key_team_id"] == "team-1" + assert forwarded[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "autorouter_classifier" + assert "user_api_key_budget_reservation" not in forwarded + assert router.aembedding_kwargs["litellm_session_id"] == "session-1" + assert router.aembedding_kwargs["proxy_server_request"] == { + "body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]} + } From 792852b4f362bcb9db9462b54caf7be62eb29edd Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 21:04:55 +0000 Subject: [PATCH 144/167] test(ui): name the mocked user list body to stay within the inline-object lint budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/components/networking.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 8dcd8c39d9d..7a220dd4711 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -824,7 +824,8 @@ describe("userListCall search serialization", () => { }); const mockOkFetch = () => { - const body = JSON.stringify({ users: [], total: 0, page: 1, page_size: 25, total_pages: 0 }); + const emptyPage = { users: [], total: 0, page: 1, page_size: 25, total_pages: 0 }; + const body = JSON.stringify(emptyPage); const mockFetch = vi.fn().mockResolvedValue({ ok: true, text: vi.fn().mockResolvedValue(body) } as any); global.fetch = mockFetch as any; return mockFetch; From e6e5be0989bfce24345eb62166bfdde9db4fa69c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 3 Sep 2026 14:37:48 -0700 Subject: [PATCH 145/167] fix(cache): use sync Redis batch reads (#39358) * fix(cache): use sync Redis batch reads * fix(cache): type sync circuit breaker decorator * test(cache): isolate sync Redis breaker coverage * fix(cache): keep batch result merge budget compliant * style(cache): format batch read * style(cache): satisfy type-discipline budget * test(cache): mock Redis before sync breaker setup * style(cache): avoid mutable batch placeholder * test(cache): document sync breaker patch target * fix(types): widen batch result params to Sequence * fix(cache): report real callers through breaker guards The sync guard's lambda and runner frames replaced the actual caller in _get_call_stack_info, so Redis service logs attributed every guarded call to the guard machinery. Skip guard-internal frames when walking the stack and ratchet the lint budgets this branch lowered * style(imports): import Sequence from collections.abc * test(cache): cover concurrent sync and async Redis batch reads * refactor: build sync batch_get_cache results as tuples to satisfy the LIT002 gate * chore: ratchet budgets after staging merge * fix: preserve DualCache batch list contract * style: format DualCache batch result * fix: satisfy mutable collection lint gate * fix(caching): keep breaker guard-frame skipping in bytecode-only deploys * chore: preserve staging budget ratchets * test(cache): isolate sync Redis batch reads * fix(cache): isolate service hook failures * fix(cache): preserve sync batch fallback on open breaker --- litellm/_service_logger.py | 115 +++++++---- litellm/caching/dual_cache.py | 58 +++--- litellm/caching/redis_cache.py | 90 +++++++-- litellm/router_strategy/lowest_tpm_rpm_v2.py | 5 +- tests/local_testing/test_dual_cache.py | 32 +++ tests/test_litellm/caching/test_dual_cache.py | 131 +++++++++++++ .../test_litellm/caching/test_redis_cache.py | 185 +++++++++++++++++- 7 files changed, 526 insertions(+), 90 deletions(-) diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 42a86763b6d..fc27d3a118a 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Callable, Coroutine from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final @@ -83,6 +84,47 @@ class ServiceLogging(CustomLogger): return open_telemetry_logger return None + @staticmethod + def _sync_dispatch_loop() -> asyncio.AbstractEventLoop | None: + """The event loop a blocking caller can dispatch on, or ``None`` if it has none.""" + try: + loop: Final = asyncio.get_event_loop() + except RuntimeError: + return None + return None if loop.is_closed() else loop + + @staticmethod + async def _emit_guarded(hook: Callable[[], Coroutine[object, object, None]]) -> None: + """Emit one service event, absorbing anything the callbacks raise. + + Monitoring must not break the call it monitors. Sync callers are the ones that + swallow their own service failures (a Redis batch read returns an empty dict), + so an exception from a misconfigured callback would replace a Redis outage with + a callback error and skip the caller's fallback handling. + """ + try: + await hook() + except Exception as e: + verbose_logger.exception("Error emitting service event - %s", e) + + @staticmethod + def _dispatch_from_sync(hook: Callable[[], Coroutine[object, object, None]]) -> None: + """Run an async service hook from a blocking caller, whatever event loop it holds. + + Takes a factory rather than a coroutine so the hook is built on the path that + runs it, and only ever once. + """ + loop: Final = ServiceLogging._sync_dispatch_loop() + try: + if loop is None: + asyncio.run(ServiceLogging._emit_guarded(hook)) + elif loop.is_running(): + loop.create_task(ServiceLogging._emit_guarded(hook)) + else: + loop.run_until_complete(ServiceLogging._emit_guarded(hook)) + except Exception as e: + verbose_logger.exception("Error dispatching service event - %s", e) + def service_success_hook( self, service: ServiceTypes, @@ -99,54 +141,45 @@ class ServiceLogging(CustomLogger): if self.mock_testing: self.mock_testing_sync_success_hook += 1 - try: - # Try to get the current event loop - loop: Final = asyncio.get_event_loop() - # Check if the loop is running - if loop.is_running(): - # If we're in a running loop, create a task - loop.create_task( - self.async_service_success_hook( - service=service, - duration=duration, - call_type=call_type, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - else: - # Loop exists but not running, we can use run_until_complete - loop.run_until_complete( - self.async_service_success_hook( - service=service, - duration=duration, - call_type=call_type, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - except RuntimeError: - # No event loop exists, create a new one and run - asyncio.run( - self.async_service_success_hook( - service=service, - duration=duration, - call_type=call_type, - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) + self._dispatch_from_sync( + lambda: self.async_service_success_hook( + service=service, + duration=duration, + call_type=call_type, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, ) + ) - def service_failure_hook(self, service: ServiceTypes, duration: float, error: Exception, call_type: str): + def service_failure_hook( + self, + service: ServiceTypes, + duration: float, + error: Exception, + call_type: str, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: float | datetime | None = None, + ): """ - [TODO] Not implemented for sync calls yet. V0 is focused on async monitoring (used by proxy). + Handles both sync and async monitoring by checking for existing event loop. """ if self.mock_testing: self.mock_testing_sync_failure_hook += 1 + self._dispatch_from_sync( + lambda: self.async_service_failure_hook( + service=service, + duration=duration, + error=error, + call_type=call_type, + parent_otel_span=parent_otel_span, + start_time=start_time, + end_time=end_time, + ) + ) + async def async_service_success_hook( self, service: ServiceTypes, diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 598c9e67faf..df67ba08416 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -8,10 +8,9 @@ Has 4 primary methods: - async_get_cache """ -import asyncio import time import traceback -from concurrent.futures import ThreadPoolExecutor +from collections.abc import Sequence from threading import Lock from typing import TYPE_CHECKING, Any, Final @@ -188,31 +187,38 @@ class DualCache(BaseCache): local_only: bool = False, **kwargs, ): - received_args: Final = locals() - received_args.pop("self") - - def run_in_new_loop(): - """Run the coroutine in a new event loop within this thread.""" - new_loop: Final = asyncio.new_event_loop() - try: - asyncio.set_event_loop(new_loop) - return new_loop.run_until_complete(self.async_batch_get_cache(**received_args)) - finally: - new_loop.close() - asyncio.set_event_loop(None) - try: - # First, try to get the current event loop - _ = asyncio.get_running_loop() - # If we're already in an event loop, run in a separate thread - # to avoid nested event loop issues - with ThreadPoolExecutor(max_workers=1) as executor: - future: Final = executor.submit(run_in_new_loop) - return future.result() + in_memory_result: Final = ( + self.in_memory_cache.batch_get_cache(keys, **kwargs) if self.in_memory_cache is not None else None + ) + result: Final = in_memory_result if in_memory_result is not None else tuple(None for _ in keys) - except RuntimeError: - # No running event loop, we can safely run in this thread - return run_in_new_loop() + if None not in result or self.redis_cache is None or local_only: + return result + + sublist_keys, previous_access_times = self._reserve_redis_batch_keys(time.time(), keys, result) + if len(sublist_keys) == 0: + return result + + try: + redis_result: Final = self.redis_cache.batch_get_cache( + key_list=sublist_keys, parent_otel_span=parent_otel_span + ) + except Exception: + # Do not throttle subsequent callers if the Redis read fails. + self._rollback_redis_batch_key_reservations(previous_access_times) + raise + + if self.in_memory_cache is not None: + for key, value in redis_result.items(): + if value is not None: + self.in_memory_cache.set_cache(key, value, **self._backfill_kwargs(kwargs)) + + return list( # mutable-ok: public list contract + redis_result.get(key) if value is None else value for key, value in zip(keys, result) + ) + except Exception: + verbose_logger.error(traceback.format_exc()) async def async_get_cache( self, @@ -251,7 +257,7 @@ class DualCache(BaseCache): self, current_time: float, keys: list[str], - result: list[Any], + result: Sequence[Any], ) -> tuple[list[str], dict[str, float | None]]: """ Atomically choose keys to fetch from Redis and reserve their access time. diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 2b04a075114..58733b384b9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -78,10 +78,18 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... +_BREAKER_GUARD_FRAME_NAMES: Final = frozenset( + {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} +) + + def _get_call_stack_info(num_frames: int = 2) -> str: """ Get the function names from the previous 1-2 functions in the call stack. + Frames belonging to this module's circuit-breaker guards are skipped so the + reported callers stay the real ones even on guarded methods. + Args: num_frames: Number of previous frames to include (default: 2) @@ -102,11 +110,11 @@ def _get_call_stack_info(num_frames: int = 2) -> str: return "unknown" function_names: Final = [] - for _ in range(num_frames): - if frame is None: - break - func_name = frame.f_code.co_name - function_names.append(func_name) + while frame is not None and len(function_names) < num_frames: + if frame.f_code.co_name in _BREAKER_GUARD_FRAME_NAMES and frame.f_globals.get("__name__") == __name__: + frame = frame.f_back + continue + function_names.append(frame.f_code.co_name) frame = frame.f_back if not function_names: @@ -241,6 +249,23 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int: + """Reject the call if the breaker is open, else return the swallowed-failure count to compare against.""" + if breaker.is_open(): + raise Exception(f"Redis circuit breaker is open — skipping {name}") + return _swallowed_redis_failures.get() + + +def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None: + """Record success only when nothing failed while the call ran. + + Several Redis methods catch their own connection errors and return a default, so a + method that returned is not on its own proof of a healthy Redis. + """ + if _swallowed_redis_failures.get() == swallowed_before: + breaker.record_success() + + async def _run_under_circuit_breaker( breaker: RedisCircuitBreaker, name: str, @@ -249,20 +274,33 @@ async def _run_under_circuit_breaker( """Run one Redis coroutine under a circuit breaker. Shared by the method decorator and the Lua script executor so both feed the same - health signal. Success is recorded only when nothing failed while ``call`` ran, - because several Redis methods catch their own connection errors and return a default. + health signal. """ - if breaker.is_open(): - raise Exception(f"Redis circuit breaker is open — skipping {name}") - swallowed_before: Final = _swallowed_redis_failures.get() + swallowed_before: Final = _enter_circuit_breaker(breaker, name) try: result: Final = await call() except Exception as e: if _is_redis_health_failure(e): breaker.record_failure() raise - if _swallowed_redis_failures.get() == swallowed_before: - breaker.record_success() + _exit_circuit_breaker(breaker, swallowed_before) + return result + + +def _run_under_circuit_breaker_sync( + breaker: RedisCircuitBreaker, + name: str, + call: Callable[[], _RedisCallResult], +) -> _RedisCallResult: + """Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path.""" + swallowed_before: Final = _enter_circuit_breaker(breaker, name) + try: + result: Final = call() + except Exception as e: + if _is_redis_health_failure(e): + breaker.record_failure() + raise + _exit_circuit_breaker(breaker, swallowed_before) return result @@ -288,6 +326,14 @@ def _redis_circuit_breaker_guard(method): return wrapper +def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) -> Callable[..., _RedisCallResult]: + return functools.wraps(method)( + lambda self, *args, **kwargs: _run_under_circuit_breaker_sync( + self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs) + ) + ) + + class RedisCache(BaseCache): # if users don't provider one, use the default litellm cache @@ -1146,14 +1192,13 @@ class RedisCache(BaseCache): """ key_value_dict = {} _key_list: Final = [key for key in key_list if key is not None] + start_time: Final = time.time() try: - _keys: Final = [] - for cache_key in _key_list: - cache_key = self.check_and_fix_namespace(key=cache_key or "") - _keys.append(cache_key) - start_time: Final = time.time() + swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") + _keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list] results: Final = self._run_redis_mget_operation(keys=_keys) + _exit_circuit_breaker(self._circuit_breaker, swallowed_before) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( @@ -1178,7 +1223,18 @@ class RedisCache(BaseCache): return decoded_results except Exception as e: + failed_at: Final = time.time() + self.service_logger_obj.service_failure_hook( + service=ServiceTypes.REDIS, + duration=failed_at - start_time, + error=e, + call_type=f"batch_get_cache <- {_get_call_stack_info()}", + start_time=start_time, + end_time=failed_at, + parent_otel_span=parent_otel_span, + ) verbose_logger.error("Error occurred in batch get cache - %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @_redis_circuit_breaker_guard diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 6deba5aa1cf..665ff69ab47 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -1,6 +1,7 @@ #### What this does #### # identifies lowest tpm deployment import random +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final import httpx @@ -350,9 +351,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): model_group: str, healthy_deployments: list, tpm_keys: list, - tpm_values: list | None, + tpm_values: Sequence | None, rpm_keys: list, - rpm_values: list | None, + rpm_values: Sequence | None, messages: list[dict[str, str]] | None = None, input: str | list | None = None, ) -> dict | None: diff --git a/tests/local_testing/test_dual_cache.py b/tests/local_testing/test_dual_cache.py index e60fa5f3746..43b10a9557a 100644 --- a/tests/local_testing/test_dual_cache.py +++ b/tests/local_testing/test_dual_cache.py @@ -240,3 +240,35 @@ async def test_dual_cache_delete(is_async): result = dual_cache.get_cache(test_key) assert result is None + + +@pytest.mark.asyncio +async def test_dual_cache_concurrent_sync_and_async_redis_reads(): + """Sync and async batch reads share one Redis backend in one process, and sync reads never open an async connection""" + redis_cache = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + dual_cache = DualCache(redis_cache=redis_cache) + + run_id = str(uuid.uuid4()) + sync_keys = [f"sync_{run_id}_{index}" for index in range(5)] + async_keys = [f"async_{run_id}_{index}" for index in range(5)] + in_loop_keys = [f"in_loop_{run_id}_{index}" for index in range(3)] + survivor_key = f"survivor_{run_id}" + expected = {key: {"key": key} for key in [*sync_keys, *async_keys, *in_loop_keys, survivor_key]} + for key, value in expected.items(): + await redis_cache.async_set_cache(key, value, ttl=60) + + concurrent_results = await asyncio.gather( + *(asyncio.to_thread(dual_cache.batch_get_cache, keys=[key]) for key in sync_keys), + *(dual_cache.async_batch_get_cache(keys=[key]) for key in async_keys), + ) + assert list(concurrent_results) == [[expected[key]] for key in [*sync_keys, *async_keys]] + + with patch.object( + redis_cache, + "async_batch_get_cache", + side_effect=AssertionError("sync batch reads must not call async Redis"), + ): + in_loop_results = [dual_cache.batch_get_cache(keys=[key]) for key in in_loop_keys] + + assert in_loop_results == [[expected[key]] for key in in_loop_keys] + assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]] diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 47be139eb5e..ded3be26630 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -61,6 +61,137 @@ async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_ assert "shared_b" not in dual_cache.last_redis_batch_access_time +def _redis_mock_for_sync_batch(redis_result: dict) -> MagicMock: + mock_redis = MagicMock(spec=RedisCache) + mock_redis.batch_get_cache.return_value = redis_result + return mock_redis + + +def _assert_sync_batch_used_blocking_client(dual_cache: DualCache, mock_redis: MagicMock) -> None: + with patch("asyncio.new_event_loop", side_effect=AssertionError("sync path must not create an event loop")): + result = dual_cache.batch_get_cache(keys=["lit6729_key"]) + + assert result == ["redis_value"] + mock_redis.batch_get_cache.assert_called_once_with(key_list=["lit6729_key"], parent_otel_span=None) + mock_redis.async_batch_get_cache.assert_not_called() + mock_redis.init_async_client.assert_not_called() + assert dual_cache.in_memory_cache.get_cache("lit6729_key") == "redis_value" + + +@pytest.mark.asyncio +async def test_dual_cache_batch_get_cache_uses_sync_redis_client_inside_running_loop(): + """ + Regression test for LIT-6729: sync batch_get_cache ran async_batch_get_cache on a + throwaway event loop, reusing an async Redis client created on another loop and + corrupting its connection pool. The sync path must use the blocking client, never + the async one, and never create an event loop, even when called from a coroutine + (e.g. async_raise_no_deployment_exception -> get_min_cooldown). + """ + mock_redis = _redis_mock_for_sync_batch({"lit6729_key": "redis_value"}) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis) + + _assert_sync_batch_used_blocking_client(dual_cache, mock_redis) + + +def test_dual_cache_batch_get_cache_uses_sync_redis_client_without_running_loop(): + mock_redis = _redis_mock_for_sync_batch({"lit6729_key": "redis_value"}) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis) + + _assert_sync_batch_used_blocking_client(dual_cache, mock_redis) + + +def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): + mock_redis = _redis_mock_for_sync_batch({"miss_key": "from_redis"}) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis) + dual_cache.in_memory_cache.set_cache("hit_key", "from_memory") + + result = dual_cache.batch_get_cache(keys=["hit_key", "miss_key"]) + + assert result == ["from_memory", "from_redis"] + mock_redis.batch_get_cache.assert_called_once_with(key_list=["miss_key"], parent_otel_span=None) + + +def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): + mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) + + first = dual_cache.batch_get_cache(keys=["absent_key"]) + second = dual_cache.batch_get_cache(keys=["absent_key"]) + + assert first == [None] + assert second == [None] + mock_redis.batch_get_cache.assert_called_once() + + +def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): + mock_redis = MagicMock(spec=RedisCache) + mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) + + first_result = dual_cache.batch_get_cache(keys=["shared_a"]) + second_result = dual_cache.batch_get_cache(keys=["shared_a"]) + + assert first_result is None + assert second_result is None + assert mock_redis.batch_get_cache.call_count == 2 + assert "shared_a" not in dual_cache.last_redis_batch_access_time + + +def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): + mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) + dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() + + result = dual_cache.batch_get_cache(keys=["throttled_key"]) + + assert result == [None] + mock_redis.batch_get_cache.assert_not_called() + + +def test_dual_cache_sync_batch_redis_backfill_injects_default_in_memory_ttl(): + """Sync batch_get_cache's Redis-to-memory backfill must honor + default_in_memory_ttl, same as the async path.""" + in_memory_cache = InMemoryCache(default_ttl=600) + mock_redis = _redis_mock_for_sync_batch({"batch_backfill_key": "redis_value"}) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + redis_cache=mock_redis, + default_in_memory_ttl=60, + ) + + before = time.time() + result = dual_cache.batch_get_cache(keys=["batch_backfill_key"]) + after = time.time() + + assert result == ["redis_value"] + expiry = in_memory_cache.ttl_dict["batch_backfill_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +def test_dual_cache_batch_get_cache_forwards_explicit_ttl_to_backfill(): + """An explicit ttl kwarg must reach the in-memory backfill flat, not nested + under a 'kwargs' key the way the old locals()-forwarding path sent it.""" + in_memory_cache = InMemoryCache(default_ttl=600) + mock_redis = _redis_mock_for_sync_batch({"explicit_ttl_key": "redis_value"}) + dual_cache = DualCache(in_memory_cache=in_memory_cache, redis_cache=mock_redis) + + before = time.time() + result = dual_cache.batch_get_cache(keys=["explicit_ttl_key"], ttl=5) + after = time.time() + + assert result == ["redis_value"] + expiry = in_memory_cache.ttl_dict["explicit_ttl_key"] + assert expiry >= before + 5 + assert expiry <= after + 5 + + @pytest.mark.asyncio async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl(): """ diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 487a64797d1..71be8730df1 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1,10 +1,10 @@ import asyncio -from unittest.mock import MagicMock, patch +from collections.abc import Iterator +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unittest.mock import AsyncMock - +from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import RedisCache @@ -17,6 +17,17 @@ def redis_no_ping(): yield +@pytest.fixture +def sync_batch_redis_cache(redis_no_ping): + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ) as get_client: + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.side_effect = OSError("redis unavailable") + get_client.assert_called_once() + yield cache + + @pytest.mark.parametrize( ("namespace", "key", "expected"), [ @@ -504,6 +515,173 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no await call_method(cache) +def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_redis_cache): + """An open breaker must preserve the sync batch read's dictionary fallback.""" + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} + + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} + + +@pytest.fixture +def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: + service_logger = ServiceLogging(mock_testing=True) + failing_client = MagicMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=failing_client + ): + cache = RedisCache(host="127.0.0.1", port=6379, service_logger_obj=service_logger) + yield cache, service_logger + + +@pytest.mark.asyncio +async def test_sync_batch_get_cache_reports_a_failed_read_from_a_running_loop( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], +): + """A swallowed Redis failure must still be reported as a service failure event. + + The routing strategies call this blocking read from inside the request's event loop, + and the read hides the Redis error by returning an empty dict. Without an emitted + failure event, litellm_redis_failed_requests_total stops moving during a Redis + outage while the success path keeps reporting, so the dashboards read healthy. + """ + cache, service_logger = sync_batch_cache_with_service_logger + + assert cache.batch_get_cache(key_list=["lit6729"]) == {} + await asyncio.sleep(0.05) + + assert service_logger.mock_testing_sync_failure_hook == 1 + assert service_logger.mock_testing_async_failure_hook == 1 + + +def test_sync_batch_get_cache_reports_a_failed_read_from_a_worker_thread( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], +): + """The same report must reach the async hook when the caller has no event loop at all.""" + from concurrent.futures import ThreadPoolExecutor + + cache, service_logger = sync_batch_cache_with_service_logger + + with ThreadPoolExecutor(max_workers=1) as pool: + assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {} + + assert service_logger.mock_testing_async_failure_hook == 1 + + +def test_sync_batch_get_cache_reports_a_failed_read_on_an_idle_event_loop( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], +): + """The report must also go out when the caller holds an open loop that is not running.""" + cache, service_logger = sync_batch_cache_with_service_logger + loop = asyncio.new_event_loop() + try: + asyncio.set_event_loop(loop) + assert cache.batch_get_cache(key_list=["lit6729"]) == {} + finally: + asyncio.set_event_loop(None) + loop.close() + + assert service_logger.mock_testing_async_failure_hook == 1 + + +def test_sync_batch_get_cache_survives_a_service_callback_that_raises( + sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging], + monkeypatch: pytest.MonkeyPatch, +): + """A failing service callback must not replace the swallowed Redis failure. + + A misconfigured callback raises while emitting (a datadog callback with no + DD_API_KEY raises at construction), and the failure event is emitted from inside + the except block that swallows the Redis error. If that exception escapes, a Redis + outage surfaces to routing as a callback error and the circuit breaker never + records the failed read. + """ + from concurrent.futures import ThreadPoolExecutor + + import litellm + + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + + cache, service_logger = sync_batch_cache_with_service_logger + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + monkeypatch.setattr( + service_logger, + "init_prometheus_services_logger_if_none", + AsyncMock(side_effect=Exception("callback is misconfigured")), + ) + + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + with ThreadPoolExecutor(max_workers=1) as pool: + assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {} + + assert cache.batch_get_cache(key_list=["lit6729"]) == {} + + +def test_call_stack_info_skips_breaker_guard_frames(): + """Guarded methods must still report their real callers in service-log call_type. + + The breaker guards put their own frames between a method body and its caller, so + without skipping them every guarded method logged the guard machinery instead of + who actually issued the Redis call. + """ + from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _get_call_stack_info, + _redis_circuit_breaker_guard_sync, + ) + + class Guarded: + _circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + + @_redis_circuit_breaker_guard_sync + def probe(self): + return _get_call_stack_info() + + def caller_one(): + return Guarded().probe() + + def caller_two(): + return caller_one() + + assert caller_two() == "caller_one <- caller_two" + + +def test_call_stack_info_skips_guard_frames_when_deployed_without_sources(monkeypatch): + """Guard-frame skipping must survive a bytecode-only deployment. + + Shipping `.pyc` files without their `.py` sources leaves the module's `__file__` pointing + at the compiled file while every frame still carries the compile-time source path, so a + check comparing those two paths stops skipping and the service log then names the guard + machinery instead of the real caller. + """ + from litellm.caching import redis_cache as redis_cache_module + from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _get_call_stack_info, + _redis_circuit_breaker_guard_sync, + ) + + monkeypatch.setattr(redis_cache_module, "__file__", redis_cache_module.__file__ + "c") + + class Guarded: + _circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + + @_redis_circuit_breaker_guard_sync + def probe(self): + return _get_call_stack_info() + + def caller_one(): + return Guarded().probe() + + def caller_two(): + return caller_one() + + assert caller_two() == "caller_one <- caller_two" + + @pytest.mark.asyncio async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ping): """A reachable Redis must keep the breaker closed, however many earlier calls failed. @@ -580,7 +758,6 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): async def swallows_a_failure(): await asyncio.sleep(0.02) _record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable")) - return None async def succeeds_while_the_other_fails(): await asyncio.sleep(0.05) From 19da217167766a0dbb7d1e4ef6ea4a8685f0f16e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:38:47 -0700 Subject: [PATCH 146/167] fix(openai): mint workload identity tokens for PrivateLink and regional api.openai.com hosts --- litellm/llms/openai/workload_identity.py | 5 +-- .../openai/test_openai_workload_identity.py | 43 +++++++++++++++++++ 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index ecec161ed46..283fdfb92c2 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -8,7 +8,7 @@ from urllib.parse import urlparse import litellm from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str -from .common_utils import OpenAIError +from .common_utils import OpenAIError, is_openai_backed_api_base if TYPE_CHECKING: from collections.abc import Callable @@ -16,7 +16,6 @@ if TYPE_CHECKING: from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth OPENAI_WIF_CLIENT_ID: Final = "litellm" -_OPENAI_API_HOST: Final = "api.openai.com" _SDK_UPGRADE_MESSAGE: Final = ( "OpenAI workload identity federation requires openai>=2.32.0. " "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / " @@ -75,7 +74,7 @@ def _targets_openai_api(api_base: str | None) -> bool: if api_base is None: return True parsed: Final = urlparse(api_base) - return parsed.scheme == "https" and parsed.hostname == _OPENAI_API_HOST + return parsed.scheme == "https" and is_openai_backed_api_base(api_base) @lru_cache(maxsize=16) diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index d8d9936e9a1..db107e00df0 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -85,6 +85,31 @@ class TestResolveConfig: def test_plaintext_http_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="http://api.openai.com/v1") is None + @pytest.mark.parametrize( + "api_base", + ( + "https://southcentralus.privatelink.api.openai.com/v1", + "https://eu.api.openai.com/v1", + "https://us.api.openai.com/v1", + ), + ) + def test_openai_backed_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig, api_base: str) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base=api_base) == wif_env + + @pytest.mark.parametrize( + "api_base", + ( + "https://api.openai.com.evil.example/v1", + "https://openai.com/v1", + "https://euapi.openai.com/v1", + "http://southcentralus.privatelink.api.openai.com/v1", + ), + ) + def test_lookalike_or_plaintext_api_base_disables( + self, wif_env: OpenAIWorkloadIdentityConfig, api_base: str + ) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base=api_base) is None + def test_foreign_env_base_url_disables( self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -158,6 +183,14 @@ class TestClientConstruction: assert client.api_key == "workload-identity-auth" assert client._workload_identity_auth is not None + def test_privatelink_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client( + is_async=False, api_key=None, api_base="https://southcentralus.privatelink.api.openai.com/v1" + ) + assert isinstance(client, OpenAI) + assert client.api_key == "workload-identity-auth" + assert client._workload_identity_auth is not None + def test_static_key_client_unaffected(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key="sk-static", api_base=None) assert isinstance(client, OpenAI) @@ -231,6 +264,16 @@ class TestResponsesValidateEnvironment: ) assert headers["Authorization"] == "Bearer None" + @respx.mock + def test_privatelink_api_base_mints_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + mock_token_exchange() + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, + model="gpt-4o-mini", + litellm_params=GenericLiteLLMParams(api_base="https://southcentralus.privatelink.api.openai.com/v1"), + ) + assert headers["Authorization"] == "Bearer exchanged-bearer-token" + def test_litellm_proxy_subclass_never_mints_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: headers: Final = LiteLLMProxyResponsesAPIConfig().validate_environment( headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() From 959e730d55b087d786562ef5b4c27096d8d337c8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:40:36 -0700 Subject: [PATCH 147/167] fix(agents): hide agents from non-admins who were never granted them (#39636) Listing agents (GET /v1/agents and MCP agent_search) treated the absence of any agent grant on the key or team as permission to see every agent. Non-admin keys now list only the union of explicit grants, and dashboard sessions resolve that union through the user's real teams and user row instead of the shared dashboard team. Proxy admins still see everything and direct access to a named agent is unchanged. Resolves LIT-6862 Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auth/agent_permission_handler.py | 53 +++++++-- .../auth/test_agent_permission_handler.py | 101 +++++++++++++++++- .../proxy/agent_endpoints/test_endpoints.py | 5 +- 3 files changed, 146 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index d0ac94d3710..e4dd77e2f82 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -5,10 +5,13 @@ Handles agent permission checking for keys and teams using object_permission_id. Follows the same pattern as MCP permission handling. """ +import asyncio +from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from typing import Final, TypeAlias from litellm._logging import verbose_logger +from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts from litellm.proxy._types import ( UI_TEAM_ID, LiteLLM_ObjectPermissionTable, @@ -443,15 +446,47 @@ class AgentRequestHandler: return [] -async def accessible_agents(user_api_key_auth: UserAPIKeyAuth) -> tuple[AgentResponse, ...]: - """Every registry agent for proxy admins, else the agents the key's and team's grants reach.""" +def _granted_ids(access: AgentAccess) -> frozenset[str]: + match access: + case UnrestrictedAgentAccess(): + return frozenset() + case RestrictedAgentAccess(agent_ids): + return agent_ids + + +ResolveAgentAccess: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[AgentAccess]] +EffectiveAuthContexts: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[Sequence[UserAPIKeyAuth]]] + + +async def _granted_agent_ids( + user_api_key_auth: UserAPIKeyAuth, + resolve_access: ResolveAgentAccess, + effective_contexts: EffectiveAuthContexts, +) -> frozenset[str]: + """Union of the explicit grants reachable from the key, its team, or (for a dashboard session) + the user's real teams and user row. No grant anywhere yields the empty set, unlike the + open-by-default ``resolve_agent_access`` that guards direct access.""" + accesses: Final = await asyncio.gather( + *(resolve_access(auth_context) for auth_context in await effective_contexts(user_api_key_auth)) + ) + return frozenset().union(*(_granted_ids(access) for access in accesses)) + + +async def accessible_agents( + user_api_key_auth: UserAPIKeyAuth, + all_agents: tuple[AgentResponse, ...] | None = None, + resolve_access: ResolveAgentAccess | None = None, + effective_contexts: EffectiveAuthContexts = build_effective_auth_contexts, +) -> tuple[AgentResponse, ...]: + """Every registry agent for proxy admins, else only the agents the caller was granted.""" from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - all_agents: Final = global_agent_registry.get_agent_list() + agents: Final = global_agent_registry.get_agent_list() if all_agents is None else all_agents if user_api_key_auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value): - return all_agents - match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_auth): - case UnrestrictedAgentAccess(): - return all_agents - case RestrictedAgentAccess(allowed_agent_ids): - return tuple(agent for agent in all_agents if agent.agent_id in allowed_agent_ids) + return agents + allowed_agent_ids: Final = await _granted_agent_ids( + user_api_key_auth, + AgentRequestHandler.resolve_agent_access if resolve_access is None else resolve_access, + effective_contexts, + ) + return tuple(agent for agent in agents if agent.agent_id in allowed_agent_ids) diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 82528c58ae0..383b72e5c58 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -10,15 +10,42 @@ from unittest.mock import AsyncMock, patch import pytest -from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import UI_SESSION_TOKEN_TEAM_ID +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentAccess, AgentRequestHandler, RestrictedAgentAccess, UnrestrictedAgentAccess, + accessible_agents, ) +def _registry_with(*agent_names: str) -> AgentRegistry: + registry: Final = AgentRegistry() + registry.load_agents_from_config( + [ + { + "agent_name": name, + "agent_card_params": {"name": name, "url": "http://localhost", "version": "1.0.0"}, + } + for name in agent_names + ] + ) + return registry + + +def _agent_id(registry: AgentRegistry, agent_name: str) -> str: + agent: Final = registry.get_agent_by_name(agent_name) + assert agent is not None + return agent.agent_id + + +async def _single_context(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + return [user_api_key_auth] + + @pytest.mark.asyncio class TestAgentRequestHandler: """ @@ -265,6 +292,78 @@ class TestAgentRequestHandler: ) assert result == UnrestrictedAgentAccess() + async def test_accessible_agents_hides_ungranted_agents_from_non_admins(self): + """LIT-6862: a key with no agent grant on itself or its team must list nothing, + while a proxy admin with the same lack of grants still lists every agent.""" + registry: Final = _registry_with("alpha", "beta") + internal_user: Final = UserAPIKeyAuth( + api_key="test-key", user_id="alice", team_id="team-no-perms", user_role=LitellmUserRoles.INTERNAL_USER + ) + proxy_admin: Final = UserAPIKeyAuth( + api_key="admin-key", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + async def no_grant_anywhere(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess: + return UnrestrictedAgentAccess() + + assert ( + await accessible_agents(internal_user, registry.get_agent_list(), no_grant_anywhere, _single_context) == () + ) + assert { + agent.agent_name + for agent in await accessible_agents( + proxy_admin, registry.get_agent_list(), no_grant_anywhere, _single_context + ) + } == {"alpha", "beta"} + + async def test_accessible_agents_lists_only_granted_agents(self): + """A grant for one agent lists that agent and hides the ungranted one.""" + registry: Final = _registry_with("alpha", "beta") + granted_user: Final = UserAPIKeyAuth( + api_key="test-key", user_id="bob", team_id="team-granted", user_role=LitellmUserRoles.INTERNAL_USER + ) + + async def alpha_only(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess: + return RestrictedAgentAccess(frozenset({_agent_id(registry, "alpha")})) + + listed: Final = await accessible_agents(granted_user, registry.get_agent_list(), alpha_only, _single_context) + assert [agent.agent_name for agent in listed] == ["alpha"] + + async def test_accessible_agents_resolves_dashboard_session_through_real_teams_and_user(self): + """LIT-6862: a dashboard session carries the shared litellm-dashboard team id, which holds no + grants. Listing must union the grants of the user's real teams and of the user row instead + of treating the session as ungranted or as unrestricted.""" + registry: Final = _registry_with("alpha", "beta", "gamma") + session: Final = UserAPIKeyAuth( + api_key="session-key", + user_id="alice", + team_id=UI_SESSION_TOKEN_TEAM_ID, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + admitted_user: Final = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + grants: Final = { + "team-granted": RestrictedAgentAccess(frozenset({_agent_id(registry, "alpha")})), + "team-no-perms": UnrestrictedAgentAccess(), + UI_SESSION_TOKEN_TEAM_ID: UnrestrictedAgentAccess(), + } + + async def effective_contexts(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + assert user_api_key_auth is session + return [ + session.model_copy(update={"team_id": "team-granted"}), + session.model_copy(update={"team_id": "team-no-perms"}), + admitted_user, + ] + + async def resolve_access(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess: + if user_api_key_auth is admitted_user: + return RestrictedAgentAccess(frozenset({_agent_id(registry, "beta")})) + assert user_api_key_auth.team_id is not None + return grants[user_api_key_auth.team_id] + + listed: Final = await accessible_agents(session, registry.get_agent_list(), resolve_access, effective_contexts) + assert {agent.agent_name for agent in listed} == {"alpha", "beta"} + async def test_get_allowed_agents_for_key_via_access_group_ids(self): """ Test that _get_allowed_agents_for_key includes agents from key's access_group_ids diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 067f5a9f64c..13d6cd8a68c 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -11,7 +11,6 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( RestrictedAgentAccess, - UnrestrictedAgentAccess, ) from litellm.proxy.agent_endpoints.endpoints import ( _attach_keys_to_agents, @@ -550,9 +549,9 @@ class TestAgentRBACProxyAdminViewOnly: self.allowed_agents_spy.assert_awaited_once() def test_should_still_redact_secrets_for_view_only_admin(self): - """An unrestricted viewer sees the same agents as an admin but with keys + """A viewer granted every agent sees the same agents as an admin but with keys stripped; litellm_params secrets never appear in either response.""" - self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() + self.allowed_agents_spy.return_value = RestrictedAgentAccess(frozenset({"agent-1", "agent-2"})) viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) From ab44e8d60222726fcab4562422bda9161ddf002a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:44:09 -0700 Subject: [PATCH 148/167] fix(team_endpoints): stop partial /team/update from wiping team metadata (#36328) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/team_endpoints.py | 21 +++++ .../test_team_endpoints.py | 81 +++++++++++++++++-- .../src/components/team/TeamInfo.test.tsx | 29 +++++++ .../src/components/team/TeamInfo.tsx | 2 +- 4 files changed, 127 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 72cdc75c29e..90d7539b38d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -10,6 +10,7 @@ All /team management endpoints """ import asyncio +import copy import json import math import traceback @@ -2189,6 +2190,26 @@ async def update_team( if field in updated_kv } + _writes_metadata_backed_field: Final = any( + field in updated_kv + for field in ( + *LiteLLM_ManagementEndpoint_MetadataFields, + *LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) + ) + if isinstance(existing_team_row.metadata, dict): + if "metadata" not in updated_kv and (_team_member_fields_in_request or _writes_metadata_backed_field): + updated_kv["metadata"] = copy.deepcopy(existing_team_row.metadata) + elif isinstance(updated_kv.get("metadata"), dict): + updated_kv["metadata"] = { + **updated_kv["metadata"], + **{ + key: existing_team_row.metadata[key] + for key in TeamMemberBudgetHandler.SYSTEM_MANAGED_METADATA_KEYS + if key in existing_team_row.metadata + }, + } + if _team_member_fields_in_request and TeamMemberBudgetHandler.should_create_budget( team_member_budget=data.team_member_budget, team_member_rpm_limit=data.team_member_rpm_limit, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 166fcb2863c..019ebc9807c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -2910,6 +2910,7 @@ async def test_update_team_with_team_member_budget_duration( "metadata": {"team_member_budget_id": "budget_123"}, } mock_existing_team.metadata = {"team_member_budget_id": "budget_123"} + mock_existing_team.members_with_roles = [] mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) @@ -11290,6 +11291,78 @@ async def test_patch_preserves_required_metadata_key_that_post_would_wipe(): assert patch_meta == {"cost_center": "FINOPS-1", "team_notes": "edited"} # preserved by PATCH +_STORED_METADATA_WITH_BUDGET: Final = { + "team_member_budget_id": "budget-existing-123", + "team_member_key_duration": "30d", + "logging": [{"callback_name": "langfuse", "callback_type": "success"}], + "cost_center": "cc-1234", +} + + +async def _written_metadata_with_budget(kind, body): + """Like ``_written_metadata`` but the team already owns a member budget row.""" + from litellm.proxy._types import LiteLLM_BudgetTable + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( # test-quality-ok: update_team imports update_budget at call time; the module attribute is its only seam + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + AsyncMock(return_value=LiteLLM_BudgetTable(budget_id="budget-existing-123")), + ), + ): + return await _written_metadata(kind, dict(_STORED_METADATA_WITH_BUDGET), body) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +@pytest.mark.parametrize( + "body", + [ + {"team_member_budget": 50.0}, + {"team_member_budget_duration": "1d"}, + {"team_member_tpm_limit": 500}, + {"team_member_rpm_limit": 5}, + ], + ids=lambda body: next(iter(body)), +) +async def test_team_member_budget_only_update_preserves_stored_metadata(kind, body): + """LIT-5150: a budget-only update that omits ``metadata`` must not replace the + stored metadata JSON with just ``{"team_member_budget_id": ...}``.""" + assert await _written_metadata_with_budget(kind, body) == _STORED_METADATA_WITH_BUDGET + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_team_member_key_duration_only_update_preserves_stored_metadata(kind): + """LIT-5150: a metadata-backed field sent alone is merged into the stored + metadata instead of becoming the whole metadata JSON.""" + written = await _written_metadata_with_budget(kind, {"team_member_key_duration": "7d"}) + + assert written == {**_STORED_METADATA_WITH_BUDGET, "team_member_key_duration": "7d"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["post", "patch"]) +async def test_explicit_null_metadata_with_budget_field_still_clears_metadata(kind): + """``metadata: null`` is an explicit clear, so only the server-owned budget link survives.""" + written = await _written_metadata_with_budget(kind, {"metadata": None, "team_member_budget": 7.0}) + + assert written == {"team_member_budget_id": "budget-existing-123"} + + +@pytest.mark.asyncio +async def test_metadata_only_update_keeps_team_member_budget_link(): + """LIT-5150: rewriting metadata without any team member field must not drop the + server-owned ``team_member_budget_id``, or the member budget silently resets.""" + body = {"metadata": {"cost_center": "cc-9999"}} + + post_meta = await _written_metadata_with_budget("post", body) + patch_meta = await _written_metadata_with_budget("patch", body) + + assert post_meta == {"cost_center": "cc-9999", "team_member_budget_id": "budget-existing-123"} + assert patch_meta == {**_STORED_METADATA_WITH_BUDGET, "cost_center": "cc-9999"} + + @pytest.mark.asyncio @pytest.mark.parametrize( "body, field, expected", @@ -11318,17 +11391,15 @@ async def test_top_level_fields_identical_post_and_patch(body, field, expected): @pytest.mark.asyncio async def test_patch_strips_system_managed_metadata_key_like_post(): """A caller cannot inject/overwrite server-owned keys via PATCH any more than - via POST: team_member_budget_id is stripped from the write in both.""" + via POST: the stored team_member_budget_id wins over the caller's value in both.""" existing = {"team_member_budget_id": "budget-123", "cost_center": "1234"} body = {"metadata": {"team_member_budget_id": "HACKED", "cost_center": "9999"}} post_meta = await _written_metadata("post", existing, body) patch_meta = await _written_metadata("patch", existing, body) - assert "team_member_budget_id" not in post_meta - assert "team_member_budget_id" not in patch_meta - assert post_meta == {"cost_center": "9999"} - assert patch_meta == {"cost_center": "9999"} + assert post_meta == {"cost_center": "9999", "team_member_budget_id": "budget-123"} + assert patch_meta == {"cost_center": "9999", "team_member_budget_id": "budget-123"} @pytest.mark.parametrize( diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index aedc04283f5..a993c1f6bd6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1538,6 +1538,35 @@ describe("TeamInfoView", () => { }); }); + describe("team member settings", () => { + it("should populate Default Key Duration from the team's stored metadata", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ metadata: { team_member_key_duration: "30d" } }), + ); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await user.click(await screen.findByRole("button", { name: /team member settings/i })); + + await waitFor(() => { + expect(screen.getByLabelText(/^Default Key Duration/)).toHaveValue("30d"); + }); + }); + }); + describe("guardrails dropdown grouping", () => { const guardrail = (name: string, defaultOn: boolean) => ({ guardrail_name: name, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index c2b8cd3cc56..23ecd19fdd7 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -335,7 +335,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): default_team_member_models: info.default_team_member_models || [], team_member_budget: info.team_member_budget_table?.max_budget, team_member_budget_duration: info.team_member_budget_table?.budget_duration, - team_member_key_duration: info.team_member_key_duration, + team_member_key_duration: info.metadata?.team_member_key_duration, team_member_tpm_limit: info.team_member_budget_table?.tpm_limit, team_member_rpm_limit: info.team_member_budget_table?.rpm_limit, budget_duration: info.budget_duration, From f2f65a6e8b244a364380cfada2498b54eaa45aab Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:04:45 -0700 Subject: [PATCH 149/167] fix(mcp): resolve OAuth broker endpoints by server_id with IP access checks (#39432) * fix(mcp): resolve OAuth broker endpoints by server_id with IP access checks Resolve named OAuth lookups through server IDs while retaining client IP checks\n\nCo-authored-by: KK291860 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: retrigger e2e pipeline Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 43 ++--- .../mcp_server/mcp_server_manager.py | 8 +- .../mcp_server/test_discoverable_endpoints.py | 163 ++++++++++++++++++ .../mcp_server/test_mcp_server_manager.py | 26 +++ 4 files changed, 209 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 94bca9460dd..81dd9057c9b 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -448,6 +448,17 @@ def _append_query_params(url: str, params: dict[str, str]) -> str: return urlunparse(parsed._replace(query=urlencode(query_params))) +def _resolve_mcp_server_by_name_or_id(lookup: str, client_ip: str | None) -> MCPServer | None: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + by_name: Final = global_mcp_server_manager.get_mcp_server_by_name(lookup, client_ip=client_ip) + if by_name is not None: + return by_name + return global_mcp_server_manager.get_mcp_server_by_id(lookup, client_ip=client_ip) + + def _resolve_oauth2_server_for_root_endpoints( client_ip: str | None = None, ) -> MCPServer | None: @@ -1766,10 +1777,6 @@ async def authorize( resource: str | None = None, ): # Redirect to real OAuth provider with PKCE support - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id): if is_proxy_api_resource(request, resource): return await native_client_authorize( @@ -1797,9 +1804,7 @@ async def authorize( lookup_name: Final[str | None] = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = ( - global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) if lookup_name else None - ) + mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if lookup_name else None if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: @@ -1855,10 +1860,6 @@ async def token_endpoint( 3. Return the token 4. Return a virtual key in this response """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - if mcp_server_name is None and is_gateway_dcr_client_id(client_id): from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load master_key, @@ -1882,7 +1883,7 @@ async def token_endpoint( lookup_name: Final = mcp_server_name or client_id client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip) + mcp_server = _resolve_mcp_server_by_name_or_id(lookup_name, client_ip) if mcp_server is None and mcp_server_name is None: mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: @@ -2288,10 +2289,6 @@ async def _build_oauth_protected_resource_response( Returns: OAuth protected resource metadata dict """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - request_base_url: Final = get_request_base_url(request) client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) explicitly_named: Final = mcp_server_name is not None @@ -2304,7 +2301,7 @@ async def _build_oauth_protected_resource_response( mcp_server: MCPServer | None = None if mcp_server_name: - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) # Build resource URL based on the pattern if mcp_server_name: @@ -2562,10 +2559,6 @@ def _build_oauth_authorization_server_response( registry lookups; unlike :func:`_build_oauth_protected_resource_response` it does not need to await any upstream IO. """ - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - request_base_url: Final = get_request_base_url(request) client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) explicitly_named: Final = mcp_server_name is not None @@ -2583,7 +2576,7 @@ def _build_oauth_authorization_server_response( mcp_server: MCPServer | None = None if mcp_server_name: - mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth authorization server") @@ -2709,10 +2702,6 @@ async def oauth_authorization_server_legacy(request: Request, mcp_server_name: s @router.post("/{mcp_server_name}/register") @router.post("/register") async def register_client(request: Request, mcp_server_name: str | None = None): - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - # Get the correct base URL considering X-Forwarded-* headers request_base_url: Final = get_request_base_url(request) @@ -2748,7 +2737,7 @@ async def register_client(request: Request, mcp_server_name: str | None = None): ) return dummy_return - mcp_server: Final = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name, client_ip=client_ip) + mcp_server: Final = _resolve_mcp_server_by_name_or_id(mcp_server_name, client_ip) if mcp_server is None: return dummy_return return await register_client_with_server( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index a772c569bfa..36e1078dc76 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6147,13 +6147,13 @@ class MCPServerManager: internal_networks = IPAddressUtils.parse_internal_networks(general_settings.get("mcp_internal_ip_ranges")) return IPAddressUtils.is_internal_ip(client_ip, internal_networks) - def get_mcp_server_by_id(self, server_id: str) -> MCPServer | None: - """ - Get the MCP Server from the server id - """ + def get_mcp_server_by_id(self, server_id: str, client_ip: str | None = None) -> MCPServer | None: + """Get the MCP Server from the server id.""" registry: Final = self.get_registry() for server in registry.values(): if server.server_id == server_id: + if not self._is_server_accessible_from_ip(server, client_ip): + return None return server return None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 588eba4adb7..775f6e5f3b8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -3508,6 +3508,169 @@ def _create_oauth2_server( ) +def _create_id_lookup_oauth2_server(): + return _create_oauth2_server( + server_id="oauth-server-id", + name="oauth-server-name", + server_name="oauth-server-name", + alias="oauth-server-alias", + ) + + +@pytest.mark.asyncio +async def test_authorize_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object(discoverable_endpoints, "encrypt_value_helper", return_value="encrypted-state"), # test-quality-ok: flow seam + ): + response = await discoverable_endpoints.authorize( + request=request, + client_id=server.client_id, + mcp_server_name=server.server_id, + redirect_uri="http://localhost:62646/callback", + state="test_state", + ) + + assert response.status_code == 307 + assert "https://provider.com/oauth/authorize" in response.headers["location"] + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +@pytest.mark.asyncio +async def test_token_endpoint_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"access_token": "token", "token_type": "Bearer"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object(discoverable_endpoints, "get_async_httpx_client", return_value=client), # test-quality-ok: HTTP seam + ): + result = await discoverable_endpoints.token_endpoint( + request=request, + grant_type="authorization_code", + code="test_code", + redirect_uri="http://localhost:62646/callback", + client_id=server.client_id, + mcp_server_name=server.server_id, + client_secret=server.client_secret, + ) + + assert json.loads(result.body)["access_token"] == "token" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +@pytest.mark.asyncio +async def test_register_client_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server().model_copy( + update={"client_id": None, "client_secret": None, "registration_url": "https://provider.com/oauth/register"} + ) + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + response = MagicMock() + response.json.return_value = {"client_id": "registered-client", "client_secret": "registered-secret"} + response.raise_for_status = MagicMock() + client = MagicMock() + client.post = AsyncMock(return_value=response) + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + patch.object(discoverable_endpoints, "_read_request_body", new=AsyncMock(return_value={})), # test-quality-ok: request seam + patch.object(discoverable_endpoints, "get_async_httpx_client", return_value=client), # test-quality-ok: HTTP seam + ): + result = await discoverable_endpoints.register_client(request=request, mcp_server_name=server.server_id) + + assert json.loads(result.body)["client_id"] == "registered-client" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +@pytest.mark.asyncio +async def test_protected_resource_metadata_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + ): + result = await discoverable_endpoints._build_oauth_protected_resource_response( + request=request, + mcp_server_name=server.server_id, + use_standard_pattern=True, + ) + + assert result["authorization_servers"] == ["https://llm.example.com/mcp"] + assert result["resource"] == f"https://llm.example.com/mcp/{server.server_id}" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + +def test_authorization_server_metadata_resolves_server_by_id_when_name_lookup_fails(): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + server = _create_id_lookup_oauth2_server() + request = MagicMock(spec=Request) + request.base_url = "https://llm.example.com/" + request.headers = {} + + with ( + patch.object(global_mcp_server_manager, "get_mcp_server_by_name", return_value=None) as by_name, # test-quality-ok: resolver seam + patch.object(global_mcp_server_manager, "get_mcp_server_by_id", return_value=server) as by_id, # test-quality-ok: resolver seam + ): + result = discoverable_endpoints._build_oauth_authorization_server_response( + request=request, + mcp_server_name=server.server_id, + ) + + assert result["scopes_supported"] == server.scopes + assert result["issuer"] == f"https://llm.example.com/{server.server_id}" + by_name.assert_called_once_with(server.server_id, client_ip=None) + by_id.assert_called_once_with(server.server_id, client_ip=None) + + @pytest.mark.asyncio async def test_authorize_root_resolves_single_oauth2_server(): """When /authorize is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 91e870d2d95..3321da83007 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -129,6 +129,32 @@ class TestMCPServerManager: assert added_server.args == ["-m", "server"] assert added_server.env == {"DEBUG": "1", "TEST": "1"} + def test_get_mcp_server_by_id_allows_internal_or_unspecified_client_ip(self): + manager = MCPServerManager() + server = MCPServer( + server_id="private-server", + name="private-server", + transport=MCPTransport.http, + available_on_public_internet=False, + ) + manager.registry[server.server_id] = server + + assert manager.get_mcp_server_by_id(server.server_id) is server + assert manager.get_mcp_server_by_id(server.server_id, client_ip="10.0.0.1") is server + + def test_get_mcp_server_by_id_rejects_private_server_for_public_ip(self): + manager = MCPServerManager() + server = MCPServer( + server_id="private-server", + name="private-server", + transport=MCPTransport.http, + available_on_public_internet=False, + ) + manager.registry[server.server_id] = server + + with patch.object(manager, "_get_general_settings", return_value={}): + assert manager.get_mcp_server_by_id(server.server_id, client_ip="8.8.8.8") is None + async def test_create_mcp_client_stdio(self): """Test creating MCP client for stdio transport""" manager = MCPServerManager() From 6da516e6f3d343db1a4301659a9a82d52de4c9dd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:08:02 -0700 Subject: [PATCH 150/167] fix(mcp): strip inbound auth scheme case-insensitively before token exchange (#39346) * fix(mcp): strip inbound auth scheme case-insensitively before token exchange Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp): type the fake credential provider params in token exchange scheme tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 9 ++- .../mcp_server/mcp_server_manager.py | 4 +- .../test_mcp_client.py | 2 + .../mcp_server/test_mcp_server_manager.py | 63 +++++++++++++++++++ 4 files changed, 70 insertions(+), 8 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ea81e323da4..6dff3976231 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -70,7 +70,7 @@ def to_basic_auth(auth_value: str) -> str: def strip_auth_scheme(auth_value: str, scheme: str) -> str: - """Return ``auth_value`` with a leading `` `` removed, or unchanged when absent. + """Return ``auth_value`` with a leading ```` and separator removed, or unchanged when absent. Callers supply both a bare credential and a complete header value, so prefixing unconditionally yields ``Bearer Bearer ``. Scheme names are case-insensitive per @@ -78,10 +78,9 @@ def strip_auth_scheme(auth_value: str, scheme: str) -> str: with the scheme text and a scheme with nothing behind it are returned untouched. Surrounding whitespace is left to ``_strip_header_whitespace`` at header-build time. """ - scheme_name, _, remainder = auth_value.lstrip().partition(" ") - credential: Final = remainder.lstrip() - if credential and scheme_name.lower() == scheme.lower(): - return credential + parts: Final = auth_value.split(None, 1) + if len(parts) == 2 and parts[0].lower() == scheme.lower(): + return parts[1] return auth_value diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 36e1078dc76..daaa5b2b322 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3346,9 +3346,7 @@ class MCPServerManager: normalized: Final = {k.lower(): v for k, v in raw_headers.items()} auth_value = normalized.get("authorization") if auth_value: - if auth_value.startswith("Bearer "): - return auth_value[len("Bearer ") :] - return auth_value + return strip_auth_scheme(auth_value, "Bearer") return None @staticmethod diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index fd7ab3afdab..4db131da62c 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -1014,6 +1014,8 @@ class TestAuthSchemeNormalization: [ ("Bearer abc", "Bearer", "abc"), ("bearer abc", "Bearer", "abc"), + ("Bearer\tabc", "Bearer", "abc"), + ("bearer\t\tabc", "Bearer", "abc"), (" Bearer abc ", "Bearer", "abc "), ("abc", "Bearer", "abc"), ("Bearerabc", "Bearer", "Bearerabc"), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 3321da83007..bd35719dd52 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -2616,6 +2616,69 @@ class TestMCPServerManager: await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None) assert resolved == ["good-subject"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "authorization", + [ + "Bearer subj-jwt", + "bearer subj-jwt", + "BEARER subj-jwt", + "Bearer\tsubj-jwt", + "Bearer subj-jwt", + ], + ) + async def test_preflight_token_exchange_strips_inbound_authorization_scheme(self, authorization): + """The resolver posts inbound_token verbatim as subject_token.""" + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError, ServerSpec, Subject + + resolved: Final[list[str | None]] = [] + + class _FakeProvider: + async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]: + resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) + return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server(f"te-preflight-auth-{authorization!r}") + + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": authorization}, + user_api_key_auth=None, + ) + + assert resolved == ["subj-jwt"] + + @pytest.mark.asyncio + async def test_preflight_token_exchange_preserves_authorization_without_separator(self): + from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import ( + StaticHeaderAuth, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok + from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError, ServerSpec, Subject + + resolved: Final[list[str | None]] = [] + + class _FakeProvider: + async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]: + resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None) + return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization")) + + manager = MCPServerManager(cred_provider=_FakeProvider()) + server = self._token_exchange_server("te-preflight-auth-no-separator") + + await manager.preflight_token_exchange( + server=server, + oauth2_headers={"Authorization": "Bearersubj-jwt"}, + user_api_key_auth=None, + ) + + assert resolved == ["Bearersubj-jwt"] + @pytest.mark.asyncio async def test_preflight_token_exchange_skips_discovery_for_other_auth_modes(self): """Preflight must not make unrelated auth modes depend on OAuth discovery.""" From 35d20468cd69dba41e7f13d315163cb73a67ff61 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:09:08 -0700 Subject: [PATCH 151/167] fix(mcp): normalize a schemed authentication_token on the v2 and OpenAPI static paths (#39345) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 10 +-- .../outbound_credentials/adapter.py | 5 +- .../outbound_credentials/test_adapter.py | 44 +++++++++++ .../mcp_server/test_mcp_server_manager.py | 75 +++++++++++++++++++ 4 files changed, 128 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index daaa5b2b322..0c3932d2ba1 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -47,7 +47,7 @@ from litellm.constants import ( MCP_TOOL_LISTING_TIMEOUT, ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException -from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth, strip_auth_scheme +from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth, strip_auth_scheme, to_basic_credentials from litellm.integrations.custom_guardrail import ( _sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic ) @@ -2299,13 +2299,13 @@ class MCPServerManager: from litellm.types.mcp import MCPAuth if server.auth_type == MCPAuth.bearer_token: - headers["Authorization"] = f"Bearer {server.authentication_token}" + headers["Authorization"] = f"Bearer {strip_auth_scheme(server.authentication_token, 'Bearer')}" elif server.auth_type == MCPAuth.api_key: - headers["Authorization"] = f"ApiKey {server.authentication_token}" + headers["Authorization"] = f"ApiKey {strip_auth_scheme(server.authentication_token, 'ApiKey')}" elif server.auth_type == MCPAuth.basic: - headers["Authorization"] = f"Basic {server.authentication_token}" + headers["Authorization"] = f"Basic {to_basic_credentials(server.authentication_token)}" elif server.auth_type == MCPAuth.token: - headers["Authorization"] = f"token {server.authentication_token}" + headers["Authorization"] = f"token {strip_auth_scheme(server.authentication_token, 'token')}" # Add any static headers from server config. # diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index 4458ac7f190..6a95a93a2a8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -18,6 +18,7 @@ from fastapi import HTTPException from pydantic import SecretStr from typing_extensions import assert_never +from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, @@ -213,7 +214,9 @@ def _shared_key_spec( token: Final = server.authentication_token if not token: return None # no key configured -> defer to v1 (parity-safe) - value: Final = base64.b64encode(token.encode("utf-8")).decode() if encode else token + value: Final = ( + to_basic_credentials(token) if encode else strip_auth_scheme(token, value_prefix) if value_prefix else token + ) return ServerSpec( server_id=server.server_id, resource=resource, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py index c667db7f07c..c6f3b9cb1f4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py @@ -12,6 +12,7 @@ import pytest from fastapi import HTTPException from pydantic import ValidationError +from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( oauth_protected_resource_path, raise_public, @@ -94,6 +95,49 @@ def test_basic_scheme_base64_encodes_the_token(): assert spec.config.key_source.value.get_secret_value() == expected +@pytest.mark.parametrize( + "auth_type, authentication_token, expected_value, expected_header", + [ + (MCPAuth.bearer_token, "Bearer abc", "abc", ("Authorization", "Bearer abc")), + (MCPAuth.token, "token abc", "abc", ("Authorization", "token abc")), + (MCPAuth.basic, "user:pass", "dXNlcjpwYXNz", ("Authorization", "Basic dXNlcjpwYXNz")), + (MCPAuth.basic, "Basic dXNlcjpwYXNz", "dXNlcjpwYXNz", ("Authorization", "Basic dXNlcjpwYXNz")), + (MCPAuth.basic, "Basic user:pass", "dXNlcjpwYXNz", ("Authorization", "Basic dXNlcjpwYXNz")), + ], +) +def test_shared_key_normalizes_schemed_authentication_token( + auth_type, authentication_token, expected_value, expected_header +): + spec = to_server_spec(_server(auth_type=auth_type, authentication_token=authentication_token)) + assert spec is not None and isinstance(spec.config, ApiKeyConfig) + assert spec.config.key_source.value.get_secret_value() == expected_value + assert spec.config.header(expected_value) == expected_header + + +@pytest.mark.parametrize( + "auth_type, authentication_token", + [ + (MCPAuth.bearer_token, "Bearer abc"), + (MCPAuth.bearer_token, "abc"), + (MCPAuth.token, "token abc"), + (MCPAuth.token, "abc"), + (MCPAuth.basic, "user:pass"), + (MCPAuth.basic, "Basic dXNlcjpwYXNz"), + (MCPAuth.basic, "Basic user:pass"), + ], +) +def test_shared_key_authorization_matches_v1(auth_type, authentication_token): + spec = to_server_spec(_server(auth_type=auth_type, authentication_token=authentication_token)) + assert spec is not None and isinstance(spec.config, ApiKeyConfig) + + client = MCPClient(server_url="https://x", auth_type=auth_type) + client.update_auth_value(authentication_token) + + assert spec.config.header(spec.config.key_source.value.get_secret_value())[1] == client._get_auth_headers()[ + "Authorization" + ] + + @pytest.mark.parametrize( "oauth2_flow", [None, "authorization_code"], diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index bd35719dd52..482f779bcb8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -4634,6 +4634,81 @@ class TestMCPServerManager: # auth_type is none here, so a 401 from this upstream must not be dressed up as a re-auth signal assert captured["relays_upstream_auth"] is False + @pytest.mark.asyncio + @pytest.mark.parametrize( + "auth_type, authentication_token, expected_authorization", + [ + (MCPAuth.bearer_token, "Bearer abc", "Bearer abc"), + (MCPAuth.bearer_token, "abc", "Bearer abc"), + (MCPAuth.api_key, "ApiKey abc", "ApiKey abc"), + (MCPAuth.token, "token abc", "token abc"), + (MCPAuth.basic, "user:pass", "Basic dXNlcjpwYXNz"), + (MCPAuth.basic, "Basic dXNlcjpwYXNz", "Basic dXNlcjpwYXNz"), + (MCPAuth.basic, "Basic user:pass", "Basic dXNlcjpwYXNz"), + ], + ) + async def test_register_openapi_tools_normalizes_authentication_token( + self, tmp_path, monkeypatch, auth_type, authentication_token, expected_authorization + ): + manager = MCPServerManager() + spec_path = tmp_path / "openapi.json" + spec_path.write_text( + json.dumps( + { + "openapi": "3.0.0", + "info": {"title": "Demo", "version": "1.0.0"}, + "paths": { + "/health": { + "get": { + "operationId": "health_check", + "summary": "health", + } + } + }, + } + ) + ) + server = MCPServer( + server_id="openapi-server", + name="openapi-server", + server_name="openapi-server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=auth_type, + authentication_token=authentication_token, + ) + captured: dict = {} + + def fake_create_tool_function( + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + ): + captured["headers"] = headers + + async def tool_func(**kwargs): + return "ok" + + return tool_func + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.create_tool_function", + fake_create_tool_function, + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.build_input_schema", + lambda *args, **kwargs: {"type": "object", "properties": {}, "required": []}, + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.tool_registry.global_mcp_tool_registry.register_tool", + lambda *args, **kwargs: None, + ) + await manager._register_openapi_tools( + spec_path=str(spec_path), + server=server, + base_url="https://example.com", + ) + + assert captured["headers"]["Authorization"] == expected_authorization + @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): """Test pre_call_tool_check allows tool when it's in allowed_tools list""" From 24666d8dbfc5ab377c8cbc5a62edb53d785f3540 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:10:33 -0700 Subject: [PATCH 152/167] fix(docker): match USE_DDTRACE case-insensitively and route build_from_pip through prod_entrypoint.sh (#39344) * fix(docker): match USE_DDTRACE case-insensitively and route build_from_pip through prod_entrypoint.sh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(docker): run build_from_pip ENTRYPOINT and CMD through the shipped prod_entrypoint.sh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../build_from_pip/Dockerfile.build_from_pip | 5 +- docker/component_entrypoint.sh | 10 ++- docker/prod_entrypoint.sh | 14 ++-- terraform/litellm/aws/ecs.tf | 4 +- terraform/litellm/gcp/cloudrun.tf | 4 +- .../test_litellm/test_component_entrypoint.py | 84 ++++++++++++++++--- 6 files changed, 95 insertions(+), 26 deletions(-) diff --git a/docker/build_from_pip/Dockerfile.build_from_pip b/docker/build_from_pip/Dockerfile.build_from_pip index e27e8f26cc2..a5733f0e1a0 100644 --- a/docker/build_from_pip/Dockerfile.build_from_pip +++ b/docker/build_from_pip/Dockerfile.build_from_pip @@ -55,7 +55,10 @@ RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/ ENV PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries +COPY docker/prod_entrypoint.sh /app/docker/prod_entrypoint.sh +RUN sed -i 's/\r$//' /app/docker/prod_entrypoint.sh && chmod +x /app/docker/prod_entrypoint.sh + EXPOSE 4000/tcp -ENTRYPOINT ["litellm"] +ENTRYPOINT ["/app/docker/prod_entrypoint.sh"] CMD ["--port", "4000"] diff --git a/docker/component_entrypoint.sh b/docker/component_entrypoint.sh index 1748f1e13a7..413957b9929 100755 --- a/docker/component_entrypoint.sh +++ b/docker/component_entrypoint.sh @@ -1,8 +1,10 @@ #!/bin/sh -if [ "$USE_DDTRACE" = "true" ]; then - export DD_TRACE_OPENAI_ENABLED="False" - exec ddtrace-run "$@" -fi +case "$USE_DDTRACE" in + [Tt][Rr][Uu][Ee]) + export DD_TRACE_OPENAI_ENABLED="False" + exec ddtrace-run "$@" + ;; +esac exec "$@" diff --git a/docker/prod_entrypoint.sh b/docker/prod_entrypoint.sh index bd78bf6687b..630eb6b065b 100644 --- a/docker/prod_entrypoint.sh +++ b/docker/prod_entrypoint.sh @@ -1,8 +1,10 @@ #!/bin/sh -if [ "$USE_DDTRACE" = "true" ]; then - export DD_TRACE_OPENAI_ENABLED="False" - exec ddtrace-run litellm "$@" -else - exec litellm "$@" -fi +case "$USE_DDTRACE" in + [Tt][Rr][Uu][Ee]) + export DD_TRACE_OPENAI_ENABLED="False" + exec ddtrace-run litellm "$@" + ;; +esac + +exec litellm "$@" diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 01b730dac65..2fb1ca08205 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -214,8 +214,8 @@ locals { gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" backend_uvicorn_args = "--host 0.0.0.0 --port 4001" - gateway_launch_cmd = "if [ \"$USE_DDTRACE\" = \"true\" ]; then export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args}; else exec uvicorn gateway.main:app ${local.gateway_uvicorn_args}; fi" - backend_launch_cmd = "if [ \"$USE_DDTRACE\" = \"true\" ]; then export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args}; else exec uvicorn backend.main:app ${local.backend_uvicorn_args}; fi" + gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" + backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac" gateway_proxy_overrides = local.proxy_config_enabled ? { entryPoint = ["sh", "-c"] diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 3340522734f..5a5c361b832 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -138,8 +138,8 @@ locals { gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" backend_uvicorn_args = "--host 0.0.0.0 --port 4001" - gateway_launch_cmd = "if [ \"$USE_DDTRACE\" = \"true\" ]; then export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args}; else exec uvicorn gateway.main:app ${local.gateway_uvicorn_args}; fi" - backend_launch_cmd = "if [ \"$USE_DDTRACE\" = \"true\" ]; then export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args}; else exec uvicorn backend.main:app ${local.backend_uvicorn_args}; fi" + gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" + backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac" gateway_args = join(" && ", concat( local.redis_ca_fragment, diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/test_litellm/test_component_entrypoint.py index 2536b76e91c..b0969e2c694 100644 --- a/tests/test_litellm/test_component_entrypoint.py +++ b/tests/test_litellm/test_component_entrypoint.py @@ -7,7 +7,6 @@ import re import stat import subprocess from pathlib import Path -from typing import Optional import pytest @@ -16,11 +15,15 @@ COMPONENT_ENTRYPOINT = REPO_ROOT / "docker" / "component_entrypoint.sh" PROD_ENTRYPOINT = REPO_ROOT / "docker" / "prod_entrypoint.sh" GATEWAY_DOCKERFILE = REPO_ROOT / "gateway" / "Dockerfile" BACKEND_DOCKERFILE = REPO_ROOT / "backend" / "Dockerfile" +BUILD_FROM_PIP_DOCKERFILE = REPO_ROOT / "docker" / "build_from_pip" / "Dockerfile.build_from_pip" TERRAFORM_ECS = REPO_ROOT / "terraform" / "litellm" / "aws" / "ecs.tf" TERRAFORM_CLOUDRUN = REPO_ROOT / "terraform" / "litellm" / "gcp" / "cloudrun.tf" IMAGE_ENTRYPOINT_PATH = "/app/docker/component_entrypoint.sh" +TRUTHY_USE_DDTRACE = ("true", "True", "TRUE", "tRuE") +FALSY_USE_DDTRACE = (None, "", "false", "False", "1", "yes", "on", "truex") + PYTHONPATH_SENTINEL = "/lit-entrypoint-sentinel:/app" _STUB_TEMPLATE = """#!/bin/sh @@ -34,6 +37,7 @@ _STUB_TEMPLATE = """#!/bin/sh _ENTRYPOINT_RE = re.compile(r"^ENTRYPOINT\s+(\[.*\])\s*$", re.MULTILINE) _CMD_RE = re.compile(r"^CMD\s+(\[.*\])\s*$", re.MULTILINE) +_COPY_RE = re.compile(r"^COPY\s+(?!--from)(\S+)\s+(\S+)\s*$", re.MULTILINE) _APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app") _TF_STRING_LOCAL_RE = re.compile(r'^\s*(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"\s*$', re.MULTILINE) _TF_INTERPOLATION_RE = re.compile(r"\$\{(local|var)\.(\w+)\}") @@ -53,7 +57,7 @@ def _write_stubs(bin_dir: Path, names: tuple[str, ...]) -> None: def _run_entrypoint( script: Path, argv: tuple[str, ...], - use_ddtrace: Optional[str], + use_ddtrace: str | None, tmp_path: Path, ) -> tuple[str, ...]: """Run `script` with stubbed executables on PATH and return the recorded lines.""" @@ -84,7 +88,7 @@ def _run_entrypoint( return tuple(record.read_text().splitlines()) if record.exists() else () -def _run_shell_command(command: str, bin_dir: Path, record: Path, use_ddtrace: Optional[str]) -> tuple[str, ...]: +def _run_shell_command(command: str, bin_dir: Path, record: Path, use_ddtrace: str | None) -> tuple[str, ...]: """Run a resolved Terraform launch command through `sh -c` and return the recorded lines.""" env = { **os.environ, @@ -184,9 +188,19 @@ def test_ddtrace_disabled_execs_the_command_directly(tmp_path: Path) -> None: ) -@pytest.mark.parametrize("use_ddtrace", [None, "", "false", "True", "TRUE", "1", "yes"]) -def test_gating_matches_the_monolithic_entrypoint(use_ddtrace: Optional[str], tmp_path: Path) -> None: - """The componentized images must honor `USE_DDTRACE` exactly as the monolith does.""" +@pytest.mark.parametrize( + "use_ddtrace, traced", + [*((v, True) for v in TRUTHY_USE_DDTRACE), *((v, False) for v in FALSY_USE_DDTRACE)], +) +def test_gating_matches_the_monolithic_entrypoint_and_get_secret_bool( + use_ddtrace: str | None, traced: bool, tmp_path: Path +) -> None: + """Both entrypoints must accept exactly the spellings `get_secret_bool` accepts. + + `ProxyStartupEvent._init_dd_tracer` reads `USE_DDTRACE` through `get_secret_bool`, which + matches `true` case-insensitively. If the shell gate were stricter, `USE_DDTRACE=True` would + give in-process LLM spans without `ddtrace-run` HTTP spans, a half-enabled state. + """ component = _run_entrypoint( COMPONENT_ENTRYPOINT, ("uvicorn", "gateway.main:app"), @@ -200,8 +214,56 @@ def test_gating_matches_the_monolithic_entrypoint(use_ddtrace: Optional[str], tm tmp_path=tmp_path / "monolith", ) - assert component[0].startswith("exec=") and monolith[0].startswith("exec=") - assert (component[0] == "exec=ddtrace-run") == (monolith[0] == "exec=ddtrace-run") + expected_exec = "exec=ddtrace-run" if traced else "exec=uvicorn" + expected_openai = "DD_TRACE_OPENAI_ENABLED=False" if traced else "DD_TRACE_OPENAI_ENABLED=" + assert component[0] == expected_exec + assert component[2] == expected_openai + assert monolith[0] == ("exec=ddtrace-run" if traced else "exec=litellm") + assert monolith[2] == expected_openai + assert monolith[1] == ("args=litellm --port 4000" if traced else "args=--port 4000") + + +def _copied_script(dockerfile: Path, image_path: str) -> Path: + """Resolve the repo file a Dockerfile `COPY`s to `image_path`, so tests run what the image ships.""" + matches = _COPY_RE.findall(dockerfile.read_text()) + sources = tuple(src for src, dst in matches if dst == image_path) + assert sources, f"{dockerfile} never COPYs anything to {image_path}" + source = REPO_ROOT / sources[-1] + assert source.is_file(), f"{dockerfile} COPYs {sources[-1]}, which does not exist in the build context" + return source + + +@pytest.mark.parametrize( + "use_ddtrace, traced", + [*((v, True) for v in TRUTHY_USE_DDTRACE), *((v, False) for v in FALSY_USE_DDTRACE)], +) +def test_build_from_pip_image_launches_litellm_through_the_prod_entrypoint( + use_ddtrace: str | None, traced: bool, tmp_path: Path +) -> None: + """Run the build_from_pip image's ENTRYPOINT + CMD through the script it actually COPYs. + + The image used to `ENTRYPOINT ["litellm"]`, so `USE_DDTRACE` was inert there at any spelling. + Resolving the ENTRYPOINT path back to its COPY source and executing it with the Dockerfile's + CMD checks the launch the container performs, not just that the Dockerfile mentions the script. + """ + entrypoint = _entrypoint_argv(BUILD_FROM_PIP_DOCKERFILE) + assert len(entrypoint) == 1, ( + f"{BUILD_FROM_PIP_DOCKERFILE} ENTRYPOINT must be the bare script so CMD reaches litellm" + ) + script = _copied_script(BUILD_FROM_PIP_DOCKERFILE, entrypoint[0]) + assert script == PROD_ENTRYPOINT, f"{BUILD_FROM_PIP_DOCKERFILE} bypasses the ddtrace-aware entrypoint" + assert f"chmod +x {entrypoint[0]}" in BUILD_FROM_PIP_DOCKERFILE.read_text() + + cmd = _cmd_argv(BUILD_FROM_PIP_DOCKERFILE) + recorded = _run_entrypoint(script, cmd, use_ddtrace=use_ddtrace, tmp_path=tmp_path) + + cmd_str = " ".join(cmd) + assert recorded == ( + "exec=ddtrace-run" if traced else "exec=litellm", + f"args=litellm {cmd_str}" if traced else f"args={cmd_str}", + "DD_TRACE_OPENAI_ENABLED=False" if traced else "DD_TRACE_OPENAI_ENABLED=", + f"PYTHONPATH={PYTHONPATH_SENTINEL}", + ) def test_entrypoint_script_is_executable() -> None: @@ -239,9 +301,9 @@ def test_component_images_make_the_entrypoint_executable(dockerfile: Path) -> No @pytest.mark.parametrize("terraform_file", TERRAFORM_LAUNCH_SITES, ids=lambda p: p.parent.name) @pytest.mark.parametrize("component", ["gateway", "backend"]) -@pytest.mark.parametrize("use_ddtrace", [None, "true", "false", "True"]) +@pytest.mark.parametrize("use_ddtrace", [*TRUTHY_USE_DDTRACE, *FALSY_USE_DDTRACE]) def test_terraform_launch_command_matches_the_script_contract( - terraform_file: Path, component: str, use_ddtrace: Optional[str], tmp_path: Path + terraform_file: Path, component: str, use_ddtrace: str | None, tmp_path: Path ) -> None: """The Terraform command and `docker/component_entrypoint.sh` must decide identically. @@ -273,7 +335,7 @@ def test_terraform_launch_command_matches_the_script_contract( assert from_terraform[2] == from_script[2], f"{terraform_file} disagrees with the script on the openai integration" assert app_target in from_terraform[1] - if use_ddtrace == "true": + if use_ddtrace in TRUTHY_USE_DDTRACE: assert from_terraform[0] == "exec=ddtrace-run" assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=False" else: From a0958d5c217321242ebc5227850eceea7276b855 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:12:06 -0700 Subject: [PATCH 153/167] perf(auth): skip object permission DB lookup when no vector stores requested (#39347) * perf(auth): skip object permission DB lookup when no vector stores requested Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(auth): justify module patches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- .../proxy/auth/test_auth_checks.py | 32 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0db4cbc3bf2..2b328b455dc 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5788,7 +5788,7 @@ async def vector_store_access_check( vector_store_ids_to_run: Final = litellm.vector_store_registry.get_vector_store_ids_to_run( non_default_params=request_body, tools=request_body.get("tools", None) ) - if vector_store_ids_to_run is None: + if not vector_store_ids_to_run: verbose_proxy_logger.debug("Vector store to run not found, skipping vector store access check") return True diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ffe3e0fee12..be83ca57e76 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1349,6 +1349,38 @@ async def test_vector_store_access_check_early_returns( assert result == expected_result +@pytest.mark.asyncio +async def test_vector_store_access_check_skips_db_lookup_when_no_vector_stores_requested(): + """Registry returns [] (not None) for plain requests; no object permission DB lookup should happen.""" + valid_token = UserAPIKeyAuth(token="test-token", object_permission_id="perm-123") + team_object = MagicMock() + team_object.object_permission_id = "team-permission" + + mock_prisma_client = MagicMock() + find_unique = AsyncMock() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = find_unique + + mock_vector_store_registry = MagicMock() + mock_vector_store_registry.get_vector_store_ids_to_run.return_value = [] + + with ( + patch( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ), + patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.vector_store_registry", mock_vector_store_registry + ), + ): + result = await vector_store_access_check( + request_body={"messages": [{"role": "user", "content": "test"}]}, + team_object=team_object, + valid_token=valid_token, + ) + + assert result is True + find_unique.assert_not_awaited() + + @pytest.mark.parametrize( "object_permissions,vector_store_ids,should_raise,error_type", [ From 942a46ffd7793ca1c535bf73ae1c8f3366976617 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:18:20 +0000 Subject: [PATCH 154/167] fix(caching): don't trip redis circuit breaker on short timeout bursts (#38999) * fix(caching): don't trip redis circuit breaker on short timeout bursts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): scope timeout duration gate to timeout failures and count breaker states per label Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): reset the timeout streak on hard failures so stale timeouts cannot pre-age the duration gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 142 +++++++++++++++-- litellm/constants.py | 3 + .../test_litellm/caching/test_redis_cache.py | 150 +++++++++++++++++- 3 files changed, 284 insertions(+), 11 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 58733b384b9..aaee7188d86 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -27,6 +27,7 @@ from litellm.constants import ( REDIS_CIRCUIT_BREAKER_ENABLED, REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, + REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, ) from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker @@ -41,6 +42,8 @@ from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prometheus_client import Counter as _PromCounter + from prometheus_client import Gauge as _PromGauge from redis.asyncio import Redis, RedisCluster from redis.asyncio.client import Pipeline from redis.asyncio.cluster import ClusterPipeline @@ -135,10 +138,20 @@ class RedisCircuitBreaker: HALF_OPEN - recovery probe: allow one request through Transitions: - CLOSED -> OPEN after failure_threshold consecutive failures + CLOSED -> OPEN after failure_threshold consecutive hard connectivity + failures, or after an unbroken run of timeout failures + (no success or hard failure in between) that reaches + failure_threshold and spans timeout_min_duration seconds OPEN -> HALF_OPEN after recovery_timeout seconds HALF_OPEN -> CLOSED on success HALF_OPEN -> OPEN on failure (resets timer) + + Timeouts are accounted separately from hard connectivity failures because the async + Redis timeout includes time waiting for the worker event loop to resume: one loop + stall makes every in-flight operation time out together, which satisfies a purely + consecutive threshold instantly even though Redis is healthy. Requiring a + timeout-only streak to also span timeout_min_duration filters such bursts while a + real outage that surfaces as timeouts still opens the breaker after that duration. """ CLOSED = "closed" @@ -150,13 +163,19 @@ class RedisCircuitBreaker: failure_threshold: int, recovery_timeout: int, enabled: bool = True, + timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, ) -> None: self.failure_threshold = failure_threshold self.recovery_timeout = recovery_timeout self.enabled = enabled + self.timeout_min_duration = timeout_min_duration self._failure_count = 0 + self._hard_failure_count = 0 + self._timeout_count = 0 + self._timeout_streak_started_at: float | None = None self._opened_at: float | None = None self._state = self.CLOSED + _breaker_metrics().record_state_change(None, self._state) def is_open(self) -> bool: """Returns True if Redis calls should be skipped.""" @@ -169,24 +188,45 @@ class RedisCircuitBreaker: return True if self._state == self.OPEN: if time.time() - (self._opened_at or 0) > self.recovery_timeout: - self._state = self.HALF_OPEN + self._set_state(self.HALF_OPEN) return False # this caller is the designated probe return True return False - def record_failure(self) -> None: + def _should_open(self, now: float) -> bool: + if self._state == self.HALF_OPEN: + return True + if self._hard_failure_count >= self.failure_threshold: + return True + if self._timeout_count < self.failure_threshold: + return False + return now - (self._timeout_streak_started_at or now) >= self.timeout_min_duration + + def record_failure(self, is_timeout: bool = False) -> None: if not self.enabled: return + now: Final = time.time() self._failure_count += 1 - self._opened_at = time.time() - if self._failure_count >= self.failure_threshold: + if is_timeout: + self._timeout_count += 1 + if self._timeout_streak_started_at is None: + self._timeout_streak_started_at = now + else: + self._hard_failure_count += 1 + self._timeout_count = 0 + self._timeout_streak_started_at = None + self._opened_at = now + _breaker_metrics().record_failure("timeout" if is_timeout else "connectivity") + if self._should_open(now): if self._state != self.OPEN: verbose_logger.warning( - "Redis circuit breaker OPENED after %d consecutive failures — fast-failing Redis calls for %ds", + "Redis circuit breaker OPENED after %d consecutive failures" + " (%d hard connectivity) — fast-failing Redis calls for %ds", self._failure_count, + self._hard_failure_count, self.recovery_timeout, ) - self._state = self.OPEN + self._set_state(self.OPEN) def record_success(self) -> None: if not self.enabled: @@ -194,7 +234,17 @@ class RedisCircuitBreaker: if self._state == self.HALF_OPEN: verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered") self._failure_count = 0 - self._state = self.CLOSED + self._hard_failure_count = 0 + self._timeout_count = 0 + self._timeout_streak_started_at = None + self._set_state(self.CLOSED) + + def _set_state(self, state: str) -> None: + if state == self._state: + return + _breaker_metrics().record_transition(state) + _breaker_metrics().record_state_change(self._state, state) + self._state = state _RedisCallResult = TypeVar("_RedisCallResult") @@ -234,6 +284,78 @@ def _is_redis_health_failure(exc: BaseException) -> bool: return True +@functools.lru_cache(maxsize=1) +def _redis_timeout_error_types() -> tuple[type, ...]: + """Health failures that are timeouts rather than unambiguous connectivity errors. + + ``builtins.TimeoutError`` covers ``asyncio.TimeoutError`` and ``socket.timeout`` + (aliases since py3.11 / py3.10). ``redis.exceptions.TimeoutError`` does not subclass + either, so it is listed explicitly. + """ + try: + from redis.exceptions import TimeoutError as RedisTimeoutError + except ImportError: + return (TimeoutError,) + return (RedisTimeoutError, TimeoutError) + + +def _is_redis_timeout_failure(exc: BaseException) -> bool: + return isinstance(exc, _redis_timeout_error_types()) + + +class _BreakerMetrics: + """Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent. + + Registered lazily on the default registry (which /metrics serves) via the module-level + ``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers. + """ + + def __init__(self) -> None: + self._state_gauge: _PromGauge | None = None + self._transitions: _PromCounter | None = None + self._failures: _PromCounter | None = None + try: + from prometheus_client import Counter as PromCounter + from prometheus_client import Gauge + except ImportError: + return + self._state_gauge = Gauge( + "litellm_redis_circuit_breaker_state", + "Number of Redis circuit breakers currently in each state", + labelnames=("state",), + ) + self._transitions = PromCounter( + "litellm_redis_circuit_breaker_transitions", + "Redis circuit breaker state transitions", + labelnames=("state",), + ) + self._failures = PromCounter( + "litellm_redis_circuit_breaker_failures", + "Redis health failures counted by the circuit breaker", + labelnames=("failure_class",), + ) + + def record_state_change(self, old_state: str | None, new_state: str) -> None: + if self._state_gauge is None: + return + if old_state is not None: + self._state_gauge.labels(old_state).dec() + self._state_gauge.labels(new_state).inc() + + def record_transition(self, state: str) -> None: + if self._transitions is not None: + self._transitions.labels(state).inc() + + def record_failure(self, failure_class: str) -> None: + if self._failures is not None: + self._failures.labels(failure_class).inc() + + +@functools.lru_cache(maxsize=1) +def _breaker_metrics() -> _BreakerMetrics: + return _BreakerMetrics() + + def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None: """Record a Redis failure that the calling method is about to swallow. @@ -245,7 +367,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep """ if not _is_redis_health_failure(exc): return - breaker.record_failure() + breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc)) _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) @@ -281,7 +403,7 @@ async def _run_under_circuit_breaker( result: Final = await call() except Exception as e: if _is_redis_health_failure(e): - breaker.record_failure() + breaker.record_failure(is_timeout=_is_redis_timeout_failure(e)) raise _exit_circuit_breaker(breaker, swallowed_before) return result diff --git a/litellm/constants.py b/litellm/constants.py index f5acadc32ab..cd72adc3db5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -432,6 +432,9 @@ REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIME REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)) REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)) REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" +# minimum seconds a timeout-only failure streak must span before it can open the breaker, +# so one event-loop stall timing out many queued calls at once does not trip it +REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0)) # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart # (e.g. ElastiCache Serverless maintenance) is not reused while broken diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 71be8730df1..e4724ff8705 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -779,7 +779,7 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): "error, opens_breaker", [ pytest.param("ConnectionError", True, id="connection_refused_is_unhealthy"), - pytest.param("TimeoutError", True, id="timeout_is_unhealthy"), + pytest.param("TimeoutError", False, id="timeout_burst_is_ambiguous"), pytest.param("BusyLoadingError", True, id="loading_is_unhealthy"), pytest.param("ResponseError", False, id="wrong_type_command_is_not"), pytest.param("DataError", False, id="bad_data_is_not"), @@ -791,6 +791,10 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker) They say nothing about connectivity, and a caller able to provoke them (an INCR against a non-numeric value, say) could otherwise trip the shared breaker on demand and drop rate limiting to per-process counters, which spreading traffic across replicas outruns. + + A rapid burst of timeouts is ambiguous too: the async timeout includes event-loop + scheduling delay, so a loop stall times out every queued call at once against a + healthy Redis. It must not open the breaker until the streak spans a minimum duration. """ import redis.exceptions @@ -810,3 +814,147 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker) await _run_under_circuit_breaker(breaker, "op", failing_call) assert breaker.is_open() is opens_breaker + + +@pytest.mark.asyncio +async def test_event_loop_stall_timeout_burst_keeps_breaker_closed(): + """One blocking stall of the worker event loop must not trip the breaker. + + Every operation already waiting on the loop times out together when the loop resumes, + so a purely consecutive threshold is satisfied instantly even though the Redis on the + other end (here an in-process fake that answers immediately) is healthy. + """ + import time as time_mod + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + async def healthy_redis_call_with_client_timeout(): + return await asyncio.wait_for(asyncio.sleep(0.001, result="ok"), timeout=0.05) + + async def stall_the_loop(): + await asyncio.sleep(0) + time_mod.sleep(0.2) + + results = await asyncio.gather( + *(_run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) for _ in range(8)), + stall_the_loop(), + return_exceptions=True, + ) + timeouts = [r for r in results if isinstance(r, asyncio.TimeoutError)] + assert len(timeouts) >= breaker.failure_threshold, "the stall must time out a full burst" + + assert breaker.is_open() is False, "a healthy Redis behind one loop stall must stay in the pool" + assert await _run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) == "ok" + + +@pytest.mark.asyncio +async def test_persistent_timeouts_still_open_the_breaker(): + """A real outage that surfaces only as timeouts must still open the breaker. + + Once the timeout-only streak spans the minimum duration with no success in between, + Redis is genuinely unusable from this worker and protection has to kick in. + """ + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.1) + + async def timing_out_call(): + raise RedisTimeoutError("read timed out") + + for _ in range(breaker.failure_threshold): + with pytest.raises(RedisTimeoutError): + await _run_under_circuit_breaker(breaker, "op", timing_out_call) + assert breaker.is_open() is False, "the burst has not spanned the minimum duration yet" + + await asyncio.sleep(0.12) + with pytest.raises(RedisTimeoutError): + await _run_under_circuit_breaker(breaker, "op", timing_out_call) + + assert breaker.is_open() is True + + +@pytest.mark.asyncio +async def test_stale_timeout_does_not_let_sub_threshold_hard_failures_open_the_breaker(): + """Hard connectivity failures below the threshold must not open the breaker just + because an old timeout already started the streak and the duration has elapsed. + + Each class has to earn the open on its own terms: hard failures by reaching the + threshold, timeouts by reaching the threshold and spanning the minimum duration. + """ + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) + + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + await asyncio.sleep(0.06) + for _ in range(breaker.failure_threshold - 1): + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + assert breaker.is_open() is False, "2 hard failures and 1 stale timeout are below both thresholds" + + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + assert breaker.is_open() is True, "the threshold-th hard failure must still open it" + + +@pytest.mark.asyncio +async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its_own_duration(): + """A stale timeout followed by hard failures must not pre-age the duration gate: + a later short timeout burst has to span timeout_min_duration on its own. + """ + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) + + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + await asyncio.sleep(0.06) + for _ in range(breaker.failure_threshold): + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + assert breaker.is_open() is False, "the burst is instantaneous, so the duration gate must hold it closed" + + await asyncio.sleep(0.06) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out"))) + assert breaker.is_open() is True, "the same run of timeouts persisting past the duration must open it" + + +@pytest.mark.asyncio +async def test_breaker_metrics_track_state_and_failure_class(): + """Breaker accounting must be observable: failure class, transitions, and state.""" + from prometheus_client import REGISTRY + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure + + def sample(name, labels=None): + return REGISTRY.get_sample_value(name, labels) or 0.0 + + timeout_before = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) + hard_before = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) + opened_before = sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) + open_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "open"}) + closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) + + breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("t"))) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused"))) + + assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) == timeout_before + 1 + assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) == hard_before + 2 + assert sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) == opened_before + 1 + assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + 1 + assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + + breaker.record_success() + assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1 From 8e4397d5f1358d8465ce80c28a56cb60f3faa997 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:19:49 -0700 Subject: [PATCH 155/167] chore(ci): rebuild the PR merge ref against staging's router coverage fix From 4811041048fb8bca24bf00938d0111e060a8b686 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 15:20:13 -0700 Subject: [PATCH 156/167] feat(proxy): add a search param to key, memory, audit, and spend log listings GET /key/list?search= matches the key hash (a raw sk- key is hashed first) or a case-insensitive alias substring, and key_hash= now hashes a raw sk- value too. GET /v1/memory?search= matches a key prefix or an exact memory_id. GET /audit?search= matches id, object_id, changed_by, or changed_by_api_key. GET /spend/logs/ui?search= matches request_id across all time and api_key, team_id, user, end_user, session_id, or model_id inside the date window; session grouping is skipped while a search is active. Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf --- .../proxy/audit_logging_endpoints.py | 28 +- .../key_management_endpoints.py | 30 +- litellm/proxy/memory/memory_endpoints.py | 56 +++- .../spend_management_endpoints.py | 62 +++- .../key_management_endpoints.py | 17 ++ .../proxy/test_audit_logging_endpoints.py | 66 +++- .../test_key_management_endpoints.py | 152 +++++++++ .../proxy/memory/test_memory_endpoints.py | 90 ++++++ .../test_spend_management_endpoints.py | 288 +++++++++++++++++- 9 files changed, 767 insertions(+), 22 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index b6f8bf2dc5b..72f7a66a420 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -18,6 +18,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import ( from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import _hash_token_if_needed from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import AuditLogRepository @@ -48,6 +49,19 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec } +def _build_search_condition(search: str) -> dict[str, object]: + """Match any id column; a raw sk- key is hashed for the two columns that store key hashes.""" + hashed: Final = _hash_token_if_needed(search) + return { + "OR": ( + {"id": search}, + {"changed_by": search}, + {"object_id": hashed}, + {"changed_by_api_key": hashed}, + ) + } + + @router.get( "/audit", tags=["Audit Logging"], @@ -83,6 +97,13 @@ async def get_audit_logs( None, description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)", ), + search: str | None = Query( + None, + description=( + "Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value " + "(a raw sk- virtual key is hashed first)" + ), + ), # Sorting parameters sort_by: str | None = Query( None, @@ -118,6 +139,11 @@ async def get_audit_logs( *([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []), ] + and_conditions: Final[tuple[dict[str, object], ...]] = ( + *json_field_conditions, + *((_build_search_condition(search),) if search else ()), + ) + # Build filter conditions where_conditions: Final[dict[str, object]] = { **({"changed_by": changed_by} if changed_by else {}), @@ -126,7 +152,7 @@ async def get_audit_logs( **({"table_name": table_name} if table_name else {}), **({"object_id": object_id} if object_id else {}), **({"updated_at": date_filter} if start_date or end_date else {}), - **({"AND": json_field_conditions} if json_field_conditions else {}), + **({"AND": and_conditions} if and_conditions else {}), } order_by: Final[dict[str, str]] = ( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d7d20d168b5..758644ff01b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -147,6 +147,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyResponse, BulkUpdateTeamKeysRequest, FailedKeyUpdate, + KeySearchWhere, SuccessfulKeyUpdate, ) from litellm.types.router import Deployment @@ -5800,6 +5801,10 @@ async def list_keys( None, description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.", ), + search: str | None = Query( + None, + description="Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive).", + ), return_full_object: bool = Query(False, description="Return full key object"), include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."), include_created_by_keys: bool = Query(False, description="Include keys created by the user"), @@ -5862,13 +5867,17 @@ async def list_keys( detail={"error": "Invalid expires value. Supported: 'active', 'expired'."}, ) + hashed_key_hash: Final[str | None] = ( + _hash_token_if_needed(token=key_hash) if isinstance(key_hash, str) else None + ) + complete_user_info: Final = await validate_key_list_check( user_api_key_dict=user_api_key_dict, user_id=user_id, team_id=team_id, organization_id=organization_id, key_alias=key_alias, - key_hash=key_hash, + key_hash=hashed_key_hash, prisma_client=prisma_client, ) @@ -5928,7 +5937,7 @@ async def list_keys( user_id=user_id, team_id=team_id, key_alias=key_alias, - key_hash=key_hash, + key_hash=hashed_key_hash, return_full_object=return_full_object, organization_id=organization_id, admin_team_ids=admin_team_ids, @@ -5943,6 +5952,7 @@ async def list_keys( agent_id=agent_id, use_substring_matching=use_substring_matching, expires_filter=expires if isinstance(expires, str) else None, + search=search, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -6162,6 +6172,16 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, return {"OR": [{"expires": None}, {"expires": {"gte": now}}]} +def _build_key_search_where(search: str) -> KeySearchWhere: + search_where: Final[KeySearchWhere] = { + "OR": ( + {"token": _hash_token_if_needed(token=search)}, + {"key_alias": {"contains": search, "mode": "insensitive"}}, + ) + } + return search_where + + def _build_key_filter_conditions( user_id: str | None, team_id: str | None, @@ -6177,6 +6197,7 @@ def _build_key_filter_conditions( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, + search: str | None = None, ) -> Mapping[str, object]: """Build filter conditions for key listing. @@ -6266,7 +6287,7 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) - global_filters: Final[tuple[dict[str, object], ...]] = ( + global_filters: Final[tuple[Mapping[str, object], ...]] = ( *( ( {"key_alias": {"contains": key_alias, "mode": "insensitive"}} @@ -6277,6 +6298,7 @@ def _build_key_filter_conditions( else () ), *(({"token": key_hash},) if key_hash and isinstance(key_hash, str) else ()), + *((_build_key_search_where(search),) if isinstance(search, str) and search else ()), *(({"team_id": team_id},) if team_id and isinstance(team_id, str) else ()), *(({"project_id": project_id},) if project_id else ()), *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()), @@ -6316,6 +6338,7 @@ async def _list_key_helper( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, + search: str | None = None, ) -> KeyListResponseObject: """ Helper function to list keys @@ -6354,6 +6377,7 @@ async def _list_key_helper( agent_id=agent_id, use_substring_matching=use_substring_matching, expires_filter=expires_filter, + search=search, ) # Calculate skip for pagination diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 98c5fdd198c..d8f72d200c7 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -22,6 +22,7 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Final from fastapi import APIRouter, Depends, HTTPException, Query +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -91,6 +92,36 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object return {"OR": ors} +class _StartsWith(TypedDict): + startsWith: ReadOnly[str] + + +class _MemoryKeyWhere(TypedDict): + key: ReadOnly[str | _StartsWith] + + +class _MemoryIdWhere(TypedDict): + memory_id: ReadOnly[str] + + +class _MemorySearchWhere(TypedDict): + OR: ReadOnly[tuple[_MemoryKeyWhere, _MemoryIdWhere]] + + +def _key_filter(search: str | None, key_prefix: str | None, key: str | None) -> Mapping[str, object] | None: + """`search` matches a key prefix or an exact memory_id; otherwise `key_prefix` wins over `key`.""" + if search is not None: + search_where: Final[_MemorySearchWhere] = {"OR": ({"key": {"startsWith": search}}, {"memory_id": search})} + return search_where + if key_prefix is not None: + prefix_where: Final[_MemoryKeyWhere] = {"key": {"startsWith": key_prefix}} + return prefix_where + if key is not None: + exact_where: Final[_MemoryKeyWhere] = {"key": key} + return exact_where + return None + + def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow: return LiteLLM_MemoryRow( memory_id=row.memory_id, @@ -326,6 +357,13 @@ async def list_memory( "Mutually exclusive with `key`; if both are provided, `key_prefix` wins." ), ), + search: str | None = Query( + None, + description=( + "Match entries whose key starts with this value or whose memory_id equals it. " + "Takes precedence over `key_prefix` and `key` when provided." + ), + ), page: int = Query(1, ge=1), page_size: int = Query(50, ge=1, le=500), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -333,22 +371,16 @@ async def list_memory( """List memory entries visible to the caller.""" prisma_client: Final = _require_prisma() - # Build the key filter first (prefix wins if both `key` and `key_prefix` - # are passed). Then AND it with the visibility filter via an explicit - # top-level "AND" — safer than `dict.update` since future visibility - # filters could grow an "OR" key that would clobber this one if merged - # by key. - key_filter: Final[dict[str, object]] = {} - if key_prefix is not None: - key_filter["key"] = {"startsWith": key_prefix} - elif key is not None: - key_filter["key"] = key + # AND the key filter with the visibility filter via an explicit top-level + # "AND": both sides can carry an "OR" key (`search`, non-admin visibility), + # so merging them by key would let one clobber the other and leak rows. + key_filter: Final = _key_filter(search=search, key_prefix=key_prefix, key=key) vis: Final = _visibility_filter(user_api_key_dict) - where: Mapping[str, object] + where: Mapping[str, object] | None if vis is None: where = key_filter - elif not key_filter: + elif key_filter is None: where = vis else: where = {"AND": [key_filter, vis]} diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 2a50d5170f0..9ec8dd205a6 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2229,6 +2229,33 @@ async def calculate_spend(request: SpendCalculateRequest): ) +class _SpendLogSearchCondition(NamedTuple): + sql: str + params: tuple[object, ...] + + +def _build_spend_log_search_condition( + search: str, + start_date: datetime, + end_date: datetime, + next_param_index: int, +) -> _SpendLogSearchCondition: + """request_id (indexed) matches across all time; the unindexed id columns only inside the window (sk- keys hashed).""" + raw: Final = f"${next_param_index}" + hashed: Final = f"${next_param_index + 1}" + window_start: Final = f"${next_param_index + 2}" + window_end: Final = f"${next_param_index + 3}" + sql: Final = ( + f"(request_id = {raw} OR (" + f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') " + f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') " + f'AND (api_key = {hashed} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} ' + f"OR session_id = {raw} OR model_id = {raw})))" + ) + hashed_search: Final = hash_token(token=search) if search.startswith("sk-") else search + return _SpendLogSearchCondition(sql=sql, params=(search, hashed_search, start_date, end_date)) + + @router.get( "/spend/logs/v2", tags=["Budget & Spend Tracking"], @@ -2329,6 +2356,14 @@ async def ui_view_spend_logs( "UI route only, honored when sorting by startTime" ), ), + search: str | None = fastapi.Query( + default=None, + description=( + "Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, " + "session_id, or model_id equals this value. request_id matches across all time; the other columns " + "match inside start_date/end_date, which stay required" + ), + ), ): """ View spend logs with pagination support. @@ -2392,8 +2427,10 @@ async def ui_view_spend_logs( try: is_admin_view: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) is_request_id_lookup: Final = request_id is not None and not is_v2 + is_search_lookup: Final = search is not None + search_owns_window: Final = is_search_lookup and not is_v2 - if is_request_id_lookup: + if is_request_id_lookup and not is_search_lookup: # request_id is the @id primary key: it identifies a single row, so a # time window is meaningless. The dashboard always sends a default 24h # window, which hid ids copied from an older page (LIT-3981). Drop the @@ -2576,7 +2613,7 @@ async def ui_view_spend_logs( # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison # against the plain `timestamp` column does not depend on the DB session # timezone (see #22529). Absent for a request_id-only lookup (see above). - if start_date_obj is not None and end_date_obj is not None: + if start_date_obj is not None and end_date_obj is not None and not search_owns_window: sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')") sql_params.append(start_date_obj) p += 1 @@ -2584,6 +2621,17 @@ async def ui_view_spend_logs( sql_params.append(end_date_obj) p += 1 + if search is not None and start_date_obj is not None and end_date_obj is not None: + search_condition: Final = _build_spend_log_search_condition( + search=search, + start_date=start_date_obj, + end_date=end_date_obj, + next_param_index=p, + ) + sql_conditions.append(search_condition.sql) + sql_params.extend(search_condition.params) + p += len(search_condition.params) # rebind-ok: advances the file's shared $N placeholder counter + # Equality filters - read effective values from where_conditions (post-authorization) for sql_col, wc_key in [ ("team_id", "team_id"), @@ -2662,7 +2710,13 @@ async def ui_view_spend_logs( sql_params.append(f"%{error_message}%") p += 1 - if group_by_session is True and not is_v2 and not is_request_id_lookup and sort_by == "startTime": + if ( + group_by_session is True + and not is_v2 + and not is_request_id_lookup + and not is_search_lookup + and sort_by == "startTime" + ): return await _ui_session_grouped_spend_logs( prisma_client=prisma_client, sql_conditions=sql_conditions, @@ -2696,7 +2750,7 @@ async def ui_view_spend_logs( _order_expr = order_column joined_conditions: Final = " AND ".join(sql_conditions) - session_grouping: Final = group_by_session is True + session_grouping: Final = group_by_session is True and not is_search_lookup count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else "" count_query: Final = f""" SELECT COUNT(*) AS total_count diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 0f17f2f23ab..5d410d7b55b 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -2,6 +2,23 @@ from datetime import datetime from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, model_validator +from typing_extensions import ReadOnly, TypedDict + +from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains + + +class KeyTokenWhere(TypedDict): + token: ReadOnly[str] + + +class KeyAliasContainsWhere(TypedDict): + key_alias: ReadOnly[InsensitiveContains] + + +class KeySearchWhere(TypedDict): + """Prisma filter behind `/key/list?search=`: exact token (sk- keys hashed) or alias substring, case-insensitive.""" + + OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]] class BulkUpdateKeyRequestItem(BaseModel): diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py index a0a26c089eb..cd2c8b0b904 100644 --- a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py +++ b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py @@ -1,5 +1,7 @@ +import hashlib from datetime import datetime, timedelta -from unittest.mock import AsyncMock, patch +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import FastAPI @@ -8,10 +10,12 @@ from litellm_enterprise.proxy.audit_logging_endpoints import router as audit_rou from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth # Create an app with just the audit router for testing app = FastAPI() app.include_router(audit_router) +app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin") client = TestClient(app) # Mock data for testing @@ -130,3 +134,63 @@ async def test_get_audit_log_by_id_not_found(mock_prisma_client): data = response.json() assert "message" in data["detail"] assert "not found" in data["detail"]["message"].lower() + + +def _list_audit_logs_where(mock_prisma_client: MagicMock, query: str) -> dict[str, object]: + mock_prisma_client.db.litellm_auditlog.find_many.return_value = [] + mock_prisma_client.db.litellm_auditlog.count.return_value = 0 + + response: Final = client.get(f"/audit?{query}") + + assert response.status_code == 200, response.text + find_many_where: Final = mock_prisma_client.db.litellm_auditlog.find_many.call_args.kwargs["where"] + assert mock_prisma_client.db.litellm_auditlog.count.call_args.kwargs["where"] == find_many_where + return find_many_where + + +def test_search_matches_any_id_column_alongside_the_other_filters(mock_prisma_client): + where: Final = _list_audit_logs_where(mock_prisma_client, "search=abc-123&action=create&object_team_id=team-1") + + assert where == { + "action": "create", + "AND": ( + { + "OR": [ + {"before_value": {"path": ["team_id"], "string_contains": "team-1"}}, + {"updated_values": {"path": ["team_id"], "string_contains": "team-1"}}, + ] + }, + { + "OR": ( + {"id": "abc-123"}, + {"changed_by": "abc-123"}, + {"object_id": "abc-123"}, + {"changed_by_api_key": "abc-123"}, + ) + }, + ), + } + + +def test_search_hashes_a_raw_virtual_key_for_the_hashed_columns(mock_prisma_client): + where: Final = _list_audit_logs_where(mock_prisma_client, "search=sk-raw") + + hashed: Final = hashlib.sha256(b"sk-raw").hexdigest() + assert where == { + "AND": ( + { + "OR": ( + {"id": "sk-raw"}, + {"changed_by": "sk-raw"}, + {"object_id": hashed}, + {"changed_by_api_key": hashed}, + ) + }, + ) + } + + +def test_an_empty_search_leaves_the_where_clause_unchanged(mock_prisma_client): + where: Final = _list_audit_logs_where(mock_prisma_client, "action=create&search=") + + assert where == {"action": "create"} diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 7954a4693cc..1e399e4fb58 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6347,6 +6347,114 @@ def test_build_key_filter_conditions_key_hash_narrows_team_admin_visibility(): assert {"token": "hashed-token-123"} in where["AND"], f"key_hash not ANDed: {where}" +def _search_clause(search: str, token: str) -> dict: + return {"OR": [{"token": token}, {"key_alias": {"contains": search, "mode": "insensitive"}}]} + + +def test_build_key_filter_conditions_search_hashes_raw_key_and_ors_alias_contains(): + """ + LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by + its ID. A pasted raw sk- key is hashed to its token first; an already-hashed + value is used verbatim. + """ + from litellm.proxy._types import hash_token + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + raw_where = json.loads( + json.dumps( + _build_key_filter_conditions( + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + search="sk-raw", + ) + ) + ) + assert _search_clause("sk-raw", hash_token("sk-raw")) in raw_where["AND"], f"raw search not ANDed: {raw_where}" + + hashed_where = json.loads( + json.dumps( + _build_key_filter_conditions( + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + search="already-hashed-token", + ) + ) + ) + assert _search_clause("already-hashed-token", "already-hashed-token") in hashed_where["AND"], ( + f"hashed search not used verbatim: {hashed_where}" + ) + + +def test_build_key_filter_conditions_search_narrows_team_admin_visibility(): + """ + LIT-4741, same class as LIT-3243: `search` must be a top-level AND so it + narrows a team admin's admin-team branch instead of being bypassed by it. + """ + from litellm.proxy._types import hash_token + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + where = json.loads( + json.dumps( + _build_key_filter_conditions( + user_id="team-admin-user", + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=["team-a"], + member_team_ids=["team-a"], + include_created_by_keys=False, + search="sk-member", + ) + ) + ) + + assert where.get("AND"), f"expected top-level AND, got: {where}" + assert _search_clause("sk-member", hash_token("sk-member")) in where["AND"], f"search not ANDed: {where}" + assert json.dumps({"team_id": {"in": ["team-a"]}}) in json.dumps(where) + + +@pytest.mark.asyncio +async def test_list_key_helper_applies_search_to_prisma_where(): + """LIT-4741: `search` given to _list_key_helper must reach the Prisma where clause.""" + from litellm.proxy._types import hash_token + + mock_prisma_client = AsyncMock() + mock_find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + + await _list_key_helper( + prisma_client=mock_prisma_client, + page=1, + size=50, + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + search="sk-raw", + ) + + where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"])) + assert _search_clause("sk-raw", hash_token("sk-raw")) in where["AND"], f"search not in Prisma where: {where}" + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ @@ -14870,6 +14978,50 @@ async def test_list_keys_non_admin_cannot_opt_into_substring(): assert kwargs["user_id"] == "alice" +@pytest.mark.asyncio +async def test_list_keys_hashes_raw_key_hash_before_validation(): + """LIT-4741: a raw sk- key pasted as key_hash is hashed before the ownership + check and the query, so a non-admin filtering by their own raw key gets the + row instead of the 'Key Hash not found.' 403.""" + from litellm.proxy._types import hash_token + + user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice") + validate = AsyncMock( + return_value=LiteLLM_UserTable( + user_id="alice", user_email="alice@example.com", teams=[], organization_memberships=[] + ) + ) + helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0}) + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check", + validate, + ), + patch("litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper", helper), + ): + await list_keys( + request=MagicMock(), + user_api_key_dict=user, + status=None, + user_id=None, + key_hash="sk-raw", + ) + + assert validate.call_args.kwargs["key_hash"] == hash_token("sk-raw") + assert helper.call_args.kwargs["key_hash"] == hash_token("sk-raw") + + +@pytest.mark.asyncio +async def test_list_keys_search_is_honored_for_non_admin(): + """LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's + search reaches the helper while their own-user scoping stays in place.""" + user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice") + kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="sk-raw") + assert kwargs["search"] == "sk-raw" + assert kwargs["user_id"] == "alice" + + @pytest.mark.asyncio async def test_cli_session_token_delegation_ceiling_blocked_by_team_budget(): team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0) diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py index be75d980d9d..dff0e80fa77 100644 --- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py +++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py @@ -615,6 +615,96 @@ class TestMemoryEndpoints: assert keys == {"user:profile"} assert body["total"] == 1 + def test_list_memory_search_matches_key_prefix_or_memory_id_within_scope(self): + """ + `search` matches a key prefix OR an exact memory_id, and stays ANDed + with the visibility filter so a pasted foreign id cannot leak a row. + """ + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None), + _make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None), + _make_row(memory_id="mem-foreign", key="user:secret", user_id="user-b", team_id=None), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + by_id = client.get("/v1/memory?search=mem-target") + by_prefix = client.get("/v1/memory?search=user:") + foreign_id = client.get("/v1/memory?search=mem-foreign") + + assert by_id.status_code == 200, by_id.text + assert [m["memory_id"] for m in by_id.json()["memories"]] == ["mem-target"] + assert by_id.json()["total"] == 1 + + assert by_prefix.status_code == 200, by_prefix.text + assert {m["key"] for m in by_prefix.json()["memories"]} == {"user:profile"} + assert by_prefix.json()["total"] == 1 + + assert foreign_id.status_code == 200, foreign_id.text + assert foreign_id.json()["memories"] == [] + assert foreign_id.json()["total"] == 0 + + def test_list_memory_search_by_memory_id_for_admin_sees_any_scope(self): + """Admins have no visibility filter, so an id search returns the row whoever owns it.""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="mem-a", key="a", user_id="user-a", team_id=None), + _make_row(memory_id="mem-b", key="b", user_id="user-b", team_id=None), + ] + ) + client = _make_client(_admin_auth()) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?search=mem-b") + assert resp.status_code == 200, resp.text + assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-b"] + assert resp.json()["total"] == 1 + + def test_list_memory_search_wins_over_key_prefix(self): + """When both are sent, `search` decides the match and `key_prefix` is ignored.""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None), + _make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?search=mem-target&key_prefix=user:") + assert resp.status_code == 200, resp.text + assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-target"] + assert resp.json()["total"] == 1 + + def test_list_memory_key_prefix_never_matches_memory_id(self): + """`key_prefix` stays a pure key-prefix match; only `search` consults memory_id.""" + table = self.prisma.db.litellm_memorytable + table.rows.append(_make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None)) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?key_prefix=mem-target") + assert resp.status_code == 200, resp.text + assert resp.json()["memories"] == [] + assert resp.json()["total"] == 0 + + def test_list_memory_key_exact_filter(self): + """`key` is an exact match, never a prefix.""" + table = self.prisma.db.litellm_memorytable + table.rows.extend( + [ + _make_row(memory_id="m1", key="user:profile", user_id="user-a", team_id=None), + _make_row(memory_id="m2", key="user:profile:archived", user_id="user-a", team_id=None), + ] + ) + client = _make_client(_user_auth("user-a", "team-a")) + with _patch_prisma(self.prisma): + resp = client.get("/v1/memory?key=user:profile") + assert resp.status_code == 200, resp.text + assert [m["memory_id"] for m in resp.json()["memories"]] == ["m1"] + assert resp.json()["total"] == 1 + def test_list_memory_admin_sees_all(self): table = self.prisma.db.litellm_memorytable table.rows.extend( diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index f4cd8814bc1..f869f3ffba2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -58,6 +58,25 @@ def _filter_logs_by_date_range(logs, where): return filtered +_SEARCH_CLAUSE_RE = re.compile( + r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' + r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' + r'AND \(api_key = \$(\d+) OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 ' + r"OR session_id = \$\1 OR model_id = \$\1\)\)\)" +) + + +def _matches_spend_log_search(log, search): + """Mirror the search clause: request_id across all time, the other id columns inside the window.""" + if log.get("request_id") == search["value"]: + return True + if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}): + return False + if log.get("api_key") == search["api_key"]: + return True + return any(log.get(col) == search["value"] for col in ("team_id", "user", "end_user", "session_id", "model_id")) + + def _reconstruct_ui_where_from_sql(sql_query, params): """ Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the @@ -77,6 +96,17 @@ def _reconstruct_ui_where_from_sql(sql_query, params): def _iso(value): return value.isoformat() if hasattr(value, "isoformat") else str(value) + search_clause = _SEARCH_CLAUSE_RE.search(clause.group(1)) + if search_clause: + raw_index, start_index, end_index, hashed_index = (int(g) for g in search_clause.groups()) + where["search"] = { + "value": params[raw_index - 1], + "api_key": params[hashed_index - 1], + "gte": _iso(params[start_index - 1]), + "lte": _iso(params[end_index - 1]), + } + remaining = clause.group(1) if search_clause is None else clause.group(1).replace(search_clause.group(0), "") + eq_cols = { "team_id": "team_id", '"user"': "user", @@ -89,7 +119,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params): } date_bounds: dict = {} metadata_conds: list = [] - for cond in (c.strip() for c in clause.group(1).split(" AND ")): + for cond in (c.strip() for c in remaining.split(" AND ")): gte = re.search(r'"startTime" >= \(\$(\d+)', cond) lte = re.search(r'"startTime" <= \(\$(\d+)', cond) alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond) @@ -2352,6 +2382,219 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( app.dependency_overrides.pop(ps.user_api_key_auth, None) +def test_build_spend_log_search_condition_windows_every_branch_except_request_id(): + """LIT-4741: request_id matches across all time; the six other id columns only inside the window, + and a raw sk- key is hashed for the api_key branch alone.""" + start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc) + end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc) + + condition = spend_management_endpoints._build_spend_log_search_condition( + search="sk-raw-key", start_date=start, end_date=end, next_param_index=3 + ) + + assert condition.sql == ( + "(request_id = $3 OR (\"startTime\" >= ($5::timestamptz AT TIME ZONE 'UTC') " + "AND \"startTime\" <= ($6::timestamptz AT TIME ZONE 'UTC') " + 'AND (api_key = $4 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))' + ) + assert condition.params == ("sk-raw-key", hashlib.sha256(b"sk-raw-key").hexdigest(), start, end) + + +def test_build_spend_log_search_condition_leaves_non_key_values_unhashed(): + start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc) + end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc) + + condition = spend_management_endpoints._build_spend_log_search_condition( + search="sess-42", start_date=start, end_date=end, next_param_index=1 + ) + + assert condition.params == ("sess-42", "sess-42", start, end) + + +def _search_fixture_logs(today): + recent = (today - datetime.timedelta(days=1)).isoformat() + old = (today - datetime.timedelta(days=90)).isoformat() + base = { + "api_key": "hashed-other", + "user": "user-x", + "team_id": "team-x", + "end_user": "cust-x", + "session_id": "sess-x", + "model_id": "mdl-x", + "spend": 0.01, + "model": "gpt-4", + } + return [ + {**base, "request_id": "req-session", "session_id": "sess-42", "startTime": recent}, + {**base, "request_id": "req-session-old", "session_id": "sess-42", "startTime": old}, + {**base, "request_id": "req-key", "api_key": hashlib.sha256(b"sk-raw-key").hexdigest(), "startTime": recent}, + {**base, "request_id": "req-team", "team_id": "team-7", "startTime": recent}, + {**base, "request_id": "req-user", "user": "user-7", "startTime": recent}, + {**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent}, + {**base, "request_id": "req-model", "model_id": "mdl-7", "startTime": recent}, + ] + + +def _search_filter_fn(logs, captured): + def filter_fn(where): + captured["where"] = where + rows = _filter_logs_by_date_range(logs, where) + if "user" in where: + rows = [row for row in rows if row["user"] == where["user"]] + if "search" in where: + rows = [row for row in rows if _matches_spend_log_search(row, where["search"])] + return rows + + return filter_fn + + +def _five_day_window(today): + return { + "start_date": (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": today.strftime("%Y-%m-%d %H:%M:%S"), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "search,expected_request_ids", + [ + ("req-session-old", {"req-session-old"}), + ("sess-42", {"req-session"}), + ("sk-raw-key", {"req-key"}), + ("team-7", {"req-team"}), + ("user-7", {"req-user"}), + ("cust-7", {"req-end-user"}), + ("mdl-7", {"req-model"}), + ("no-such-id", set()), + ], +) +async def test_ui_view_spend_logs_search_matches_any_id(client, monkeypatch, search, expected_request_ids): + """LIT-4741: one box matches any id column. A request_id is found across all time (the 5-day + window excludes the 90-day-old row), every other column only inside the window, and a raw + sk- key is hashed before it is compared with api_key. The window is not applied globally.""" + today = datetime.datetime.now(timezone.utc) + logs = _search_fixture_logs(today) + captured = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + try: + response = client.get( + "/spend/logs/ui", + params={"search": search, **_five_day_window(today)}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert {row["request_id"] for row in data["data"]} == expected_request_ids + assert data["total"] == len(expected_request_ids) + assert "startTime" not in captured["where"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_spend_logs_v2_search_keeps_global_window(client, monkeypatch): + """The public route keeps the caller's window on the whole query, so a search only finds rows + inside it even by request_id; the windowless request_id branch is a dashboard-only relaxation.""" + today = datetime.datetime.now(timezone.utc) + logs = _search_fixture_logs(today) + captured = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + try: + response = client.get( + "/spend/logs/v2", + params={"search": "req-session-old", **_five_day_window(today)}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["data"] == [] + assert data["total"] == 0 + assert "startTime" in captured["where"] + assert captured["where"]["search"]["value"] == "req-session-old" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "params", + [ + {"search": "req-old"}, + {"search": "req-old", "request_id": "req-old"}, + ], +) +async def test_ui_view_spend_logs_search_requires_dates(client, monkeypatch, params): + """A search needs the window for its non-request_id branches, so it stays required even + alongside a request_id, which on its own may drop the window.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma([], lambda where: []), + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + try: + response = client.get("/spend/logs/ui", params=params, headers={"Authorization": "Bearer sk-test"}) + assert response.status_code == 400 + assert "date" in response.text.lower() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "search,expected_request_ids", + [("sess-9", {"req-own"}), ("req-foreign", set())], +) +async def test_ui_view_spend_logs_search_keeps_non_admin_scope(client, monkeypatch, search, expected_request_ids): + """A search is scoped like any other listing: an internal user only sees their own rows even + when the id is on someone else's row, and the request_id ownership shortcut is not used.""" + yesterday = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=1)).isoformat() + base = {"api_key": "hashed-key", "team_id": None, "spend": 0.01, "startTime": yesterday, "model": "gpt-4"} + logs = [ + {**base, "request_id": "req-own", "user": "internal_user_1", "session_id": "sess-9"}, + {**base, "request_id": "req-own-other", "user": "internal_user_1", "session_id": "sess-other"}, + {**base, "request_id": "req-foreign", "user": "internal_user_2", "session_id": "sess-9"}, + ] + captured = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)), + ) + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs", + AsyncMock(return_value=[]), + ) + ownership_check = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.spend_tracking.spend_management_endpoints._assert_user_can_view_request_id", + ownership_check, + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={"search": search, "start_date": start_date, "end_date": end_date}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + assert {row["request_id"] for row in response.json()["data"]} == expected_request_ids + assert captured["where"]["user"] == "internal_user_1" + ownership_check.assert_not_awaited() + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + @pytest.mark.asyncio async def test_ui_view_spend_logs_unauthorized(client): # Test without authorization header @@ -6351,3 +6594,46 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort assert "OFFSET" in emitted_sql[1] finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_session(client, monkeypatch): + """The dashboard lists sessions by default; a search for an id lists every matching row instead, + so both calls of a session show up rather than one representative, and no session cursor is returned.""" + rows = [_session_representative_row("req-1", "sess-1"), _session_representative_row("req-2", "sess-1")] + + async def mock_query_raw(sql_query, *params): + if "mcp_tool_call_count" in sql_query: + return [] + grouped = "DISTINCT ON" in sql_query or "GROUP BY" in sql_query + visible = rows[:1] if grouped else rows + if "COUNT(*)" in sql_query: + return [{"total_count": len(visible)}] + return visible + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + try: + start_date, end_date = _default_date_range() + response = client.get( + "/spend/logs/ui", + params={ + "search": "sess-1", + "group_by_session": "true", + "start_date": start_date, + "end_date": end_date, + }, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert [row["request_id"] for row in data["data"]] == ["req-1", "req-2"] + assert data["total"] == 2 + assert "next_session_cursor" not in data + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) From ea12510f1a5799e7d9fd2cf87bd5c94263440f46 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:20:42 -0700 Subject: [PATCH 157/167] test(router): drop duplicate get_configured_mode test failing ruff F811 PRs #39630 and #39634 both added test_get_configured_mode_reads_deployment_model_info to tests/test_litellm/test_router.py, so the staging tip defines it twice and the required lint check fails with F811 on every PR synced past 321636ef5d. Keep the four tests from #39634 (mode read, None for unset or unknown, no wildcard pattern routing, malformed values treated as absent), which subsume the #39630 pair, and delete that pair. Five hand-applied mutations of Router.get_configured_mode are all still killed by the surviving tests. --- tests/test_litellm/test_router.py | 30 ------------------------------ 1 file changed, 30 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index e4236afc586..417c58b95f6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12701,33 +12701,3 @@ async def test_prompt_management_factory_marks_injection_for_every_deployment(mo bucket = captured.get("litellm_metadata") or captured["metadata"] assert captured["model_info"]["id"] == "provisional-dep" assert bucket["litellm_gateway_injected_cache"] == "" - - -def test_get_configured_mode_reads_deployment_model_info(): - router = Router( - model_list=[ - { - "model_name": "my-tts", - "litellm_params": {"model": "openai/some-unmapped-mode-model"}, - "model_info": {"mode": "audio_speech"}, - } - ] - ) - - assert router.get_configured_mode("my-tts") == "audio_speech" - - -@pytest.mark.parametrize("model_info", [{}, {"mode": ""}, {"mode": " "}, {"mode": 123}]) -def test_get_configured_mode_returns_none_for_unset_blank_or_unknown(model_info): - router = Router( - model_list=[ - { - "model_name": "plain-model", - "litellm_params": {"model": "openai/some-unmapped-mode-model"}, - "model_info": model_info, - } - ] - ) - - assert router.get_configured_mode("plain-model") is None - assert router.get_configured_mode("unknown-model") is None From a6b7ef6abed774bb6461f9b09b3621b0533cd8c6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 15:20:46 -0700 Subject: [PATCH 158/167] feat(ui): find a row by its pasted ID on every list page Virtual Keys and Team Virtual Keys send the search box to the new key/list search param so a key hash matches. Agents matches agent_id client-side. Memory sends the box as search so a memory_id matches. Audit Logs gains a search box. Request Logs sends the box as search so a session, team, user, key hash, or model id matches without opening the filter drawer. Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf --- .../agents/_components/AgentsTable.test.tsx | 29 +++- .../agents/_components/AgentsTable.tsx | 9 +- .../(dashboard)/hooks/keys/useKeys.test.ts | 18 +++ .../src/app/(dashboard)/hooks/keys/useKeys.ts | 2 + .../memory/_components/MemoryTable.test.tsx | 2 + .../memory/_components/MemoryTable.tsx | 4 +- .../memory/_components/MemoryView.test.tsx | 39 ++++- .../memory/_components/MemoryView.tsx | 4 +- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 17 +- .../VirtualKeysPage/VirtualKeysTable.tsx | 4 +- .../src/components/networking.test.ts | 45 ++++++ .../src/components/networking.tsx | 9 +- .../team/TeamVirtualKeysTable.test.tsx | 39 ++++- .../components/team/TeamVirtualKeysTable.tsx | 41 +++-- .../view_logs/AuditLogsPanel.test.tsx | 145 ++++++++++++++++++ .../components/view_logs/AuditLogsPanel.tsx | 16 +- .../view_logs/AuditLogsTable.test.tsx | 18 +++ .../components/view_logs/AuditLogsTable.tsx | 10 +- .../view_logs/RequestLogsPanel.test.tsx | 61 +++++++- .../components/view_logs/RequestLogsPanel.tsx | 23 ++- .../components/view_logs/RequestLogsTable.tsx | 2 +- .../view_logs/log_filter_logic.test.tsx | 1 + .../components/view_logs/log_filter_logic.tsx | 3 + 23 files changed, 489 insertions(+), 52 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx index 4d18ec2ef5f..bef938cd31c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx @@ -90,7 +90,7 @@ describe("AgentsTable", () => { />, ); - const search = screen.getByPlaceholderText("Search agent names or descriptions..."); + const search = screen.getByPlaceholderText("Search agents by name, ID, or description..."); await user.type(search, "billing"); expect(screen.getByText("Billing Router")).toBeInTheDocument(); expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); @@ -101,11 +101,36 @@ describe("AgentsTable", () => { expect(screen.queryByText("Billing Router")).not.toBeInTheDocument(); }); + it("filters agents by a pasted agent_id so only that agent's row survives", async () => { + const user = userEvent.setup(); + render( + , + ); + + const search = screen.getByPlaceholderText("Search agents by name, ID, or description..."); + await user.click(search); + await user.paste("5f3c2a1b-9d8e-4f7a-b6c5-d4e3f2a1b0c9"); + expect(screen.getByText("Billing Router")).toBeInTheDocument(); + expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); + + await user.clear(search); + await user.paste("ffffffff-0000-4000-8000-000000000000"); + expect(screen.queryByText("Billing Router")).not.toBeInTheDocument(); + expect(screen.queryByText("Second Agent")).not.toBeInTheDocument(); + expect(screen.getByText("No matching agents")).toBeInTheDocument(); + }); + it("shows the no-match empty state when the search matches nothing", async () => { const user = userEvent.setup(); render(); - await user.type(screen.getByPlaceholderText("Search agent names or descriptions..."), "zzzz"); + await user.type(screen.getByPlaceholderText("Search agents by name, ID, or description..."), "zzzz"); expect(screen.queryByText("Test Agent")).not.toBeInTheDocument(); expect(screen.getByText("No matching agents")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx index 35ed6b66425..aceb07e2e9a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx @@ -55,7 +55,12 @@ const AgentsTable: React.FC = ({ const [sorting, setSorting] = useState(DEFAULT_SORTING); const [searchTerm, setSearchTerm] = useState(""); const filteredAgents = useMemo( - () => filterBySearchTerm(agents, searchTerm, (agent) => [agent.agent_name, agent.agent_card_params?.description]), + () => + filterBySearchTerm(agents, searchTerm, (agent) => [ + agent.agent_name, + agent.agent_id, + agent.agent_card_params?.description, + ]), [agents, searchTerm], ); @@ -83,7 +88,7 @@ const AgentsTable: React.FC = ({ setSearchTerm(e.target.value)} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts index 8c9b33f2c3e..f23fcf811f2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -518,6 +518,24 @@ describe("useKeys", () => { const callUrl = mockFetch.mock.calls[0][0]; expect(callUrl).not.toContain("agent_id"); }); + + it("sends the combined alias-or-ID search as the search param, separate from key_alias and key_hash", async () => { + mockFetch.mockResolvedValueOnce({ + ok: true, + json: async () => mockKeysResponse, + }); + + const { result } = renderHook(() => useKeys(1, 10, { search: "sk-pasted-key" }), { wrapper }); + + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + const callUrl = new URL(mockFetch.mock.calls[0][0], "http://localhost"); + expect(callUrl.searchParams.get("search")).toBe("sk-pasted-key"); + expect(callUrl.searchParams.has("key_alias")).toBe(false); + expect(callUrl.searchParams.has("key_hash")).toBe(false); + }); }); describe("useDeletedKeys", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts index 94ded01679d..7e7089e685f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -40,6 +40,7 @@ export interface KeyListCallOptions { selectedKeyAlias?: string | null; userID?: string | null; keyHash?: string | null; + search?: string | null; sortBy?: string | null; sortOrder?: string | null; expand?: string | null; @@ -61,6 +62,7 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number, organization_id: options.organizationID, key_alias: options.selectedKeyAlias, key_hash: options.keyHash, + search: options.search, user_id: options.userID, page, size: pageSize, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx index 5100b998b80..984b8135466 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx @@ -95,6 +95,7 @@ describe("MemoryTable", () => { it("shows the filtered-empty copy when a search is active", () => { render(); expect(screen.getByText("No matching memories")).toBeInTheDocument(); + expect(screen.getByText("No memories match your search.")).toBeInTheDocument(); expect(screen.queryByText("No memories stored yet")).not.toBeInTheDocument(); }); @@ -128,6 +129,7 @@ describe("MemoryTable", () => { const onRefresh = vi.fn(); render(); + expect(screen.getByPlaceholderText("Search by key prefix or memory ID…")).toBeInTheDocument(); fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "u" } }); expect(onSearchChange).toHaveBeenCalledWith("u"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx index 50dd04ee14c..3e37faafe15 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx @@ -36,7 +36,7 @@ function MemoryEmptyState({ hasActiveSearch }: { hasActiveSearch: boolean }) {
{hasActiveSearch - ? "No memories have keys starting with your search." + ? "No memories match your search." : "Memories your agents store under /v1/memory will appear here."}
@@ -81,7 +81,7 @@ export function MemoryTable({ table={table} searchValue={searchValue} onSearchChange={onSearchChange} - searchPlaceholder='Filter by key prefix, e.g. "user:"' + searchPlaceholder="Search by key prefix or memory ID…" onRefresh={onRefresh} isRefreshing={isRefreshing} showViewOptions={false} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx index 9ccef5357b9..b703df652c2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx @@ -1,8 +1,9 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, render, screen } from "@testing-library/react"; +import type { PaginationState } from "@tanstack/react-table"; +import { act, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { MemoryRow } from "@/components/networking"; @@ -13,10 +14,13 @@ interface CapturedTableProps { rowCount: number; data: MemoryRow[]; hasActiveSearch: boolean; + onSearchChange: (value: string) => void; + onPaginationChange: (state: PaginationState) => void; onViewClick: (row: MemoryRow) => void; } const captured = vi.hoisted(() => ({ current: null as CapturedTableProps | null })); +const fetchMemoryListMock = vi.hoisted(() => vi.fn()); vi.mock("./MemoryTable", () => ({ MemoryTable: function MemoryTableMock(props: CapturedTableProps) { @@ -25,6 +29,15 @@ vi.mock("./MemoryTable", () => ({ }, })); +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + fetchMemoryList: fetchMemoryListMock, +})); + +vi.mock("@tanstack/react-pacer/debouncer", () => ({ + useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }], +})); + const renderView = (accessToken: string | null) => { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); return render( @@ -35,6 +48,28 @@ const renderView = (accessToken: string | null) => { }; describe("MemoryView", () => { + beforeEach(() => { + fetchMemoryListMock.mockReset(); + fetchMemoryListMock.mockResolvedValue({ memories: [], total: 0 }); + }); + + it("queries the server with the search box value as `search` and resets to page 1", async () => { + renderView("token"); + await waitFor(() => expect(fetchMemoryListMock).toHaveBeenCalled()); + + act(() => captured.current?.onPaginationChange({ pageIndex: 2, pageSize: 50 })); + await waitFor(() => + expect(fetchMemoryListMock).toHaveBeenLastCalledWith("token", expect.objectContaining({ page: 3 })), + ); + + act(() => captured.current?.onSearchChange("mem-abc123")); + + await waitFor(() => + expect(fetchMemoryListMock).toHaveBeenLastCalledWith("token", { search: "mem-abc123", page: 1, pageSize: 50 }), + ); + expect(captured.current?.hasActiveSearch).toBe(true); + }); + it("keeps the table out of the skeleton state when the token is null (disabled query)", () => { renderView(null); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx index 1d2e5150a62..58d4e42aa96 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx @@ -43,10 +43,8 @@ export const MemoryView: React.FC = ({ accessToken }) => { queryKey: [MEMORY_LIST_KEY, debouncedSearch, pagination.pageIndex, pagination.pageSize], queryFn: () => { if (!accessToken) throw new Error("Access token required"); - // Prefix search matches the Redis-style mental model (namespace scan): - // typing "user:" finds "user:profile", "user:prefs", etc. return fetchMemoryList(accessToken, { - keyPrefix: debouncedSearch || undefined, + search: debouncedSearch || undefined, page: pagination.pageIndex + 1, pageSize: pagination.pageSize, }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index c69662b69dc..881b0b93ff9 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -542,6 +542,19 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { expect((lastCall[2] ?? {}).userID).toBeUndefined(); }); }); + + it("sends the search box as the combined alias-or-ID search rather than the key-alias filter", async () => { + renderWithProviders(); + + fireEvent.change(screen.getByPlaceholderText(/Search by key alias or ID/), { target: { value: mockKey.token } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: mockKey.token })); + }); + const lastOptions = mockUseKeys.mock.calls.at(-1)?.[2]; + expect(lastOptions?.selectedKeyAlias).toBeUndefined(); + expect(lastOptions?.keyHash).toBeUndefined(); + }); }); describe("pagination display – total count comes from useKeys", () => { @@ -663,7 +676,7 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(mockUseKeys).toHaveBeenLastCalledWith( 3, 25, - expect.objectContaining({ selectedKeyAlias: "prod", sortBy: "spend", sortOrder: "asc" }), + expect.objectContaining({ search: "prod", sortBy: "spend", sortOrder: "asc" }), ); }); expect(screen.getByPlaceholderText(/Search by key alias/)).toHaveValue("prod"); @@ -736,7 +749,7 @@ describe("table state lives in the URL so it survives leaving and returning to t fireEvent.change(screen.getByPlaceholderText(/Search by key alias/), { target: { value: "prod" } }); await waitFor(() => { - expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "prod" })); + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: "prod" })); }); await waitFor(() => { expect(lastSearchParam(onUrlUpdate, "page")).toBeNull(); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index c424966a0a3..ebedf57af45 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -118,7 +118,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const keyListOptions = { teamID: appliedFilters.team_id || undefined, organizationID: appliedFilters.org_id || undefined, - selectedKeyAlias: searchQuery.trim() || undefined, + search: searchQuery.trim() || undefined, userID: appliedFilters.user_id || undefined, keyHash: appliedFilters.key_hash || undefined, sortBy, @@ -291,7 +291,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { table={table} searchValue={searchInput} onSearchChange={handleSearchChange} - searchPlaceholder="Search by key alias…" + searchPlaceholder="Search by key alias or ID…" onRefresh={() => refetch?.()} isRefreshing={isFetching} onOpenFilters={() => setFiltersOpen(true)} diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index 7a220dd4711..9df9e9a9209 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -854,3 +854,48 @@ describe("userListCall search serialization", () => { expect(lastParams(mockFetch).get("user_email")).toBe("ada@example.com"); }); }); + +describe("fetchMemoryList search serialization", () => { + const originalFetch = global.fetch; + + afterEach(() => { + global.fetch = originalFetch; + }); + + const mockOkFetch = () => { + const emptyPage = { memories: [], total: 0 }; + const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue(emptyPage) } as any); + global.fetch = mockFetch as any; + return mockFetch; + }; + + const lastParams = (mockFetch: ReturnType) => { + const [url] = mockFetch.mock.calls.at(-1) ?? []; + return new URL(url as string, "http://example.com").searchParams; + }; + + it("sends the search box value as search and omits key_prefix and key", async () => { + const mockFetch = mockOkFetch(); + + await Networking.fetchMemoryList("token", { search: "mem-abc123", page: 1, pageSize: 50 }); + + const params = lastParams(mockFetch); + expect(params.get("search")).toBe("mem-abc123"); + expect(params.has("key_prefix")).toBe(false); + expect(params.has("key")).toBe(false); + expect(params.get("page")).toBe("1"); + expect(params.get("page_size")).toBe("50"); + }); + + it("keeps key_prefix and key working when no search is given", async () => { + const mockFetch = mockOkFetch(); + + await Networking.fetchMemoryList("token", { keyPrefix: "user:" }); + expect(lastParams(mockFetch).get("key_prefix")).toBe("user:"); + expect(lastParams(mockFetch).has("search")).toBe(false); + + await Networking.fetchMemoryList("token", { key: "user:profile" }); + expect(lastParams(mockFetch).get("key")).toBe("user:profile"); + expect(lastParams(mockFetch).has("search")).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index d8762565e08..1384679a88a 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2056,6 +2056,7 @@ interface UiSpendLogsParams { exclude_internal_health_checks?: boolean; group_by_session?: boolean; session_cursor?: string; + search?: string; } interface UiSpendLogsCallOptions { @@ -6563,6 +6564,7 @@ interface UiAuditLogsParams { changed_by_api_key?: string; object_team_id?: string; object_key_hash?: string; + search?: string | null; sort_by?: string; sort_order?: "asc" | "desc"; } @@ -8061,15 +8063,18 @@ export const fetchMemoryList = async ( options: { key?: string; keyPrefix?: string; + search?: string; page?: number; pageSize?: number; } = {}, ): Promise => { const base = proxyBaseUrl ? `${proxyBaseUrl}/v1/memory` : `/v1/memory`; const params = new URLSearchParams(); - // keyPrefix takes precedence — backend also does, but we omit `key` + // Backend precedence is search > key_prefix > key; only the winner is sent // to keep the URL clean and intent obvious. - if (options.keyPrefix) { + if (options.search) { + params.append("search", options.search); + } else if (options.keyPrefix) { params.append("key_prefix", options.keyPrefix); } else if (options.key) { params.append("key", options.key); diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx index 8bf5d639d6c..0d9d0988aa1 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx @@ -30,6 +30,8 @@ vi.mock("@tanstack/react-pacer/debouncer", () => ({ const mockUseKeys = useKeys as MockedFunction; +const KEY_HASH = "88a145505dd6e87e2ea166fcef1e4b53948dbdb32af6431dfd05ec06b571ee52"; + const createMockKey = (overrides: Partial = {}): KeyResponse => ({ token: "sk-test123", @@ -277,7 +279,7 @@ describe("TeamVirtualKeysTable", () => { ); }); - it("maps the search box to a server-side key-alias query", async () => { + it("maps the Key ID drawer filter to a server-side useKeys query and clears it", async () => { const user = userEvent.setup(); mockUseKeys.mockReturnValue({ data: { keys: [createMockKey()], total_count: 1, current_page: 1, total_pages: 1 }, @@ -288,11 +290,42 @@ describe("TeamVirtualKeysTable", () => { renderWithProviders(); - fireEvent.change(await screen.findByTestId("datatable-search"), { target: { value: "check-002" } }); + await user.click(await screen.findByTestId("datatable-filters-trigger")); + const drawerBody = await screen.findByTestId("filter-drawer-body"); + fireEvent.change(within(drawerBody).getByPlaceholderText("Enter Key ID…"), { target: { value: KEY_HASH } }); + await user.click(screen.getByTestId("filter-drawer-apply")); await waitFor(() => - expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "check-002" })), + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ keyHash: KEY_HASH })), ); + expect(screen.getByTestId("filter-chip-key_hash")).toHaveTextContent("Key ID"); + + await user.click(screen.getByTestId("datatable-clear-filters")); + await waitFor(() => + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ keyHash: undefined })), + ); + }); + + it("maps the search box to the combined alias-or-ID search rather than the key-alias filter", async () => { + mockUseKeys.mockReturnValue({ + data: { keys: [createMockKey()], total_count: 1, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as unknown as ReturnType); + + renderWithProviders(); + + const searchBox = await screen.findByTestId("datatable-search"); + expect(searchBox).toHaveAttribute("placeholder", "Search by key alias or ID…"); + fireEvent.change(searchBox, { target: { value: KEY_HASH } }); + + await waitFor(() => + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: KEY_HASH })), + ); + const lastOptions = mockUseKeys.mock.calls.at(-1)?.[2]; + expect(lastOptions?.selectedKeyAlias).toBeUndefined(); + expect(lastOptions?.keyHash).toBeUndefined(); }); it("should show Loading keys when isPending", async () => { diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx index bd7faa41ee9..aa9df4a0319 100644 --- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx @@ -68,19 +68,17 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi const pageIndex = tablePagination.pageIndex; const pageSize = tablePagination.pageSize; - const { - data: keys, - isPending: isLoading, - isFetching, - refetch, - } = useKeys(pageIndex + 1, pageSize, { + const keyListOptions = { teamID: teamId, - selectedKeyAlias: searchQuery.trim() || undefined, + search: searchQuery.trim() || undefined, userID: getFilterValue("user_id"), + keyHash: getFilterValue("key_hash"), sortBy: sortBy || undefined, sortOrder: sortOrder || undefined, expand: "user", - }); + }; + + const { data: keys, isPending: isLoading, isFetching, refetch } = useKeys(pageIndex + 1, pageSize, keyListOptions); const displayKeys = useMemo(() => { const kList = keys?.keys || []; @@ -481,11 +479,11 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi table={table} searchValue={searchInput} onSearchChange={handleSearchChange} - searchPlaceholder="Search by key alias…" + searchPlaceholder="Search by key alias or ID…" onRefresh={() => refetch?.()} isRefreshing={isFetching} onOpenFilters={() => setFiltersOpen(true)} - filterLabels={{ user_id: "User ID" }} + filterLabels={{ user_id: "User ID", key_hash: "Key ID" }} /> {({ get, set }) => ( - - set("user_id", event.target.value)} - placeholder="Filter by user ID…" - /> - + <> + + set("user_id", event.target.value)} + placeholder="Filter by user ID…" + /> + + + set("key_hash", event.target.value)} + placeholder="Enter Key ID…" + /> + + )} diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx new file mode 100644 index 00000000000..3b27f663b8e --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx @@ -0,0 +1,145 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { chooseSelectOption } from "../../../tests/test-utils"; +import AuditLogsPanel from "./AuditLogsPanel"; + +vi.mock("../networking", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, uiAuditLogsCall: vi.fn() }; +}); + +// Resolve the debounced search synchronously so typed input reaches the query within the test tick. +vi.mock("@tanstack/react-pacer/debouncer", () => ({ + useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }], +})); + +import { uiAuditLogsCall } from "../networking"; + +type AuditLogsParams = NonNullable[0]["params"]>; + +const PAGE_SIZE = 50; + +const ID_PARAM_KEYS = [ + "search", + "object_id", + "changed_by", + "object_team_id", + "object_key_hash", + "action", + "table_name", +] as const satisfies readonly (keyof AuditLogsParams)[]; + +const respondWith = (total: number) => { + const response = { audit_logs: [], total, page: 1, page_size: PAGE_SIZE, total_pages: Math.ceil(total / PAGE_SIZE) }; + return vi.mocked(uiAuditLogsCall).mockResolvedValue(response); +}; + +const lastCall = () => vi.mocked(uiAuditLogsCall).mock.calls.at(-1)?.[0]; +const sentIdParams = () => ID_PARAM_KEYS.filter((key) => lastCall()?.params?.[key] !== undefined); + +const defaultProps = { + accessToken: "sk-test", + token: "jwt-test", + userRole: "Admin", + userID: "user-1", + isActive: true, + premiumUser: true, +}; + +const renderPanel = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; + +const TEXT_FILTERS: { filterId: string; placeholder: string; paramKey: keyof AuditLogsParams }[] = [ + { filterId: "object_id", placeholder: "Enter object ID…", paramKey: "object_id" }, + { filterId: "changed_by", placeholder: "Enter user ID…", paramKey: "changed_by" }, + { filterId: "team_id", placeholder: "Enter team ID…", paramKey: "object_team_id" }, + { filterId: "key_hash", placeholder: "Enter key hash…", paramKey: "object_key_hash" }, +]; + +const SELECT_FILTERS: { + label: string; + comboboxIndex: number; + option: string; + paramKey: keyof AuditLogsParams; + value: string; +}[] = [ + { label: "Action", comboboxIndex: 0, option: "Created", paramKey: "action", value: "created" }, + { label: "Table", comboboxIndex: 1, option: "Teams", paramKey: "table_name", value: "LiteLLM_TeamTable" }, +]; + +describe("AuditLogsPanel", () => { + beforeEach(() => { + vi.clearAllMocks(); + respondWith(0); + }); + + it("sends the typed search as params.search and returns to the first page", async () => { + const user = userEvent.setup(); + respondWith(120); + renderPanel(); + await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled()); + expect(lastCall()?.params?.search).toBeUndefined(); + + await user.click(screen.getByTestId("pagination-next")); + await waitFor(() => expect(lastCall()?.page).toBe(2)); + + await user.type(screen.getByTestId("datatable-search"), "team-abc"); + + await waitFor(() => expect(lastCall()?.params?.search).toBe("team-abc")); + expect(lastCall()?.page).toBe(1); + expect(sentIdParams()).toEqual(["search"]); + }); + + it("trims the search and drops params.search once the box is cleared", async () => { + const user = userEvent.setup(); + renderPanel(); + const input = await screen.findByTestId("datatable-search"); + + await user.type(input, " abc"); + await waitFor(() => expect(lastCall()?.params?.search).toBe("abc")); + + await user.clear(input); + + await waitFor(() => expect(lastCall()?.params?.search).toBeUndefined()); + expect(sentIdParams()).toEqual([]); + }); + + it.each(TEXT_FILTERS)("maps the $filterId drawer filter to params.$paramKey", async ({ placeholder, paramKey }) => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled()); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + fireEvent.change(await screen.findByPlaceholderText(placeholder), { target: { value: "val-1" } }); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastCall()?.params?.[paramKey]).toBe("val-1")); + expect(sentIdParams()).toEqual([paramKey]); + }); + + it.each(SELECT_FILTERS)( + "maps the $label drawer select to params.$paramKey", + async ({ comboboxIndex, option, paramKey, value }) => { + const user = userEvent.setup(); + renderPanel(); + await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled()); + + await user.click(screen.getByTestId("datatable-filters-trigger")); + const triggers = await screen.findAllByRole("combobox"); + await chooseSelectOption(user, triggers[comboboxIndex], option); + await user.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => expect(lastCall()?.params?.[paramKey]).toBe(value)); + expect(sentIdParams()).toEqual([paramKey]); + }, + ); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx index 81bd4a19f76..5ce4ca6053f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx @@ -1,7 +1,9 @@ import { useCallback, useState } from "react"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { useQuery, keepPreviousData } from "@tanstack/react-query"; import { ColumnFiltersState, OnChangeFn, PaginationState } from "@tanstack/react-table"; import { resolveLogoSrc } from "@/lib/assetPaths"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { uiAuditLogsCall } from "../networking"; import { AuditLogEntry } from "./AuditLogsTableColumns"; import { AuditLogsTable } from "./AuditLogsTable"; @@ -39,9 +41,13 @@ export default function AuditLogsPanel({ }: AuditLogsProps) { const [pagination, setPagination] = useState({ pageIndex: 0, pageSize: PAGE_SIZE }); const [columnFilters, setColumnFilters] = useState([]); + const [searchInput, setSearchInput] = useState(""); + const [debouncedSearch] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS }); const [selectedLog, setSelectedLog] = useState(null); const [drawerOpen, setDrawerOpen] = useState(false); + const searchTerm = debouncedSearch.trim(); + const getFilterValue = (columnId: string): string | undefined => { const entry = columnFilters.find((filter) => filter.id === columnId); return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined; @@ -50,7 +56,7 @@ export default function AuditLogsPanel({ const canQueryAuditLogs = !!accessToken && !!token && !!userRole && !!userID && isActive && premiumUser; const query = useQuery({ - queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters], + queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters, searchTerm], queryFn: async () => { if (!accessToken) { return { audit_logs: [], total: 0, page: 1, page_size: pagination.pageSize, total_pages: 0 }; @@ -60,6 +66,7 @@ export default function AuditLogsPanel({ page: pagination.pageIndex + 1, page_size: pagination.pageSize, params: { + search: searchTerm || undefined, object_id: getFilterValue("object_id"), changed_by: getFilterValue("changed_by"), object_key_hash: getFilterValue("key_hash"), @@ -80,6 +87,11 @@ export default function AuditLogsPanel({ setPagination((prev) => ({ ...prev, pageIndex: 0 })); }, []); + const handleSearchChange = useCallback((value: string) => { + setSearchInput(value); + setPagination((prev) => ({ ...prev, pageIndex: 0 })); + }, []); + const handleViewLog = useCallback((log: AuditLogEntry) => { setSelectedLog(log); setDrawerOpen(true); @@ -128,6 +140,8 @@ export default function AuditLogsPanel({ onPaginationChange={setPagination} columnFilters={columnFilters} onColumnFiltersChange={handleColumnFiltersChange} + searchValue={searchInput} + onSearchChange={handleSearchChange} onRefresh={() => query.refetch()} onViewLog={handleViewLog} /> diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx index 7349c3019ae..d8e549715af 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx @@ -120,6 +120,24 @@ describe("AuditLogsTable", () => { expect(screen.getByText("No matching audit logs")).toBeInTheDocument(); }); + it("renders the toolbar search box from the search props and forwards typed input", () => { + const onSearchChange = vi.fn(); + renderTable({ searchValue: "team-", onSearchChange }); + + const input = screen.getByPlaceholderText("Search audit logs by ID…"); + expect(input).toHaveValue("team-"); + + fireEvent.change(input, { target: { value: "team-7" } }); + expect(onSearchChange).toHaveBeenCalledWith("team-7"); + }); + + it("treats an active search as a filter for the empty state", () => { + const emptySearchResult = { data: [], rowCount: 0, searchValue: "zzz", onSearchChange: vi.fn() }; + renderTable(emptySearchResult); + + expect(screen.getByText("No matching audit logs")).toBeInTheDocument(); + }); + it("renders active filter chips with human-readable labels", () => { const filters: ColumnFiltersState = [{ id: "action", value: "created" }]; renderTable({ columnFilters: filters }); diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx index 799505ed07c..cef828838e2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx @@ -24,6 +24,8 @@ interface AuditLogsTableProps { onPaginationChange: OnChangeFn; columnFilters: ColumnFiltersState; onColumnFiltersChange: OnChangeFn; + searchValue?: string; + onSearchChange?: (value: string) => void; onRefresh: () => void; onViewLog: (log: AuditLogEntry) => void; } @@ -102,11 +104,14 @@ export function AuditLogsTable({ onPaginationChange, columnFilters, onColumnFiltersChange, + searchValue, + onSearchChange, onRefresh, onViewLog, }: AuditLogsTableProps) { const [filtersOpen, setFiltersOpen] = useState(false); const columns = useMemo(() => getAuditLogsTableColumns({ onViewLog }), [onViewLog]); + const hasActiveSearch = Boolean(searchValue?.trim()); return ( 0} />} + noDataMessage={ 0 || hasActiveSearch} />} size="compact" toolbar={(table) => ( <> setFiltersOpen(true)} diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index 93066c106ee..be0d0049c13 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -54,6 +54,14 @@ vi.mock("./LogDetailsDrawer", () => ({ }, })); +const debounce = vi.hoisted(() => ({ settled: null as string | null })); + +vi.mock("@tanstack/react-pacer/debouncer", () => ({ + useDebouncedValue: vi.fn((value: unknown) => [debounce.settled ?? value, { cancel: vi.fn(), flush: vi.fn() }]), +})); + +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { uiSpendLogsCall } from "../networking"; const logEntry = (overrides: Partial): LogEntry => ({ @@ -136,6 +144,7 @@ describe("RequestLogsPanel", () => { sessionStorage.clear(); testQueryClient.clear(); respondWith([]); + debounce.settled = null; }); describe("server-grouped session pagination (#38060)", () => { @@ -322,9 +331,8 @@ describe("RequestLogsPanel", () => { }); }); - describe("search by request id (LIT-3981)", () => { - it("sends the typed request id to the server on the first page instead of filtering the loaded rows", async () => { - const user = userEvent.setup(); + describe("search by any id (LIT-3981, LIT-4741)", () => { + it("sends the typed id to the server as search on the first page instead of filtering the loaded rows", async () => { renderPanel(); await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); @@ -334,9 +342,54 @@ describe("RequestLogsPanel", () => { await waitFor(() => { const call = lastCall(); if (!call) throw new Error("uiSpendLogsCall was not called"); - expect(call.params?.request_id).toBe("req-on-another-page"); + expect(call.params?.search).toBe("req-on-another-page"); expect(call.page).toBe(1); }); + expect(lastCall()?.params?.request_id).toBeUndefined(); + expect(lastCall()?.params?.session_cursor).toBeUndefined(); + }); + + it("sends the debounced value to the server while the box shows what is being typed", async () => { + debounce.settled = "settled-id"; + renderPanel(); + + await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled()); + + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "still-typing" } }); + + expect(screen.getByTestId("datatable-search")).toHaveValue("still-typing"); + await waitFor(() => + expect(useDebouncedValue).toHaveBeenLastCalledWith("still-typing", { wait: DEBOUNCE_WAIT_MS }), + ); + await waitFor(() => expect(lastCall()?.params?.search).toBe("settled-id")); + const sentLiveValue = vi + .mocked(uiSpendLogsCall) + .mock.calls.some(([options]) => options.params?.search === "still-typing"); + expect(sentLiveValue).toBe(false); + }); + + it("shows a Search chip whose remove button clears the box and restores the unsearched listing", async () => { + const user = userEvent.setup(); + vi.mocked(uiSpendLogsCall).mockImplementation(async ({ params }) => { + const data = + params?.search === "sess-42" + ? [logEntry({ request_id: "req-sess", session_id: "sess-42" })] + : [logEntry({ request_id: "req-initial" })]; + return { data, total: data.length, page: 1, page_size: 50, total_pages: 1 }; + }); + renderPanel(); + + await waitFor(() => expect(row("req-initial")).not.toBeNull()); + fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "sess-42" } }); + await waitFor(() => expect(row("req-sess")).not.toBeNull()); + expect(row("req-initial")).toBeNull(); + expect(screen.getByTestId("filter-chip-search")).toHaveTextContent("Search:sess-42"); + + await user.click(screen.getByRole("button", { name: "Remove Search filter" })); + + expect(screen.getByTestId("datatable-search")).toHaveValue(""); + await waitFor(() => expect(row("req-initial")).not.toBeNull()); + expect(row("req-sess")).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 96cf2bd5dde..9b99c6af923 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -1,11 +1,13 @@ "use client"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { useQuery, type UseQueryOptions } from "@tanstack/react-query"; import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import moment from "moment"; import { useCallback, useEffect, useMemo, useState } from "react"; import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells"; +import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import type { KeyResponse } from "../key_team_helpers/key_list"; import { keyInfoV1Call, uiSpendLogsCall } from "../networking"; import KeyInfoView from "../templates/key_info_view"; @@ -75,12 +77,22 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, sessionStorage.setItem("excludeInternalHealthChecks", JSON.stringify(excludeInternalHealthChecks)); }, [excludeInternalHealthChecks]); + const searchTerm = useMemo(() => { + const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.SEARCH); + return typeof entry?.value === "string" ? entry.value : ""; + }, [columnFilters]); + const [debouncedSearch] = useDebouncedValue(searchTerm, { wait: DEBOUNCE_WAIT_MS }); + const queryColumnFilters = useMemo(() => { + const others = columnFilters.filter((filter) => filter.id !== LOG_FILTER_IDS.SEARCH); + return debouncedSearch === "" ? others : [...others, { id: LOG_FILTER_IDS.SEARCH, value: debouncedSearch }]; + }, [columnFilters, debouncedSearch]); + const { logsQuery, filteredLogs, allTeams, usesSessionCursor } = useLogFilterLogic({ accessToken, token, userRole, userID, - columnFilters, + columnFilters: queryColumnFilters, activeTab: isActive ? "request logs" : "inactive", isLiveTail, excludeInternalHealthChecks, @@ -155,15 +167,10 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const rows: LogEntry[] = filteredLogs.data; - const searchTerm = useMemo(() => { - const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.REQUEST_ID); - return typeof entry?.value === "string" ? entry.value : ""; - }, [columnFilters]); - const handleSearchChange = useCallback((value: string) => { setColumnFilters((previous) => { - const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.REQUEST_ID); - return value === "" ? others : [...others, { id: LOG_FILTER_IDS.REQUEST_ID, value }]; + const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.SEARCH); + return value === "" ? others : [...others, { id: LOG_FILTER_IDS.SEARCH, value }]; }); setSessionCursors({}); setPagination((previous) => ({ ...previous, pageIndex: 0 })); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx index 4159b3b699b..17caa4466fa 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx @@ -108,7 +108,7 @@ export function RequestLogsTable({ table={table} searchValue={searchValue} onSearchChange={onSearchChange} - searchPlaceholder="Search by Request ID" + searchPlaceholder="Search logs by ID…" onRefresh={onRefresh} isRefreshing={isRefreshing} onOpenFilters={() => setFiltersOpen(true)} diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx index 1b738db097d..6af791cc50e 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx @@ -91,6 +91,7 @@ describe("useLogFilterLogic", () => { { id: LOG_FILTER_IDS.ERROR_CODE, value: "429", param: "error_code" }, { id: LOG_FILTER_IDS.ERROR_MESSAGE, value: "rate limited", param: "error_message" }, { id: LOG_FILTER_IDS.USER_ID, value: "user-9", param: "user_id" }, + { id: LOG_FILTER_IDS.SEARCH, value: "any-id", param: "search" }, ]; it.each(cases)("sends $id as $param", async ({ id, value, param }) => { diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index 90f0f0a60f1..3d368527ad9 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -33,6 +33,7 @@ export const LOG_FILTER_IDS = { PUBLIC_MODEL_OR_SEARCH_TOOL: "model", REQUEST_ID: "request_id", USER_ID: "user_id", + SEARCH: "search", } as const; export const LOG_FILTER_LABELS: Record = { @@ -48,6 +49,7 @@ export const LOG_FILTER_LABELS: Record = { [LOG_FILTER_IDS.SESSION_ID]: "Session ID", [LOG_FILTER_IDS.MODEL_ID]: "Model", [LOG_FILTER_IDS.PUBLIC_MODEL_OR_SEARCH_TOOL]: "Public model / search tool", + [LOG_FILTER_IDS.SEARCH]: "Search", }; export interface LogsWindow { @@ -175,6 +177,7 @@ export function useLogFilterLogic({ api_key: getFilterValue(columnFilters, LOG_FILTER_IDS.KEY_HASH), team_id: getFilterValue(columnFilters, LOG_FILTER_IDS.TEAM_ID), request_id: getFilterValue(columnFilters, LOG_FILTER_IDS.REQUEST_ID), + search: getFilterValue(columnFilters, LOG_FILTER_IDS.SEARCH), session_id: getFilterValue(columnFilters, LOG_FILTER_IDS.SESSION_ID), user_id: userIdFilter, end_user: getFilterValue(columnFilters, LOG_FILTER_IDS.END_USER), From e504477a696ca0b8c82c2083510e4ae0c9a391f6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 15:20:55 -0700 Subject: [PATCH 159/167] chore(ui): regenerate schema.d.ts for the new search params Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6742076fa78..48a8cfd54d6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -40904,6 +40904,8 @@ export interface operations { object_team_id?: string | null; /** @description Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only) */ object_key_hash?: string | null; + /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value (a raw sk- virtual key is hashed first) */ + search?: string | null; /** @description Column to sort by (e.g. 'updated_at', 'action', 'table_name') */ sort_by?: string | null; /** @description Sort order ('asc' or 'desc') */ @@ -49609,6 +49611,8 @@ export interface operations { key_hash?: string | null; /** @description Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */ key_alias?: string | null; + /** @description Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive). */ + search?: string | null; /** @description Return full key object */ return_full_object?: boolean; /** @description Include all keys for teams that user is an admin of. */ @@ -56867,6 +56871,8 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; + /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + search?: string | null; }; header?: never; path?: never; @@ -56983,6 +56989,8 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; + /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + search?: string | null; }; header?: never; path?: never; @@ -63282,6 +63290,8 @@ export interface operations { key?: string | null; /** @description Filter by key prefix (Redis-style namespace scan). Mutually exclusive with `key`; if both are provided, `key_prefix` wins. */ key_prefix?: string | null; + /** @description Match entries whose key starts with this value or whose memory_id equals it. Takes precedence over `key_prefix` and `key` when provided. */ + search?: string | null; page?: number; page_size?: number; }; From 24531ee576b4a5c535f46e840d27e2d89e990878 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:37:56 -0700 Subject: [PATCH 160/167] refactor(cost): drop the docstrings that restate TokenRates and the new tests --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 5 ----- .../llm_cost_calc/test_llm_cost_calc_utils.py | 13 ------------- .../dashscope/test_dashscope_cost_calculator.py | 5 ----- 3 files changed, 23 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index e03c2c93c26..8e24302b440 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -417,11 +417,6 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = @dataclass(frozen=True, slots=True) class TokenRates: - """The per-token rates one request bills at. reasoning_rate is None when reasoning bills at - output_rate: the model has no dedicated reasoning rate, or the caller resolves reasoning on - its own. - """ - input_rate: float output_rate: float cache_read_rate: float diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 7bc02145841..d5d7b6a47f4 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -815,8 +815,6 @@ def _off_peak_reasoning_usage() -> Usage: def test_generic_cost_per_token_off_peak_reasoning_rate(): - """Regression (LIT-6887): the block's output_cost_per_reasoning_token used to be ignored, so - reasoning tokens billed at the model's standard reasoning rate all through the window.""" from datetime import datetime, timezone model_name = "litellm-test-off-peak-reasoning" @@ -843,8 +841,6 @@ def test_generic_cost_per_token_off_peak_reasoning_rate(): def test_generic_cost_per_token_off_peak_block_without_reasoning_rate(): - """A block that leaves output_cost_per_reasoning_token unset keeps the model's own reasoning - rate, and a model with no reasoning rate at all follows the off-peak output rate.""" from datetime import datetime, timezone inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc) @@ -870,7 +866,6 @@ def test_generic_cost_per_token_off_peak_block_without_reasoning_rate(): def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier(): - """Tiered models resolve reasoning on their own path, so the block has to win there too.""" from datetime import datetime, timezone model_name = "litellm-test-off-peak-tiered-reasoning" @@ -914,8 +909,6 @@ def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier(): def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_tier(): - """A priority request bills its service-tier reasoning rate outside the window and the block's - rate inside it.""" from datetime import datetime, timezone model_name = "litellm-test-off-peak-reasoning-service-tier" @@ -946,7 +939,6 @@ def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_ti def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings(): - """A YAML true never turns into a rate of 1.0, and a quoted number still counts.""" from datetime import datetime, timezone model_name = "litellm-test-off-peak-odd-values" @@ -972,9 +964,6 @@ def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings(): def test_get_token_base_cost_off_peak_cache_creation_rate(): - """Regression (LIT-6887): the block's cache_creation_input_token_cost used to be ignored. It - replaces the five-minute cache-creation rate inside the window; the one-hour rate, and a - block without the key, keep the standard rate.""" from datetime import datetime, timezone from typing import cast @@ -1008,8 +997,6 @@ def test_get_token_base_cost_off_peak_cache_creation_rate(): def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates(): - """The per-token-type breakdown feeds the spend logs, so it has to bill the new keys the same - way the total does.""" from datetime import datetime, timezone model_name = "litellm-test-off-peak-breakdown" diff --git a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py index f0949ce041a..a30d35d46f2 100644 --- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py +++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py @@ -650,8 +650,6 @@ class TestDashscopeCostCalculator: assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10) def test_dashscope_off_peak_reasoning_rate_replaces_the_dedicated_reasoning_rate(self): - """Regression (LIT-6887): a block carrying output_cost_per_reasoning_token bills reasoning - tokens at it inside the window, over the model's own reasoning rate, which returns outside.""" self._register_off_peak_flat_model( "dashscope/qwen-reasoning-rate-off-peak-test", { @@ -678,8 +676,6 @@ class TestDashscopeCostCalculator: assert math.isclose(peak_completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10) def test_dashscope_off_peak_cache_creation_rate_replaces_the_standard_rate(self): - """Regression (LIT-6887): a block carrying cache_creation_input_token_cost bills cache-creation - tokens at it inside the window, while the cache-read rate it leaves unset stays standard.""" self._register_off_peak_flat_model( "dashscope/qwen-cache-creation-off-peak-test", {"hours_utc": self.OFF_PEAK_WINDOW, "cache_creation_input_token_cost": 1.5e-06}, @@ -701,7 +697,6 @@ class TestDashscopeCostCalculator: assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10) def test_dashscope_off_peak_reasoning_and_cache_creation_rates_override_the_selected_tier(self): - """The new keys override the selected tier the way the input and output rates already do.""" self._register_tiered_model( "dashscope/qwen-tiered-reasoning-off-peak-test", [ From 1d71e306cc13008f551d8964623ca813ca44bccb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:38:15 -0700 Subject: [PATCH 161/167] test(router): assert a non-chat configured mode round-trips With every kept get_configured_mode test using mode "chat", a Router that answered "chat" for any non-blank configured mode passed all four of them (the deleted #39630 pair's audio_speech case was the only test catching it). Read the mode back as audio_speech on an unmapped model so the configured value itself is what the test checks. Six hand-applied mutations of Router.get_configured_mode, including that hardcoded-chat one, are now all killed by the four surviving tests. --- tests/test_litellm/test_router.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 417c58b95f6..228588d974f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7362,14 +7362,14 @@ def test_get_configured_mode_reads_deployment_model_info(): router = litellm.Router( model_list=[ { - "model_name": "chat-model", - "litellm_params": {"model": "openai/some-unmapped-model"}, - "model_info": {"mode": "chat"}, + "model_name": "tts-model", + "litellm_params": {"model": "openai/some-unmapped-tts-model"}, + "model_info": {"mode": "audio_speech"}, } ] ) - assert router.get_configured_mode("chat-model") == "chat" + assert router.get_configured_mode("tts-model") == "audio_speech" def test_get_configured_mode_returns_none_for_unset_or_unknown(): From 9baa19c7d13c4d5937d7c6f337659b94b514f773 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 15:54:42 -0700 Subject: [PATCH 162/167] fix(proxy): stop hashing raw sk- values in list searches The search= param on /key/list, /audit, and /spend/logs/ui, plus key_hash= on /key/list, now compare the pasted value verbatim. Only a copied key ID (the hash) matches, so a raw virtual key never needs to travel in a GET query string Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf --- .../proxy/audit_logging_endpoints.py | 17 ++--- .../key_management_endpoints.py | 12 ++-- .../spend_management_endpoints.py | 14 ++-- .../key_management_endpoints.py | 2 +- .../proxy/test_audit_logging_endpoints.py | 19 ----- .../test_key_management_endpoints.py | 71 +++---------------- .../test_spend_management_endpoints.py | 37 ++++------ .../(dashboard)/hooks/keys/useKeys.test.ts | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 +-- 9 files changed, 43 insertions(+), 141 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index 72f7a66a420..7df14565c3f 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -18,7 +18,6 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import ( from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.utils import _hash_token_if_needed from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import AuditLogRepository @@ -50,14 +49,13 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec def _build_search_condition(search: str) -> dict[str, object]: - """Match any id column; a raw sk- key is hashed for the two columns that store key hashes.""" - hashed: Final = _hash_token_if_needed(search) + """Match a row whose id, changed_by, object_id, or changed_by_api_key equals the search value.""" return { "OR": ( {"id": search}, {"changed_by": search}, - {"object_id": hashed}, - {"changed_by_api_key": hashed}, + {"object_id": search}, + {"changed_by_api_key": search}, ) } @@ -99,10 +97,7 @@ async def get_audit_logs( ), search: str | None = Query( None, - description=( - "Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value " - "(a raw sk- virtual key is hashed first)" - ), + description="Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value", ), # Sorting parameters sort_by: str | None = Query( @@ -159,7 +154,7 @@ async def get_audit_logs( {sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order} ) - audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table + audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table # Get paginated results audit_logs: Final = await audit_log_table.find_many( @@ -221,7 +216,7 @@ async def get_audit_log_by_id( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) - audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table + audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table # Get the audit log by ID audit_log: Final = await audit_log_table.find_unique(where={"id": id}) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 758644ff01b..324d380b85b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -5803,7 +5803,7 @@ async def list_keys( ), search: str | None = Query( None, - description="Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive).", + description="Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive).", ), return_full_object: bool = Query(False, description="Return full key object"), include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."), @@ -5867,17 +5867,13 @@ async def list_keys( detail={"error": "Invalid expires value. Supported: 'active', 'expired'."}, ) - hashed_key_hash: Final[str | None] = ( - _hash_token_if_needed(token=key_hash) if isinstance(key_hash, str) else None - ) - complete_user_info: Final = await validate_key_list_check( user_api_key_dict=user_api_key_dict, user_id=user_id, team_id=team_id, organization_id=organization_id, key_alias=key_alias, - key_hash=hashed_key_hash, + key_hash=key_hash, prisma_client=prisma_client, ) @@ -5937,7 +5933,7 @@ async def list_keys( user_id=user_id, team_id=team_id, key_alias=key_alias, - key_hash=hashed_key_hash, + key_hash=key_hash, return_full_object=return_full_object, organization_id=organization_id, admin_team_ids=admin_team_ids, @@ -6175,7 +6171,7 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, def _build_key_search_where(search: str) -> KeySearchWhere: search_where: Final[KeySearchWhere] = { "OR": ( - {"token": _hash_token_if_needed(token=search)}, + {"token": search}, {"key_alias": {"contains": search, "mode": "insensitive"}}, ) } diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 9ec8dd205a6..b86a877e8f9 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2240,20 +2240,18 @@ def _build_spend_log_search_condition( end_date: datetime, next_param_index: int, ) -> _SpendLogSearchCondition: - """request_id (indexed) matches across all time; the unindexed id columns only inside the window (sk- keys hashed).""" + """request_id (indexed) matches across all time; the unindexed id columns only inside the window.""" raw: Final = f"${next_param_index}" - hashed: Final = f"${next_param_index + 1}" - window_start: Final = f"${next_param_index + 2}" - window_end: Final = f"${next_param_index + 3}" + window_start: Final = f"${next_param_index + 1}" + window_end: Final = f"${next_param_index + 2}" sql: Final = ( f"(request_id = {raw} OR (" f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') " f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') " - f'AND (api_key = {hashed} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} ' + f'AND (api_key = {raw} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} ' f"OR session_id = {raw} OR model_id = {raw})))" ) - hashed_search: Final = hash_token(token=search) if search.startswith("sk-") else search - return _SpendLogSearchCondition(sql=sql, params=(search, hashed_search, start_date, end_date)) + return _SpendLogSearchCondition(sql=sql, params=(search, start_date, end_date)) @router.get( @@ -2359,7 +2357,7 @@ async def ui_view_spend_logs( search: str | None = fastapi.Query( default=None, description=( - "Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, " + "Match a log whose request_id, api_key (hash), team_id, user, end_user, " "session_id, or model_id equals this value. request_id matches across all time; the other columns " "match inside start_date/end_date, which stay required" ), diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 5d410d7b55b..9fb5bea81e3 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -16,7 +16,7 @@ class KeyAliasContainsWhere(TypedDict): class KeySearchWhere(TypedDict): - """Prisma filter behind `/key/list?search=`: exact token (sk- keys hashed) or alias substring, case-insensitive.""" + """Prisma filter behind `/key/list?search=`: exact token or case-insensitive alias substring.""" OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]] diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py index cd2c8b0b904..fd1b05ff060 100644 --- a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py +++ b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py @@ -1,4 +1,3 @@ -import hashlib from datetime import datetime, timedelta from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -172,24 +171,6 @@ def test_search_matches_any_id_column_alongside_the_other_filters(mock_prisma_cl } -def test_search_hashes_a_raw_virtual_key_for_the_hashed_columns(mock_prisma_client): - where: Final = _list_audit_logs_where(mock_prisma_client, "search=sk-raw") - - hashed: Final = hashlib.sha256(b"sk-raw").hexdigest() - assert where == { - "AND": ( - { - "OR": ( - {"id": "sk-raw"}, - {"changed_by": "sk-raw"}, - {"object_id": hashed}, - {"changed_by_api_key": hashed}, - ) - }, - ) - } - - def test_an_empty_search_leaves_the_where_clause_unchanged(mock_prisma_client): where: Final = _list_audit_logs_where(mock_prisma_client, "action=create&search=") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 1e399e4fb58..0e4af9f75a5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6351,33 +6351,15 @@ def _search_clause(search: str, token: str) -> dict: return {"OR": [{"token": token}, {"key_alias": {"contains": search, "mode": "insensitive"}}]} -def test_build_key_filter_conditions_search_hashes_raw_key_and_ors_alias_contains(): +def test_build_key_filter_conditions_search_ors_token_and_alias_contains(): """ LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by - its ID. A pasted raw sk- key is hashed to its token first; an already-hashed - value is used verbatim. + its ID (the token column), with the pasted value used verbatim. """ - from litellm.proxy._types import hash_token from litellm.proxy.management_endpoints.key_management_endpoints import ( _build_key_filter_conditions, ) - raw_where = json.loads( - json.dumps( - _build_key_filter_conditions( - user_id=None, - team_id=None, - organization_id=None, - key_alias=None, - key_hash=None, - exclude_team_id=None, - admin_team_ids=None, - search="sk-raw", - ) - ) - ) - assert _search_clause("sk-raw", hash_token("sk-raw")) in raw_where["AND"], f"raw search not ANDed: {raw_where}" - hashed_where = json.loads( json.dumps( _build_key_filter_conditions( @@ -6402,7 +6384,6 @@ def test_build_key_filter_conditions_search_narrows_team_admin_visibility(): LIT-4741, same class as LIT-3243: `search` must be a top-level AND so it narrows a team admin's admin-team branch instead of being bypassed by it. """ - from litellm.proxy._types import hash_token from litellm.proxy.management_endpoints.key_management_endpoints import ( _build_key_filter_conditions, ) @@ -6419,21 +6400,19 @@ def test_build_key_filter_conditions_search_narrows_team_admin_visibility(): admin_team_ids=["team-a"], member_team_ids=["team-a"], include_created_by_keys=False, - search="sk-member", + search="member-key-id", ) ) ) assert where.get("AND"), f"expected top-level AND, got: {where}" - assert _search_clause("sk-member", hash_token("sk-member")) in where["AND"], f"search not ANDed: {where}" + assert _search_clause("member-key-id", "member-key-id") in where["AND"], f"search not ANDed: {where}" assert json.dumps({"team_id": {"in": ["team-a"]}}) in json.dumps(where) @pytest.mark.asyncio async def test_list_key_helper_applies_search_to_prisma_where(): """LIT-4741: `search` given to _list_key_helper must reach the Prisma where clause.""" - from litellm.proxy._types import hash_token - mock_prisma_client = AsyncMock() mock_find_many = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many @@ -6448,11 +6427,11 @@ async def test_list_key_helper_applies_search_to_prisma_where(): organization_id=None, key_alias=None, key_hash=None, - search="sk-raw", + search="key-id-123", ) where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"])) - assert _search_clause("sk-raw", hash_token("sk-raw")) in where["AND"], f"search not in Prisma where: {where}" + assert _search_clause("key-id-123", "key-id-123") in where["AND"], f"search not in Prisma where: {where}" @pytest.mark.asyncio @@ -14978,47 +14957,13 @@ async def test_list_keys_non_admin_cannot_opt_into_substring(): assert kwargs["user_id"] == "alice" -@pytest.mark.asyncio -async def test_list_keys_hashes_raw_key_hash_before_validation(): - """LIT-4741: a raw sk- key pasted as key_hash is hashed before the ownership - check and the query, so a non-admin filtering by their own raw key gets the - row instead of the 'Key Hash not found.' 403.""" - from litellm.proxy._types import hash_token - - user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice") - validate = AsyncMock( - return_value=LiteLLM_UserTable( - user_id="alice", user_email="alice@example.com", teams=[], organization_memberships=[] - ) - ) - helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0}) - with ( - patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), - patch( - "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check", - validate, - ), - patch("litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper", helper), - ): - await list_keys( - request=MagicMock(), - user_api_key_dict=user, - status=None, - user_id=None, - key_hash="sk-raw", - ) - - assert validate.call_args.kwargs["key_hash"] == hash_token("sk-raw") - assert helper.call_args.kwargs["key_hash"] == hash_token("sk-raw") - - @pytest.mark.asyncio async def test_list_keys_search_is_honored_for_non_admin(): """LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's search reaches the helper while their own-user scoping stays in place.""" user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice") - kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="sk-raw") - assert kwargs["search"] == "sk-raw" + kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="key-id-123") + assert kwargs["search"] == "key-id-123" assert kwargs["user_id"] == "alice" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index f869f3ffba2..73a29afd9b9 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -61,7 +61,7 @@ def _filter_logs_by_date_range(logs, where): _SEARCH_CLAUSE_RE = re.compile( r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) ' - r'AND \(api_key = \$(\d+) OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 ' + r'AND \(api_key = \$\1 OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 ' r"OR session_id = \$\1 OR model_id = \$\1\)\)\)" ) @@ -72,9 +72,8 @@ def _matches_spend_log_search(log, search): return True if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}): return False - if log.get("api_key") == search["api_key"]: - return True - return any(log.get(col) == search["value"] for col in ("team_id", "user", "end_user", "session_id", "model_id")) + columns = ("api_key", "team_id", "user", "end_user", "session_id", "model_id") + return any(log.get(col) == search["value"] for col in columns) def _reconstruct_ui_where_from_sql(sql_query, params): @@ -98,10 +97,9 @@ def _reconstruct_ui_where_from_sql(sql_query, params): search_clause = _SEARCH_CLAUSE_RE.search(clause.group(1)) if search_clause: - raw_index, start_index, end_index, hashed_index = (int(g) for g in search_clause.groups()) + raw_index, start_index, end_index = (int(g) for g in search_clause.groups()) where["search"] = { "value": params[raw_index - 1], - "api_key": params[hashed_index - 1], "gte": _iso(params[start_index - 1]), "lte": _iso(params[end_index - 1]), } @@ -2384,31 +2382,20 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( def test_build_spend_log_search_condition_windows_every_branch_except_request_id(): """LIT-4741: request_id matches across all time; the six other id columns only inside the window, - and a raw sk- key is hashed for the api_key branch alone.""" + all comparing the pasted value verbatim.""" start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc) end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc) condition = spend_management_endpoints._build_spend_log_search_condition( - search="sk-raw-key", start_date=start, end_date=end, next_param_index=3 + search="key-hash-7", start_date=start, end_date=end, next_param_index=3 ) assert condition.sql == ( - "(request_id = $3 OR (\"startTime\" >= ($5::timestamptz AT TIME ZONE 'UTC') " - "AND \"startTime\" <= ($6::timestamptz AT TIME ZONE 'UTC') " - 'AND (api_key = $4 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))' + "(request_id = $3 OR (\"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC') " + "AND \"startTime\" <= ($5::timestamptz AT TIME ZONE 'UTC') " + 'AND (api_key = $3 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))' ) - assert condition.params == ("sk-raw-key", hashlib.sha256(b"sk-raw-key").hexdigest(), start, end) - - -def test_build_spend_log_search_condition_leaves_non_key_values_unhashed(): - start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc) - end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc) - - condition = spend_management_endpoints._build_spend_log_search_condition( - search="sess-42", start_date=start, end_date=end, next_param_index=1 - ) - - assert condition.params == ("sess-42", "sess-42", start, end) + assert condition.params == ("key-hash-7", start, end) def _search_fixture_logs(today): @@ -2427,7 +2414,7 @@ def _search_fixture_logs(today): return [ {**base, "request_id": "req-session", "session_id": "sess-42", "startTime": recent}, {**base, "request_id": "req-session-old", "session_id": "sess-42", "startTime": old}, - {**base, "request_id": "req-key", "api_key": hashlib.sha256(b"sk-raw-key").hexdigest(), "startTime": recent}, + {**base, "request_id": "req-key", "api_key": "hashed-7", "startTime": recent}, {**base, "request_id": "req-team", "team_id": "team-7", "startTime": recent}, {**base, "request_id": "req-user", "user": "user-7", "startTime": recent}, {**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent}, @@ -2461,7 +2448,7 @@ def _five_day_window(today): [ ("req-session-old", {"req-session-old"}), ("sess-42", {"req-session"}), - ("sk-raw-key", {"req-key"}), + ("hashed-7", {"req-key"}), ("team-7", {"req-team"}), ("user-7", {"req-user"}), ("cust-7", {"req-end-user"}), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts index f23fcf811f2..84be7e2ef49 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -525,14 +525,14 @@ describe("useKeys", () => { json: async () => mockKeysResponse, }); - const { result } = renderHook(() => useKeys(1, 10, { search: "sk-pasted-key" }), { wrapper }); + const { result } = renderHook(() => useKeys(1, 10, { search: "pasted-key-id" }), { wrapper }); await waitFor(() => { expect(result.current.isLoading).toBe(false); }); const callUrl = new URL(mockFetch.mock.calls[0][0], "http://localhost"); - expect(callUrl.searchParams.get("search")).toBe("sk-pasted-key"); + expect(callUrl.searchParams.get("search")).toBe("pasted-key-id"); expect(callUrl.searchParams.has("key_alias")).toBe(false); expect(callUrl.searchParams.has("key_hash")).toBe(false); }); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 48a8cfd54d6..324085fdae4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -40904,7 +40904,7 @@ export interface operations { object_team_id?: string | null; /** @description Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only) */ object_key_hash?: string | null; - /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value (a raw sk- virtual key is hashed first) */ + /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value */ search?: string | null; /** @description Column to sort by (e.g. 'updated_at', 'action', 'table_name') */ sort_by?: string | null; @@ -49611,7 +49611,7 @@ export interface operations { key_hash?: string | null; /** @description Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */ key_alias?: string | null; - /** @description Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive). */ + /** @description Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive). */ search?: string | null; /** @description Return full key object */ return_full_object?: boolean; @@ -56871,7 +56871,7 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; - /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ search?: string | null; }; header?: never; @@ -56989,7 +56989,7 @@ export interface operations { group_by_session?: boolean; /** @description Keyset cursor '||' from a previous group_by_session page. UI route only, honored when sorting by startTime */ session_cursor?: string | null; - /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ + /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */ search?: string | null; }; header?: never; From 9464888ee9064df4083eee8843424b16eacd7da0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 16:04:41 -0700 Subject: [PATCH 163/167] test(proxy): pass search=None in direct ui_view_spend_logs calls Calling the endpoint without going through FastAPI leaves the new search param set to its Query default object, which is not None, so the grouped-session and request_id lookup tests started taking the search branch Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf --- .../proxy/spend_tracking/test_spend_query_optimization.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py index 9ae932ff01f..a7de3f1d8d6 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py @@ -184,6 +184,7 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -247,6 +248,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -314,6 +316,7 @@ async def test_spend_logs_ui_caps_total_for_large_result_sets(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -359,6 +362,7 @@ async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -406,6 +410,7 @@ async def test_spend_logs_ui_out_of_range_page_keeps_total(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=99, @@ -552,6 +557,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=1, @@ -616,6 +622,7 @@ async def test_spend_logs_ui_group_by_session_offset_pages_for_other_sorts(monke api_key=None, user_id=None, request_id=None, + search=None, start_date="2026-02-16 00:00:00", end_date="2026-02-16 23:59:59", page=2, @@ -664,6 +671,7 @@ async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(m api_key=None, user_id=None, request_id="req-deep-link", + search=None, start_date=None, end_date=None, page=1, From 7bdd148f38f35e5baed4bced6fd980dd77a83bdd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 3 Sep 2026 16:10:44 -0700 Subject: [PATCH 164/167] test(proxy-extras): fake run_prisma instead of subprocess.run in the migrate deploy harness --- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py index 9ffb57924b6..3fab20a28ad 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py @@ -763,7 +763,7 @@ class _MigrateDeployHarness: "_resolve_specific_migration", staticmethod(self.resolved.append), ) - monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run) + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run) monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None) self.baseline_succeeds = True From dc98901dc1645391986e3434a72cd256617837cf Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 19:53:32 +0000 Subject: [PATCH 165/167] fix(scim): apply default_internal_user_params.teams to SCIM-created users Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/scim/scim_v2.py | 3 +- .../scim/test_scim_v2_endpoints.py | 80 ++++++++++++++++++- 2 files changed, 80 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 069f86c852c..0f0124ee9d9 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1123,7 +1123,6 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou user_id=user_id, user_email=user_id, # We don't have email from group membership user_alias=None, - teams=[], # Teams will be added separately metadata={"created_via": created_via}, auto_create_key=False, user_role=default_role, @@ -1699,7 +1698,7 @@ async def create_user( user_id=user_id, user_email=user_data["user_email"], user_alias=user_data["user_alias"], - teams=user_data["teams"], + teams=user_data["teams"] or None, metadata=metadata, auto_create_key=False, user_role=resolved_role if admin_group is not None else default_role, diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 1697b77b99a..f4627f82506 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( SCIMRosterSyncError, UserProvisionerHelpers, _apply_group_patch_updates, + _create_user_if_not_exists, _extract_group_member_ids, _extract_ids_from_path_filter, _handle_group_membership_changes, @@ -37,8 +38,8 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import ( delete_group, delete_user, get_groups, - get_users, get_service_provider_config, + get_users, merge_placeholder, patch_group, patch_team_membership, @@ -304,6 +305,83 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN +def _mock_scim_create_user_deps(mocker: MockerFixture, scim_user: SCIMUser) -> AsyncMock: + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=()) + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=scim_user), + ) + return mocker.patch( # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user + "litellm.proxy.management_endpoints.scim.scim_v2.new_user", + AsyncMock(return_value=NewUserRequest(user_id=scim_user.userName)), + ) + + +@pytest.mark.asyncio +async def test_create_user_without_groups_defers_to_default_team(mocker: MockerFixture, monkeypatch): + """IdPs omit groups on POST /Users; teams must stay unset so new_user applies default_internal_user_params.teams""" + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-user", + emails=[SCIMUserEmail(value="new@example.com")], + ) + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = _mock_scim_create_user_deps(mocker, scim_user) + + await create_user(user=scim_user) + + assert new_user_mock.call_args.kwargs["data"].teams is None + + +@pytest.mark.asyncio +async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch): + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="new-user", + emails=[SCIMUserEmail(value="new@example.com")], + groups=[SCIMUserGroup(value="idp-team")], + ) + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = _mock_scim_create_user_deps(mocker, scim_user) + + await create_user(user=scim_user) + + assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"] + + +@pytest.mark.asyncio +async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch): + monkeypatch.setattr( + "litellm.default_internal_user_params", + {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]}, + raising=False, + ) + new_user_mock = mocker.patch( # test-quality-ok: new_user is imported inside the helper, not injectable + "litellm.proxy.management_endpoints.internal_user_endpoints.new_user", + AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")), + ) + + created = await _create_user_if_not_exists(user_id="group-user") + + assert created is not None + assert new_user_mock.call_args.kwargs["data"].teams is None + + @pytest.mark.asyncio async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch): """ From 0429339204ac41f2c0420d1693f78155244b6205 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 3 Sep 2026 20:16:09 +0000 Subject: [PATCH 166/167] fix(scim): pass proxy admin auth to new_user so default team add succeeds Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/scim/scim_v2.py | 6 +++++- .../management_endpoints/scim/test_scim_v2_endpoints.py | 3 +++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 0f0124ee9d9..98770cf9c2b 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1128,7 +1128,10 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou user_role=default_role, ) - created_user: Final = await new_user(data=new_user_request) + created_user: Final = await new_user( + data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) verbose_proxy_logger.info("Created user %s via %s", user_id, created_via) return created_user @@ -1716,6 +1719,7 @@ async def create_user( created_user: Final = await new_user( data=new_user_request, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), ) scim_user: Final = await ScimTransformations.transform_litellm_user_to_scim_user(user=created_user) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index f4627f82506..88f67cfe0e4 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -19,6 +19,7 @@ from litellm.proxy._types import ( NewUserResponse, ProxyErrorTypes, ProxyException, + UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.scim.scim_v2 import ( SCIMRosterSyncError, @@ -342,6 +343,7 @@ async def test_create_user_without_groups_defers_to_default_team(mocker: MockerF await create_user(user=scim_user) assert new_user_mock.call_args.kwargs["data"].teams is None + assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) @pytest.mark.asyncio @@ -380,6 +382,7 @@ async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFi assert created is not None assert new_user_mock.call_args.kwargs["data"].teams is None + assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) @pytest.mark.asyncio From 07dd8a7e47957a020841439da35da1d287998e08 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 3 Sep 2026 15:46:44 -0700 Subject: [PATCH 167/167] fix(scim): keep team memberships when PUT /Users carries no groups Okta sends profile updates as full PUTs with no groups or groups: [], since SCIM User.groups is readOnly and membership is synced through /Groups. The PUT handler diffed that empty list against the stored teams, removed the user from every team (which also deletes their team keys) and recomputed the role from an empty group list. Treat an empty groups list on PUT as unspecified: keep the stored teams and leave the role alone. Explicit non-empty groups still replace memberships as before Claude-Session: https://claude.ai/code/session_01CqwUV4Ywnu5aUjXx1UhJrM --- .../management_endpoints/scim/scim_v2.py | 11 ++-- .../scim/test_scim_v2_endpoints.py | 61 +++++++++++++++++++ type-discipline-budget.json | 2 +- 3 files changed, 69 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 98770cf9c2b..ceb67e3eee8 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1774,22 +1774,25 @@ async def update_user( roles=user_data["roles"], ) + # SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send + # no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team" + target_teams: Final = user_data["teams"] or existing_user.teams await _handle_team_membership_changes( user_id=user_id, - existing_teams=existing_user.teams or [], - new_teams=user_data["teams"], + existing_teams=existing_user.teams, + new_teams=target_teams, ) update_data: Final = { "user_email": user_data["user_email"], "user_alias": user_data["user_alias"], "sso_user_id": user_data["sso_user_id"], - "teams": user_data["teams"], + "teams": target_teams, "metadata": safe_dumps(metadata), } admin_group: Final = await _get_scim_admin_group() - if admin_group is not None: + if admin_group is not None and user_data["teams"]: update_data["user_role"] = _resolve_scim_user_role( user.groups or [], admin_group, _default_scim_user_role() ) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 88f67cfe0e4..60f9a1a55e2 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -1257,6 +1257,67 @@ async def test_update_user_success(mocker): assert call_args[1]["data"]["teams"] == ["new-team"] +@pytest.mark.asyncio +@pytest.mark.parametrize("groups", [None, []], ids=["groups-omitted", "groups-empty"]) +async def test_update_user_without_groups_preserves_memberships_and_role(mocker, monkeypatch, groups): + """Okta profile PUTs carry no `groups` or `groups: []`; neither may drop teams (and their keys) or recompute role""" + from litellm.proxy.proxy_server import proxy_config + + async def mock_get_config(): + return {"litellm_settings": {"scim_admin_group": "litellm-admins"}} + + monkeypatch.setattr(proxy_config, "get_config", mock_get_config) + monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False) + + existing_user = mocker.MagicMock() + existing_user.teams = ["litellm-admins", "engineering"] + existing_user.metadata = {} + + scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + userName="okta-user", + name=SCIMUserName(familyName="Renamed", givenName="Okta"), + emails=[SCIMUserEmail(value="okta@example.com")], + **({} if groups is None else {"groups": groups}), + ) + response_scim_user = SCIMUser( + schemas=["urn:ietf:params:scim:schemas:core:2.0:User"], + id="okta-user", + userName="okta-user", + emails=[SCIMUserEmail(value="okta@example.com")], + ) + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "okta-user"}) + + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception", + AsyncMock(return_value=mock_prisma_client), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists", + AsyncMock(return_value=existing_user), + ) + mocker.patch( # test-quality-ok: update_user's collaborators are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=response_scim_user), + ) + patch_membership = mocker.patch( # test-quality-ok: roster writes are module-level, not injectable + "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership", + AsyncMock(), + ) + + result = await update_user(user_id="okta-user", user=scim_user) + + assert result == response_scim_user + patch_membership.assert_not_awaited() + update_data = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"] + assert update_data["teams"] == ["litellm-admins", "engineering"] + assert "user_role" not in update_data + + @pytest.mark.asyncio async def test_update_user_not_found(mocker): """Should raise 404 when user doesn't exist""" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2f85128b4b6..78090779109 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22328 }, "LIT002": { - "limit": 26760 + "limit": 26758 }, "LIT003": { "limit": 261