diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 71a598a6fe7..4749ac6ac11 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -130,6 +130,7 @@ from litellm.types.utils import ( EmbeddingResponse, FileTypes, LiteLLMBatch, + LlmProviders, TranscriptionResponse, ) from litellm.types.vector_store_files import ( @@ -242,6 +243,19 @@ def _responses_api_optional_request_param_names() -> frozenset[str]: return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams).keys()) +def _ensure_openrouter_responses_usage( + data: dict, # mutable-ok: request payload accepts extra_body fields + responses_api_provider_config: BaseResponsesAPIConfig, +) -> None: + """Require OpenRouter usage data after callers merge an extra request body.""" + if responses_api_provider_config.custom_llm_provider != LlmProviders.OPENROUTER: + return + + usage: Final = data.get("usage") + normalized_usage = {**usage, "include": True} if isinstance(usage, Mapping) else {"include": True} + data["usage"] = normalized_usage # rebind-ok: OpenRouter usage must override extra_body + + def _custom_logger_callbacks(logging_obj: LiteLLMLoggingObj) -> list["CustomLogger"]: from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import ( @@ -2661,6 +2675,7 @@ class BaseLLMHTTPHandler: if extra_body: data.update(extra_body) + _ensure_openrouter_responses_usage(data, responses_api_provider_config) stream = bool(stream or data.get("stream")) # Preserve the OpenAI-style request context (not sent to the provider) for streaming @@ -2839,6 +2854,7 @@ class BaseLLMHTTPHandler: if extra_body: data.update(extra_body) + _ensure_openrouter_responses_usage(data, responses_api_provider_config) stream = bool(stream or data.get("stream")) # Preserve the OpenAI-style request context (not sent to the provider) for streaming diff --git a/litellm/llms/openrouter/responses/transformation.py b/litellm/llms/openrouter/responses/transformation.py index 9e32356c2fa..43a3d626a4d 100644 --- a/litellm/llms/openrouter/responses/transformation.py +++ b/litellm/llms/openrouter/responses/transformation.py @@ -8,14 +8,26 @@ encrypted_content for multi-turn stateless workflows. Docs: https://openrouter.ai/docs/api/reference/responses/overview """ -from typing import Final +from collections.abc import Mapping +from json import JSONDecodeError +from typing import TYPE_CHECKING, Final + +import httpx import litellm from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = object + class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): """ @@ -75,3 +87,51 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): def supports_native_websocket(self) -> bool: """OpenRouter does not support native WebSocket for Responses API""" return False + + def transform_responses_api_request( + self, + model: str, + input: str | ResponseInputParam, + response_api_optional_request_params: dict, # mutable-ok: base method requires mutable payload + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: base method requires mutable payload + ) -> dict: # mutable-ok: base method returns mutable JSON payload + """Request usage details so OpenRouter returns the response cost.""" + request: Final = super().transform_responses_api_request( + model=model, + input=input, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + usage: Final = request.get("usage") + request["usage"] = {**usage, "include": True} if isinstance(usage, Mapping) else {"include": True} + return request + + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + """Store the OpenRouter response cost for LiteLLM cost calculation.""" + response: Final = super().transform_response_api_response( + model=model, + raw_response=raw_response, + logging_obj=logging_obj, + ) + try: + response_json: Final = raw_response.json() + except JSONDecodeError: + return response + if not isinstance(response_json, Mapping): + return response + usage: Final = response_json.get("usage") + cost: Final = usage.get("cost") if isinstance(usage, Mapping) else None + try: + cost_value: Final = float(cost) if cost is not None else None + except (TypeError, ValueError): + return response + if cost_value is not None: + response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = cost_value + return response diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py index df883d5a208..cea72cde75d 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/envelope.py @@ -471,7 +471,7 @@ def _open_claims( def _envelope_ttl_seconds(upstream_expires_in: int | None) -> int: if upstream_expires_in is None: return MAX_ENVELOPE_TTL_SECONDS - return upstream_expires_in + return min(upstream_expires_in, MAX_ENVELOPE_TTL_SECONDS) def _refresh_ttl_seconds(upstream_refresh_expires_in: int | None) -> int: diff --git a/litellm/proxy/search_endpoints/search_tool_management.py b/litellm/proxy/search_endpoints/search_tool_management.py index 81a008cf4c8..f8161981e0e 100644 --- a/litellm/proxy/search_endpoints/search_tool_management.py +++ b/litellm/proxy/search_endpoints/search_tool_management.py @@ -17,6 +17,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.search_endpoints.search_tool_registry import SearchToolRegistry from litellm.types.search import ( ListSearchToolsResponse, @@ -31,6 +32,12 @@ router: Final = APIRouter() SEARCH_TOOL_REGISTRY: Final = SearchToolRegistry() +async def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)) -> None: + """Allow search tool changes only for proxy administrators.""" + if not is_proxy_admin(user_api_key_dict): + raise HTTPException(status_code=403, detail="Only proxy administrators can manage search tools") + + def _convert_datetime_to_str(value: datetime | str | None) -> str | None: """ Convert datetime object to ISO format string. @@ -267,7 +274,7 @@ class CreateSearchToolRequest(BaseModel): @router.post( "/search_tools", tags=["Search Tools"], - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_require_proxy_admin)], ) async def create_search_tool(request: CreateSearchToolRequest): """ @@ -339,7 +346,7 @@ class UpdateSearchToolRequest(BaseModel): @router.put( "/search_tools/{search_tool_id}", tags=["Search Tools"], - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_require_proxy_admin)], ) async def update_search_tool(search_tool_id: str, request: UpdateSearchToolRequest): """ @@ -422,7 +429,7 @@ async def update_search_tool(search_tool_id: str, request: UpdateSearchToolReque @router.delete( "/search_tools/{search_tool_id}", tags=["Search Tools"], - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_require_proxy_admin)], ) async def delete_search_tool(search_tool_id: str): """ diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f012ec8f07b..7475f1cb34e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -632,7 +632,11 @@ def _resolve_prompt_swapped_provider( prompt_id: str | None, ) -> str: swapped_provider: Final = litellm.get_llm_provider(model=swapped_model)[1] - if kwargs.get("api_key") is None and kwargs.get("api_base") is None: + extra_headers: Final = kwargs.get("extra_headers") + has_authorization_header: Final = isinstance(extra_headers, Mapping) and any( + str(header_name).lower() == "authorization" for header_name in extra_headers + ) + if kwargs.get("api_key") is None and kwargs.get("api_base") is None and not has_authorization_header: return swapped_provider try: original_provider: Final = custom_llm_provider or litellm.get_llm_provider(model=original_model)[1] diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index 63bad6d0e22..9faaaf492e8 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -376,7 +376,6 @@ async def test_bedrock_kb_request_body_has_transformed_filters( timeout=None, client=None, _is_async=False, - embedding_executor=None, ): litellm_params_dict = ( litellm_params.model_dump(exclude_none=False) diff --git a/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py b/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py index d3ea8d5b907..8c2b2322b7a 100644 --- a/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py +++ b/tests/test_litellm/llms/openrouter/responses/test_openrouter_responses_transformation.py @@ -9,6 +9,9 @@ reasoning.encrypted_content for multi-turn stateless workflows. Related issue: https://github.com/BerriAI/litellm/issues/22189 """ +from unittest.mock import MagicMock + +import httpx import pytest import litellm @@ -19,6 +22,21 @@ from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +def _successful_response_body(cost: float | None = None) -> dict: + usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15} + if cost is not None: + usage["cost"] = cost + return { + "id": "resp_123", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "openai/gpt-4o-mini", + "output": [], + "usage": usage, + } + + class TestOpenRouterResponsesAPIConfig: """Test OpenRouter Responses API configuration.""" @@ -57,9 +75,7 @@ class TestOpenRouterResponsesAPIConfig: from litellm.types.router import GenericLiteLLMParams params = GenericLiteLLMParams(api_key="sk-or-test-key") - headers = config.validate_environment( - headers={}, model="openai/o4-mini", litellm_params=params - ) + headers = config.validate_environment(headers={}, model="openai/o4-mini", litellm_params=params) assert headers["Authorization"] == "Bearer sk-or-test-key" def test_validate_environment_raises_without_key(self, monkeypatch): @@ -81,6 +97,53 @@ class TestOpenRouterResponsesAPIConfig: e = exc_info.value assert "OpenRouter API key is required" in str(e) + def test_transform_request_includes_usage(self): + """OpenRouter must return usage details for response cost accounting.""" + from litellm.types.router import GenericLiteLLMParams + + request = OpenRouterResponsesAPIConfig().transform_responses_api_request( + model="openai/gpt-4o-mini", + input="hello", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert request["usage"] == {"include": True} + + def test_transform_response_stores_usage_cost(self): + """OpenRouter usage.cost must reach the LiteLLM cost calculator.""" + raw_response = httpx.Response(200, json=_successful_response_body(cost=0.0125)) + + response = OpenRouterResponsesAPIConfig().transform_response_api_response( + model="openai/gpt-4o-mini", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0125 + + def test_transform_response_allows_missing_usage_cost(self): + """A response without cost data must remain usable.""" + raw_response = httpx.Response(200, json=_successful_response_body()) + + response = OpenRouterResponsesAPIConfig().transform_response_api_response( + model="openai/gpt-4o-mini", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert "llm_provider-x-litellm-response-cost" not in response._hidden_params["additional_headers"] + + def test_extra_body_usage_cannot_disable_openrouter_cost_tracking(self): + """The final request merge must force usage.include to true.""" + from litellm.llms.custom_httpx.llm_http_handler import _ensure_openrouter_responses_usage + + request = {"usage": {"include": False, "detail": "all"}} + _ensure_openrouter_responses_usage(request, OpenRouterResponsesAPIConfig()) + + assert request["usage"] == {"include": True, "detail": "all"} + class TestOpenRouterResponsesAPIRegistration: """Test that OpenRouter is properly registered as a native Responses API provider.""" @@ -97,8 +160,7 @@ class TestOpenRouterResponsesAPIRegistration: provider=LlmProviders.OPENROUTER, ) assert config is not None, ( - "OpenRouter must be registered as a native Responses API provider " - "to preserve reasoning.encrypted_content" + "OpenRouter must be registered as a native Responses API provider to preserve reasoning.encrypted_content" ) assert isinstance(config, OpenRouterResponsesAPIConfig) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py index bd310339a1d..de597a11413 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_envelope.py @@ -30,7 +30,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import DecryptFailed, EnvelopeIdentity, EnvelopeKeys, - EnvelopeLifetimeUnrepresentable, EnvelopeMintError, EnvelopeTooLarge, Expired, @@ -161,7 +160,7 @@ def test_claim_layout_and_no_plaintext_token_in_envelope(): assert _REFRESH_TOKEN not in json.dumps(claims) -def test_unrepresentable_access_lifetime_is_a_typed_mint_error(): +def test_large_access_lifetime_is_capped(): grant = UpstreamTokenGrant( access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", @@ -170,9 +169,8 @@ def test_unrepresentable_access_lifetime_is_a_typed_mint_error(): result = mint_envelope(_IDENTITY, grant, _KEYS, _NOW) - assert isinstance(result, EnvelopeLifetimeUnrepresentable) - assert result.tag == "envelope_lifetime_unrepresentable" - assert result.expires_in == 10**30 + assert isinstance(result, SealedEnvelope) + assert result.expires_at == _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS) def _refresh_credential() -> RefreshCredential: @@ -259,7 +257,7 @@ def test_refresh_envelope_never_leaks_the_refresh_token_in_plaintext(): "expires_in, expected_ttl", [ (600, 600), - (MAX_ENVELOPE_TTL_SECONDS + 82800, MAX_ENVELOPE_TTL_SECONDS + 82800), + (MAX_ENVELOPE_TTL_SECONDS + 82800, MAX_ENVELOPE_TTL_SECONDS), (None, MAX_ENVELOPE_TTL_SECONDS), ], ) @@ -277,11 +275,11 @@ def test_expiry_honored_against_injected_clock(): assert isinstance(open_envelope(token, _KEYS, _NOW + timedelta(seconds=601)), Expired) -def test_upstream_token_lifetime_is_enforced_on_open(): +def test_envelope_lifetime_is_capped_on_open(): grant = UpstreamTokenGrant(access_token=SecretStr(_ACCESS_TOKEN), token_type="Bearer", expires_in=86400) token = _sealed_token(grant) - just_before_expiry = _NOW + timedelta(seconds=86399) - at_expiry = _NOW + timedelta(seconds=86400) + just_before_expiry = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS - 1) + at_expiry = _NOW + timedelta(seconds=MAX_ENVELOPE_TTL_SECONDS) assert isinstance(open_envelope(token, _KEYS, just_before_expiry), OpenedEnvelope) assert isinstance(open_envelope(token, _KEYS, at_expiry), Expired) 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..95f80ae33c2 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 @@ -5690,7 +5690,7 @@ async def test_bridge_envelope_too_large_upstream_token_is_502(): @pytest.mark.asyncio -async def test_bridge_envelope_unrepresentable_upstream_lifetime_is_502(): +async def test_bridge_envelope_large_upstream_lifetime_is_capped(): from litellm.types.mcp import MCPAuth server = _bridge_server(auth_type=MCPAuth.oauth_delegate) @@ -5706,11 +5706,10 @@ async def test_bridge_envelope_unrepresentable_upstream_lifetime_is_502(): key_hash="hashed-litellm-key-77", ) - assert response.status_code == 502 - assert json.loads(response.body) == { - "error": "server_error", - "error_description": "the upstream token response reports an unrepresentable lifetime", - } + assert response.status_code == 200 + body = json.loads(response.body) + assert body["access_token"].startswith("llm_env_") + assert body["expires_in"] <= 3600 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index 70e9a96b316..1d4be321a3b 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -114,9 +114,7 @@ async def test_list_search_tools_config_only(monkeypatch): with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): # Mock proxy_config mock_proxy_config = MagicMock() - mock_proxy_config.get_config = AsyncMock( - return_value={"search_tools": config_tools} - ) + mock_proxy_config.get_config = AsyncMock(return_value={"search_tools": config_tools}) mock_proxy_config.parse_search_tools = MagicMock(return_value=config_tools) with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config): # Mock auth @@ -190,9 +188,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch): with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): # Mock proxy_config mock_proxy_config = MagicMock() - mock_proxy_config.get_config = AsyncMock( - return_value={"search_tools": config_tools} - ) + mock_proxy_config.get_config = AsyncMock(return_value={"search_tools": config_tools}) mock_proxy_config.parse_search_tools = MagicMock(return_value=config_tools) with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config): # Mock auth @@ -213,11 +209,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch): # Verify DB tool is present db_tool = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "existing-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "existing-tool"), None, ) assert db_tool is not None @@ -230,11 +222,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch): # Verify unique config tool is present config_tool = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "unique-config-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "unique-config-tool"), None, ) assert config_tool is not None @@ -245,8 +233,7 @@ async def test_list_search_tools_filters_duplicate_config_tools(monkeypatch): ( t for t in data["search_tools"] - if t["search_tool_name"] == "existing-tool" - and t["is_from_config"] is True + if t["search_tool_name"] == "existing-tool" and t["is_from_config"] is True ), None, ) @@ -321,11 +308,7 @@ async def test_list_search_tools_datetime_conversion(monkeypatch): # Test datetime conversion for tool 1 tool1 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "datetime-test-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "datetime-test-tool"), None, ) assert tool1 is not None @@ -339,11 +322,7 @@ async def test_list_search_tools_datetime_conversion(monkeypatch): # Test None handling for tool 2 tool2 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "null-datetime-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "null-datetime-tool"), None, ) assert tool2 is not None @@ -355,11 +334,7 @@ async def test_list_search_tools_datetime_conversion(monkeypatch): # Test string passthrough for tool 3 tool3 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "string-datetime-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "string-datetime-tool"), None, ) assert tool3 is not None @@ -399,9 +374,7 @@ async def test_list_search_tools_config_error_handling(monkeypatch): with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): # Mock proxy_config to raise an error mock_proxy_config = MagicMock() - mock_proxy_config.get_config = AsyncMock( - side_effect=Exception("Config error") - ) + mock_proxy_config.get_config = AsyncMock(side_effect=Exception("Config error")) with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config): # Mock auth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -420,13 +393,8 @@ async def test_list_search_tools_config_error_handling(monkeypatch): assert len(data["search_tools"]) == 1 assert data["search_tools"][0]["search_tool_name"] == "db-tool-1" # Verify masking of sensitive values - assert ( - data["search_tools"][0]["litellm_params"]["api_key"] - != "sk-test" - ) - assert ( - "****" in data["search_tools"][0]["litellm_params"]["api_key"] - ) + assert data["search_tools"][0]["litellm_params"]["api_key"] != "sk-test" + assert "****" in data["search_tools"][0]["litellm_params"]["api_key"] finally: app.dependency_overrides.pop(user_api_key_auth, None) @@ -541,31 +509,18 @@ async def test_list_search_tools_db_masking_sensitive_values(monkeypatch): # Test tool 1: api_key should be masked tool1 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "perplexity-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "perplexity-tool"), None, ) assert tool1 is not None - assert ( - tool1["litellm_params"]["api_key"] != "pplx-sk-1234567890abcdef" - ) + assert tool1["litellm_params"]["api_key"] != "pplx-sk-1234567890abcdef" assert "****" in tool1["litellm_params"]["api_key"] assert tool1["litellm_params"]["search_provider"] == "perplexity" - assert ( - tool1["litellm_params"]["api_base"] - == "https://api.perplexity.ai" - ) + assert tool1["litellm_params"]["api_base"] == "https://api.perplexity.ai" # Test tool 2: api_key should be masked tool2 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "tavily-tool" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "tavily-tool"), None, ) assert tool2 is not None @@ -575,29 +530,18 @@ async def test_list_search_tools_db_masking_sensitive_values(monkeypatch): # Test tool 3: access_token and secret_key should be masked tool3 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "tool-with-token" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "tool-with-token"), None, ) assert tool3 is not None - assert ( - tool3["litellm_params"]["access_token"] - != "token-abcdefghijklmnop" - ) + assert tool3["litellm_params"]["access_token"] != "token-abcdefghijklmnop" assert "****" in tool3["litellm_params"]["access_token"] assert tool3["litellm_params"]["secret_key"] != "secret-xyz123" assert "****" in tool3["litellm_params"]["secret_key"] # Test tool 4: non-sensitive fields should remain unmasked tool4 = next( - ( - t - for t in data["search_tools"] - if t["search_tool_name"] == "tool-with-non-sensitive" - ), + (t for t in data["search_tools"] if t["search_tool_name"] == "tool-with-non-sensitive"), None, ) assert tool4 is not None @@ -626,25 +570,18 @@ async def test_get_all_search_tools_from_db_retries_on_transport_error(): return [] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_searchtoolstable.find_many = AsyncMock( - side_effect=_flaky_find_many - ) + mock_prisma_client.db.litellm_searchtoolstable.find_many = AsyncMock(side_effect=_flaky_find_many) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 - result = await SearchToolRegistry.get_all_search_tools_from_db( - prisma_client=mock_prisma_client - ) + result = await SearchToolRegistry.get_all_search_tools_from_db(prisma_client=mock_prisma_client) assert result == [] assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "get_all_search_tools_from_db_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "get_all_search_tools_from_db_lookup_failure" @contextlib.contextmanager @@ -715,6 +652,46 @@ def _override_auth(user): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.mark.parametrize( + ("method", "path", "body", "registry_method"), + [ + ( + "post", + "/search_tools", + {"search_tool": {"search_tool_name": "blocked", "litellm_params": {"search_provider": "tavily"}}}, + "add_search_tool_to_db", + ), + ( + "put", + "/search_tools/blocked-id", + {"search_tool": {"search_tool_name": "blocked", "litellm_params": {"search_provider": "tavily"}}}, + "update_search_tool_in_db", + ), + ("delete", "/search_tools/blocked-id", None, "delete_search_tool_from_db"), + ], +) +def test_search_tool_mutations_require_proxy_admin(method, path, body, registry_method): + """Internal users cannot configure proxy-originated search requests.""" + registry = _fake_registry([]) + internal_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user") + + with ( + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY", registry + ), + _override_auth(internal_user), + ): + client = TestClient(app) + request_kwargs = {} if body is None else {"json": body} + response = getattr(client, method)(path, **request_kwargs) + + assert response.status_code == 403 + getattr(registry, registry_method).assert_not_awaited() + + @pytest.mark.asyncio async def test_list_search_tools_scoped_to_key_object_permission(): """ @@ -747,9 +724,7 @@ async def test_list_search_tools_scoped_to_key_object_permission(): @pytest.mark.asyncio async def test_list_search_tools_unrestricted_internal_user_sees_all(): """An internal user with no search_tools allowlist is unrestricted and sees every tool.""" - unrestricted_user = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user" - ) + unrestricted_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user") with ( _mock_search_tool_backend(_scoping_db_tools()), @@ -789,9 +764,7 @@ async def test_list_search_tools_scoped_to_team_object_permission(): response = TestClient(app).get("/search_tools/list") assert response.status_code == 200 - assert [t["search_tool_name"] for t in response.json()["search_tools"]] == [ - "db-tool-2" - ] + assert [t["search_tool_name"] for t in response.json()["search_tools"]] == ["db-tool-2"] @pytest.mark.asyncio @@ -1059,9 +1032,21 @@ def _live_router_and_db(db_rows: list): fake_router.search_tools = list(db_rows) with contextlib.ExitStack() as stack: - stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client", MagicMock())) # test-quality-ok: proxy globals are the only seam; see the module note above - stack.enter_context(patch("litellm.proxy.proxy_server.proxy_config", proxy_config)) # test-quality-ok: proxy globals are the only seam; see the module note above - stack.enter_context(patch("litellm.proxy.proxy_server.llm_router", fake_router)) # test-quality-ok: proxy globals are the only seam; see the module note above + stack.enter_context( + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ) + ) + stack.enter_context( + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.proxy_server.proxy_config", proxy_config + ) + ) + stack.enter_context( + patch( # test-quality-ok: proxy globals are the only seam; see the module note above + "litellm.proxy.proxy_server.llm_router", fake_router + ) + ) stack.enter_context( patch( # test-quality-ok: proxy globals are the only seam; see the module note above "litellm.proxy.search_endpoints.search_tool_management.SEARCH_TOOL_REGISTRY", diff --git a/tests/test_litellm/responses/test_responses_prompt_management.py b/tests/test_litellm/responses/test_responses_prompt_management.py index 204b4d00f01..1c6bb6a2652 100644 --- a/tests/test_litellm/responses/test_responses_prompt_management.py +++ b/tests/test_litellm/responses/test_responses_prompt_management.py @@ -67,18 +67,15 @@ def _patch_responses_dispatch(): side_effect=_provider_by_model, ), patch( - "litellm.responses.mcp.litellm_proxy_mcp_handler." - "LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", + "litellm.responses.mcp.litellm_proxy_mcp_handler.LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway", return_value=False, ), patch( - "litellm.responses.main.ProviderConfigManager" - ".get_provider_responses_api_config", + "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config", return_value=None, ), patch( - "litellm.responses.main.litellm_completion_transformation_handler" - ".response_api_handler", + "litellm.responses.main.litellm_completion_transformation_handler.response_api_handler", return_value=MagicMock(), ), ] @@ -140,7 +137,6 @@ def _make_cache_control_case() -> tuple[ class TestResponsesAPIPromptManagement: - def test_str_input_coerced_and_merged(self): """[A] str input is wrapped into a message list before being passed to the hook.""" template_messages: List[AllMessageValues] = [ @@ -171,9 +167,7 @@ class TestResponsesAPIPromptManagement: logging_obj.get_chat_completion_prompt.assert_called_once() call_kwargs = logging_obj.get_chat_completion_prompt.call_args.kwargs # str was coerced to a single user message before being passed to the hook - assert call_kwargs["messages"] == [ - {"role": "user", "content": "Tell me about AI."} - ] + assert call_kwargs["messages"] == [{"role": "user", "content": "Tell me about AI."}] assert call_kwargs["prompt_id"] == "summariser-prompt" def test_list_input_merged_with_template(self): @@ -579,6 +573,36 @@ def test_resolve_prompt_swapped_provider_allows_swap_without_credentials(): ) +@pytest.mark.parametrize("authorization_header", ["Authorization", "authorization"]) +def test_resolve_prompt_swapped_provider_rejects_authorization_extra_headers(authorization_header: str): + import litellm + from litellm.responses.main import _resolve_prompt_swapped_provider + + with pytest.raises(litellm.BadRequestError, match="Refusing to send"): + _resolve_prompt_swapped_provider( + original_model="anthropic/claude-haiku-4-5", + swapped_model="gpt-4o-mini", + custom_llm_provider="anthropic", + kwargs={"extra_headers": {authorization_header: "Bearer secret"}}, + prompt_id="p1", + ) + + +def test_resolve_prompt_swapped_provider_allows_harmless_extra_headers(): + from litellm.responses.main import _resolve_prompt_swapped_provider + + assert ( + _resolve_prompt_swapped_provider( + original_model="anthropic/claude-haiku-4-5", + swapped_model="gpt-4o-mini", + custom_llm_provider="anthropic", + kwargs={"extra_headers": {"X-Request-ID": "request-1"}}, + prompt_id="p1", + ) + == "openai" + ) + + def test_resolve_prompt_swapped_provider_allows_same_provider_swap_with_credentials(): from litellm.responses.main import _resolve_prompt_swapped_provider