diff --git a/.github/workflows/check-lazy-openapi-snapshot.yml b/.github/workflows/check-lazy-openapi-snapshot.yml deleted file mode 100644 index 2e4ed3637f..0000000000 --- a/.github/workflows/check-lazy-openapi-snapshot.yml +++ /dev/null @@ -1,75 +0,0 @@ -name: Check Lazy OpenAPI Snapshot - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - "litellm_**" - -permissions: - contents: read - checks: write - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -jobs: - verify: - runs-on: ubuntu-latest - timeout-minutes: 10 - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 - with: - version: "0.10.9" - - - name: Cache uv dependencies - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: | - ~/.cache/uv - .venv - key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }} - restore-keys: | - ${{ runner.os }}-uv- - - - name: Install dependencies - run: uv sync --frozen --all-groups --all-extras - - - name: Regenerate snapshot to /tmp - id: regen - run: | - cp litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.committed.json - uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot - mv litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.fresh.json - mv /tmp/snapshot.committed.json litellm/proxy/_lazy_openapi_snapshot.json - - - name: Compare - id: diff - continue-on-error: true - run: | - diff -q /tmp/snapshot.fresh.json litellm/proxy/_lazy_openapi_snapshot.json - - - name: Mark neutral if drift - if: steps.diff.outcome == 'failure' - uses: LouisBrunner/checks-action@6b626ffbad7cc56fd58627f774b9067e6118af23 # v2.0.0 - with: - token: ${{ secrets.GITHUB_TOKEN }} - name: lazy-openapi-snapshot - conclusion: neutral - output: | - { - "title": "Lazy openapi snapshot is stale", - "summary": "Run `python -m litellm.proxy._lazy_openapi_snapshot` and commit the regenerated `litellm/proxy/_lazy_openapi_snapshot.json`. Not blocking — the snapshot will regenerate at release if not committed." - } diff --git a/.gitignore b/.gitignore index 38bf9554b5..59812ed6ed 100644 --- a/.gitignore +++ b/.gitignore @@ -90,7 +90,6 @@ test.py litellm_config.yaml !.github/observatory/litellm_config.yaml .cursor -.vscode/launch.json litellm/proxy/to_delete_loadtest_work/* update_model_cost_map.py tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -100,4 +99,5 @@ STABILIZATION_TODO.md **/test-results **/playwright-report **/*.storageState.json -**/coverage \ No newline at end of file +**/coverage +test-config \ No newline at end of file diff --git a/Makefile b/Makefile index b6b674ff3b..5dbd308a3e 100644 --- a/Makefile +++ b/Makefile @@ -185,3 +185,6 @@ test-llm-translation-single: install-test-deps $(UV_RUN) pytest tests/llm_translation/$(FILE) \ --junitxml=test-results/junit.xml \ -v --tb=short --maxfail=100 --timeout=300 + +test-llm-translation-flush-vcr-cache: + $(UV_RUN) python tests/_flush_vcr_cache.py diff --git a/README.md b/README.md index d72fb746ed..72fd43925c 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,7 @@ Managing LLM calls across providers gets complicated fast — different SDKs, au Stripe image Google ADK - Greptile + Greptile OpenHands

Netflix

OpenAI Agents SDK diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index f6ed7767c4..4bfe9d3187 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -857,10 +857,16 @@ async def project_info( where={"team_id": project.team_id} ) if team: - is_team_member = ( - user_api_key_dict.user_id in team.admins - or user_api_key_dict.user_id in team.members - ) + caller_user_id = user_api_key_dict.user_id + for m in team.members_with_roles or []: + m_user_id = ( + m.get("user_id") + if isinstance(m, dict) + else getattr(m, "user_id", None) + ) + if m_user_id == caller_user_id: + is_team_member = True + break if not (is_admin or is_team_member): raise HTTPException( @@ -911,20 +917,20 @@ async def list_projects( include={"litellm_budget_table": True, "object_permission": True} ) else: - # Get projects for teams the user belongs to - user_teams = await prisma_client.db.litellm_teamtable.find_many( - where={ - "OR": [ - {"members": {"has": user_api_key_dict.user_id}}, - {"admins": {"has": user_api_key_dict.user_id}}, - ] - } + # Look up the user's team memberships via the reverse-index on + # LiteLLM_UserTable.teams (maintained by team_member_add alongside + # members_with_roles). This avoids a full scan of all team rows. + user_record = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + ) + user_team_ids = ( + user_record.teams + if user_record is not None and user_record.teams + else [] ) - team_ids = [team.team_id for team in user_teams] - projects = await prisma_client.db.litellm_projecttable.find_many( - where={"team_id": {"in": team_ids}}, + where={"team_id": {"in": user_team_ids}}, include={"litellm_budget_table": True, "object_permission": True}, ) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 41c78296fb..b8710da343 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.69" +version = "0.4.70" 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.69" +version = "0.4.70" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index ce1bc26c5e..11733ce4ce 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -432,9 +432,10 @@ class Cache: str: The final hashed cache key with the redis namespace. """ dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {}) + metadata = kwargs.get("metadata") or {} namespace = ( dynamic_cache_control.get("namespace") - or kwargs.get("metadata", {}).get("redis_namespace") + or metadata.get("redis_namespace") or self.namespace ) if namespace: diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7d514e648f..3cf1d911d7 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -87,6 +87,18 @@ class CachingHandlerResponse(BaseModel): in_memory_cache_obj = InMemoryCache() +def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool: + """ + When stream=True, do not run success callbacks at cache-hit time. + + Cached chat/text completion replay uses CustomStreamWrapper; cached Responses + replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success + handlers when the stream finishes; firing them here too would double-count + spend and callback records. + """ + return kwargs.get("stream", False) is True + + class LLMCachingHandler: def __init__( self, @@ -99,6 +111,7 @@ class LLMCachingHandler: self.async_streaming_chunks: List[ModelResponse] = [] self.sync_streaming_chunks: List[ModelResponse] = [] self.request_kwargs = request_kwargs + self.preset_cache_key: Optional[str] = None self.original_function = original_function self.start_time = start_time if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache): @@ -206,7 +219,7 @@ class LLMCachingHandler: custom_llm_provider=kwargs.get("custom_llm_provider", None), args=args, ) - if kwargs.get("stream", False) is False: + if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs): # LOG SUCCESS self._async_log_cache_hit_on_callbacks( logging_obj=logging_obj, @@ -215,11 +228,12 @@ class LLMCachingHandler: end_time=end_time, cache_hit=cache_hit, ) - cache_key = litellm.cache.get_cache_key(**kwargs) - if ( - isinstance(cached_result, BaseModel) - or isinstance(cached_result, CustomStreamWrapper) - ) and hasattr(cached_result, "_hidden_params"): + cache_key = ( + self.preset_cache_key + or self.request_kwargs.get("cache_key") + or litellm.cache.get_cache_key(**self.request_kwargs) + ) + if hasattr(cached_result, "_hidden_params"): cached_result._hidden_params["cache_key"] = cache_key # type: ignore return CachingHandlerResponse(cached_result=cached_result) elif ( @@ -265,8 +279,6 @@ class LLMCachingHandler: kwargs: Dict[str, Any], args: Optional[Tuple[Any, ...]] = None, ) -> CachingHandlerResponse: - from litellm.utils import CustomStreamWrapper - cached_result: Optional[Any] = None # Check if caching should be performed BEFORE doing expensive kwargs copy @@ -282,6 +294,11 @@ class LLMCachingHandler: args, ) ) + if new_kwargs.get("metadata") is None: + new_kwargs.pop("metadata", None) + if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: + new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) + self.request_kwargs = new_kwargs print_verbose("Checking Sync Cache") cached_result = litellm.cache.get_cache(**new_kwargs) if cached_result is not None: @@ -322,17 +339,19 @@ class LLMCachingHandler: is_async=False, ) - logging_obj.handle_sync_success_callbacks_for_async_calls( - result=cached_result, - start_time=start_time, - end_time=end_time, - cache_hit=cache_hit, + if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs): + logging_obj.handle_sync_success_callbacks_for_async_calls( + result=cached_result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + ) + cache_key = ( + self.preset_cache_key + or self.request_kwargs.get("cache_key") + or litellm.cache.get_cache_key(**self.request_kwargs) ) - cache_key = litellm.cache.get_cache_key(**kwargs) - if ( - isinstance(cached_result, BaseModel) - or isinstance(cached_result, CustomStreamWrapper) - ) and hasattr(cached_result, "_hidden_params"): + if hasattr(cached_result, "_hidden_params"): cached_result._hidden_params["cache_key"] = cache_key # type: ignore return CachingHandlerResponse(cached_result=cached_result) return CachingHandlerResponse(cached_result=cached_result) @@ -686,6 +705,11 @@ class LLMCachingHandler: args, ) ) + if new_kwargs.get("metadata") is None: + new_kwargs.pop("metadata", None) + if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: + new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) + self.request_kwargs = new_kwargs cached_result: Optional[Any] = None if call_type == CallTypes.aembedding.value: if isinstance(new_kwargs["input"], str): @@ -710,14 +734,26 @@ class LLMCachingHandler: if all(result is None for result in cached_result): cached_result = None else: + request_kwargs = new_kwargs.copy() + request_cache_key = request_kwargs.pop("cache_key", None) if litellm.cache._supports_async() is True: ## check if dual cache is supported ## + self.preset_cache_key = ( + request_cache_key or litellm.cache.get_cache_key(**request_kwargs) + ) cached_result = await litellm.cache.async_get_cache( - dynamic_cache_object=self.dual_cache, **new_kwargs + dynamic_cache_object=self.dual_cache, + cache_key=self.preset_cache_key, + **request_kwargs, ) else: # fallback for caches that don't support async + self.preset_cache_key = ( + request_cache_key or litellm.cache.get_cache_key(**request_kwargs) + ) cached_result = litellm.cache.get_cache( - dynamic_cache_object=self.dual_cache, **new_kwargs + dynamic_cache_object=self.dual_cache, + cache_key=self.preset_cache_key, + **request_kwargs, ) return cached_result @@ -825,8 +861,27 @@ class LLMCachingHandler: elif (call_type == "aresponses" or call_type == "responses") and isinstance( cached_result, dict ): - # Convert cached dict back to ResponsesAPIResponse object - cached_result = ResponsesAPIResponse(**cached_result) + from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + ) + + response_obj = ResponsesAPIResponse(**cached_result) + if ( + hasattr(response_obj, "_hidden_params") + and response_obj._hidden_params is not None + and isinstance(response_obj._hidden_params, dict) + ): + response_obj._hidden_params["cache_hit"] = True + + if kwargs.get("stream", False) is True: + cached_result = CachedResponsesAPIStreamingIterator( + response=response_obj, + logging_obj=logging_obj, + request_data=kwargs, + call_type=call_type, + ) + else: + cached_result = response_obj if ( hasattr(cached_result, "_hidden_params") diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 34ae3638a5..8060a65b78 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -92,6 +92,25 @@ class DualCache(BaseCache): if default_redis_ttl is not None: self.default_redis_ttl = default_redis_ttl + def attach_redis_cache( + self, + redis_cache: Optional[RedisCache] = None, + *, + default_redis_ttl: Optional[float] = None, + ) -> None: + """ + Attach a Redis backend if this DualCache does not already have one. + + No-op when ``redis_cache`` is None or when Redis was already set (constructor + or a prior attach). Use this for lazy wiring after a shared Redis client exists. + Does not backfill in-memory-only keys to Redis. + """ + if redis_cache is None or self.redis_cache is not None: + return + self.redis_cache = redis_cache + if default_redis_ttl is not None: + self.default_redis_ttl = default_redis_ttl + def set_cache(self, key, value, local_only: bool = False, **kwargs): # Update both Redis and in-memory cache try: @@ -392,6 +411,7 @@ class DualCache(BaseCache): value: float, parent_otel_span: Optional[Span] = None, local_only: bool = False, + refresh_ttl: bool = False, **kwargs, ) -> Optional[float]: """ @@ -399,6 +419,9 @@ class DualCache(BaseCache): Value - float - the value you want to increment by + Refresh_ttl - bool - if True, resets the Redis TTL on every write. + Default False preserves window-style semantics. + Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ @@ -415,6 +438,7 @@ class DualCache(BaseCache): value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), + refresh_ttl=refresh_ttl, ) return result diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 84a2887f52..cb9ce475d3 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -551,6 +551,13 @@ class RedisCache(BaseCache): async def async_set_cache(self, key, value, **kwargs): from redis.asyncio import Redis + if key is None: + verbose_logger.debug( + "LiteLLM Redis Caching: async set() skipped — key is None, value=%r", + value, + ) + return None + start_time = time.time() try: _redis_client: Redis = self.init_async_client() # type: ignore @@ -569,8 +576,9 @@ class RedisCache(BaseCache): ) ) verbose_logger.error( - "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s", + "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r", str(e), + key, value, ) raise e @@ -824,6 +832,7 @@ class RedisCache(BaseCache): value: float, ttl: Optional[int] = None, parent_otel_span: Optional[Span] = None, + refresh_ttl: bool = False, ) -> float: from redis.asyncio import Redis @@ -834,11 +843,12 @@ class RedisCache(BaseCache): try: result = await _redis_client.incrbyfloat(name=key, amount=value) if _used_ttl is not None: - # check if key already has ttl, if not -> set ttl - current_ttl = await _redis_client.ttl(key) - if current_ttl == -1: - # Key has no expiration + if refresh_ttl: await _redis_client.expire(key, _used_ttl) + else: + current_ttl = await _redis_client.ttl(key) + if current_ttl == -1: + await _redis_client.expire(key, _used_ttl) ## LOGGING ## end_time = time.time() diff --git a/litellm/constants.py b/litellm/constants.py index d78c124d71..6c889a317b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1425,6 +1425,7 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id" LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli" LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token" CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session" +CLI_SSO_SESSION_TTL_SECONDS = 600 CLI_JWT_TOKEN_NAME = "cli-jwt-token" # Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility CLI_JWT_EXPIRATION_HOURS = int( diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py index 3c83517bb5..8c3c2a5ff0 100644 --- a/litellm/integrations/arize/arize_phoenix_client.py +++ b/litellm/integrations/arize/arize_phoenix_client.py @@ -2,11 +2,23 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix. """ +import urllib.parse from typing import Any, Dict, Optional from litellm.llms.custom_httpx.http_handler import HTTPHandler +def _sanitize_id(identifier: str) -> str: + """Reject path traversal characters and URL-encode the identifier.""" + if any(c in identifier for c in ("/", "\\", "#", "?")): + raise ValueError( + f"Invalid identifier {identifier!r}: contains disallowed characters" + ) + if ".." in identifier: + raise ValueError(f"Invalid identifier {identifier!r}: path traversal detected") + return urllib.parse.quote(identifier, safe="") + + class ArizePhoenixClient: """ Client for interacting with Arize Phoenix API to fetch prompt versions. @@ -53,7 +65,8 @@ class ArizePhoenixClient: Returns: Dictionary containing prompt version data, or None if not found """ - url = f"{self.api_base}/v1/prompt_versions/{prompt_version_id}" + safe_id = _sanitize_id(prompt_version_id) + url = f"{self.api_base}/v1/prompt_versions/{safe_id}" try: # Use the underlying httpx client directly to avoid query param extraction diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index 0502422cf8..e742cc14b7 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -3,11 +3,27 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. """ import base64 +import urllib.parse from typing import Any, Dict, List, Optional from litellm.llms.custom_httpx.http_handler import HTTPHandler +def _sanitize_file_path(file_path: str) -> str: + """Reject path traversal and URL-encode each path segment.""" + if "#" in file_path or "?" in file_path: + raise ValueError( + f"Invalid file path {file_path!r}: contains URL special characters" + ) + parts = file_path.split("/") + for part in parts: + if part == "..": + raise ValueError( + f"Invalid file path {file_path!r}: path traversal detected" + ) + return "/".join(urllib.parse.quote(part, safe="") for part in parts) + + class BitBucketClient: """ Client for interacting with BitBucket API to fetch .prompt files. @@ -72,7 +88,8 @@ class BitBucketClient: Returns: File content as string, or None if file not found """ - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}" + safe_path = _sanitize_file_path(file_path) + url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}" try: response = self.http_handler.get(url, headers=self.headers) @@ -119,7 +136,8 @@ class BitBucketClient: Returns: List of file paths """ - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{directory_path}" + safe_dir = _sanitize_file_path(directory_path) if directory_path else "" + url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_dir}" try: response = self.http_handler.get(url, headers=self.headers) @@ -211,7 +229,8 @@ class BitBucketClient: Returns: Dictionary containing file metadata, or None if file not found """ - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}" + safe_path = _sanitize_file_path(file_path) + url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}" try: # Use GET with Range header to get just the headers (HEAD equivalent) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 723b142dfa..d9e57ee7ce 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -265,6 +265,7 @@ class PrometheusLogger(CustomLogger): ######################################## # LiteLLM Virtual API KEY metrics ######################################## + # Remaining MODEL RPM limit for API Key self.litellm_remaining_api_key_requests_for_model = self._gauge_factory( "litellm_remaining_api_key_requests_for_model", diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index e2e304931a..3776d27691 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -31,15 +31,23 @@ def load_cli_token() -> Optional[dict]: return None -def get_litellm_gateway_api_key() -> Optional[str]: +def get_litellm_gateway_api_key( + expected_base_url: Optional[str] = None, +) -> Optional[str]: """ Get the stored CLI API key for use with LiteLLM SDK. This function reads the token file created by `litellm-proxy login` and returns the API key for use in Python scripts. + Args: + expected_base_url: When provided, the key is only returned if it was + originally issued for this URL. Pass the target server URL to + prevent credential leakage when the client is pointed at a + different (possibly malicious) server. + Returns: - str: The API key if found, None otherwise + str: The API key if found (and origin matches), None otherwise Example: >>> import litellm @@ -53,6 +61,10 @@ def get_litellm_gateway_api_key() -> Optional[str]: >>> ) """ token_data = load_cli_token() - if token_data and "key" in token_data: - return token_data["key"] - return None + if not token_data or "key" not in token_data: + return None + if expected_base_url is not None: + stored_url = token_data.get("base_url") + if stored_url != expected_base_url.rstrip("/"): + return None + return token_data["key"] diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index f5f28822ca..7be7085297 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -77,8 +77,8 @@ def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: if litellm_params is None: return {} - proxy_request_headers = ( - litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} - ) + proxy_request_headers = (litellm_params.get("proxy_server_request") or {}).get( + "headers" + ) or {} return proxy_request_headers diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3a83162fb2..ba840bc3d8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4582,6 +4582,11 @@ class BedrockConverseMessagesProcessor: message=cast(ChatCompletionFileObject, element) ) _parts.append(_part) + elif element["type"] == "document": + _part = BedrockConverseMessagesProcessor._process_document_message( + element + ) + _parts.append(_part) _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( message_block=cast( @@ -4864,6 +4869,44 @@ class BedrockConverseMessagesProcessor: image_url=cast(str, file_id or file_data), format=format ) + @staticmethod + def _process_document_message(element: dict) -> BedrockContentBlock: + """Convert a document content block to a Bedrock DocumentBlock. + + Handles the Anthropic-style document format: + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "..."}} + """ + source = element["source"] + source_type = source.get("type") + if source_type != "base64": + raise ValueError( + f"Bedrock Converse only supports base64-encoded document sources, got '{source_type}'. " + "Please convert the document to base64 before sending to Bedrock." + ) + media_type: str = source["media_type"] + data: str = source["data"] + doc_format = BedrockImageProcessor._validate_format( + mime_type=media_type, image_format=media_type.split("/")[1] + ) + + # Deterministic name using the same hashing pattern as _create_bedrock_block + HASH_SAMPLE_BYTES = 64 * 1024 + normalized = "".join(data.split()).encode("utf-8") + sample = normalized[:HASH_SAMPLE_BYTES] + hasher = hashlib.sha256() + hasher.update(sample) + hasher.update(str(len(normalized)).encode("utf-8")) + content_hash = hasher.hexdigest()[:16] + document_name = f"Document_{content_hash}_{doc_format}" + + return BedrockContentBlock( + document=BedrockDocumentBlock( + source=BedrockSourceBlock(bytes=data), + format=doc_format, + name=document_name, + ) + ) + @staticmethod def add_thinking_blocks_to_assistant_content( thinking_blocks: List[BedrockContentBlock], @@ -4961,6 +5004,11 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 ) ) _parts.append(_part) + elif element["type"] == "document": + _part = BedrockConverseMessagesProcessor._process_document_message( + element + ) + _parts.append(_part) _cache_point_block = ( litellm.AmazonConverseConfig()._get_cache_point_block( message_block=cast( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index e281b17268..fa7faf3035 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2244,7 +2244,7 @@ class CustomStreamWrapper: asyncio.create_task( self.logging_obj.async_failure_handler(e, traceback_exception) ) - raise e + self._handle_stream_fallback_error(e) except Exception as e: traceback_exception = traceback.format_exc() if self.logging_obj is not None: diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index d1649452f7..ff53b5808b 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -22,7 +22,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): import socket from ipaddress import ip_address, ip_network from typing import Any, List, Optional, Set, Tuple -from urllib.parse import urlparse, urlunparse +from urllib.parse import quote, urlparse, urlunparse import httpx @@ -46,6 +46,46 @@ class SSRFError(ValueError): pass +def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: + """Percent-encode one user-controlled URL path segment. + + ``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986 + unreserved characters such as ``.`` unescaped, so reject standalone dot + segments before they can be appended to an upstream URL and normalized by + the HTTP client. + """ + if value is None: + raise ValueError(f"{field_name} is required") + + value_str = str(value) + if value_str == "": + raise ValueError(f"{field_name} is required") + if value_str in {".", ".."}: + raise ValueError(f"{field_name} cannot be a dot path segment") + + return quote(value_str, safe="") + + +def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: + """Percent-encode a user-controlled URL path made of multiple segments. + + Empty segments are rejected, so leading, trailing, or consecutive slashes + fail closed instead of being normalized by the HTTP client. + """ + if value is None: + raise ValueError(f"{field_name} is required") + + value_str = str(value) + if value_str == "": + raise ValueError(f"{field_name} is required") + + encoded_segments = [] + for segment in value_str.split("/"): + encoded_segments.append(encode_url_path_segment(segment, field_name=field_name)) + + return "/".join(encoded_segments) + + def _is_blocked_ip(addr: str) -> bool: """Return True for any IP not safe to reach from a user-supplied URL. @@ -278,6 +318,47 @@ def validate_url(url: str) -> Tuple[str, str]: return rewritten, host_header +def assert_same_origin(candidate_url: str, expected_url: str) -> None: + """Verify ``candidate_url`` shares scheme, host, and port with ``expected_url``. + + Use when an upstream API returns a URL meant for follow-up requests + (e.g. an async-job polling URL that will be hit with the operator's + API key in the headers). The upstream is trusted because the operator + configured ``api_base``, but the URL it hands back must actually point + back at the same origin or we'd be blindly forwarding credentials + wherever the upstream told us to. + + Hostnames are compared case-insensitively. Default ports are made + explicit (HTTP→80, HTTPS→443) so ``https://api.example.com:443/...`` + and ``https://api.example.com/...`` are treated as the same origin. + + Error messages identify *which* component mismatched but never echo + the operator's ``expected`` host or the candidate's hostname back to + the caller — in the SSRF threat model the caller is the attacker, + and reflecting host info would be a secondary leak of operator + infrastructure details. + """ + candidate = urlparse(candidate_url) + expected = urlparse(expected_url) + + if candidate.scheme not in _ALLOWED_SCHEMES: + raise SSRFError("URL scheme is not allowed") + + if candidate.scheme != expected.scheme: + raise SSRFError("Origin mismatch on scheme") + + candidate_host = _normalize_host(candidate.hostname or "") + expected_host = _normalize_host(expected.hostname or "") + if not candidate_host or candidate_host != expected_host: + raise SSRFError("Origin mismatch on host") + + default_port = 443 if candidate.scheme == "https" else 80 + candidate_port = candidate.port if candidate.port is not None else default_port + expected_port = expected.port if expected.port is not None else default_port + if candidate_port != expected_port: + raise SSRFError("Origin mismatch on port") + + _MAX_REDIRECTS = 10 diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 3f03c744ef..fd67a7fbaf 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cas import httpx from httpx import Headers, Response +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest @@ -122,7 +123,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig): Complete URL for Anthropic batch retrieval: {api_base}/v1/messages/batches/{batch_id} """ api_base = api_base or self.anthropic_model_info.get_api_base(api_base) - return f"{api_base.rstrip('/')}/v1/messages/batches/{batch_id}" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + return f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}" def transform_retrieve_batch_request( self, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 31b9bad339..61ddf801a2 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1553,25 +1553,43 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) data["output_config"] = output_config - def _transform_response_for_json_mode( + def _resolve_json_mode_non_streaming( self, json_mode: Optional[bool], tool_calls: List[ChatCompletionToolCallChunk], - ) -> Optional[LitellmMessage]: - _message: Optional[LitellmMessage] = None - if json_mode is True and len(tool_calls) == 1: - # check if tool name is the default tool name - json_mode_content_str: Optional[str] = None - if ( - "name" in tool_calls[0]["function"] - and tool_calls[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME - ): - json_mode_content_str = tool_calls[0]["function"].get("arguments") - if json_mode_content_str is not None: - _message = AnthropicConfig._convert_tool_response_to_message( - tool_calls=tool_calls, - ) - return _message + ) -> Tuple[ + Optional[LitellmMessage], + List[ChatCompletionToolCallChunk], + Optional[str], + ]: + """Strip internal response_format tool calls; merge payload into content when mixed with user tools.""" + if json_mode is not True or not tool_calls: + return None, tool_calls, None + + json_indices = [ + i + for i, t in enumerate(tool_calls) + if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME + ] + if not json_indices: + return None, tool_calls, None + + if len(json_indices) == len(tool_calls): + json_tool = tool_calls[json_indices[0]] + if json_tool.get("function", {}).get("arguments") is None: + return None, tool_calls, None + _message = AnthropicConfig._convert_tool_response_to_message( + tool_calls=[json_tool] + ) + return _message, [], None + + first_json = tool_calls[json_indices[0]] + json_msg = AnthropicConfig._convert_tool_response_to_message([first_json]) + extra_content: Optional[str] = ( + json_msg.content if json_msg is not None else None + ) + filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices] + return None, filtered_tools, extra_content def extract_response_content(self, completion_response: dict) -> Tuple[ str, @@ -1931,19 +1949,27 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): tool_calls, ) + json_mode_message, tool_calls_for_message, json_extra_content = ( + self._resolve_json_mode_non_streaming( + json_mode=json_mode, + tool_calls=tool_calls, + ) + ) + merged_text = text_content or "" + if json_extra_content: + merged_text = ( + merged_text + json_extra_content if merged_text else json_extra_content + ) + _message = litellm.Message( - tool_calls=tool_calls, - content=text_content or None, + tool_calls=tool_calls_for_message, + content=merged_text or None, provider_specific_fields=provider_specific_fields, thinking_blocks=thinking_blocks, reasoning_content=reasoning_content, ) _message.provider_specific_fields = provider_specific_fields - json_mode_message = self._transform_response_for_json_mode( - json_mode=json_mode, - tool_calls=tool_calls, - ) if json_mode_message is not None: completion_response["stop_reason"] = "stop" _message = json_mode_message diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index c56799f30c..56296df94a 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( FileContentRequest, @@ -89,7 +90,10 @@ class AnthropicFilesHandler: raise ValueError("Missing Anthropic API Key") # Construct the Anthropic batch results URL - results_url = f"{api_base.rstrip('/')}/v1/messages/batches/{batch_id}/results" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + results_url = ( + f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}/results" + ) # Prepare headers headers = { diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index aeaab4e57b..ea9bf00f50 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -19,6 +19,7 @@ from typing import Any, Dict, List, Optional, Union, cast import httpx from openai.types.file_deleted import FileDeleted +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( @@ -185,7 +186,8 @@ class AnthropicFilesConfig(BaseFilesConfig): AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE ) - return f"{api_base.rstrip('/')}/v1/files/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} def transform_retrieve_file_response( self, @@ -206,7 +208,8 @@ class AnthropicFilesConfig(BaseFilesConfig): AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE ) - return f"{api_base.rstrip('/')}/v1/files/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} def transform_delete_file_response( self, @@ -268,7 +271,8 @@ class AnthropicFilesConfig(BaseFilesConfig): AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE ) - return f"{api_base.rstrip('/')}/v1/files/{file_id}/content", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}/content", {} def transform_file_content_response( self, diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index a992d84d45..4ea768b02a 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Tuple import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.skills.transformation import ( BaseSkillsAPIConfig, LiteLLMLoggingObj, @@ -81,7 +82,8 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): api_base = AnthropicModelInfo.get_api_base() if skill_id: - return f"{api_base}/v1/skills/{skill_id}" + encoded_skill_id = encode_url_path_segment(skill_id, field_name="skill_id") + return f"{api_base}/v1/skills/{encoded_skill_id}" return f"{api_base}/v1/{endpoint}" def transform_create_skill_request( diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 61cfd54b56..877a7d3c84 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -16,6 +16,7 @@ import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -898,6 +899,17 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): operation_location_url = response.headers["operation-location"] else: raise AzureOpenAIError(status_code=500, message=response.text) + # Reject polling URLs that don't share an origin with ``api_base``. + # Without this an upstream-controlled or attacker-controlled + # value would receive the operator's Azure API key in the + # request headers below. VERIA-51. + try: + assert_same_origin(operation_location_url, api_base) + except SSRFError as ssrf_err: + raise AzureOpenAIError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) response = await async_handler.get( url=operation_location_url, headers=headers, @@ -908,8 +920,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): timeout_secs: int = AZURE_OPERATION_POLLING_TIMEOUT start_time = time.time() if "status" not in response.json(): - raise Exception( - "Expected 'status' in response. Got={}".format(response.json()) + # Don't reflect the raw response body — when the polling + # URL points at an internal JSON API (cloud metadata + # service etc.) reflecting it here turns Blind SSRF into + # Full-Read SSRF. VERIA-51. + raise AzureOpenAIError( + status_code=502, + message="Polling response missing 'status' field", ) while response.json()["status"] not in ["succeeded", "failed"]: if time.time() - start_time > timeout_secs: @@ -1009,6 +1026,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): operation_location_url = response.headers["operation-location"] else: raise AzureOpenAIError(status_code=500, message=response.text) + try: + assert_same_origin(operation_location_url, api_base) + except SSRFError as ssrf_err: + raise AzureOpenAIError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) response = sync_handler.get( url=operation_location_url, headers=headers, @@ -1019,8 +1043,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): timeout_secs: int = AZURE_OPERATION_POLLING_TIMEOUT start_time = time.time() if "status" not in response.json(): - raise Exception( - "Expected 'status' in response. Got={}".format(response.json()) + raise AzureOpenAIError( + status_code=502, + message="Polling response missing 'status' field", ) while response.json()["status"] not in ["succeeded", "failed"]: if time.time() - start_time > timeout_secs: diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 76a6d485bc..ca9293325f 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -5,6 +5,7 @@ import httpx from openai.types.responses import ResponseReasoningItem from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import * @@ -201,7 +202,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Insert the response_id at the end of the path component # Remove trailing slash if present to avoid double slashes path = parsed_url.path.rstrip("/") - new_path = f"{path}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + new_path = f"{path}/{encoded_response_id}" # Reconstruct the URL with all original components but with the modified path constructed_url = urlunparse( @@ -322,7 +326,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Insert the response_id and /cancel at the end of the path component # Remove trailing slash if present to avoid double slashes path = parsed_url.path.rstrip("/") - new_path = f"{path}/{response_id}/cancel" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + new_path = f"{path}/{encoded_response_id}/cancel" # Reconstruct the URL with all original components but with the modified path cancel_url = urlunparse( diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index c3cd06ab4d..9bae8abce8 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -36,6 +36,7 @@ from typing import ( import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, @@ -75,20 +76,29 @@ class AzureAIAgentsHandler: def _build_messages_url( self, api_base: str, thread_id: str, api_version: str ) -> str: - return f"{api_base}/threads/{thread_id}/messages?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + return ( + f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + ) def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: - return f"{api_base}/threads/{thread_id}/runs?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + return f"{api_base}/threads/{encoded_thread_id}/runs?api-version={api_version}" def _build_run_status_url( self, api_base: str, thread_id: str, run_id: str, api_version: str ) -> str: - return f"{api_base}/threads/{thread_id}/runs/{run_id}?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + return f"{api_base}/threads/{encoded_thread_id}/runs/{encoded_run_id}?api-version={api_version}" def _build_list_messages_url( self, api_base: str, thread_id: str, api_version: str ) -> str: - return f"{api_base}/threads/{thread_id}/messages?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + return ( + f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + ) def _build_create_thread_and_run_url(self, api_base: str, api_version: str) -> str: """URL for the create-thread-and-run endpoint (supports streaming).""" diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 76c247aea8..d4144a7571 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -17,11 +17,13 @@ from urllib.parse import quote import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.constants import ( AZURE_DOCUMENT_INTELLIGENCE_API_VERSION, AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, AZURE_OPERATION_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.ocr.transformation import ( BaseOCRConfig, DocumentType, @@ -217,11 +219,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): if "/" in model: # Extract the last part after the last slash model_id = model.split("/")[-1] + encoded_model_id = encode_url_path_segment(model_id, field_name="model_id") # Azure Document Intelligence analyze endpoint # Note: API version 2024-11-30+ uses /documentintelligence/ (not /formrecognizer/) url = ( - f"{api_base}/documentintelligence/documentModels/{model_id}:analyze" + f"{api_base}/documentintelligence/documentModels/{encoded_model_id}:analyze" f"?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" ) @@ -599,6 +602,16 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "Azure Document Intelligence returned 202 but no Operation-Location header found" ) + # Reject cross-origin polling URLs — the auth headers + # below would otherwise leak to whatever URL the upstream + # (or an attacker-controlled upstream) returns. VERIA-51. + try: + assert_same_origin(operation_url, str(raw_response.request.url)) + except SSRFError as ssrf_err: + raise ValueError( + f"Azure Document Intelligence: rejected polling URL ({ssrf_err})" + ) + # Get headers for polling (need auth) poll_headers = { "Ocp-Apim-Subscription-Key": raw_response.request.headers.get( @@ -711,6 +724,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "Azure Document Intelligence returned 202 but no Operation-Location header found" ) + # Reject cross-origin polling URLs (see sync path). VERIA-51. + try: + assert_same_origin(operation_url, str(raw_response.request.url)) + except SSRFError as ssrf_err: + raise ValueError( + f"Azure Document Intelligence: rejected polling URL ({ssrf_err})" + ) + # Get headers for polling (need auth) poll_headers = { "Ocp-Apim-Subscription-Key": raw_response.request.headers.get( diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 7874201f7f..166f876ba0 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -33,6 +33,7 @@ class BaseRerankConfig(ABC): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: return {} diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index 2c7135f4d8..4c667b0ce3 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -12,6 +12,7 @@ import httpx from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -97,8 +98,15 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): agent_id, agent_alias_id = self._get_agent_id_and_alias_id(model) session_id = self._get_session_id(optional_params) + encoded_agent_id = encode_url_path_segment(agent_id, field_name="agent_id") + encoded_agent_alias_id = encode_url_path_segment( + agent_alias_id, field_name="agent_alias_id" + ) + encoded_session_id = encode_url_path_segment( + session_id, field_name="session_id" + ) - endpoint_url = f"{endpoint_url}/agents/{agent_id}/agentAliases/{agent_alias_id}/sessions/{session_id}/text" + endpoint_url = f"{endpoint_url}/agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text" return endpoint_url diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index a37af13162..c967fd334b 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -201,13 +201,14 @@ class BedrockCountTokensConfig(BaseAWSLLM): # Remove bedrock/ prefix if present if model_id.startswith("bedrock/"): model_id = model_id[8:] # Remove "bedrock/" prefix + encoded_model_id = self.encode_model_id(model_id=model_id) base_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=aws_region_name, ) - endpoint = f"{base_url}/model/{model_id}/count-tokens" + endpoint = f"{base_url}/model/{encoded_model_id}/count-tokens" return endpoint diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index f028503c6a..ec20d76102 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -5,6 +5,7 @@ from urllib.parse import urlparse import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( @@ -209,7 +210,10 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if isinstance(query, list): query = " ".join(query) - url = f"{api_base}/{vector_store_id}/retrieve" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}/retrieve" request_body: Dict[str, Any] = { "retrievalQuery": BedrockKBRetrievalQuery(text=query), diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py index dea2683a04..f5784e0836 100644 --- a/litellm/llms/black_forest_labs/image_edit/handler.py +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -15,6 +15,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -331,6 +332,17 @@ class BlackForestLabsImageEdit: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} @@ -416,6 +428,17 @@ class BlackForestLabsImageEdit: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py index 5a1d885e52..8af4a236fd 100644 --- a/litellm/llms/black_forest_labs/image_generation/handler.py +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -15,6 +15,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -317,6 +318,17 @@ class BlackForestLabsImageGeneration: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} @@ -402,6 +414,17 @@ class BlackForestLabsImageGeneration: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index a72f732a30..5b08670f9f 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import httpx +from litellm.litellm_core_utils.url_utils import encode_url_path_segments from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -149,7 +150,8 @@ class BytezChatConfig(BaseConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - return f"{API_BASE}/{model}" + encoded_model = encode_url_path_segments(model, field_name="model") + return f"{API_BASE}/{encoded_model}" def transform_request( self, diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 9e59782bf7..b9e219f5cb 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -5,6 +5,7 @@ from typing import AsyncIterator, Iterator, List, Optional, Union import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segments from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import ( BaseConfig, @@ -89,7 +90,8 @@ class CloudflareChatConfig(BaseConfig): api_base = ( f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/" ) - return api_base + model + encoded_model = encode_url_path_segments(model, field_name="model") + return api_base + encoded_model def get_supported_openai_params(self, model: str) -> List[str]: return [ diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 531b94d180..64ae8e8ffa 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -111,6 +111,7 @@ class CohereRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for Cohere rerank") diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index 60d22ff4be..4c800d6455 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -71,6 +71,7 @@ class CohereRerankV2Config(CohereRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for Cohere rerank") diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index afdd7bc6a8..599cd705eb 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Coroutine, Dict, Optional, Type, Union import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -72,7 +73,8 @@ def _build_url( # Substitute path parameters for param, value in path_params.items(): - path_template = path_template.replace(f"{{{param}}}", value) + encoded_value = encode_url_path_segment(value, field_name=param) + path_template = path_template.replace(f"{{{param}}}", encoded_value) # Parse the api_base to extract existing query params parsed_base = httpx.URL(api_base) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a34b73b531..0c4816fcda 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -26,6 +26,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -1007,6 +1008,7 @@ class BaseLLMHTTPHandler: api_key: Optional[str] = None, api_base: Optional[str] = None, client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + litellm_params: Optional[Dict[str, Any]] = None, ) -> RerankResponse: # get config from model, custom llm provider headers = provider_config.validate_environment( @@ -1026,6 +1028,7 @@ class BaseLLMHTTPHandler: model=model, optional_rerank_params=optional_rerank_params, headers=headers, + litellm_params=litellm_params, ) ## LOGGING @@ -2535,10 +2538,16 @@ class BaseLLMHTTPHandler: }, ) + delete_kwargs: Dict[str, Any] = { + "url": url, + "headers": headers, + "timeout": timeout, + } + if data: + delete_kwargs["json"] = data + try: - response = await async_httpx_client.delete( - url=url, headers=headers, json=data, timeout=timeout - ) + response = await async_httpx_client.delete(**delete_kwargs) except Exception as e: raise self._handle_error( @@ -2619,10 +2628,16 @@ class BaseLLMHTTPHandler: }, ) + delete_kwargs: Dict[str, Any] = { + "url": url, + "headers": headers, + "timeout": timeout, + } + if data: + delete_kwargs["json"] = data + try: - response = sync_httpx_client.delete( - url=url, headers=headers, json=data, timeout=timeout - ) + response = sync_httpx_client.delete(**delete_kwargs) except Exception as e: raise self._handle_error( @@ -8934,7 +8949,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", @@ -9001,7 +9019,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", @@ -9200,7 +9221,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" request_body: Dict[str, Any] = dict(vector_store_update_optional_params) @@ -9283,7 +9307,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" request_body: Dict[str, Any] = dict(vector_store_update_optional_params) @@ -9349,7 +9376,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", @@ -9414,7 +9444,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 71e300d258..276735f475 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -132,6 +132,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: # Convert OptionalRerankParams to dict as expected by parent class if optional_rerank_params is None: diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 4dac2b8ba9..6a59911701 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -11,13 +11,14 @@ import httpx from httpx import Headers import litellm -from litellm.types.utils import all_litellm_params +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, ) from litellm.secret_managers.main import get_secret_str +from litellm.types.utils import all_litellm_params from ..common_utils import ElevenLabsException @@ -321,7 +322,8 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): "ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`." ) - url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{voice_id}" + encoded_voice_id = encode_url_path_segment(voice_id, field_name="voice_id") + url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}" query_params = litellm_params.get(self.ELEVENLABS_QUERY_PARAMS_KEY, {}) if query_params: diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index eb92399a05..4a7b64b9b7 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -127,6 +127,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request to Fireworks AI rerank format diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 3804145eb5..00822da43e 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -6,14 +6,17 @@ For vertex ai, check out the vertex_ai/files/handler.py file. import time from typing import Any, List, Literal, Optional -from urllib.parse import unquote, urlparse +from urllib.parse import urlparse import httpx from openai.types.file_deleted import FileDeleted from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host +from litellm.litellm_core_utils.url_utils import ( + encode_url_path_segment, + is_url_destination_allowed_by_host, +) from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, @@ -268,11 +271,14 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): normalized_file_id = file_id normalized_file_id = normalized_file_id.strip("/") - if not normalized_file_id.startswith("files/"): - normalized_file_id = f"files/{normalized_file_id}" - self._validate_gemini_file_name(normalized_file_id) + if normalized_file_id.startswith("files/"): + normalized_file_id = normalized_file_id.removeprefix("files/") - return normalized_file_id + encoded_file_id = encode_url_path_segment( + normalized_file_id, field_name="file_id" + ) + + return f"files/{encoded_file_id}" @staticmethod def _is_allowed_gemini_file_url( @@ -288,26 +294,6 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): ) return is_url_destination_allowed_by_host(file_url, allowed_hosts) - @staticmethod - def _validate_gemini_file_name(file_name: str) -> None: - parts = file_name.split("/") - decoded_file_id = "" - if len(parts) == 2: - decoded_file_id = parts[1] - while True: - next_decoded_file_id = unquote(decoded_file_id) - if next_decoded_file_id == decoded_file_id: - break - decoded_file_id = next_decoded_file_id - if ( - len(parts) != 2 - or parts[0] != "files" - or not parts[1] - or decoded_file_id in {".", ".."} - or any(char in decoded_file_id for char in ("/", "\\", "?", "#")) - ): - raise ValueError("Invalid Gemini file name") - def transform_retrieve_file_response( self, raw_response: httpx.Response, diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index c34da83cb8..593cbf7c2c 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -15,6 +15,7 @@ import httpx from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.interactions import ( @@ -205,8 +206,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") + encoded_interaction_id = encode_url_path_segment( + interaction_id, field_name="interaction_id" + ) return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", {}, ) @@ -238,8 +242,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") + encoded_interaction_id = encode_url_path_segment( + interaction_id, field_name="interaction_id" + ) return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", {}, ) @@ -268,8 +275,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") + encoded_interaction_id = encode_url_path_segment( + interaction_id, field_name="interaction_id" + ) return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}:cancel", {}, ) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 8066e53afc..60b6dc7d23 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -121,6 +121,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for Hosted VLLM rerank") diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 3f83b8e422..2c847b617e 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -146,6 +146,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Union[OptionalRerankParams, dict], headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: if "query" not in optional_rerank_params: raise ValueError("query is required for HuggingFace rerank") diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 48d876f8ea..ad4416925a 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -74,7 +74,11 @@ class JinaAIRerankConfig(BaseRerankConfig): return cleaned_base def transform_rerank_request( - self, model: str, optional_rerank_params: Dict, headers: Dict + self, + model: str, + optional_rerank_params: Dict, + headers: Dict, + litellm_params: Optional[dict] = None, ) -> Dict: return {"model": model, **optional_rerank_params} diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index 3381a5327e..3416616139 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -18,6 +18,7 @@ from openai.types.file_deleted import FileDeleted import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( @@ -306,7 +307,8 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_retrieve_file_response( self, @@ -336,7 +338,8 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_delete_file_response( self, @@ -422,7 +425,8 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - return f"{api_base}/{file_id}/content", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}/content", {} def transform_file_content_response( self, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index 510c41304a..b3a0073a5c 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -6,6 +6,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -270,7 +271,10 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): Reference: https://open.manus.im/docs/openai-compatibility """ - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 757d874bf3..b9a46b8ac2 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -66,6 +66,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request, using clean model name without 'ranking/' prefix. @@ -75,4 +76,5 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): model=clean_model, optional_rerank_params=optional_rerank_params, headers=headers, + litellm_params=litellm_params, ) diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index bd5abac60c..fc317293ac 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -177,6 +177,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request to Nvidia NIM format. diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 955b9f760d..7f874ffd3b 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -6,6 +6,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, @@ -198,7 +199,10 @@ class OpenAIContainerConfig(BaseContainerConfig): ) -> Tuple[str, Dict]: """Transform the OpenAI container retrieve request.""" # For container retrieve, we just need to construct the URL - url = join_container_api_base_path(api_base, f"/{container_id}") + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request data: Dict[str, Any] = {} @@ -230,7 +234,10 @@ class OpenAIContainerConfig(BaseContainerConfig): - DELETE /v1/containers/{container_id} """ # Construct the URL for container delete - url = join_container_api_base_path(api_base, f"/{container_id}") + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request data: Dict[str, Any] = {} @@ -267,7 +274,10 @@ class OpenAIContainerConfig(BaseContainerConfig): - GET /v1/containers/{container_id}/files """ # Construct the URL for container files - url = join_container_api_base_path(api_base, f"/{container_id}/files") + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters params: Dict[str, Any] = {} @@ -311,8 +321,12 @@ class OpenAIContainerConfig(BaseContainerConfig): - GET /v1/containers/{container_id}/files/{file_id}/content """ # Construct the URL for container file content + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") url = join_container_api_base_path( - api_base, f"/{container_id}/files/{file_id}/content" + api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content" ) # No query parameters needed diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index c24dbf8637..66537e56a6 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Tuple import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.evals.transformation import ( BaseEvalsAPIConfig, LiteLLMLoggingObj, @@ -76,7 +77,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base = "https://api.openai.com" if eval_id: - return f"{api_base}/v1/evals/{eval_id}" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + return f"{api_base}/v1/evals/{encoded_eval_id}" return f"{api_base}/v1/{endpoint}" def transform_create_eval_request( @@ -276,7 +278,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if litellm_params and litellm_params.api_base: api_base = litellm_params.api_base - url = f"{api_base}/v1/evals/{eval_id}/runs" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs" # Build request body request_body = {k: v for k, v in create_request.items() if v is not None} @@ -310,7 +313,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if litellm_params and litellm_params.api_base: api_base = litellm_params.api_base - url = f"{api_base}/v1/evals/{eval_id}/runs" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs" # Build query parameters query_params: Dict[str, Any] = {} @@ -350,7 +354,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform get run request for OpenAI""" - url = f"{api_base}/v1/evals/{eval_id}/runs/{run_id}" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" verbose_logger.debug("Get run request - URL: %s", url) @@ -376,7 +382,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict, Dict]: """Transform cancel run request for OpenAI""" - url = f"{api_base}/v1/evals/{eval_id}/runs/{run_id}/cancel" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel" # Empty body for cancel request request_body: Dict[str, Any] = {} @@ -405,7 +413,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict, Dict]: """Transform delete run request for OpenAI""" - url = f"{api_base}/v1/evals/{eval_id}/runs/{run_id}" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" # Empty body for delete request request_body: Dict[str, Any] = {} diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 87c502032c..b7d5340d8d 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -7,6 +7,7 @@ from pydantic import BaseModel, ValidationError import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -421,7 +422,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - DELETE /v1/responses/{response_id} """ - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -457,7 +461,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - GET /v1/responses/{response_id} """ - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -498,7 +505,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}/input_items" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/input_items" params: Dict[str, Any] = {} if after is not None: params["after"] = after @@ -540,7 +550,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - POST /v1/responses/{response_id}/cancel """ - url = f"{api_base}/{response_id}/cancel" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/cancel" data: Dict = {} return url, data diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index cd5f10251b..52202f57fd 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -3,6 +3,7 @@ from typing import Any, Dict, Optional, Tuple, cast import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) @@ -98,7 +99,10 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): or "https://api.openai.com/v1" ) base_url = base_url.rstrip("/") - return f"{base_url}/vector_stores/{vector_store_id}/files" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + return f"{base_url}/vector_stores/{encoded_vector_store_id}/files" def transform_create_vector_store_file_request( self, @@ -163,7 +167,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, api_base: str, ) -> Tuple[str, Dict[str, Any]]: - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_retrieve_vector_store_file_response( self, @@ -186,7 +191,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, api_base: str, ) -> Tuple[str, Dict[str, Any]]: - return f"{api_base}/{file_id}/content", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}/content", {} def transform_retrieve_vector_store_file_content_response( self, @@ -218,7 +224,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): payload["attributes"] = filtered_attributes else: payload.pop("attributes", None) - return f"{api_base}/{file_id}", payload + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", payload def transform_update_vector_store_file_response( self, @@ -241,7 +248,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, api_base: str, ) -> Tuple[str, Dict[str, Any]]: - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_delete_vector_store_file_response( self, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 2c11d13748..bd095a0a1b 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -108,7 +109,10 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: - url = f"{api_base}/{vector_store_id}/search" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( query=query, filters=vector_store_search_optional_params.get("filters", None), diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 61baa56949..2d165a7d7d 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -1,11 +1,13 @@ import mimetypes from io import BufferedReader, BytesIO from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from urllib.parse import quote import httpx from httpx._types import RequestFiles import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils from litellm.secret_managers.main import get_secret_str @@ -220,11 +222,18 @@ class OpenAIVideoConfig(BaseVideoConfig): - GET /v1/videos/{video_id}/content?variant=thumbnail """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for video content download - url = f"{api_base.rstrip('/')}/{original_video_id}/content" + url = f"{api_base.rstrip('/')}/{encoded_video_id}/content" if variant is not None: - url = f"{url}?variant={variant}" + # Encode the user-controlled ``variant`` so a value like + # ``thumbnail&extra=1`` cannot inject additional query params + # into the upstream request — same hardening rationale as the + # path-segment encoding above. + url = f"{url}?variant={quote(variant, safe='')}" # No additional data needed for GET content request data: Dict[str, Any] = {} @@ -247,9 +256,12 @@ class OpenAIVideoConfig(BaseVideoConfig): - POST /v1/videos/{video_id}/remix """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for video remix - url = f"{api_base.rstrip('/')}/{original_video_id}/remix" + url = f"{api_base.rstrip('/')}/{encoded_video_id}/remix" # Prepare the request data data = {"prompt": prompt} @@ -391,9 +403,12 @@ class OpenAIVideoConfig(BaseVideoConfig): - DELETE /v1/videos/{video_id} """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for video delete - url = f"{api_base.rstrip('/')}/{original_video_id}" + url = f"{api_base.rstrip('/')}/{encoded_video_id}" # No data needed for DELETE request data: Dict[str, Any] = {} @@ -427,9 +442,12 @@ class OpenAIVideoConfig(BaseVideoConfig): """ # Extract the original video_id (remove provider encoding if present) original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # For video retrieve, we just need to construct the URL - url = f"{api_base.rstrip('/')}/{original_video_id}" + url = f"{api_base.rstrip('/')}/{encoded_video_id}" # No additional data needed for GET request data: Dict[str, Any] = {} @@ -494,7 +512,11 @@ class OpenAIVideoConfig(BaseVideoConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base.rstrip('/')}/characters/{character_id}" + original_character_id = extract_original_character_id(character_id) + encoded_character_id = encode_url_path_segment( + original_character_id, field_name="character_id" + ) + url = f"{api_base.rstrip('/')}/characters/{encoded_character_id}" return url, {} def transform_video_get_character_response( diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 7b22edd867..fc4cfc7b08 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -1,5 +1,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -82,7 +83,10 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: - url = f"{api_base}/{vector_store_id}/search" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py index d49a5fd370..990fc2b2e6 100644 --- a/litellm/llms/ragflow/chat/transformation.py +++ b/litellm/llms/ragflow/chat/transformation.py @@ -13,6 +13,7 @@ Model name format: from typing import List, Optional, Tuple import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.openai.openai import OpenAIConfig from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -126,10 +127,11 @@ class RAGFlowConfig(OpenAIConfig): api_base = api_base[:-3] # Remove /v1 # Construct the RAGFlow-specific path + encoded_entity_id = encode_url_path_segment(entity_id, field_name="entity_id") if endpoint_type == "chat": - path = f"/api/v1/chats_openai/{entity_id}/chat/completions" + path = f"/api/v1/chats_openai/{encoded_entity_id}/chat/completions" else: # agent - path = f"/api/v1/agents_openai/{entity_id}/chat/completions" + path = f"/api/v1/agents_openai/{encoded_entity_id}/chat/completions" # Ensure path starts with / if not path.startswith("/"): diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 8377dea952..4f84816a2b 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -6,6 +6,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -334,9 +335,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): We'll retrieve the task and extract the video URL. """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Get task status to retrieve video URL - url = f"{api_base}/tasks/{original_video_id}" + url = f"{api_base}/tasks/{encoded_video_id}" params: Dict[str, Any] = {} @@ -495,9 +499,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): RunwayML uses task cancellation. """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for task cancellation - url = f"{api_base}/tasks/{original_video_id}/cancel" + url = f"{api_base}/tasks/{encoded_video_id}/cancel" data: Dict[str, Any] = {} @@ -533,9 +540,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): RunwayML uses GET /v1/tasks/{task_id} to retrieve task status. """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the full URL for task status retrieval - url = f"{api_base}/tasks/{original_video_id}" + url = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) data: Dict[str, Any] = {} diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7436bfef58..c627599da8 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -4,7 +4,11 @@ from typing import Any, Coroutine, Dict, Optional, Union import httpx import litellm -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import ( + async_safe_get, + encode_url_path_segment, + safe_get, +) from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -170,7 +174,8 @@ class VertexAIBatchPrediction(VertexLLM): ) # Append batch_id to the URL - default_api_base = f"{default_api_base}/{batch_id}" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + default_api_base = f"{default_api_base}/{encoded_batch_id}" if len(default_api_base.split(":")) > 1: endpoint = default_api_base.split(":")[-1] @@ -413,7 +418,8 @@ class VertexAIBatchPrediction(VertexLLM): vertex_project=vertex_project or project_id, ) - retrieve_api_base_default = f"{default_api_base}/{batch_id}" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + retrieve_api_base_default = f"{default_api_base}/{encoded_batch_id}" cancel_api_base_default = f"{retrieve_api_base_default}:cancel" _, api_base = self._check_custom_proxy( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index b4bfde5f54..c72160f7d0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -27,6 +27,53 @@ class VertexAIError(BaseLLMException): super().__init__(message=message, status_code=status_code, headers=headers) +def vertex_request_labels_from_litellm_params( + litellm_params: Optional[dict], +) -> Optional[Dict[str, str]]: + """ + Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``: + ``metadata`` (``completion(..., metadata=...)``) or ``litellm_metadata``, + using ``requester_metadata`` string key-value pairs (same convention as Gemini). + ``metadata`` is tried first when both are present. + """ + if not litellm_params: + return None + for key in ("metadata", "litellm_metadata"): + if key not in litellm_params: + continue + metadata = litellm_params[key] + if metadata is None or not isinstance(metadata, dict): + continue + if "requester_metadata" not in metadata: + continue + rm = metadata["requester_metadata"] + if not isinstance(rm, dict): + continue + labels = {k: v for k, v in rm.items() if isinstance(v, str)} + if labels: + return labels + return None + + +def pop_vertex_request_labels( + optional_params: Optional[dict], + litellm_params: Optional[dict], +) -> Optional[Dict[str, str]]: + """ + Resolve labels from optional ``labels`` (Gemini-style) and/or + ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]`` + (``requester_metadata``). Pops ``labels`` from optional_params when present. + """ + labels: Optional[Dict[str, str]] = None + if optional_params is not None and "labels" in optional_params: + raw = optional_params.pop("labels") + if isinstance(raw, dict): + labels = {k: v for k, v in raw.items() if isinstance(v, str)} + if not labels: + labels = vertex_request_labels_from_litellm_params(litellm_params) + return labels if labels else None + + class VertexAIModelRoute(str, Enum): """Enum for Vertex AI model routing""" @@ -50,7 +97,7 @@ def get_vertex_ai_model_route( Determine which handler to use for a Vertex AI model based on the model name. Args: - model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "openai/gpt-oss-120b") + model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "xai/grok-4.1-fast-non-reasoning") litellm_params: Optional litellm parameters dict that may contain base_model for routing Returns: @@ -66,7 +113,7 @@ def get_vertex_ai_model_route( >>> get_vertex_ai_model_route("gemma/gemma-3-12b-it") VertexAIModelRoute.GEMMA - >>> get_vertex_ai_model_route("openai/gpt-oss-120b") + >>> get_vertex_ai_model_route("xai/grok-4.1-fast-non-reasoning") VertexAIModelRoute.MODEL_GARDEN >>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"}) @@ -102,8 +149,11 @@ def get_vertex_ai_model_route( if "gemma/" in model: return VertexAIModelRoute.GEMMA - # Check for model garden openai models - if "openai" in model: + # Check for model garden OpenAI-compatible publisher models. + # Examples: + # - openai/gpt-oss-120b-maas + # - xai/grok-4.1-fast-non-reasoning + if "openai" in model or model.startswith("xai/"): return VertexAIModelRoute.MODEL_GARDEN # Check for gemini models @@ -209,8 +259,8 @@ def get_vertex_base_model_name(model: str) -> str: >>> get_vertex_base_model_name("gemma/gemma-3-12b-it") "gemma-3-12b-it" - >>> get_vertex_base_model_name("openai/gpt-oss-120b") - "gpt-oss-120b" + >>> get_vertex_base_model_name("xai/grok-4.1-fast-non-reasoning") + "grok-4.1-fast-non-reasoning" >>> get_vertex_base_model_name("1234567890") "1234567890" diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 533bd06d2d..87bd484382 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( response_schema_prompt, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.files import ( get_file_mime_type_for_file_type, get_file_type_from_extension, @@ -714,16 +715,8 @@ def _transform_request_body( # noqa: PLR0915 optional_params.pop("output_config", None) config_fields = GenerationConfig.__annotations__.keys() - # If the LiteLLM client sends Gemini-supported parameter "labels", add it - # as "labels" field to the request sent to the Gemini backend. - labels: Optional[dict[str, str]] = optional_params.pop("labels", None) - # If the LiteLLM client sends OpenAI-supported parameter "metadata", add it - # as "labels" field to the request sent to the Gemini backend. - if labels is None and "metadata" in litellm_params: - metadata = litellm_params["metadata"] - if metadata is not None and "requester_metadata" in metadata: - rm = metadata["requester_metadata"] - labels = {k: v for k, v in rm.items() if isinstance(v, str)} + # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata) + labels = pop_vertex_request_labels(optional_params, litellm_params) filtered_params = { k: v diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 2371bc4865..99165c37c9 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import Any, Dict, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Tuple, Union import httpx @@ -13,8 +13,8 @@ from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, ) -from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( + GeminiEmbeddingInput, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, ) @@ -23,7 +23,6 @@ from litellm.types.utils import EmbeddingResponse from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM from .batch_embed_content_transformation import ( _is_file_reference, - _is_multimodal_input, process_embed_content_response, process_response, transform_openai_input_gemini_content, @@ -32,9 +31,24 @@ from .batch_embed_content_transformation import ( class GoogleBatchEmbeddings(VertexLLM): + @staticmethod + def _flatten_and_detect_file_refs( + input: GeminiEmbeddingInput, + ) -> Tuple[List[str], bool]: + """Flatten nested input lists and detect file references.""" + input_list = [input] if isinstance(input, str) else input + flat_elements = [ + e + for item in input_list + for e in (item if isinstance(item, list) else [item]) + if isinstance(e, str) + ] + has_file_refs = any(_is_file_reference(e) for e in flat_elements) + return flat_elements, has_file_refs + def _resolve_file_references( self, - input: EmbeddingInput, + input: GeminiEmbeddingInput, api_key: str, sync_handler: HTTPHandler, ) -> Dict[str, Dict[str, str]]: @@ -42,7 +56,7 @@ class GoogleBatchEmbeddings(VertexLLM): Resolve Gemini file references (files/...) to get mime_type and uri. Args: - input: EmbeddingInput that may contain file references + input: GeminiEmbeddingInput that may contain file references api_key: Gemini API key sync_handler: HTTP client @@ -73,7 +87,7 @@ class GoogleBatchEmbeddings(VertexLLM): async def _async_resolve_file_references( self, - input: EmbeddingInput, + input: GeminiEmbeddingInput, api_key: str, async_handler: AsyncHTTPHandler, ) -> Dict[str, Dict[str, str]]: @@ -81,7 +95,7 @@ class GoogleBatchEmbeddings(VertexLLM): Async version of _resolve_file_references. Args: - input: EmbeddingInput that may contain file references + input: GeminiEmbeddingInput that may contain file references api_key: Gemini API key async_handler: Async HTTP client @@ -110,10 +124,10 @@ class GoogleBatchEmbeddings(VertexLLM): return resolved_files - def batch_embeddings( + def batch_embeddings( # noqa: PLR0915 self, model: str, - input: EmbeddingInput, + input: GeminiEmbeddingInput, print_verbose, model_response: EmbeddingResponse, custom_llm_provider: Literal["gemini", "vertex_ai"], @@ -151,8 +165,7 @@ class GoogleBatchEmbeddings(VertexLLM): optional_params = optional_params or {} - is_multimodal = _is_multimodal_input(input) - use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai") + use_embed_content = custom_llm_provider == "vertex_ai" mode: Literal["embedding", "batch_embedding"] if use_embed_content: mode = "embedding" @@ -215,8 +228,22 @@ class GoogleBatchEmbeddings(VertexLLM): resolved_files=resolved_files, ) else: + flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input) + if has_file_refs and not api_key: + raise ValueError( + "An API key is required to resolve Gemini file references (files/...). " + "Pass api_key= or set GEMINI_API_KEY." + ) + resolved_files = {} + if api_key and has_file_refs: + resolved_files = self._resolve_file_references( + input=flat_elements, api_key=api_key, sync_handler=sync_handler + ) request_data = transform_openai_input_gemini_content( - input=input, model=model, optional_params=optional_params + input=input, + model=model, + optional_params=optional_params, + resolved_files=resolved_files, ) ## LOGGING @@ -264,7 +291,7 @@ class GoogleBatchEmbeddings(VertexLLM): url: str, data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]], model_response: EmbeddingResponse, - input: EmbeddingInput, + input: GeminiEmbeddingInput, timeout: Optional[Union[float, httpx.Timeout]], headers={}, client: Optional[AsyncHTTPHandler] = None, @@ -303,8 +330,22 @@ class GoogleBatchEmbeddings(VertexLLM): resolved_files=resolved_files, ) else: + flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input) + if has_file_refs and not api_key: + raise ValueError( + "An API key is required to resolve Gemini file references (files/...). " + "Pass api_key= or set GEMINI_API_KEY." + ) + resolved_files = {} + if api_key and has_file_refs: + resolved_files = await self._async_resolve_file_references( + input=flat_elements, api_key=api_key, async_handler=async_handler + ) data = transform_openai_input_gemini_content( - input=input, model=model, optional_params=optional_params or {} + input=input, + model=model, + optional_params=optional_params or {}, + resolved_files=resolved_files, ) ## LOGGING diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 34fc95e0af..e1b365c9f4 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -6,12 +6,12 @@ Why separate file? Make it easy to see how transformation works from typing import Dict, List, Optional, Tuple -from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( BlobType, ContentType, EmbedContentRequest, FileDataType, + GeminiEmbeddingInput, PartType, VertexAIBatchEmbeddingsRequestBody, VertexAIBatchEmbeddingsResponseObject, @@ -114,33 +114,77 @@ def _parse_data_url(data_url: str) -> Tuple[str, str]: return media_type, base64_data -def _is_multimodal_input(input: EmbeddingInput) -> bool: +def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool: """ - Check if the input contains multimodal data (data URIs, file references, or GCS URLs). + Check if the input contains multimodal data (data URIs, file references, + GCS URLs, or nested lists for combined embeddings). Args: - input: EmbeddingInput (str or List[str]) + input: GeminiEmbeddingInput — str, List[str], or List[List[str]] for combined embeddings Returns: - bool: True if any element is a data URI, file reference, or GCS URL + bool: True if any element is multimodal or a nested list """ if isinstance(input, str): - input_list = [input] - else: - input_list = input + return _is_multimodal_element(input) - for element in input_list: - if isinstance(element, str): - if element.startswith("data:") and ";base64," in element: - return True - if _is_file_reference(element): - return True - if _is_gcs_url(element): + for element in input: + if isinstance(element, list): + if any( + _is_multimodal_element(sub) for sub in element if isinstance(sub, str) + ): return True + elif isinstance(element, str) and _is_multimodal_element(element): + return True return False +def _is_multimodal_element(element: str) -> bool: + """Check if a single string element is multimodal.""" + if element.startswith("data:") and ";base64," in element: + return True + if _is_file_reference(element): + return True + if _is_gcs_url(element): + return True + return False + + +def _build_part_for_input( + element: str, + resolved_files: Optional[Dict[str, Dict[str, str]]] = None, +) -> PartType: + """ + Build a single PartType for an input element, handling text, data URIs, + file references, and GCS URLs. + """ + resolved_files = resolved_files or {} + + if element.startswith("data:") and ";base64," in element: + mime_type, base64_data = _parse_data_url(element) + blob: BlobType = {"mime_type": mime_type, "data": base64_data} + return PartType(inline_data=blob) + elif _is_gcs_url(element): + mime_type = _infer_mime_type_from_gcs_url(element) + file_data: FileDataType = { + "mime_type": mime_type, + "file_uri": element, + } + return PartType(file_data=file_data) + elif _is_file_reference(element): + if element not in resolved_files: + raise ValueError(f"File reference {element} not resolved") + file_info = resolved_files[element] + file_data_ref: FileDataType = { + "mime_type": file_info["mime_type"], + "file_uri": file_info["uri"], + } + return PartType(file_data=file_data_ref) + else: + return PartType(text=element) + + _SUPPORTED_EMBED_PARAMS = {"outputDimensionality", "taskType", "title"} @@ -155,37 +199,60 @@ def _filter_embed_params(optional_params: dict) -> dict: def transform_openai_input_gemini_content( - input: EmbeddingInput, model: str, optional_params: dict + input: GeminiEmbeddingInput, + model: str, + optional_params: dict, + resolved_files: Optional[Dict[str, Dict[str, str]]] = None, ) -> VertexAIBatchEmbeddingsRequestBody: """ - The content to embed. Only the parts.text fields will be counted. + Transform OpenAI embedding input to Gemini batchEmbedContents format. + + Each input element becomes a separate EmbedContentRequest, supporting + text, data URIs, file references, and GCS URLs. + + If an element is a list (nested input), all sub-elements are combined + into a single content with multiple parts, producing one combined + embedding for the group. + + Examples: + input=["text", "image"] → 2 separate embeddings + input=[["text", "image"]] → 1 combined embedding + input=[["text", "image"], "x"] → 2 embeddings (1 combined + 1 separate) """ gemini_model_name = "models/{}".format(model) gemini_params = _filter_embed_params(optional_params) + input_list = [input] if isinstance(input, str) else input requests: List[EmbedContentRequest] = [] - if isinstance(input, str): + + for element in input_list: + if isinstance(element, list): + if not element: + raise ValueError("Nested input list must not be empty") + for sub in element: + if not isinstance(sub, str): + raise ValueError( + f"Elements inside a nested input list must be strings, got {type(sub)}" + ) + parts = [ + _build_part_for_input(sub, resolved_files=resolved_files) + for sub in element + ] + else: + parts = [_build_part_for_input(element, resolved_files=resolved_files)] request = EmbedContentRequest( model=gemini_model_name, - content=ContentType(parts=[PartType(text=input)]), + content=ContentType(parts=parts), **gemini_params, ) requests.append(request) - else: - for i in input: - request = EmbedContentRequest( - model=gemini_model_name, - content=ContentType(parts=[PartType(text=i)]), - **gemini_params, - ) - requests.append(request) return VertexAIBatchEmbeddingsRequestBody(requests=requests) def transform_openai_input_gemini_embed_content( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model: str, optional_params: dict, resolved_files: Optional[Dict[str, Dict[str, str]]] = None, @@ -194,7 +261,7 @@ def transform_openai_input_gemini_embed_content( Transform OpenAI embedding input to Gemini embedContent format (multimodal). Args: - input: EmbeddingInput (str or List[str]) with text, data URIs, or file references + input: GeminiEmbeddingInput with text, data URIs, or file references model: Model name optional_params: Additional parameters (taskType, outputDimensionality, etc.) resolved_files: Dict mapping file names (files/abc) to {mime_type, uri} @@ -210,31 +277,14 @@ def transform_openai_input_gemini_embed_content( parts: List[PartType] = [] for element in input_list: + if isinstance(element, list): + raise ValueError( + "Nested (combined) embeddings are not supported on the embedContent path. " + "Use the batchEmbedContents path or pass a flat list instead." + ) if not isinstance(element, str): raise ValueError(f"Unsupported input type: {type(element)}") - - if element.startswith("data:") and ";base64," in element: - mime_type, base64_data = _parse_data_url(element) - blob: BlobType = {"mime_type": mime_type, "data": base64_data} - parts.append(PartType(inline_data=blob)) - elif _is_gcs_url(element): - mime_type = _infer_mime_type_from_gcs_url(element) - file_data: FileDataType = { - "mime_type": mime_type, - "file_uri": element, - } - parts.append(PartType(file_data=file_data)) - elif _is_file_reference(element): - if element not in resolved_files: - raise ValueError(f"File reference {element} not resolved") - file_info = resolved_files[element] - file_data_ref: FileDataType = { - "mime_type": file_info["mime_type"], - "file_uri": file_info["uri"], - } - parts.append(PartType(file_data=file_data_ref)) - else: - parts.append(PartType(text=element)) + parts.append(_build_part_for_input(element, resolved_files=resolved_files)) request_body: dict = { "content": ContentType(parts=parts), @@ -245,7 +295,7 @@ def transform_openai_input_gemini_embed_content( def process_embed_content_response( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, response_json: dict, @@ -291,7 +341,7 @@ def process_embed_content_response( def process_response( - input: EmbeddingInput, + input: GeminiEmbeddingInput, model_response: EmbeddingResponse, model: str, _predictions: VertexAIBatchEmbeddingsResponseObject, @@ -308,8 +358,29 @@ def process_response( model_response.data = openai_embeddings model_response.model = model - input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") - prompt_tokens = token_counter(model=model, text=input_text) + has_nested = isinstance(input, list) and any(isinstance(e, list) for e in input) + if _is_multimodal_input(input) or has_nested: + input_list = input if isinstance(input, list) else [input] + text_elements: List[str] = [] + for e in input_list: + if isinstance(e, list): + text_elements.extend( + sub + for sub in e + if isinstance(sub, str) and not _is_multimodal_element(sub) + ) + elif isinstance(e, str) and not _is_multimodal_element(e): + text_elements.append(e) + if text_elements: + input_text = get_formatted_prompt( + data={"input": text_elements}, call_type="embedding" + ) + prompt_tokens = token_counter(model=model, text=input_text) + else: + prompt_tokens = 0 + else: + input_text = get_formatted_prompt(data={"input": input}, call_type="embedding") + prompt_tokens = token_counter(model=model, text=input_text) model_response.usage = Usage( prompt_tokens=prompt_tokens, total_tokens=prompt_tokens ) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 1c7696d55a..05ebd685d9 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -7,7 +7,10 @@ import litellm from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import ( + get_vertex_base_url, + pop_vertex_request_labels, +) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -203,13 +206,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): "sampleCount": 1, } - # Merge with optional params + labels = pop_vertex_request_labels(optional_params, litellm_params) + # Merge with optional params (after popping labels so they are not sent as Imagen parameters) parameters = {**default_params, **optional_params} - request_body = { + request_body: dict = { "instances": [{"prompt": prompt}], "parameters": parameters, } + if labels: + request_body["labels"] = labels return request_body diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 5365183967..3b84972e94 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -11,12 +11,15 @@ import httpx import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.llms.vertex_ai.common_utils import ( + vertex_request_labels_from_litellm_params, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( + RerankBilledUnits, RerankResponse, RerankResponseMeta, - RerankBilledUnits, RerankResponseResult, ) @@ -109,6 +112,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform the request from Cohere format to Vertex AI Discovery Engine format @@ -145,6 +149,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): # When return_documents is False, we want to ignore record details (return only IDs) request_data["ignoreRecordDetailsInResponse"] = not return_documents + user_labels = vertex_request_labels_from_litellm_params(litellm_params) + if user_labels: + request_data["userLabels"] = user_labels + return request_data def transform_rerank_response( diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 6cb7a86bea..61fb848b40 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import httpx from litellm import get_model_info +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.router import GenericLiteLLMParams @@ -91,12 +92,18 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): raise ValueError("vector_store_id is required") if api_base: return api_base.rstrip("/") + encoded_collection_id = encode_url_path_segment( + collection_id, field_name="vertex_collection_id" + ) + encoded_datastore_id = encode_url_path_segment( + datastore_id, field_name="vector_store_id" + ) # Vertex AI Search API endpoint for search return ( f"https://discoveryengine.googleapis.com/v1/" f"projects/{vertex_project}/locations/{vertex_location}/" - f"collections/{collection_id}/dataStores/{datastore_id}/servingConfigs/default_config" + f"collections/{encoded_collection_id}/dataStores/{encoded_datastore_id}/servingConfigs/default_config" ) def transform_search_vector_store_request( diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index 18c5ec3d83..696341598e 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional, Union +from typing import Dict, Literal, Optional, Union import httpx @@ -44,6 +44,7 @@ class VertexEmbedding(VertexBase): vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = None, gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, + litellm_params: Optional[Dict] = None, ) -> EmbeddingResponse: if aembedding is True: return self.async_embedding( # type: ignore @@ -61,6 +62,7 @@ class VertexEmbedding(VertexBase): vertex_credentials=vertex_credentials, gemini_api_key=gemini_api_key, extra_headers=extra_headers, + litellm_params=litellm_params, ) should_use_v1beta1_features = self.is_using_v1beta1_features( @@ -92,7 +94,10 @@ class VertexEmbedding(VertexBase): headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) vertex_request: VertexEmbeddingRequest = ( litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, optional_params=optional_params, model=model + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, ) ) @@ -156,6 +161,7 @@ class VertexEmbedding(VertexBase): gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, encoding=None, + litellm_params: Optional[Dict] = None, ) -> EmbeddingResponse: """ Async embedding implementation @@ -188,7 +194,10 @@ class VertexEmbedding(VertexBase): headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers) vertex_request: VertexEmbeddingRequest = ( litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( - input=input, optional_params=optional_params, model=model + input=input, + optional_params=optional_params, + model=model, + litellm_params=litellm_params, ) ) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 132f29987a..24396628db 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -3,6 +3,7 @@ from typing import List, Literal, Optional, Union from pydantic import BaseModel +from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.utils import EmbeddingResponse, Usage from .types import * @@ -100,7 +101,11 @@ class VertexAITextEmbeddingConfig(BaseModel): return optional_params def transform_openai_request_to_vertex_embedding_request( - self, input: Union[list, str], optional_params: dict, model: str + self, + input: Union[list, str], + optional_params: dict, + model: str, + litellm_params: Optional[dict] = None, ) -> VertexEmbeddingRequest: """ Transforms an openai request to a vertex embedding request. @@ -108,16 +113,26 @@ class VertexAITextEmbeddingConfig(BaseModel): # Import here to avoid circular import issues with litellm.__init__ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + labels = pop_vertex_request_labels(optional_params, litellm_params) + if model.isdigit(): - return self._transform_openai_request_to_fine_tuned_embedding_request( - input, optional_params, model + vertex_request = ( + self._transform_openai_request_to_fine_tuned_embedding_request( + input, optional_params, model + ) ) + if labels: + vertex_request["labels"] = labels + return vertex_request if VertexBGEConfig.is_bge_model(model): - return VertexBGEConfig.transform_request( + vertex_request = VertexBGEConfig.transform_request( input=input, optional_params=optional_params, model=model ) + if labels: + vertex_request["labels"] = labels + return vertex_request - vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest() + vertex_request = VertexEmbeddingRequest() vertex_text_embedding_input_list: List[TextEmbeddingInput] = [] task_type: Optional[TaskType] = optional_params.get("task_type") title = optional_params.get("title") @@ -133,6 +148,8 @@ class VertexAITextEmbeddingConfig(BaseModel): vertex_request["instances"] = vertex_text_embedding_input_list vertex_request["parameters"] = EmbeddingParameters(**optional_params) + if labels: + vertex_request["labels"] = labels return vertex_request diff --git a/litellm/llms/vertex_ai/vertex_embeddings/types.py b/litellm/llms/vertex_ai/vertex_embeddings/types.py index 317b9c4fb8..bf73f4d193 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/types.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/types.py @@ -3,7 +3,7 @@ Types for Vertex Embeddings Requests """ from enum import Enum -from typing import List, Optional, Union +from typing import Dict, List, Optional, Union from typing_extensions import TypedDict @@ -56,6 +56,7 @@ class VertexEmbeddingRequest(TypedDict, total=False): List[TextEmbeddingFineTunedInput], ] parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]] + labels: Optional[Dict[str, str]] # Example usage: diff --git a/litellm/llms/vertex_ai/vertex_model_garden/main.py b/litellm/llms/vertex_ai/vertex_model_garden/main.py index c37bb449ec..7240d9dce5 100644 --- a/litellm/llms/vertex_ai/vertex_model_garden/main.py +++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py @@ -27,6 +27,17 @@ from ..common_utils import VertexAIError, get_vertex_base_model_name from ..vertex_llm_base import VertexBase +def _vertex_model_garden_model_id_in_json_body(model: str) -> bool: + """ + Vertex catalog / publisher models are addressed as publisher/model (e.g. + xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body. + + Deployed Model Garden endpoints are typically a single segment (often numeric) + and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field. + """ + return "/" in model + + def create_vertex_url( vertex_location: str, vertex_project: str, @@ -34,8 +45,13 @@ def create_vertex_url( model: str, api_base: Optional[str] = None, ) -> str: - """Return the base url for the vertex garden models""" + """Return the api base for vertex model garden (without /chat/completions).""" base_url = get_vertex_base_url(vertex_location) + if _vertex_model_garden_model_id_in_json_body(model): + return ( + f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}" + "/endpoints/openapi" + ) return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}" @@ -129,7 +145,10 @@ class VertexAIModelGardenModels(VertexBase): vertex_location=vertex_location or "us-central1", vertex_api_version="v1beta1", ) - model = "" + # Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route). + # Single-segment endpoint ids: model is encoded in the URL path; body model stays empty. + if not _vertex_model_garden_model_id_in_json_body(model): + model = "" return openai_like_chat_completions.completion( model=model, messages=messages, diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index f6dda4dd25..99e0a958ef 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -17,6 +17,7 @@ from pydantic import fields as pyd_fields import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -300,7 +301,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -333,7 +337,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -372,7 +379,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}/input_items" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/input_items" params: Dict[str, Any] = {} if after is not None: params["after"] = after @@ -408,7 +418,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}/cancel" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/cancel" data: Dict = {} return url, data diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 521dae980d..d64450a121 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -67,7 +67,11 @@ class VoyageRerankConfig(BaseRerankConfig): return api_base def transform_rerank_request( - self, model: str, optional_rerank_params: Dict, headers: Dict + self, + model: str, + optional_rerank_params: Dict, + headers: Dict, + litellm_params: Optional[dict] = None, ) -> Dict: return {"model": model, **optional_rerank_params} diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 4f8e196f25..202760f68a 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -143,6 +143,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): model: str, optional_rerank_params: Dict, headers: dict, + litellm_params: Optional[dict] = None, ) -> dict: """ Transform request to IBM watsonx.ai rerank format diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index bfa55105a6..64b4a545ac 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -43,6 +43,7 @@ class XAIChatConfig(OpenAIGPTConfig): "logprobs", "max_tokens", "n", + "parallel_tool_calls", "presence_penalty", "response_format", "seed", diff --git a/litellm/main.py b/litellm/main.py index daa0fb063a..0079bd750c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5311,6 +5311,7 @@ def embedding( # noqa: PLR0915 api_key=api_key, api_base=api_base, client=client, + litellm_params=litellm_params_dict, ) elif custom_llm_provider == "oobabooga": response = oobabooga.embedding( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 13a45fd165..a1e3e42a9c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -33337,6 +33337,72 @@ "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas", "supports_reasoning": true }, + "vertex_ai/xai/grok-4.1-fast-non-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "vertex_ai/xai/grok-4.1-fast-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "vertex_ai/xai/grok-4.20-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "vertex_ai/xai/grok-4.20-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-qwen_models", diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 21abaa3f98..a6f0d145e9 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,4 +1,5 @@ import base64 +import binascii import json from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast @@ -498,6 +499,82 @@ async def rotate_mcp_server_credentials_master_key( ) +def _decode_user_credential(stored: str) -> Optional[str]: + """Read back a value persisted in ``LiteLLM_MCPUserCredentials.credential_b64``. + + Tries nacl decryption first (current write format). Falls back to a + plain ``urlsafe_b64decode`` for rows persisted by older code that wrote + the credential without encryption. Returns ``None`` when neither path + yields a valid string. + """ + decrypted = decrypt_value_helper( + value=stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + if decrypted is not None: + return decrypted + try: + return base64.urlsafe_b64decode(stored).decode() + except (binascii.Error, UnicodeDecodeError, ValueError, TypeError): + return None + + +def _decode_oauth_payload(stored: str) -> Optional[Dict[str, Any]]: + """Return the OAuth2 payload dict if ``stored`` holds one, else ``None``. + + A row is considered an OAuth2 credential iff its decoded value parses as + a JSON object with ``"type": "oauth2"``. Plain BYOK credentials (which + share the same column) decode to a non-JSON string and return ``None``. + """ + decoded = _decode_user_credential(stored) + if decoded is None: + return None + try: + parsed = json.loads(decoded) + except (ValueError, TypeError): + return None + if isinstance(parsed, dict) and parsed.get("type") == "oauth2": + return parsed + return None + + +async def rotate_mcp_user_credentials_master_key( + prisma_client: PrismaClient, new_master_key: str +): + """Re-encrypt every ``LiteLLM_MCPUserCredentials`` row with ``new_master_key``. + + Reads each ``credential_b64`` with the current salt key (falling back to + legacy plain base64 for unmigrated rows) and writes it back encrypted + under the new master key. Rows that are unreadable under both paths + are logged and skipped so one corrupt row does not abort the rotation. + """ + rows = await prisma_client.db.litellm_mcpusercredentials.find_many() + for row in rows: + plaintext = _decode_user_credential(row.credential_b64) + if plaintext is None: + verbose_proxy_logger.warning( + "rotate_mcp_user_credentials_master_key: could not decode " + "credential for user_id=%s server_id=%s, skipping", + row.user_id, + row.server_id, + ) + continue + re_encrypted = encrypt_value_helper( + plaintext, new_encryption_key=new_master_key + ) + await prisma_client.db.litellm_mcpusercredentials.update( + where={ + "user_id_server_id": { + "user_id": row.user_id, + "server_id": row.server_id, + } + }, + data={"credential_b64": re_encrypted}, + ) + + async def store_user_credential( prisma_client: PrismaClient, user_id: str, @@ -506,7 +583,7 @@ async def store_user_credential( ) -> None: """Store a user credential for a BYOK MCP server.""" - encoded = base64.urlsafe_b64encode(credential.encode()).decode() + encoded = encrypt_value_helper(credential) await prisma_client.db.litellm_mcpusercredentials.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ @@ -532,16 +609,7 @@ async def get_user_credential( ) if row is None: return None - try: - return base64.urlsafe_b64decode(row.credential_b64).decode() - except Exception: - # Fall back to nacl decryption for credentials stored by older code - return decrypt_value_helper( - value=row.credential_b64, - key="byok_credential", - exception_type="debug", - return_original_value=False, - ) + return _decode_user_credential(row.credential_b64) async def has_user_credential( @@ -582,7 +650,7 @@ async def store_user_oauth_credential( ) -> None: """Persist an OAuth2 access token for a user+server pair. - The payload is JSON-serialised and stored base64-encoded in the same + The payload is JSON-serialised and stored encrypted in the same ``credential_b64`` column used by BYOK. A ``"type": "oauth2"`` key differentiates it from plain BYOK API keys. """ @@ -606,29 +674,27 @@ async def store_user_oauth_credential( payload["scopes"] = scopes # Guard against silently overwriting a BYOK credential with an OAuth token. - # BYOK credentials lack a "type" field (or use a non-"oauth2" type). # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - if existing is not None: - _byok_error = ValueError( - f"A non-OAuth2 credential already exists for user {user_id} " - f"and server {server_id}. Refusing to overwrite." + if ( + existing is not None + and _decode_oauth_payload(existing.credential_b64) is None + ): + # Existing row is either a BYOK secret or an OAuth2 row that no + # longer decrypts (e.g. after a salt-key rotation). In either + # case, refuse to overwrite — the caller would clobber data + # that may still be recoverable. + raise ValueError( + f"Existing credential for user {user_id} and server " + f"{server_id} could not be verified as an OAuth2 token. " + f"Refusing to overwrite." ) - try: - raw = json.loads( - base64.urlsafe_b64decode(existing.credential_b64).decode() - ) - except Exception: - # Credential is not base64+JSON — it's a plain-text BYOK key. - raise _byok_error - if raw.get("type") != "oauth2": - raise _byok_error - encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + encoded = encrypt_value_helper(json.dumps(payload)) await prisma_client.db.litellm_mcpusercredentials.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ @@ -672,15 +738,7 @@ async def get_user_oauth_credential( ) if row is None: return None - try: - decoded = base64.urlsafe_b64decode(row.credential_b64).decode() - parsed = json.loads(decoded) - if isinstance(parsed, dict) and parsed.get("type") == "oauth2": - return parsed - # Row exists but is a BYOK (plain string), not an OAuth token - return None - except Exception: - return None + return _decode_oauth_payload(row.credential_b64) async def list_user_oauth_credentials( @@ -694,14 +752,11 @@ async def list_user_oauth_credentials( ) results: List[Dict[str, Any]] = [] for row in rows: - try: - decoded = base64.urlsafe_b64decode(row.credential_b64).decode() - parsed = json.loads(decoded) - if isinstance(parsed, dict) and parsed.get("type") == "oauth2": - parsed["server_id"] = row.server_id - results.append(parsed) - except Exception: - pass # Skip non-OAuth rows (BYOK plain strings) + payload = _decode_oauth_payload(row.credential_b64) + if payload is None: + continue + payload["server_id"] = row.server_id + results.append(payload) return results diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cebd224a1a..1794cd1438 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -33,10 +33,12 @@ def get_request_base_url(request: Request) -> str: """ Get the base URL for the request, considering X-Forwarded-* headers. - When behind a proxy (like nginx), the proxy may set: - - X-Forwarded-Proto: The original protocol (http/https) - - X-Forwarded-Host: The original host (may include port) - - X-Forwarded-Port: The original port (if not in Host header) + X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured + when the request comes from a configured trusted proxy + (``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``). + Otherwise the request's literal ``base_url`` is returned, so an + untrusted caller cannot poison OAuth-discovery / redirect_uri values + by injecting headers. Args: request: FastAPI Request object @@ -47,34 +49,28 @@ def get_request_base_url(request: Request) -> str: base_url = str(request.base_url).rstrip("/") parsed = urlparse(base_url) - # Get forwarded headers + if not IPAddressUtils.is_request_from_trusted_proxy(request): + return base_url + x_forwarded_proto = request.headers.get("X-Forwarded-Proto") x_forwarded_host = request.headers.get("X-Forwarded-Host") x_forwarded_port = request.headers.get("X-Forwarded-Port") - # Start with the original scheme scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme - # Handle host and port if x_forwarded_host: # X-Forwarded-Host may already include port (e.g., "example.com:8080") if ":" in x_forwarded_host and not x_forwarded_host.startswith("["): - # Host includes port netloc = x_forwarded_host elif x_forwarded_port: - # Port is separate netloc = f"{x_forwarded_host}:{x_forwarded_port}" else: - # Just host, no explicit port netloc = x_forwarded_host else: - # No X-Forwarded-Host, use original netloc netloc = parsed.netloc if x_forwarded_port and ":" not in netloc: - # Add forwarded port if not already in netloc netloc = f"{netloc}:{x_forwarded_port}" - # Reconstruct the URL return urlunparse((scheme, netloc, parsed.path, "", "", "")) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 25c47aacf0..9923c3ce4b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -41,6 +41,7 @@ from litellm.constants import ( MCP_TOOL_LISTING_TIMEOUT, ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth 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 ( @@ -168,6 +169,37 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]: class MCPServerManager: _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$") + @staticmethod + def _resolve_oauth2_flow( + *, + auth_type: Optional[MCPAuthType], + oauth2_flow: Optional[str], + token_url: Optional[str], + authorization_url: Optional[str], + client_id: Optional[str], + client_secret: Optional[str], + ) -> Optional[Literal["client_credentials", "authorization_code"]]: + """Infer oauth2_flow for legacy records that omit the field. + + DB rows created before oauth2_flow support may have OAuth2 client + credentials + token_url but a null oauth2_flow. Treat these as M2M, + unless authorization_url is present (interactive OAuth). + """ + if oauth2_flow in ("client_credentials", "authorization_code"): + return cast( + Literal["client_credentials", "authorization_code"], oauth2_flow + ) + if oauth2_flow: + # Ignore unknown/untyped values and continue legacy inference. + return None + if auth_type != MCPAuth.oauth2: + return None + if authorization_url: + return None + if token_url and client_id and client_secret: + return "client_credentials" + return None + def __init__(self): self.registry: Dict[str, MCPServer] = {} self.config_mcp_servers: Dict[str, MCPServer] = {} @@ -341,7 +373,14 @@ class MCPServerManager: # oauth specific fields client_id=server_config.get("client_id", None), client_secret=server_config.get("client_secret", None), - oauth2_flow=server_config.get("oauth2_flow", None), + oauth2_flow=self._resolve_oauth2_flow( + auth_type=auth_type, + oauth2_flow=server_config.get("oauth2_flow", None), + token_url=resolved_token_url, + authorization_url=resolved_authorization_url, + client_id=server_config.get("client_id", None), + client_secret=server_config.get("client_secret", None), + ), scopes=resolved_scopes, authorization_url=resolved_authorization_url, token_url=resolved_token_url, @@ -678,7 +717,17 @@ class MCPServerManager: client_id=client_id_value or getattr(mcp_server, "client_id", None), client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), - oauth2_flow=getattr(mcp_server, "oauth2_flow", None), + oauth2_flow=self._resolve_oauth2_flow( + auth_type=auth_type, + oauth2_flow=getattr(mcp_server, "oauth2_flow", None), + token_url=mcp_server.token_url + or getattr(mcp_oauth_metadata, "token_url", None), + authorization_url=mcp_server.authorization_url + or getattr(mcp_oauth_metadata, "authorization_url", None), + client_id=client_id_value or getattr(mcp_server, "client_id", None), + client_secret=client_secret_value + or getattr(mcp_server, "client_secret", None), + ), scopes=resolved_scopes, authorization_url=mcp_server.authorization_url or getattr(mcp_oauth_metadata, "authorization_url", None), @@ -1499,6 +1548,47 @@ class MCPServerManager: ) return await client.get_prompt(get_prompt_request_params) + @staticmethod + def _is_same_authority_metadata_url(url: str, server_url: str) -> bool: + """ + Whether ``url`` shares scheme, host, and port with ``server_url``. + + Same-authority metadata URLs are produced by our well-known discovery + construction and by resource servers that publish protected-resource + metadata on the resource origin. These must keep working for + administrator-configured internal MCP servers, so they are fetched + directly. Cross-origin URLs are fetched through ``async_safe_get``. + """ + try: + target = urlparse(url) + base = urlparse(server_url) + except Exception: + return False + + if target.scheme not in ("http", "https") or not target.hostname: + return False + + target_port = target.port or (443 if target.scheme == "https" else 80) + base_port = base.port or (443 if base.scheme == "https" else 80) + return ( + base.scheme == target.scheme + and (base.hostname or "").lower() == target.hostname.lower() + and base_port == target_port + ) + + async def _fetch_oauth_discovery_url(self, url: str, server_url: str) -> Any: + client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.MCP, + params={"timeout": MCP_METADATA_TIMEOUT}, + ) + if self._is_same_authority_metadata_url(url, server_url): + # Same-authority URLs may point at administrator-configured + # internal MCP servers. Do not run them through user URL + # validation, but also do not follow redirects because the + # redirect target would not inherit the same-authority guarantee. + return await client.get(url, follow_redirects=False) + return await async_safe_get(client, url) + async def _descovery_metadata( self, server_url: str, @@ -1514,7 +1604,7 @@ class MCPServerManager: resource_scopes, ) = await self._attempt_well_known_discovery(server_url) metadata = await self._fetch_authorization_server_metadata( - authorization_servers + authorization_servers, server_url ) if ( metadata is None @@ -1555,7 +1645,7 @@ class MCPServerManager: authorization_servers, resource_scopes, ) = await self._fetch_oauth_metadata_from_resource( - resource_metadata_url + resource_metadata_url, server_url ) else: ( @@ -1576,7 +1666,7 @@ class MCPServerManager: if authorization_servers: metadata = await self._fetch_authorization_server_metadata( - authorization_servers + authorization_servers, server_url ) preferred_scopes = scopes or resource_scopes @@ -1616,19 +1706,26 @@ class MCPServerManager: return resource_metadata_url, scopes async def _fetch_oauth_metadata_from_resource( - self, resource_metadata_url: str + self, resource_metadata_url: str, server_url: str ) -> Tuple[List[str], Optional[List[str]]]: if not resource_metadata_url: return [], None try: - client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.MCP, - params={"timeout": MCP_METADATA_TIMEOUT}, + response = await self._fetch_oauth_discovery_url( + resource_metadata_url, server_url ) - response = await client.get(resource_metadata_url) response.raise_for_status() data = response.json() + except SSRFError as exc: + verbose_logger.warning( + "MCP OAuth discovery: refusing to fetch resource metadata from %s " + "(rejected by SSRF guard for server %s): %s", + resource_metadata_url, + server_url, + exc, + ) + return [], None except Exception as exc: # pragma: no cover - network issues verbose_logger.debug( "Failed to fetch MCP OAuth metadata from %s: %s", @@ -1677,23 +1774,25 @@ class MCPServerManager: ( authorization_servers, scopes, - ) = await self._fetch_oauth_metadata_from_resource(url) + ) = await self._fetch_oauth_metadata_from_resource(url, server_url) if authorization_servers: return authorization_servers, scopes return [], None async def _fetch_authorization_server_metadata( - self, authorization_servers: List[str] + self, authorization_servers: List[str], server_url: str ) -> Optional[MCPOAuthMetadata]: for issuer in authorization_servers: - metadata = await self._fetch_single_authorization_server_metadata(issuer) + metadata = await self._fetch_single_authorization_server_metadata( + issuer, server_url + ) if metadata is not None: return metadata return None async def _fetch_single_authorization_server_metadata( - self, issuer_url: str + self, issuer_url: str, server_url: str ) -> Optional[MCPOAuthMetadata]: try: parsed = urlparse(issuer_url) @@ -1721,13 +1820,18 @@ class MCPServerManager: for url in candidate_urls: try: - client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.MCP, - params={"timeout": MCP_METADATA_TIMEOUT}, - ) - response = await client.get(url) + response = await self._fetch_oauth_discovery_url(url, server_url) response.raise_for_status() data = response.json() + except SSRFError as exc: + verbose_logger.warning( + "MCP OAuth discovery: refusing to fetch authorization-server " + "metadata from %s (rejected by SSRF guard for server %s): %s", + url, + server_url, + exc, + ) + continue except Exception as exc: # pragma: no cover - network issues verbose_logger.debug( "Failed to fetch authorization metadata from %s: %s", @@ -2370,7 +2474,7 @@ class MCPServerManager: ) ) - async def _call_regular_mcp_tool( + async def _call_regular_mcp_tool( # noqa: PLR0915 self, mcp_server: MCPServer, original_tool_name: str, @@ -2433,7 +2537,11 @@ class MCPServerManager: # oauth2 headers extra_headers: Optional[Dict[str, str]] = None if mcp_server.auth_type == MCPAuth.oauth2: - extra_headers = oauth2_headers + if mcp_server.has_client_credentials: + # For M2M OAuth servers, Authorization must come from token fetch. + extra_headers = None + else: + extra_headers = oauth2_headers if mcp_server.extra_headers and raw_headers: if extra_headers is None: @@ -2445,6 +2553,11 @@ class MCPServerManager: for header in mcp_server.extra_headers: if not isinstance(header, str): continue + if ( + mcp_server.has_client_credentials + and header.lower() == "authorization" + ): + continue header_value = normalized_raw_headers.get(header.lower()) if header_value is None: continue @@ -2480,6 +2593,10 @@ class MCPServerManager: ) extra_headers.update(hook_extra_headers) + # Reset to None if no headers were actually added + if extra_headers is not None and len(extra_headers) == 0: + extra_headers = None + stdio_env = self._build_stdio_env(mcp_server, raw_headers) client = await self._create_mcp_client( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ae6055217b..abb4b5cfa6 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -153,6 +153,7 @@ if MCP_AVAILABLE: MCPAuthenticatedUser, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, global_mcp_server_manager, ) from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( @@ -900,6 +901,20 @@ if MCP_AVAILABLE: allowed_mcp_server_id ) if mcp_server is not None: + # Apply oauth2_flow resolution for legacy DB rows where it may be NULL + resolved_flow = MCPServerManager._resolve_oauth2_flow( + auth_type=mcp_server.auth_type, + oauth2_flow=mcp_server.oauth2_flow, + token_url=mcp_server.token_url, + authorization_url=mcp_server.authorization_url, + client_id=mcp_server.client_id, + client_secret=mcp_server.client_secret, + ) + if resolved_flow and resolved_flow != mcp_server.oauth2_flow: + # Create a new instance with the resolved flow for this request + mcp_server = mcp_server.model_copy( + update={"oauth2_flow": resolved_flow} + ) allowed_mcp_servers.append(mcp_server) if mcp_servers is not None: @@ -1100,8 +1115,13 @@ if MCP_AVAILABLE: extra_headers: Optional[Dict[str, str]] = None if server.auth_type == MCPAuth.oauth2: - # Copy to avoid mutating the original dict (important for parallel fetching) - extra_headers = oauth2_headers.copy() if oauth2_headers else None + # For OAuth2 M2M servers, upstream Authorization must come from + # client_credentials token fetch, never from caller headers. + if server.has_client_credentials: + extra_headers = None + else: + # Copy to avoid mutating the original dict (important for parallel fetching) + extra_headers = oauth2_headers.copy() if oauth2_headers else None if server.extra_headers and raw_headers: if extra_headers is None: @@ -1114,11 +1134,17 @@ if MCP_AVAILABLE: for header in server.extra_headers: if not isinstance(header, str): continue + if server.has_client_credentials and header.lower() == "authorization": + continue header_value = normalized_raw_headers.get(header.lower()) if header_value is None: continue extra_headers[header] = header_value + # Reset to None if no headers were actually added + if extra_headers is not None and len(extra_headers) == 0: + extra_headers = None + if server_auth_header is None: server_auth_header = mcp_auth_header @@ -1377,11 +1403,19 @@ if MCP_AVAILABLE: spend_meta["per_server_tool_counts"] = per_server_tool_counts end_time = datetime.now() - await litellm_logging_obj.async_success_handler( - result=all_tools, - start_time=list_tools_start_time, - end_time=end_time, - ) + try: + await litellm_logging_obj.async_success_handler( + result=all_tools, + start_time=list_tools_start_time, + end_time=end_time, + ) + except Exception as log_exc: + # list_tools responses must not be dropped due to non-blocking + # observability/serialization failures. + verbose_logger.warning( + "MCP list_tools success logging failed (continuing): %s", + log_exc, + ) verbose_logger.info( f"Successfully fetched {len(all_tools)} tools total from all MCP servers" diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/skills.html b/litellm/proxy/_experimental/out/skills/index.html similarity index 100% rename from litellm/proxy/_experimental/out/skills.html rename to litellm/proxy/_experimental/out/skills/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index dca54eb7fa..309a0276aa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -8,16 +8,12 @@ any drift as a neutral check. """ import json -import re import sys -from copy import copy from pathlib import Path -from typing import Dict, List, Optional - -from fastapi.routing import APIRoute -from starlette.routing import BaseRoute +from typing import Dict, Optional SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json" +HTTP_METHODS = {"delete", "get", "head", "options", "patch", "post", "put"} def load_snapshot() -> Optional[Dict[str, Dict]]: @@ -30,36 +26,37 @@ def load_snapshot() -> Optional[Dict[str, Dict]]: return None -def _stable_unique_id(route: APIRoute, method: str) -> str: - operation_id = f"{route.name}{route.path_format}" - operation_id = re.sub(r"\W", "_", operation_id) - return f"{operation_id}_{method.lower()}" +def _normalize_operation_ids(paths: Dict[str, Dict]) -> None: + """Make FastAPI-generated operation IDs stable for multi-method routes. - -def _routes_with_stable_unique_ids(routes: List[BaseRoute]) -> List[BaseRoute]: - stable_routes: List[BaseRoute] = [] - for route in routes: - if not isinstance(route, APIRoute) or not route.methods: - stable_routes.append(route) + FastAPI derives the default operation ID suffix from the first item in the + route's methods set. For routes registered with several HTTP methods, that + set iteration order can vary between processes, which makes the snapshot + drift even when no routes changed. + """ + for path_ops in paths.values(): + if not isinstance(path_ops, dict): continue - methods = sorted(route.methods) - has_multiple_methods = len(methods) > 1 - for method in methods: - method_route = copy(route) - method_route.methods = {method} - if route.operation_id is not None: - method_route.operation_id = ( - f"{route.operation_id}_{method.lower()}" - if has_multiple_methods - else route.operation_id - ) - method_route.unique_id = method_route.operation_id - else: - method_route.unique_id = _stable_unique_id(route, method) - stable_routes.append(method_route) + methods = {method for method in path_ops if method in HTTP_METHODS} + if not methods: + continue - return stable_routes + for method, operation in path_ops.items(): + if method not in HTTP_METHODS or not isinstance(operation, dict): + continue + + operation_id = operation.get("operationId") + if not isinstance(operation_id, str): + continue + + for suffix in methods: + suffix_token = f"_{suffix}" + if operation_id.endswith(suffix_token): + operation["operationId"] = ( + operation_id[: -len(suffix_token)] + f"_{method}" + ) + break def generate_snapshot() -> Dict[str, Dict]: @@ -88,18 +85,16 @@ def generate_snapshot() -> Dict[str, Dict]: ] if not feat_routes: continue - full = get_openapi( - title=app.title, - version=app.version, - routes=_routes_with_stable_unique_ids(feat_routes), - ) + full = get_openapi(title=app.title, version=app.version, routes=feat_routes) + paths = full.get("paths", {}) + _normalize_operation_ids(paths) # Group all of a feature's routes under one tag. - for path_ops in full.get("paths", {}).values(): + for path_ops in paths.values(): for op in path_ops.values(): if isinstance(op, dict): op["tags"] = [feat.name] fragments[feat.name] = { - "paths": full.get("paths", {}), + "paths": paths, "components": {"schemas": full.get("components", {}).get("schemas", {})}, } return fragments diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 81e972e035..8520e03f83 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -668,6 +668,8 @@ class LiteLLMRoutes(enum.Enum): "/models/{model_id}", "/guardrails/list", "/v2/guardrails/list", + "/project/list", + "/project/info", ] + spend_tracking_routes + key_management_routes @@ -692,6 +694,9 @@ class LiteLLMRoutes(enum.Enum): "/model/{model_id}/update", "/prompt/list", "/prompt/info", + # Project read routes - endpoint scopes results to caller's teams (non-admin) + "/project/list", + "/project/info", # Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges "/invitation/new", "/invitation/delete", @@ -2161,8 +2166,8 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase): description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.", ) auth: bool = Field( - default=False, - description="Whether authentication is required for the pass-through endpoint. If True, requests to the endpoint will require a valid LiteLLM API key.", + default=True, + description="Whether authentication is required for the pass-through endpoint. Defaults to True so a pass-through silently created without an explicit value still requires a valid LiteLLM API key — set to False only if the endpoint is meant to be a public forwarder (e.g. an unauthenticated webhook target).", ) guardrails: Optional[PassThroughGuardrailsConfig] = Field( default=None, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 65638ed6c1..113a8f538c 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -12,14 +12,13 @@ Run checks for: import asyncio import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast from fastapi import HTTPException, Request, status from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.constants import ( CLI_JWT_EXPIRATION_HOURS, @@ -66,6 +65,8 @@ from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names, ) +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -852,7 +853,7 @@ def get_actual_routes(allowed_routes: list) -> list: async def get_default_end_user_budget( prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, ) -> Optional[LiteLLM_BudgetTable]: """ @@ -875,9 +876,12 @@ async def get_default_end_user_budget( cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}" # Check cache first - cached_budget = await user_api_key_cache.async_get_cache(key=cache_key) + cached_budget = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_BudgetTable, + ) if cached_budget is not None: - return LiteLLM_BudgetTable(**cached_budget) + return cached_budget # Fetch from database try: @@ -891,14 +895,16 @@ async def get_default_end_user_budget( ) return None + _budget_obj = LiteLLM_BudgetTable(**budget_record.dict()) # Cache the budget for 60 seconds await user_api_key_cache.async_set_cache( key=cache_key, - value=budget_record.dict(), + value=_budget_obj, + model_type=LiteLLM_BudgetTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - return LiteLLM_BudgetTable(**budget_record.dict()) + return _budget_obj except Exception as e: verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}") @@ -909,7 +915,7 @@ async def get_default_end_user_budget( async def get_team_member_default_budget( budget_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> Optional[LiteLLM_BudgetTable]: """ Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"]. @@ -966,7 +972,7 @@ async def get_team_member_default_budget( async def _apply_default_budget_to_end_user( end_user_obj: LiteLLM_EndUserTable, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, ) -> LiteLLM_EndUserTable: """ @@ -1039,7 +1045,7 @@ def _check_end_user_budget( async def get_end_user_object( end_user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, route: str, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, @@ -1070,10 +1076,12 @@ async def get_end_user_object( _key = "end_user_id:{}".format(end_user_id) # Check cache first - cached_user_obj = await user_api_key_cache.async_get_cache(key=_key) + cached_user_obj = await user_api_key_cache.async_get_cache( + key=_key, + model_type=LiteLLM_EndUserTable, + ) if cached_user_obj is not None: - return_obj = LiteLLM_EndUserTable(**cached_user_obj) - + return_obj = cached_user_obj # Apply default budget if needed return_obj = await _apply_default_budget_to_end_user( end_user_obj=return_obj, @@ -1108,9 +1116,11 @@ async def get_end_user_object( parent_otel_span=parent_otel_span, ) - # Save to cache (always store as dict for consistency) + # Save to cache await user_api_key_cache.async_set_cache( - key="end_user_id:{}".format(end_user_id), value=_response.dict() + key="end_user_id:{}".format(end_user_id), + value=_response, + model_type=LiteLLM_EndUserTable, ) # Check budget limits @@ -1128,7 +1138,7 @@ async def get_end_user_object( async def get_tag_objects_batch( tag_names: List[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Dict[str, LiteLLM_TagTable]: @@ -1161,12 +1171,12 @@ async def get_tag_objects_batch( # Try to get all tags from cache first for tag_name in tag_names: cache_key = f"tag:{tag_name}" - cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) + cached_tag = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_TagTable, + ) if cached_tag is not None: - if isinstance(cached_tag, dict): - tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag) - else: - tag_objects[tag_name] = cached_tag + tag_objects[tag_name] = cached_tag else: uncached_tags.append(tag_name) @@ -1182,11 +1192,13 @@ async def get_tag_objects_batch( for db_tag in db_tags: tag_name = db_tag.tag_name cache_key = f"tag:{tag_name}" - # Cache with default TTL (same as end_user objects) + _tag_obj = LiteLLM_TagTable(**db_tag.dict()) await user_api_key_cache.async_set_cache( - key=cache_key, value=db_tag.dict() + key=cache_key, + value=_tag_obj, + model_type=LiteLLM_TagTable, ) - tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict()) + tag_objects[tag_name] = _tag_obj except Exception as e: verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}") @@ -1197,7 +1209,7 @@ async def get_tag_objects_batch( async def get_tag_object( tag_name: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_TagTable]: @@ -1236,7 +1248,7 @@ async def get_team_membership( user_id: str, team_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional["LiteLLM_TeamMembership"]: @@ -1256,9 +1268,12 @@ async def get_team_membership( _key = "team_membership:{}:{}".format(user_id, team_id) # check if in cache - cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key) + cached_membership_obj = await user_api_key_cache.async_get_cache( + key=_key, + model_type=LiteLLM_TeamMembership, + ) if cached_membership_obj is not None: - return LiteLLM_TeamMembership(**cached_membership_obj) + return cached_membership_obj # else, check db try: @@ -1270,10 +1285,12 @@ async def get_team_membership( if response is None: return None - # save the team membership object to cache (store as dict) - await user_api_key_cache.async_set_cache(key=_key, value=response.dict()) - _response = LiteLLM_TeamMembership(**response.dict()) + await user_api_key_cache.async_set_cache( + key=_key, + value=_response, + model_type=LiteLLM_TeamMembership, + ) return _response except Exception: @@ -1441,7 +1458,7 @@ async def _get_fuzzy_user_object( async def get_user_object( user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_id_upsert: bool, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, @@ -1460,12 +1477,12 @@ async def get_user_object( # check if in cache if not check_db_only: - cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id) + cached_user_obj = await user_api_key_cache.async_get_cache( + key=user_id, + model_type=LiteLLM_UserTable, + ) if cached_user_obj is not None: - if isinstance(cached_user_obj, dict): - return LiteLLM_UserTable(**cached_user_obj) - elif isinstance(cached_user_obj, LiteLLM_UserTable): - return cached_user_obj + return cached_user_obj # else, check db if prisma_client is None: raise Exception("No db connected") @@ -1527,7 +1544,8 @@ async def get_user_object( # save the user object to cache await user_api_key_cache.async_set_cache( key=user_id, - value=response_dict, + value=_response, + model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1548,13 +1566,21 @@ async def get_user_object( async def _cache_management_object( key: str, - value: BaseModel, - user_api_key_cache: DualCache, + value: Union[BaseModel, Dict[str, Any]], + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], + *, + model_type: Type[BaseModel], ): + """ + Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis). + + ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned. + """ await user_api_key_cache.async_set_cache( key=key, value=value, + model_type=model_type, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1562,7 +1588,7 @@ async def _cache_management_object( async def _cache_team_object( team_id: str, team_table: LiteLLM_TeamTableCachedObj, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = "team_id:{}".format(team_id) @@ -1575,13 +1601,14 @@ async def _cache_team_object( value=team_table, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + model_type=LiteLLM_TeamTableCachedObj, ) async def _cache_key_object( hashed_token: str, user_api_key_obj: UserAPIKeyAuth, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = hashed_token @@ -1594,12 +1621,13 @@ async def _cache_key_object( value=user_api_key_obj, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + model_type=UserAPIKeyAuth, ) async def _delete_cache_key_object( hashed_token: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging], ): key = hashed_token @@ -1647,7 +1675,7 @@ async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): async def _get_team_object_from_user_api_key_cache( team_id: str, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, last_db_access_time: LimitedSizeOrderedDict, db_cache_expiry: int, proxy_logging_obj: Optional[ProxyLogging], @@ -1708,38 +1736,38 @@ async def _get_team_object_from_user_api_key_cache( async def _get_team_object_from_cache( key: str, proxy_logging_obj: Optional[ProxyLogging], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], ) -> Optional[LiteLLM_TeamTableCachedObj]: - cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None - - ## CHECK REDIS CACHE ## + ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## if ( proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache ): - cached_team_obj = ( + cached_raw = ( await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span ) ) + if cached_raw is not None: + from_internal = CacheCodec.deserialize( + cached_raw, LiteLLM_TeamTableCachedObj + ) + if from_internal is not None: + return from_internal - if cached_team_obj is None: - cached_team_obj = await user_api_key_cache.async_get_cache(key=key) - - if cached_team_obj is not None: - if isinstance(cached_team_obj, dict): - return LiteLLM_TeamTableCachedObj(**cached_team_obj) - elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj): - return cached_team_obj - - return None + decoded = await user_api_key_cache.async_get_cache( + key=key, + parent_otel_span=parent_otel_span, + model_type=LiteLLM_TeamTableCachedObj, + ) + return decoded async def get_team_object( team_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, @@ -1805,20 +1833,21 @@ async def get_team_object( async def _cache_access_object( access_group_id: str, access_group_table: LiteLLM_AccessGroupTable, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) await user_api_key_cache.async_set_cache( key=key, value=access_group_table, + model_type=LiteLLM_AccessGroupTable, ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL, ) async def _delete_cache_access_object( access_group_id: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ): key = "access_group_id:{}".format(access_group_id) @@ -1836,7 +1865,7 @@ async def _delete_cache_access_object( async def get_access_object( access_group_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> LiteLLM_AccessGroupTable: """ @@ -1858,13 +1887,12 @@ async def get_access_object( key = "access_group_id:{}".format(access_group_id) - # Always check cache first - cached_access_obj = await user_api_key_cache.async_get_cache(key=key) + cached_access_obj = await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_AccessGroupTable, + ) if cached_access_obj is not None: - if isinstance(cached_access_obj, dict): - return LiteLLM_AccessGroupTable(**cached_access_obj) - elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable): - return cached_access_obj + return cached_access_obj # Not in cache - fetch from DB try: @@ -1910,7 +1938,7 @@ async def get_access_object( async def get_team_object_by_alias( team_alias: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional["Span"] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> LiteLLM_TeamTableCachedObj: @@ -1992,6 +2020,7 @@ async def get_team_object_by_alias( await user_api_key_cache.async_set_cache( key=cache_key, value=team_obj, + model_type=LiteLLM_TeamTableCachedObj, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by team_id for consistency @@ -1999,6 +2028,7 @@ async def get_team_object_by_alias( await user_api_key_cache.async_set_cache( key=team_id_cache_key, value=team_obj, + model_type=LiteLLM_TeamTableCachedObj, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2020,7 +2050,7 @@ async def get_team_object_by_alias( async def get_org_object_by_alias( org_alias: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional["Span"] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_OrganizationTable]: @@ -2047,12 +2077,12 @@ async def get_org_object_by_alias( # Check cache first (keyed by alias) cache_key = "org_alias:{}".format(org_alias) - cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) + cached_org_obj = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_OrganizationTable, + ) if cached_org_obj is not None: - if isinstance(cached_org_obj, dict): - return LiteLLM_OrganizationTable(**cached_org_obj) - elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): - return cached_org_obj + return cached_org_obj # Query database by organization_alias try: @@ -2082,13 +2112,15 @@ async def get_org_object_by_alias( # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=org_obj.model_dump(), + value=org_obj, + model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) # Also cache by org_id for consistency await user_api_key_cache.async_set_cache( key="org_id:{}".format(org_obj.organization_id), - value=org_obj.model_dump(), + value=org_obj, + model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2291,7 +2323,7 @@ async def get_jwt_key_mapping_object( async def get_key_object( hashed_token: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, check_cache_only: Optional[bool] = None, @@ -2309,15 +2341,14 @@ async def get_key_object( # check if in cache key = hashed_token - cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache( - key=key + # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth + # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB. + user_api_key_auth = await user_api_key_cache.async_get_cache( + key=key, + model_type=UserAPIKeyAuth, ) - - if cached_key_obj is not None: - if isinstance(cached_key_obj, dict): - return UserAPIKeyAuth(**cached_key_obj) - elif isinstance(cached_key_obj, UserAPIKeyAuth): - return cached_key_obj + if user_api_key_auth is not None: + return user_api_key_auth if check_cache_only: raise Exception( @@ -2374,7 +2405,7 @@ async def get_key_object( async def get_object_permission( object_permission_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_ObjectPermissionTable]: @@ -2390,12 +2421,12 @@ async def get_object_permission( # check if in cache key = "object_permission_id:{}".format(object_permission_id) - cached_obj_permission = await user_api_key_cache.async_get_cache(key=key) - if cached_obj_permission is not None: - if isinstance(cached_obj_permission, dict): - return LiteLLM_ObjectPermissionTable(**cached_obj_permission) - elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable): - return cached_obj_permission + deserialized_perm = await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_ObjectPermissionTable, + ) + if deserialized_perm is not None: + return deserialized_perm # else, check db try: @@ -2406,14 +2437,15 @@ async def get_object_permission( if response is None: return None - # save the object permission to cache + _perm_obj = LiteLLM_ObjectPermissionTable(**response.dict()) await user_api_key_cache.async_set_cache( key=key, - value=response.model_dump(), + value=_perm_obj, + model_type=LiteLLM_ObjectPermissionTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) - return LiteLLM_ObjectPermissionTable(**response.dict()) + return _perm_obj except Exception: return None @@ -2422,7 +2454,7 @@ async def get_object_permission( async def get_managed_vector_store_rows_by_uuids( uuids: List[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[LiteLLM_ManagedVectorStoresTable]: @@ -2442,14 +2474,12 @@ async def get_managed_vector_store_rows_by_uuids( for uuid in uuids: key = "managed_vector_store_id:{}".format(uuid) - cached = await user_api_key_cache.async_get_cache(key=key) - if cached is not None: - if isinstance(cached, dict): - result.append(LiteLLM_ManagedVectorStoresTable(**cached)) - elif isinstance(cached, LiteLLM_ManagedVectorStoresTable): - result.append(cached) - else: - cache_misses.append(uuid) + deserialized_vs = await user_api_key_cache.async_get_cache( + key=key, + model_type=LiteLLM_ManagedVectorStoresTable, + ) + if deserialized_vs is not None: + result.append(deserialized_vs) else: cache_misses.append(uuid) @@ -2475,7 +2505,8 @@ async def get_managed_vector_store_rows_by_uuids( key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) await user_api_key_cache.async_set_cache( key=key, - value=row_dict, + value=cached_obj, + model_type=LiteLLM_ManagedVectorStoresTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) result.append(cached_obj) @@ -2487,7 +2518,7 @@ async def get_managed_vector_store_rows_by_uuids( async def get_org_object( org_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span] = None, proxy_logging_obj: Optional[ProxyLogging] = None, include_budget_table: bool = False, @@ -2518,12 +2549,12 @@ async def get_org_object( cache_key = "org_id:{}:with_budget".format(org_id) # check if in cache - cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key) - if cached_org_obj is not None: - if isinstance(cached_org_obj, dict): - return LiteLLM_OrganizationTable(**cached_org_obj) - elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): - return cached_org_obj + deserialized_org = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_OrganizationTable, + ) + if deserialized_org is not None: + return deserialized_org # else, check db try: query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}} @@ -2537,16 +2568,16 @@ async def get_org_object( if response is None: raise Exception + _org_obj = LiteLLM_OrganizationTable(**response.model_dump()) # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=( - response.model_dump() if hasattr(response, "model_dump") else response - ), + value=_org_obj, + model_type=LiteLLM_OrganizationTable, ttl=DEFAULT_IN_MEMORY_TTL, ) - return response + return _org_obj except Exception: raise Exception( f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call." @@ -2559,7 +2590,7 @@ async def _get_resources_from_access_groups( "access_model_names", "access_mcp_server_ids", "access_agent_ids" ], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -2617,7 +2648,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -2636,7 +2667,7 @@ async def _get_models_from_access_groups( async def _get_mcp_server_ids_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -2655,7 +2686,7 @@ async def _get_mcp_server_ids_from_access_groups( async def _get_agent_ids_from_access_groups( access_group_ids: List[str], prisma_client: Optional[PrismaClient] = None, - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> List[str]: """ @@ -3379,7 +3410,7 @@ async def _check_team_member_budget( user_object: Optional[LiteLLM_UserTable], valid_token: Optional[UserAPIKeyAuth], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ): """Check if team member is over their max budget within the team.""" @@ -3447,7 +3478,7 @@ async def _check_team_member_model_access( valid_token: UserAPIKeyAuth, llm_router: Optional[Router], prisma_client: Optional["PrismaClient"], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> None: """ @@ -3754,7 +3785,7 @@ async def _project_soft_budget_check( async def get_project_object( project_id: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Optional[ProxyLogging] = None, ) -> Optional[LiteLLM_ProjectTableCachedObj]: """ @@ -3769,12 +3800,12 @@ async def get_project_object( # Check cache first cache_key = "project_id:{}".format(project_id) - cached_obj = await user_api_key_cache.async_get_cache(key=cache_key) - if cached_obj is not None: - if isinstance(cached_obj, dict): - return LiteLLM_ProjectTableCachedObj(**cached_obj) - elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj): - return cached_obj + deserialized_project = await user_api_key_cache.async_get_cache( + key=cache_key, + model_type=LiteLLM_ProjectTableCachedObj, + ) + if deserialized_project is not None: + return deserialized_project # Fetch from DB project_row = await prisma_client.db.litellm_projecttable.find_unique( @@ -3793,6 +3824,7 @@ async def get_project_object( value=project_obj, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, + model_type=LiteLLM_ProjectTableCachedObj, ) return project_obj @@ -3802,7 +3834,7 @@ async def _organization_max_budget_check( valid_token: Optional[UserAPIKeyAuth], team_object: Optional[LiteLLM_TeamTable], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ): """ @@ -3896,7 +3928,7 @@ async def _organization_max_budget_check( async def _tag_max_budget_check( request_body: dict, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, valid_token: Optional[UserAPIKeyAuth], ): diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 91c8f2dd7c..cbed34adac 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -167,6 +167,81 @@ def _allow_model_level_clientside_configurable_parameters( ) +# Config dicts whose entries are spread as ``**dict`` into outbound LLM +# API calls. ``litellm_embedding_config`` is consumed by the Milvus +# vector store transformer; future nested-config keys with the same +# threat shape should be added here. +_NESTED_CONFIG_KEYS: Tuple[str, ...] = ("litellm_embedding_config",) + +# Banned root-level params. Same list applies to every entry in +# ``_NESTED_CONFIG_KEYS`` because those dicts get spread as ``**kwargs`` +# into the same outbound calls. +_BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( + "api_base", + "base_url", + "user_config", + "aws_sts_endpoint", + "aws_web_identity_token", + "aws_role_name", + "vertex_credentials", + # Endpoint-targeting fields that retarget the outbound request or + # an observability callback. An attacker-controlled value either + # exfiltrates the request payload (incl. messages + admin-set + # tokens) to the attacker's host, or coerces the proxy into + # authenticating against the attacker's host with admin secrets. + "aws_bedrock_runtime_endpoint", + "langsmith_base_url", + "langfuse_host", + "posthog_host", + "braintrust_host", + "slack_webhook_url", + # Provider-specific endpoint overrides that flow into the outbound + # request via ``optional_params``. Same threat as ``api_base``: + # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker + # S3; ``sagemaker_base_url`` redirects all SageMaker traffic; + # ``deployment_url`` redirects SAP deployments. + "s3_endpoint_url", + "sagemaker_base_url", + "deployment_url", +) + + +def _check_banned_params( + body: dict, + general_settings: dict, + llm_router: Optional[Router], + model: str, +) -> None: + """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in. + + Shared between the root-level check and the nested-config check so a + new banned param only needs to be added in one place. + """ + for param in _BANNED_REQUEST_BODY_PARAMS: + if param not in body: + continue + if general_settings.get("allow_client_side_credentials") is True: + return + if ( + _allow_model_level_clientside_configurable_parameters( + model=model, + param=param, + request_body_value=body[param], + llm_router=llm_router, + ) + is True + ): + return + raise ValueError( + f"Rejected Request: {param} is not allowed in request body. " + "Clientside passthrough requires explicit admin opt-in via " + "either `general_settings.allow_client_side_credentials = true` " + "(proxy-wide) or `configurable_clientside_auth_params` on the " + "deployment in your proxy config.yaml. " + "Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997", + ) + + def is_request_body_safe( request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str ) -> bool: @@ -175,72 +250,31 @@ def is_request_body_safe( A malicious user can set the api_base to their own domain and invoke POST /chat/completions to intercept and steal the OpenAI API key. Relevant issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997 + + The blocklist is enforced unconditionally. Legitimate clientside + credential / endpoint passthrough goes through one of the two + explicit admin opt-ins (``general_settings.allow_client_side_credentials`` + proxy-wide or ``configurable_clientside_auth_params`` per deployment). + Historically there was a third, *implicit*, *caller-controlled* path: + ``check_complete_credentials`` returned True when the caller supplied + any non-empty ``api_key``, which made the entire blocklist a no-op. + That bypass turned every missing entry on the blocklist into an + exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp, + GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l, + b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now + has a single, predictable failure mode for missing entries (a 400), + not a credential leak. + + Iterative single-level descent into ``_NESTED_CONFIG_KEYS`` (rather + than recursion) covers nested-config attacks like Milvus's + ``litellm_embedding_config.api_base`` (VERIA-6) without exposing a + recursion-depth DoS surface. """ - banned_params = [ - "api_base", - "base_url", - "user_config", - "aws_sts_endpoint", - "aws_web_identity_token", - "aws_role_name", - "vertex_credentials", - # Endpoint-targeting fields that retarget the outbound request or - # an observability callback. An attacker-controlled value either - # exfiltrates the request payload (incl. messages + admin-set - # tokens) to the attacker's host, or coerces the proxy into - # authenticating against the attacker's host with admin secrets. - "aws_bedrock_runtime_endpoint", - "langsmith_base_url", - "langfuse_host", - "posthog_host", - "braintrust_host", - "slack_webhook_url", - # Provider-specific endpoint overrides that flow into the outbound - # request via ``optional_params``. Same threat as ``api_base``: - # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker - # S3; ``sagemaker_base_url`` redirects all SageMaker traffic; - # ``deployment_url`` redirects SAP deployments. - "s3_endpoint_url", - "sagemaker_base_url", - "deployment_url", - ] - - # The blocklist is enforced unconditionally. Legitimate clientside - # credential / endpoint passthrough goes through one of the two - # explicit admin opt-ins (``general_settings.allow_client_side_credentials`` - # proxy-wide or ``configurable_clientside_auth_params`` per deployment). - # Historically there was a third, *implicit*, *caller-controlled* path: - # ``check_complete_credentials`` returned True when the caller supplied - # any non-empty ``api_key``, which made the entire blocklist a no-op. - # That bypass turned every missing entry on the blocklist into an - # exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp, - # GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l, - # b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now - # has a single, predictable failure mode for missing entries (a 400), - # not a credential leak. - for param in banned_params: - if param in request_body: - if general_settings.get("allow_client_side_credentials") is True: - return True - elif ( - _allow_model_level_clientside_configurable_parameters( - model=model, - param=param, - request_body_value=request_body[param], - llm_router=llm_router, - ) - is True - ): - return True - raise ValueError( - f"Rejected Request: {param} is not allowed in request body. " - "Clientside passthrough requires explicit admin opt-in via " - "either `general_settings.allow_client_side_credentials = true` " - "(proxy-wide) or `configurable_clientside_auth_params` on the " - "deployment in your proxy config.yaml. " - "Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997", - ) - + _check_banned_params(request_body, general_settings, llm_router, model) + for nested_key in _NESTED_CONFIG_KEYS: + nested = request_body.get(nested_key) + if isinstance(nested, dict): + _check_banned_params(nested, general_settings, llm_router, model) return True diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index f50c950d74..71411bed7f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -6,6 +6,8 @@ Currently only supports admin. JWT token must have 'litellm_proxy_admin' in scope. """ +from __future__ import annotations + import fnmatch import hashlib import os @@ -20,7 +22,6 @@ import jwt from jwt.api_jwk import PyJWK from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.llms.custom_httpx.httpx_handler import HTTPHandler @@ -46,6 +47,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging from .auth_checks import ( @@ -73,7 +75,7 @@ class JWTHandler: """ prisma_client: Optional[PrismaClient] - user_api_key_cache: DualCache + user_api_key_cache: UserApiKeyCache # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret # the key in different ways (e.g. HS* and RS*)." @@ -99,7 +101,7 @@ class JWTHandler: def update_environment( self, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, litellm_jwtauth: LiteLLM_JWTAuth, leeway: int = 0, ) -> None: @@ -952,7 +954,7 @@ class JWTAuthManager: jwt_handler: JWTHandler, jwt_valid_token: dict, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: @@ -1045,7 +1047,7 @@ class JWTAuthManager: route: str, jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: @@ -1133,7 +1135,7 @@ class JWTAuthManager: valid_user_email: Optional[bool], jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, route: str, @@ -1349,7 +1351,7 @@ class JWTAuthManager: jwt_valid_token: dict, user_object: Optional[LiteLLM_UserTable], prisma_client: Optional[PrismaClient], - user_api_key_cache: Optional[DualCache] = None, + user_api_key_cache: Optional[UserApiKeyCache] = None, ) -> None: """ Sync user role and team memberships with JWT claims @@ -1377,7 +1379,8 @@ class JWTAuthManager: if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=user_object.model_dump(), + value=user_object, + model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) @@ -1400,7 +1403,8 @@ class JWTAuthManager: if user_api_key_cache is not None: await user_api_key_cache.async_set_cache( key=user_object.user_id, - value=user_object.model_dump(), + value=user_object, + model_type=LiteLLM_UserTable, ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, ) return None @@ -1412,7 +1416,7 @@ class JWTAuthManager: request_headers: Optional[dict], jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> None: @@ -1456,7 +1460,7 @@ class JWTAuthManager: user_object: Optional[LiteLLM_UserTable], user_id: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, team_id_upsert: Optional[bool], @@ -1514,7 +1518,7 @@ class JWTAuthManager: general_settings: dict, route: str, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, request_headers: Optional[dict] = None, diff --git a/litellm/proxy/auth/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 34fab4849e..39d3282942 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -13,6 +13,10 @@ from fastapi import Request from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.auth_utils import _get_request_ip_address +# One-shot warning so operators upgrading from the prior "always trust X-Forwarded-*" +# behaviour see an actionable message in their logs the first time it triggers. +_warned_xff_without_trusted_ranges = False + class IPAddressUtils: """Static utilities for IP-based MCP access control.""" @@ -106,6 +110,61 @@ class IPAddressUtils: return any(addr in network for network in networks) + @staticmethod + def is_request_from_trusted_proxy( + request: Request, + general_settings: Optional[Dict[str, Any]] = None, + ) -> bool: + """ + Return True if X-Forwarded-* headers on this request should be trusted. + + Trusts the headers iff both: + 1. ``use_x_forwarded_for`` is enabled in proxy settings, AND + 2. ``mcp_trusted_proxy_ranges`` is configured AND the direct + connection IP (``request.client.host``) falls inside one of + those CIDRs. + + When ``use_x_forwarded_for`` is enabled but ``mcp_trusted_proxy_ranges`` + is missing, the headers are NOT trusted: there is no way to + distinguish a trusted reverse proxy from a direct attacker, so callers + that build URLs (OAuth issuer / redirect_uri / etc.) must fall back + to the request's literal base URL instead of risking a poisoned host. + """ + if general_settings is None: + try: + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) + + general_settings = proxy_general_settings + except ImportError: + general_settings = {} + + if general_settings is None: + general_settings = {} + + if not general_settings.get("use_x_forwarded_for", False): + return False + + trusted_ranges = general_settings.get("mcp_trusted_proxy_ranges") + if not trusted_ranges: + global _warned_xff_without_trusted_ranges + if not _warned_xff_without_trusted_ranges: + verbose_proxy_logger.warning( + "use_x_forwarded_for is enabled but mcp_trusted_proxy_ranges " + "is not configured. X-Forwarded-* headers will NOT be " + "trusted, so MCP OAuth discovery URLs will use the proxy's " + "literal base URL. Set mcp_trusted_proxy_ranges in " + "general_settings to your reverse-proxy CIDR(s) to allow " + "X-Forwarded-* through." + ) + _warned_xff_without_trusted_ranges = True + return False + + direct_ip = request.client.host if request.client else None + trusted_networks = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges) + return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks) + @staticmethod def get_mcp_client_ip( request: Request, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f0c2a4514f..bfd1f2e0b3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -20,7 +20,6 @@ from fastapi.security.api_key import APIKeyHeader import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging -from litellm.caching import DualCache from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value @@ -60,6 +59,7 @@ from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -329,7 +329,7 @@ _global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEN async def _fetch_global_spend_with_event_coordination( cache_key: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> Optional[float]: """ @@ -345,14 +345,14 @@ async def _fetch_global_spend_with_event_coordination( return await _global_spend_coordinator.get_or_load( cache_key=cache_key, - cache=user_api_key_cache, + cache=user_api_key_cache, # pyright: ignore[reportArgumentType] load_fn=_load_global_spend, ) async def get_global_proxy_spend( litellm_proxy_admin_name: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: Optional[PrismaClient], token: str, proxy_logging_obj: ProxyLogging, @@ -473,7 +473,12 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints( for endpoint in pass_through_endpoints: if isinstance(endpoint, dict) and endpoint.get("path", "") == route: ## IF AUTH DISABLED - if endpoint.get("auth") is not True: + # Default to True: a config dict with no ``auth`` key + # otherwise produced an unauthenticated forwarder. The + # Pydantic ``PassThroughGenericEndpoint.auth`` default + # is also True, but raw config dicts skip that path — + # so this runtime check has to default to True too. + if endpoint.get("auth", True) is not True: return UserAPIKeyAuth() ## IF AUTH ENABLED ### IF CUSTOM PARSER REQUIRED @@ -505,7 +510,7 @@ async def _resolve_jwt_to_virtual_key( jwt_claims: dict, jwt_handler: JWTHandler, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Optional[UserAPIKeyAuth]: @@ -1107,9 +1112,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 is_master_key_valid = False ## VALIDATE MASTER KEY ## - try: - assert isinstance(master_key, str) - except Exception: + if not isinstance(master_key, str): raise HTTPException( status_code=500, detail={ @@ -1179,11 +1182,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if len(api_key) > 8 else "****" ) - assert api_key.startswith( - "sk-" - ), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( - _masked_key - ) # prevent token hashes from being used + if not api_key.startswith("sk-"): + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=( + "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( + _masked_key + ) + ), + ) # prevent token hashes from being used else: verbose_logger.warning( "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format( @@ -1291,7 +1298,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" team_member_info = await user_api_key_cache.async_get_cache( - key=_cache_key + key=_cache_key, + model_type=LiteLLM_TeamMembership, ) if team_member_info is None: # read from DB @@ -1299,18 +1307,23 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _team_id = valid_token.team_id if _user_id is not None and _team_id is not None: - team_member_info = await prisma_client.db.litellm_teammembership.find_first( + _db_member = await prisma_client.db.litellm_teammembership.find_first( where={ "user_id": _user_id, "team_id": _team_id, }, # type: ignore include={"litellm_budget_table": True}, ) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - ttl=5, - ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership( + **_db_member.dict() + ) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=team_member_info, + model_type=LiteLLM_TeamMembership, + ttl=5, + ) if ( team_member_info is not None @@ -1457,9 +1470,13 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 else: valid_token.team_object_permission = None - await user_api_key_cache.async_set_cache( - key=valid_token.team_id, value=_team_obj - ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + # Only cache when the key is a real team_id (non-team keys must not use key=None). + if valid_token.team_id is not None and _team_obj is not None: + await user_api_key_cache.async_set_cache( + key=valid_token.team_id, + value=_team_obj, + model_type=LiteLLM_TeamTableCachedObj, + ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py # Fetch project object if key belongs to a project _project_obj = None diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 5dcc88cacb..adf562d69c 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -313,23 +313,24 @@ sequenceDiagram participant Proxy as LiteLLM Proxy participant SSO as SSO Provider - CLI->>CLI: Generate key ID (sk-uuid) - CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=sk-uuid + CLI->>Proxy: POST /sso/cli/start + Proxy->>CLI: Return login_id, poll_secret, user_code + CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=login_id - Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=sk-uuid - Proxy->>Proxy: Set cli_state = litellm-session-token:sk-uuid - Proxy->>SSO: Redirect with state=litellm-session-token:sk-uuid + Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=login_id + Proxy->>Proxy: Set cli_state = litellm-session-token:login_id + Proxy->>SSO: Redirect with state=litellm-session-token:login_id SSO->>Browser: Show login page Browser->>SSO: User authenticates - SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:sk-uuid + SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:login_id Proxy->>Proxy: Check if state starts with "litellm-session-token:" - Proxy->>Proxy: Generate API key with ID=sk-uuid - Proxy->>Browser: Show success page + Proxy->>Browser: Prompt for user_code + Browser->>Proxy: POST /sso/cli/complete/login_id - CLI->>Proxy: Poll /sso/cli/poll/sk-uuid - Proxy->>CLI: Return {"status": "ready", "key": "sk-uuid"} + CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header + Proxy->>CLI: Return {"status": "ready", "key": "jwt"} CLI->>CLI: Save key to ~/.litellm/token.json ``` @@ -343,13 +344,13 @@ The CLI provides three authentication commands: ### Authentication Flow Steps -1. **Generate Session ID**: CLI generates a unique key ID (`sk-{uuid}`) -2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and key parameters -3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:sk-uuid`) as OAuth state parameter and redirects to SSO provider +1. **Start Session**: CLI creates a short-lived login session with `/sso/cli/start` +2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and login ID parameters +3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:{login_id}`) as OAuth state parameter and redirects to SSO provider 4. **User Authentication**: User completes SSO authentication in browser 5. **Callback Processing**: SSO provider redirects back to proxy with state parameter -6. **Key Generation**: Proxy detects CLI login (state starts with "litellm-session-token:") and generates API key with pre-specified ID -7. **Polling**: CLI polls `/sso/cli/poll/{key_id}` endpoint until key is ready +6. **User Code Verification**: Browser confirms the verification code shown in the CLI +7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready 8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json` ### Benefits of This Approach @@ -357,7 +358,7 @@ The CLI provides three authentication commands: - **No Local Server**: No need to run a local callback server - **Standard OAuth**: Uses OAuth 2.0 state parameter correctly - **Remote Compatible**: Works with remote proxy servers -- **Secure**: Uses UUID session identifiers +- **Secure**: Keeps the polling secret out of the browser handoff - **Simple Setup**: No additional OAuth redirect URL configuration needed ### Token Storage diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index aeb59e78a5..447837c35e 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -5,6 +5,7 @@ import time import webbrowser from pathlib import Path from typing import Any, Dict, List, Optional +from urllib.parse import urlencode import click import requests @@ -52,12 +53,16 @@ def clear_token() -> None: os.remove(token_file) -def get_stored_api_key() -> Optional[str]: - """Get the stored API key from token file""" - # Use the SDK-level utility +def get_stored_api_key(expected_base_url: Optional[str] = None) -> Optional[str]: + """Get the stored API key from token file. + + If expected_base_url is provided, the key is only returned when it was + originally issued for that URL. This prevents credential leakage when the + CLI is pointed at a different (possibly malicious) server. + """ from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key - return get_litellm_gateway_api_key() + return get_litellm_gateway_api_key(expected_base_url=expected_base_url) # Team selection utilities @@ -241,7 +246,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any def prompt_team_selection_fallback( - teams: List[Dict[str, Any]] + teams: List[Dict[str, Any]], ) -> Optional[Dict[str, Any]]: """Fallback team selection for non-interactive environments""" if not teams: @@ -279,6 +284,7 @@ def prompt_team_selection_fallback( def _poll_for_ready_data( url: str, *, + headers: Optional[Dict[str, str]] = None, total_timeout: int = 300, poll_interval: int = 2, request_timeout: int = 10, @@ -291,7 +297,10 @@ def _poll_for_ready_data( ) -> Optional[Dict[str, Any]]: for attempt in range(total_timeout // poll_interval): try: - response = requests.get(url, timeout=request_timeout) + request_kwargs: Dict[str, Any] = {"timeout": request_timeout} + if headers is not None: + request_kwargs["headers"] = headers + response = requests.get(url, **request_kwargs) if response.status_code == 200: data = response.json() status = data.get("status") @@ -346,7 +355,23 @@ def _normalize_teams(teams, team_details): return [] -def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]: +def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]: + response = requests.post(f"{base_url}/sso/cli/start", timeout=10) + response.raise_for_status() + data = response.json() + required_fields = ("login_id", "poll_secret", "user_code") + if not all(isinstance(data.get(field), str) for field in required_fields): + raise ValueError("Invalid CLI SSO start response") + return data + + +def _get_cli_sso_poll_headers(poll_secret: str) -> Dict[str, str]: + return {"x-litellm-cli-poll-secret": poll_secret} + + +def _poll_for_authentication( + base_url: str, key_id: str, poll_secret: str +) -> Optional[dict]: """ Poll the server for authentication completion and handle team selection. @@ -356,6 +381,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]: poll_url = f"{base_url}/sso/cli/poll/{key_id}" data = _poll_for_ready_data( poll_url, + headers=_get_cli_sso_poll_headers(poll_secret), pending_message="Still waiting for authentication...", ) if not data: @@ -373,6 +399,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]: jwt_with_team = _handle_team_selection_during_polling( base_url=base_url, key_id=key_id, + poll_secret=poll_secret, teams=normalized_teams, ) @@ -410,7 +437,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]: def _handle_team_selection_during_polling( - base_url: str, key_id: str, teams: List[Dict[str, Any]] + base_url: str, key_id: str, poll_secret: str, teams: List[Dict[str, Any]] ) -> Optional[str]: """ Handle team selection and re-poll with selected team_id. @@ -441,6 +468,7 @@ def _handle_team_selection_during_polling( poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}" data = _poll_for_ready_data( poll_url, + headers=_get_cli_sso_poll_headers(poll_secret), pending_message="Still waiting for team authentication...", other_status_message="Waiting for team authentication to complete...", http_error_log_every=10, @@ -514,29 +542,24 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option @click.pass_context def login(ctx: click.Context): """Login to LiteLLM proxy using SSO authentication""" - from litellm._uuid import uuid from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER from litellm.proxy.client.cli.interface import show_commands base_url = ctx.obj["base_url"] - # Check if we have an existing key to regenerate - existing_key = get_stored_api_key() - - # Generate unique key ID for this login session - key_id = f"sk-{str(uuid.uuid4())}" - try: - # Construct SSO login URL with CLI source and pre-generated key - sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}" + cli_sso_flow = _start_cli_sso_flow(base_url=base_url) + key_id = cli_sso_flow["login_id"] + poll_secret = cli_sso_flow["poll_secret"] + user_code = cli_sso_flow["user_code"] - # If we have an existing key, include it as a parameter to the login endpoint - # The server will encode it in the OAuth state parameter for the SSO flow - if existing_key: - sso_url += f"&existing_key={existing_key}" + sso_url = f"{base_url}/sso/key/generate?" + urlencode( + {"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id} + ) click.echo(f"Opening browser to: {sso_url}") click.echo("Please complete the SSO authentication in your browser...") + click.echo(f"Verification code: {user_code}") click.echo(f"Session ID: {key_id}") # Open browser @@ -545,15 +568,19 @@ def login(ctx: click.Context): # Poll for authentication completion click.echo("Waiting for authentication...") - auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id) + auth_result = _poll_for_authentication( + base_url=base_url, key_id=key_id, poll_secret=poll_secret + ) if auth_result: api_key = auth_result["api_key"] user_id = auth_result["user_id"] - # Save token data (simplified for CLI - we just need the key) + # Save token data. base_url is stored so we can verify origin + # before reusing the key on a subsequent CLI invocation. save_token( { + "base_url": base_url.rstrip("/"), "key": api_key, "user_id": user_id or "cli-user", "user_email": "unknown", diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index 22de5a7861..be55f79c06 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -74,9 +74,10 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None: """LiteLLM Proxy CLI - Manage your LiteLLM proxy server""" ctx.ensure_object(dict) - # If no API key provided via flag or environment variable, try to load from saved token + # If no API key provided via flag or environment variable, try to load from saved token. + # Pass base_url so we only use the stored key when it was issued for this server. if api_key is None: - api_key = get_stored_api_key() + api_key = get_stored_api_key(expected_base_url=base_url) ctx.obj["base_url"] = base_url ctx.obj["api_key"] = api_key diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index 12b5cd79f7..929ad46a77 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -28,12 +28,17 @@ class Client: api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. timeout: Request timeout in seconds (default: 30) """ - self._base_url = base_url.rstrip("/") # Remove trailing slash if present - self._api_key = get_litellm_gateway_api_key() or api_key + self._base_url = base_url.rstrip("/") + # Only use the stored CLI key when it was issued for this server. + self._api_key = api_key or get_litellm_gateway_api_key( + expected_base_url=self._base_url + ) # Initialize resource clients - self.http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout) + self.http = HTTPClient( + base_url=base_url, api_key=self._api_key, timeout=timeout + ) self.models = ModelsManagementClient( base_url=self._base_url, api_key=self._api_key ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 76c52f83ee..f3138f10da 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -744,6 +744,11 @@ class ProxyBaseLLMRequestProcessing: "aingest", "aretrieve_container", "adelete_container", + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", "acreate_skill", "alist_skills", "aget_skill", @@ -1001,6 +1006,11 @@ class ProxyBaseLLMRequestProcessing: "aingest", "aretrieve_container", "adelete_container", + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", "acreate_skill", "alist_skills", "aget_skill", diff --git a/litellm/proxy/common_utils/cache_coordinator.py b/litellm/proxy/common_utils/cache_coordinator.py index 24da9450ab..abb0402d3b 100644 --- a/litellm/proxy/common_utils/cache_coordinator.py +++ b/litellm/proxy/common_utils/cache_coordinator.py @@ -20,11 +20,27 @@ T = TypeVar("T") class AsyncCacheProtocol(Protocol): - """Protocol for cache backends used by EventDrivenCacheCoordinator.""" + """Protocol for cache backends used by EventDrivenCacheCoordinator. - async def async_get_cache(self, key: str, **kwargs: Any) -> Any: ... + Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params + before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate. + """ - async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> Any: ... + async def async_get_cache( + self, + key: str, + parent_otel_span: Any = None, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... + + async def async_set_cache( + self, + key: str, + value: Any, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... class EventDrivenCacheCoordinator: @@ -36,6 +52,9 @@ class EventDrivenCacheCoordinator: - Other requests: wait for the signal, then read from cache. Create one instance per resource (e.g. one for global spend, one for feature flags). + + Args: + log_prefix: Prefix for debug log messages. """ def __init__(self, log_prefix: str = "[CACHE]"): diff --git a/litellm/proxy/common_utils/cache_pydantic_utils.py b/litellm/proxy/common_utils/cache_pydantic_utils.py new file mode 100644 index 0000000000..80a8d6281a --- /dev/null +++ b/litellm/proxy/common_utils/cache_pydantic_utils.py @@ -0,0 +1,93 @@ +""" +DualCache presents a single API for reads and writes, but the two backends behave +differently: the in-memory layer can store arbitrary Python objects (including live +``BaseModel`` instances), while Redis persists strings and therefore needs JSON-safe +payloads (``json.dumps`` on the Redis side). + +Call sites therefore see cache ``value`` / ``cached`` as effectively ``Any``: the same +key may deserialize to a model on one process (memory hit) or to a ``dict`` after a +Redis round-trip. ``CacheCodec`` centralizes encode/decode at that boundary: +``CacheCodec.serialize`` before ``set``, ``CacheCodec.deserialize`` after ``get`` +when you need a typed ``BaseModel``. + +``dataclasses`` are not supported: only ``dict`` and Pydantic ``BaseModel`` inputs +are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` first. +""" + +from __future__ import annotations + +from typing import Any, Optional, Type, TypeVar + +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger + +T = TypeVar("T", bound=BaseModel) + + +class CacheCodec: + """ + Encode/decode Pydantic models for DualCache (memory vs Redis safe payloads). + + Dataclasses are not supported yet (only ``dict`` and ``BaseModel``). + + Use ``serialize`` with ``model_type`` when writing so the same schema is used + as on read (``deserialize``). Pass ``model_type`` whenever you know it + (validates ``dict`` payloads and normalizes ``BaseModel`` instances). + """ + + @staticmethod + def serialize(value: Any, model_type: Optional[Type[T]] = None) -> Any: + """ + Encode a value for DualCache / Redis (``json.dumps``-safe). + + If ``model_type`` is set, the payload is validated with that model, then + ``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``. + + If the value is already an instance of ``model_type`` (or a subclass), + ``model_validate`` is skipped to avoid an unnecessary Pydantic copy — the + value is dumped directly. + + If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other + values (e.g. plain ``dict``) are returned unchanged. + """ + if model_type is not None: + if isinstance(value, model_type): + # Already the right type: dump directly, skip re-validation copy. + return value.model_dump(mode="json", exclude_none=True) + if isinstance(value, (dict, BaseModel)): + return model_type.model_validate(value).model_dump( + mode="json", exclude_none=True + ) + return value + if isinstance(value, BaseModel): + return value.model_dump(mode="json", exclude_none=True) + return value + + @staticmethod + def deserialize(cached: Any, model_type: Type[T]) -> Optional[T]: + """ + Decode a cache entry to ``model_type``. + + - ``None`` → ``None`` + - Already an instance of ``model_type`` (including subclasses) → returned as-is + - ``dict`` → ``model_type.model_validate(...)``; on ``ValidationError``, + logs a warning and returns ``None`` (treat as cache miss; avoids serving + malformed or schema-drifted entries) + - Any other type → ``None`` (caller should treat as cache miss or log) + """ + if cached is None: + return None + if isinstance(cached, model_type): + return cached + if isinstance(cached, dict): + try: + return model_type.model_validate(cached) + except ValidationError as e: + verbose_proxy_logger.warning( + "CacheCodec.deserialize: validation failed for %s (%s)", + model_type.__name__, + e, + ) + return None + return None diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index c25d853312..67a2456746 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -8,7 +8,7 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, @@ -31,7 +31,7 @@ class ExpiredUISessionKeyCleanupManager: def __init__( self, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, pod_lock_manager=None, ): self.prisma_client = prisma_client diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index e486336cec..0928ce914d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -52,6 +52,37 @@ class ResetBudgetJob: ### RESET MULTI-WINDOW BUDGETS ### await self.reset_budget_windows() + @staticmethod + async def _invalidate_spend_counter(counter_key: str) -> None: + """Zero a spend counter so a DB-row reset takes effect immediately. + + Call AFTER the DB write commits. Clearing Redis before the DB + commit opens a window where get_current_spend reads 0 from Redis + while the DB still holds the pre-reset value, allowing bypass. + """ + try: + from litellm.proxy.proxy_server import spend_counter_cache + + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=0.0, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset spend counter %s in Redis: %s. " + "Budget may be over-enforced until counter expires.", + counter_key, + redis_err, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to reset spend counter %s: %s", counter_key, e + ) + async def reset_budget_for_litellm_team_members( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -64,46 +95,30 @@ class ResetBudgetJob: if budget.budget_id is not None ] - # Reset spend counters for affected team members. - # Reset Redis directly so a transient failure doesn't leave stale - # counters that get_current_spend would read as authoritative. try: - from litellm.proxy.proxy_server import spend_counter_cache - memberships = await self.prisma_client.db.litellm_teammembership.find_many( where={"budget_id": {"in": budget_ids}} ) - for m in memberships: - counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" - # Always reset in-memory - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 - ) - # Explicitly reset Redis with warning on failure - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset team member spend counter in Redis %s: %s. " - "Budget may be over-enforced until counter expires.", - counter_key, - redis_err, - ) except Exception as e: + memberships = [] verbose_proxy_logger.warning( - "Failed to reset team member spend counters: %s", e + "Failed to fetch team memberships for counter invalidation: %s", e ) - return await self.prisma_client.db.litellm_teammembership.update_many( + update_result = await self.prisma_client.db.litellm_teammembership.update_many( where={"budget_id": {"in": budget_ids}}, data={ "spend": 0, }, ) + for m in memberships: + await self._invalidate_spend_counter( + f"spend:team_member:{m.user_id}:{m.team_id}" + ) + + return update_result + async def reset_budget_for_keys_linked_to_budgets( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -126,17 +141,36 @@ class ResetBudgetJob: if not budget_ids: return - return await self.prisma_client.db.litellm_verificationtoken.update_many( - where={ - "budget_id": {"in": budget_ids}, - "budget_duration": None, # only keys without their own reset schedule - "spend": {"gt": 0}, # only reset keys that have accumulated spend - }, - data={ - "spend": 0, - }, + where_clause: dict = { + "budget_id": {"in": budget_ids}, + "budget_duration": None, # only keys without their own reset schedule + "spend": {"gt": 0}, # only reset keys that have accumulated spend + } + + try: + keys = await self.prisma_client.db.litellm_verificationtoken.find_many( + where=where_clause + ) + except Exception as e: + keys = [] + verbose_proxy_logger.warning( + "Failed to fetch keys for counter invalidation: %s", e + ) + + update_result = ( + await self.prisma_client.db.litellm_verificationtoken.update_many( + where=where_clause, + data={ + "spend": 0, + }, + ) ) + for k in keys: + await self._invalidate_spend_counter(f"spend:key:{k.token}") + + return update_result + async def reset_budget_for_litellm_budget_table(self): """ Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired @@ -365,6 +399,10 @@ class ResetBudgetJob: data_list=updated_keys, table_name="key", ) + for k in updated_keys: + token = getattr(k, "token", None) + if token: + await self._invalidate_spend_counter(f"spend:key:{token}") end_time = time.time() if len(failed_keys) > 0: # If any keys failed to reset @@ -450,6 +488,12 @@ class ResetBudgetJob: data_list=updated_users, table_name="user", ) + for u in updated_users: + user_id = getattr(u, "user_id", None) + if user_id: + await self._invalidate_spend_counter( + f"spend:user:{user_id}" + ) end_time = time.time() if len(failed_users) > 0: # If any users failed to reset @@ -541,6 +585,12 @@ class ResetBudgetJob: data_list=updated_teams, table_name="team", ) + for t in updated_teams: + team_id = getattr(t, "team_id", None) + if team_id: + await self._invalidate_spend_counter( + f"spend:team:{team_id}" + ) end_time = time.time() if len(failed_teams) > 0: # If any teams failed to reset diff --git a/litellm/proxy/common_utils/static_asset_utils.py b/litellm/proxy/common_utils/static_asset_utils.py new file mode 100644 index 0000000000..c108af2b47 --- /dev/null +++ b/litellm/proxy/common_utils/static_asset_utils.py @@ -0,0 +1,52 @@ +"""Helpers for unauthenticated logo / favicon endpoints.""" + +import os +from typing import Optional, Tuple + +from litellm._logging import verbose_proxy_logger + +LOCAL_IMAGE_HEADER_BYTES = 512 + + +def detect_local_image_media_type(header: bytes) -> Optional[str]: + """Return a browser image media type for supported local image signatures.""" + if header[0:8] == b"\x89PNG\r\n\x1a\n": + return "image/png" + if header[0:4] == b"GIF8" and header[5:6] == b"a": + return "image/gif" + if header[0:3] == b"\xff\xd8\xff": + return "image/jpeg" + if header[0:4] == b"RIFF" and header[8:12] == b"WEBP": + return "image/webp" + if header[0:4] in (b"\x00\x00\x01\x00", b"\x00\x00\x02\x00"): + return "image/x-icon" + return None + + +def resolve_validated_local_image_path(candidate: str) -> Optional[Tuple[str, str]]: + """Resolve ``candidate`` only when it is an existing supported image file.""" + if not candidate: + return None + try: + resolved = os.path.realpath(os.path.expanduser(candidate)) + except (OSError, ValueError): + return None + if not os.path.isfile(resolved): + return None + + try: + with open(resolved, "rb") as f: + header = f.read(LOCAL_IMAGE_HEADER_BYTES) + except OSError as exc: + verbose_proxy_logger.debug("Could not read local asset %r: %s", candidate, exc) + return None + + media_type = detect_local_image_media_type(header) + if media_type is None: + verbose_proxy_logger.warning( + "Local asset %r is not a supported image file; falling back to default.", + candidate, + ) + return None + + return resolved, media_type diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py new file mode 100644 index 0000000000..914be36457 --- /dev/null +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +from typing import Any, Optional, Type, TypeVar, Union, cast, overload + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + +T = TypeVar("T", bound=BaseModel) + + +class UserApiKeyCache(DualCache): + """ + DualCache wrapper for UserAPIKeyAuth-like payloads. + + Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid + "memory returns BaseModel, Redis returns dict" format drift. + + When ``model_type`` is provided: + - writes are serialized via ``CacheCodec.serialize(..., model_type=...)`` + - reads are deserialized via ``CacheCodec.deserialize(..., model_type)`` + and return ``Optional[T]``: the model on success, ``None`` on cache miss + **or** if the cached payload fails validation (schema drift). On + validation failure after a cache hit, an error line is emitted via + ``verbose_proxy_logger``. + + When ``model_type`` is omitted, the interface behaves like ``DualCache``: + raw cached payload is returned (dict/str/etc.). + + ``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting + ``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis). + + ``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous + (no other methods in between) so mypy resolves ``@overload`` + implementation correctly. + """ + + @overload + def get_cache( + self, + key: Any, + parent_otel_span: Any = None, + local_only: bool = False, + *, + model_type: Type[T], + **kwargs: Any, + ) -> Optional[T]: ... + + @overload + def get_cache( + self, + key: Any, + parent_otel_span: Any = None, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... + + def get_cache( # type: ignore[override] + self, + key, + parent_otel_span=None, + local_only: bool = False, + model_type: Optional[Type[BaseModel]] = None, + **kwargs, + ) -> Union[Any, Optional[BaseModel]]: + if model_type is None and "model_type" in kwargs: + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + cached = super().get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) + if model_type is None: + return cached + if cached is None: + return None + decoded = CacheCodec.deserialize(cached, model_type=model_type) + if decoded is None: + verbose_proxy_logger.error( + "UserApiKeyCache.get_cache failed to deserialize cached value for " + "key=%r model_type=%s", + key, + getattr(model_type, "__name__", str(model_type)), + ) + return None + return decoded + + @overload + async def async_get_cache( + self, + key: Any, + parent_otel_span: Any = None, + local_only: bool = False, + *, + model_type: Type[T], + **kwargs: Any, + ) -> Optional[T]: ... + + @overload + async def async_get_cache( + self, + key: Any, + parent_otel_span: Any = None, + local_only: bool = False, + **kwargs: Any, + ) -> Any: ... + + async def async_get_cache( # type: ignore[override] + self, + key, + parent_otel_span=None, + local_only: bool = False, + model_type: Optional[Type[BaseModel]] = None, + **kwargs, + ) -> Union[Any, Optional[BaseModel]]: + if model_type is None and "model_type" in kwargs: + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + cached = await super().async_get_cache( + key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs + ) + if model_type is None: + return cached + if cached is None: + return None + decoded = CacheCodec.deserialize(cached, model_type=model_type) + if decoded is None: + verbose_proxy_logger.error( + "UserApiKeyCache.async_get_cache failed to deserialize cached value for " + "key=%r model_type=%s", + key, + getattr(model_type, "__name__", str(model_type)), + ) + return None + return decoded + + def set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + payload = CacheCodec.serialize(value, model_type=model_type) + return super().set_cache( + key=key, value=payload, local_only=local_only, **kwargs + ) + + async def async_set_cache(self, key, value, local_only: bool = False, **kwargs): # type: ignore[override] + model_type = cast(Optional[Type[BaseModel]], kwargs.pop("model_type", None)) + payload = CacheCodec.serialize(value, model_type=model_type) + return await super().async_set_cache( + key=key, value=payload, local_only=local_only, **kwargs + ) + + async def async_set_cache_pipeline( # type: ignore[override] + self, cache_list: list, local_only: bool = False, **kwargs + ) -> None: + """ + Batch writes with the same Codec boundary as ``async_set_cache`` without + ``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged. + """ + normalized = [ + (key, CacheCodec.serialize(value, model_type=None)) + for key, value in cache_list + ] + return await super().async_set_cache_pipeline( + cache_list=normalized, local_only=local_only, **kwargs + ) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index fae7f939ae..794051e90f 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -19,7 +19,6 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) -from litellm.responses.utils import ResponsesAPIRequestUtils def _load_endpoints_config() -> Dict: @@ -64,10 +63,12 @@ def _create_handler_for_path_params( request: Request, container_id: str, file_id: str, + fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): return await _process_binary_request( request=request, + fastapi_response=fastapi_response, container_id=container_id, file_id=file_id, user_api_key_dict=user_api_key_dict, @@ -152,63 +153,61 @@ def _create_handler_for_path_params( async def _process_binary_request( request: Request, + fastapi_response: Response, container_id: str, file_id: str, user_api_key_dict: UserAPIKeyAuth, ): """ - Process binary content requests using the proper transformation pattern. + Process binary content requests through the standard proxy/router pipeline. - This uses the provider config transformations and llm_http_handler - to maintain consistency with the established pattern. + The router owns managed container ID decoding and deployment selection. This + handler only adapts the byte response to FastAPI. """ - from litellm.litellm_core_utils.litellm_logging import Logging - from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler - from litellm.types.router import GenericLiteLLMParams + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) - # Extract custom_llm_provider custom_llm_provider = ( get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) or "openai" ) - - # Build litellm_params - credentials are resolved by provider config from env - litellm_params = GenericLiteLLMParams() - - # Decode container ID and extract provider info - decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) - original_container_id = decoded.get("response_id", container_id) - - # If container ID has encoded provider info and user didn't explicitly set provider, use it - decoded_provider = decoded.get("custom_llm_provider") - if decoded_provider and custom_llm_provider == "openai": - custom_llm_provider = decoded_provider - - # Get the provider config - container_provider_config = _get_container_provider_config(custom_llm_provider) - - # Create logging object - logging_obj = Logging( - model="container-file-content", - messages=[], - stream=False, - call_type="container_file_content", - start_time=None, - litellm_call_id="", - function_id="", - ) - - # Use the HTTP handler to make the request - handler = BaseLLMHTTPHandler() + data: Dict[str, Any] = { + "container_id": container_id, + "file_id": file_id, + "custom_llm_provider": custom_llm_provider, + } + processor = ProxyBaseLLMRequestProcessing(data=data) try: - content = await handler.async_container_file_content_handler( - container_id=original_container_id, # Use decoded original ID - file_id=file_id, - container_provider_config=container_provider_config, - litellm_params=litellm_params, - logging_obj=logging_obj, + content = await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="aretrieve_container_file_content", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, ) # Determine content type based on common file extensions in the file_id @@ -229,13 +228,25 @@ async def _process_binary_request( elif ".pdf" in file_id_lower: content_type = "application/pdf" + if not isinstance(content, bytes): + raise TypeError( + "aretrieve_container_file_content expected bytes, got " + f"{type(content).__name__}" + ) + return Response( content=content, + headers=dict(fastapi_response.headers), media_type=content_type, ) except Exception as e: - raise e + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) async def _process_multipart_upload_request( @@ -284,16 +295,7 @@ async def _process_multipart_upload_request( or "openai" ) - # Decode container ID and extract provider info - decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) - original_container_id = decoded.get("response_id", container_id) - - # If container ID has encoded provider info and user didn't explicitly set provider, use it - decoded_provider = decoded.get("custom_llm_provider") - if decoded_provider and custom_llm_provider == "openai": - custom_llm_provider = decoded_provider - - data["container_id"] = original_container_id # Use decoded original ID + data["container_id"] = container_id data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) @@ -359,21 +361,6 @@ async def _process_request( or "openai" ) - # Decode container_id if present in path_params - if "container_id" in path_params: - decoded = ResponsesAPIRequestUtils._decode_container_id( - path_params["container_id"] - ) - original_container_id = decoded.get("response_id", path_params["container_id"]) - - # If container ID has encoded provider info and user didn't explicitly set provider, use it - decoded_provider = decoded.get("custom_llm_provider") - if decoded_provider and custom_llm_provider == "openai": - custom_llm_provider = decoded_provider - - # Update path_params with decoded original ID - data["container_id"] = original_container_id - data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index bf60a087c6..a979471dc8 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -129,7 +129,9 @@ class SpendCounterReseed: """ lock = await SpendCounterReseed._get_lock(counter_key) async with lock: - # Re-check after acquiring the lock - another waiter may have warmed it. + # Re-check after acquiring the lock. Skip in-memory on a clean + # Redis miss - in-memory is per-pod-stale. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache( @@ -137,11 +139,13 @@ class SpendCounterReseed: ) if val is not None: return float(val) + redis_clean_miss = True except Exception: pass - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key) if db_spend is None: @@ -149,7 +153,7 @@ class SpendCounterReseed: # Warm even when 0 so subsequent reads hit cache, not DB. try: await spend_counter_cache.async_increment_cache( - key=counter_key, value=db_spend + key=counter_key, value=db_spend, refresh_ttl=True ) except Exception: verbose_proxy_logger.exception( diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 6ada8f5878..967ac9f0ac 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,10 +1,6 @@ -from datetime import datetime +from fastapi import APIRouter, Depends, Request, Response +from fastapi.responses import ORJSONResponse -from fastapi import APIRouter, Depends, HTTPException, Request, Response -from fastapi.responses import ORJSONResponse, StreamingResponse - -import litellm -from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -30,12 +26,17 @@ async def google_generate_content( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ( general_settings, llm_router, proxy_config, proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, version, ) @@ -43,48 +44,33 @@ async def google_generate_content( if "model" not in data: data["model"] = model_name - # Extract generationConfig and pass it as config parameter - generation_config = data.pop("generationConfig", None) - if generation_config: - data["config"] = generation_config - - # Add user authentication metadata for cost tracking - data = await add_litellm_data_to_request( - data=data, - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - general_settings=general_settings, - version=version, - ) - - # Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id - data["litellm_call_id"] = request.headers.get( - "x-litellm-call-id", str(uuid.uuid4()) - ) - logging_obj, data = litellm.utils.function_setup( - original_function="agenerate_content", - rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), - **data, - ) - data["litellm_logging_obj"] = logging_obj - - # call router - if llm_router is None: - raise HTTPException(status_code=500, detail="Router not initialized") - response = await llm_router.agenerate_content(**data) - success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( - response=response, - request_data=data, - request=request, - user_api_key_dict=user_api_key_dict, - logging_obj=logging_obj, - version=version, - proxy_logging_obj=proxy_logging_obj, - ) - fastapi_response.headers.update(success_headers) - return response + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="agenerate_content", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=model_name, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) @router.post( @@ -101,73 +87,52 @@ async def google_stream_generate_content( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.proxy.proxy_server import ( general_settings, llm_router, proxy_config, proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, version, ) data = await _read_request_body(request=request) - if "model" not in data: data["model"] = model_name + data["stream"] = True - data["stream"] = True # enforce streaming for this endpoint - - # Extract generationConfig and pass it as config parameter - generation_config = data.pop("generationConfig", None) - if generation_config: - data["config"] = generation_config - - # Add user authentication metadata for cost tracking - data = await add_litellm_data_to_request( - data=data, - request=request, - user_api_key_dict=user_api_key_dict, - proxy_config=proxy_config, - general_settings=general_settings, - version=version, - ) - - # Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id - data["litellm_call_id"] = request.headers.get( - "x-litellm-call-id", str(uuid.uuid4()) - ) - logging_obj, data = litellm.utils.function_setup( - original_function="agenerate_content_stream", - rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), - **data, - ) - data["litellm_logging_obj"] = logging_obj - - # call router - if llm_router is None: - raise HTTPException(status_code=500, detail="Router not initialized") - response = await llm_router.agenerate_content_stream(**data) - - success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( - response=response, - request_data=data, - request=request, - user_api_key_dict=user_api_key_dict, - logging_obj=logging_obj, - version=version, - proxy_logging_obj=proxy_logging_obj, - ) - - # Check if response is an async iterator (streaming response) - if response is not None and hasattr(response, "__aiter__"): - return StreamingResponse( - content=response, - media_type="text/event-stream", - headers=success_headers, + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="agenerate_content_stream", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=model_name, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, ) - fastapi_response.headers.update(success_headers) - return response @router.post( diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 6dd0288cb0..37be832d35 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -225,10 +225,10 @@ class ToolPermissionGuardrail(CustomGuardrail): def _parse_tool_call_arguments( self, tool_call: ChatCompletionMessageToolCall - ) -> Dict[str, Any]: + ) -> tuple[Optional[Dict[str, Any]], Optional[str]]: arguments = getattr(tool_call.function, "arguments", None) if not arguments: - return {} + return None, "missing arguments" parsed_arguments: Any = {} try: @@ -236,22 +236,24 @@ class ToolPermissionGuardrail(CustomGuardrail): parsed_arguments = json.loads(arguments) elif isinstance(arguments, dict): parsed_arguments = arguments - except json.JSONDecodeError as exc: + else: + return None, "arguments must be a JSON object" + except (json.JSONDecodeError, TypeError) as exc: verbose_proxy_logger.warning( "Tool Permission Guardrail: Failed to decode arguments for tool %s: %s", tool_call.function.name, exc, ) - return {} + return None, "arguments could not be parsed" if isinstance(parsed_arguments, dict): - return parsed_arguments + return parsed_arguments, None verbose_proxy_logger.debug( - "Tool Permission Guardrail: Ignoring non-dict arguments for tool %s", + "Tool Permission Guardrail: Rejecting non-dict arguments for tool %s", tool_call.function.name, ) - return {} + return None, "arguments must be a JSON object" def _collect_argument_paths( self, @@ -331,10 +333,21 @@ class ToolPermissionGuardrail(CustomGuardrail): continue if rule.allowed_param_patterns and should_check_params: - arguments = self._parse_tool_call_arguments(tool_call) + arguments, parse_error = self._parse_tool_call_arguments(tool_call) + if parse_error: + default_message = f"Tool '{tool_identifier}' {parse_error} required by rule '{rule.id}'" + message = self.render_violation_message( + default=default_message, + context={"tool_name": tool_identifier, "rule_id": rule.id}, + ) + return False, rule.id, message if not arguments: - last_pattern_failure_msg = f"Tool '{tool_identifier}' is missing arguments required by rule '{rule.id}'" - continue + default_message = f"Tool '{tool_identifier}' is missing arguments required by rule '{rule.id}'" + message = self.render_violation_message( + default=default_message, + context={"tool_name": tool_identifier, "rule_id": rule.id}, + ) + return False, rule.id, message patterns_match, failure_message = self._patterns_match_for_rule( arguments=arguments, @@ -365,6 +378,33 @@ class ToolPermissionGuardrail(CustomGuardrail): ) return is_allowed, None, message + @staticmethod + def _get_mapping_value(item: Any, key: str) -> Any: + if isinstance(item, dict): + return item.get(key) + return getattr(item, key, None) + + @staticmethod + def _legacy_function_call_id(choice_index: int) -> str: + return f"legacy_function_call_{choice_index}" + + def _legacy_function_call_to_tool_call( + self, function_call: Any, choice_index: int + ) -> Optional[ChatCompletionMessageToolCall]: + if function_call is None: + return None + + function_name = self._get_mapping_value(function_call, "name") + arguments = self._get_mapping_value(function_call, "arguments") or "" + if not function_name: + return None + + return ChatCompletionMessageToolCall( + id=self._legacy_function_call_id(choice_index), + type="function", + function={"name": function_name, "arguments": arguments}, + ) + def _extract_tool_calls_from_response( self, response: ModelResponse ) -> List[ChatCompletionMessageToolCall]: @@ -379,13 +419,72 @@ class ToolPermissionGuardrail(CustomGuardrail): """ tool_calls = [] - for choice in response.choices: + for choice_index, choice in enumerate(response.choices): if isinstance(choice, Choices): for tool in choice.message.tool_calls or []: tool_calls.append(tool) + legacy_tool_call = self._legacy_function_call_to_tool_call( + getattr(choice.message, "function_call", None), choice_index + ) + if legacy_tool_call is not None: + tool_calls.append(legacy_tool_call) return tool_calls + def _get_request_tool_name(self, tool: Any) -> tuple[Optional[str], Optional[str]]: + tool_type = self._get_mapping_value(tool, "type") + if tool_type != "function": + return None, tool_type + + function = self._get_mapping_value(tool, "function") + tool_name = self._get_mapping_value(function, "name") + return tool_name, tool_type + + def _get_legacy_function_name(self, function: Any) -> Optional[str]: + return self._get_mapping_value(function, "name") + + def _get_named_tool_choice(self, data: dict) -> Optional[str]: + tool_choice = data.get("tool_choice") + if not tool_choice or tool_choice in ("auto", "none", "required"): + return None + if isinstance(tool_choice, str): + return tool_choice + if self._get_mapping_value(tool_choice, "type") != "function": + return None + return self._get_mapping_value( + self._get_mapping_value(tool_choice, "function"), "name" + ) + + def _get_named_function_call(self, data: dict) -> Optional[str]: + function_call = data.get("function_call") + if not function_call or function_call in ("auto", "none"): + return None + if isinstance(function_call, str): + return function_call + return self._get_mapping_value(function_call, "name") + + def _collect_request_tools(self, data: dict) -> List[tuple[str, Optional[str]]]: + request_tools: List[tuple[str, Optional[str]]] = [] + + for tool in data.get("tools") or []: + tool_name, tool_type = self._get_request_tool_name(tool) + if tool_name is not None: + request_tools.append((tool_name, tool_type)) + + for function in data.get("functions") or []: + function_name = self._get_legacy_function_name(function) + if function_name is not None: + request_tools.append((function_name, "function")) + + for forced_tool_name in ( + self._get_named_tool_choice(data), + self._get_named_function_call(data), + ): + if forced_tool_name is not None: + request_tools.append((forced_tool_name, "function")) + + return request_tools + def _modify_request_with_permission_errors( self, data: dict, @@ -410,19 +509,32 @@ class ToolPermissionGuardrail(CustomGuardrail): for tool_use in denied_tool_names: error_tool_names.add(tool_use) - # Modify the tools tools: Optional[List[ChatCompletionToolParam]] = data.get("tools") - if tools is None: - return data - - new_tools = [] - for tool in tools: - if tool["type"] != "function": - continue - tool_name: str = tool["function"]["name"] - if tool_name not in error_tool_names: + if tools is not None: + new_tools = [] + for tool in tools: + tool_name, tool_type = self._get_request_tool_name(tool) + if tool_type == "function" and tool_name in error_tool_names: + continue new_tools.append(tool) - data["tools"] = new_tools + data["tools"] = new_tools + + functions = data.get("functions") + if functions is not None: + data["functions"] = [ + function + for function in functions + if self._get_legacy_function_name(function) not in error_tool_names + ] + + named_tool_choice = self._get_named_tool_choice(data) + if named_tool_choice in error_tool_names: + data["tool_choice"] = "none" + + named_function_call = self._get_named_function_call(data) + if named_function_call in error_tool_names: + data["function_call"] = "none" + return data def _create_permission_error_result( @@ -472,7 +584,7 @@ class ToolPermissionGuardrail(CustomGuardrail): error_results[tool_use.id] = error_result # Modify the response content - for choice in response.choices: + for choice_index, choice in enumerate(response.choices): if isinstance(choice, Choices): filtered_tool_calls = [] error_messages = [] @@ -490,6 +602,15 @@ class ToolPermissionGuardrail(CustomGuardrail): filtered_tool_calls if filtered_tool_calls else None ) + legacy_tool_call = self._legacy_function_call_to_tool_call( + getattr(choice.message, "function_call", None), choice_index + ) + if legacy_tool_call is not None: + legacy_error_result = error_results.get(legacy_tool_call.id) + if legacy_error_result is not None: + choice.message.function_call = None + error_messages.append(legacy_error_result.content) + # Add error messages to content if error_messages: existing_content = choice.message.content @@ -519,21 +640,16 @@ class ToolPermissionGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - new_tools: Optional[List[ChatCompletionToolParam]] = data.get("tools") - if new_tools is None: + new_tools = self._collect_request_tools(data) + if not new_tools: verbose_proxy_logger.warning( - "Tool Permission Guardrail: not running guardrail. No tools in data" + "Tool Permission Guardrail: not running guardrail. No tools or functions in data" ) return data # Check permissions for each tool denied_tool_names = [] - for tool in new_tools: - if tool["type"] != "function": - continue - tool_name: str = tool["function"]["name"] - tool_type: Optional[str] = tool.get("type") - + for tool_name, tool_type in new_tools: is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 7d67750c78..7c340ff5df 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -29,6 +29,10 @@ ILLEGAL_DISPLAY_PARAMS = [ "exception", # internal; not JSON-serializable, never for display "litellm_metadata", # internal tracking metadata with auth objects; not for display ] +# Provider routing fields. Allowed for proxy admins so they can see which +# region/version a deployment is checking; gated at the endpoint layer for +# non-admin callers (see _strip_admin_only_fields_from_health_result). +ADMIN_ONLY_HEALTH_DISPLAY_PARAMS = ("api_base", "api_version") MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"] diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index b4b5de1746..35c9edb937 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( CallInfo, EnterpriseLicenseData, Litellm_EntityType, + LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth, @@ -28,6 +29,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.health_check import ( + ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, _update_litellm_params_for_health_check, perform_health_check, @@ -723,6 +725,129 @@ async def _save_background_health_checks_to_db( # Continue execution - don't let database save failure break health checks +_PROXY_ADMIN_ROLES = frozenset( + { + LitellmUserRoles.PROXY_ADMIN.value, + # View-only admins are operators (oncall, support); they need the + # routing fields (api_base, api_version) to diagnose health and tell + # which provider region a check is hitting. They cannot mutate config + # so granting them the read-only view is safe. + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + } +) + + +def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Return True if the caller has a proxy-admin role (full or view-only). + + user_role on UserAPIKeyAuth can be either a LitellmUserRoles enum or its + string value depending on how the auth path constructed the object, so we + compare against the raw value rather than the enum identity. + """ + role = user_api_key_dict.user_role + if role is None: + return False + role_value = role.value if hasattr(role, "value") else role + return role_value in _PROXY_ADMIN_ROLES + + +def _strip_admin_only_fields_from_health_result(result: dict) -> dict: + """ + Return a copy of the /health response with provider routing fields + (``api_base``, ``api_version``) removed from each healthy/unhealthy + endpoint entry. Used to hide those fields from non-admin callers while + still showing them which deployments they own and whether each one is + healthy. Proxy admins receive the unmodified result. + """ + out = dict(result) + drop = set(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS) + for key in ("healthy_endpoints", "unhealthy_endpoints"): + eps = out.get(key) + if isinstance(eps, list): + out[key] = [ + ( + {k: v for k, v in ep.items() if k not in drop} + if isinstance(ep, dict) + else ep + ) + for ep in eps + ] + return out + + +def _resolve_targeted_model_ids( + model_list: list, model: Optional[str], model_id: Optional[str] +) -> Optional[set]: + """ + Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of + deployment IDs the response should be scoped to. + + Mirrors the live-path semantics in ``perform_health_check()``: ``model`` + matches either the deployment's ``model_name`` alias or its + ``litellm_params.model`` provider string. ``model_id`` matches + ``model_info.id``. + + Both query params are validated against the supplied ``model_list``. + Callers pass an already-scoped list (filtered to the caller's allowed + models for non-admins, full list for admins), so a ``model_id`` that + isn't present resolves to an empty set rather than a single-element + set — preventing a non-admin from reading another deployment's cached + health entry by guessing its ID. + + Returns ``None`` when no targeting is requested — callers should treat + that as "no filter." + """ + if not model and not model_id: + return None + target_ids: set = set() + for m in model_list: + deployment_id = (m.get("model_info") or {}).get("id") + if not deployment_id: + continue + if model_id and deployment_id == model_id: + target_ids.add(deployment_id) + continue + if model: + litellm_model = (m.get("litellm_params") or {}).get("model") + if m.get("model_name") == model or litellm_model == model: + target_ids.add(deployment_id) + return target_ids + + +def _filter_health_check_results_by_model_ids( + results: dict, allowed_model_ids: set +) -> dict: + """ + Restrict a cached background health-check result dict to endpoints whose + model_id is in ``allowed_model_ids``. + + Endpoints without a model_id (e.g. CLI-model entries that predate the + model_id wiring) are dropped conservatively — we cannot prove they belong + to the caller, so they are excluded rather than leaked. + + Each retained endpoint is shallow-copied before being returned, so any + downstream transform (e.g. _strip_admin_only_fields_from_health_result) + cannot accidentally mutate the shared ``health_check_results`` cache. + """ + healthy = [ + dict(ep) + for ep in (results.get("healthy_endpoints") or []) + if ep.get("model_id") in allowed_model_ids + ] + unhealthy = [ + dict(ep) + for ep in (results.get("unhealthy_endpoints") or []) + if ep.get("model_id") in allowed_model_ids + ] + return { + "healthy_endpoints": healthy, + "unhealthy_endpoints": unhealthy, + "healthy_count": len(healthy), + "unhealthy_count": len(unhealthy), + } + + async def _perform_health_check_and_save( model_list, target_model, @@ -771,6 +896,7 @@ async def _perform_health_check_and_save( @router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def health_endpoint( + response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), model: Optional[str] = fastapi.Query( None, description="Specify the model name (optional)" @@ -838,11 +964,33 @@ async def health_endpoint( detail={"error": f"Model with ID {model_id} not found"}, ) + is_admin = _is_proxy_admin(user_api_key_dict) + model_specific_request = bool(model or model_id) + + def _post_process(result: dict) -> dict: + # api_base / api_version reveal which provider/region/internal host the + # deployment talks to; only proxy admins receive them. Non-admin keys + # still see model/model_id and the healthy/unhealthy status. We also + # set a header so non-admin clients that previously parsed those + # fields can detect the change programmatically. + # When a caller asked about a specific model/model_id and zero + # endpoints came back healthy, surface that as a 503 so monitoring + # systems can rely on the HTTP status instead of having to parse the + # body. The body shape is unchanged. + if model_specific_request and result.get("healthy_count", 0) == 0: + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE + if is_admin: + return result + response.headers["Litellm-Health-Field-Notice"] = ( + "api_base and api_version are admin-only on this endpoint" + ) + return _strip_admin_only_fields_from_health_result(result) + try: if llm_model_list is None: # if no router set, check if user set a model using litellm --model ollama/llama2 if user_model is not None: - return await _perform_health_check_and_save( + cli_result = await _perform_health_check_and_save( model_list=[], target_model=None, cli_model=user_model, @@ -853,20 +1001,81 @@ async def health_endpoint( model_id=None, # CLI model doesn't have model_id max_concurrency=health_check_concurrency, ) + return _post_process(cli_result) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) _llm_model_list = copy.deepcopy(llm_model_list) ### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ### + # Live path: scope by model_name (every deployment has one). + # Cache path: scope by model_id (the cache is keyed on model_id). + # Consequence: a deployment whose model_name the caller can access + # but which lacks model_info.id will appear in the live /health + # response but NOT in the background-cache /health response. This is + # surfaced via the "warnings" field below so operators can fix the + # missing model_info.id rather than guess at the discrepancy. if len(user_api_key_dict.models) > 0: - pass - else: - pass # + allowed_models = set(user_api_key_dict.models) + _llm_model_list = [ + m for m in _llm_model_list if m.get("model_name") in allowed_models + ] if use_background_health_checks: - return health_check_results + # The cached background result covers every model. When the + # caller targets a specific model/model_id we have to narrow the + # cache to that deployment before _post_process evaluates + # healthy_count, otherwise an unhealthy "foo" combined with any + # other healthy model would still report healthy_count > 0 and + # the targeted-503 path would never fire. + targeted_ids = _resolve_targeted_model_ids(_llm_model_list, model, model_id) + if len(user_api_key_dict.models) > 0: + allowed_model_ids = { + (m.get("model_info") or {}).get("id") + for m in _llm_model_list + if (m.get("model_info") or {}).get("id") + } + # _llm_model_list is already scoped to the caller's allowed + # model_names above, so targeted_ids is implicitly the + # intersection of "targeted" and "allowed." + filter_ids = ( + targeted_ids if targeted_ids is not None else allowed_model_ids + ) + filtered = _filter_health_check_results_by_model_ids( + health_check_results, filter_ids + ) + if targeted_ids is None and not allowed_model_ids: + # Caller has accessible model_names but none of the + # matching deployments expose a model_info.id, so the + # cache filter (which keys on model_id) drops every + # entry. Surface this both as a warning log and a + # structured "warnings" field on the response so the + # caller can distinguish "no deployments found" from + # "deployments excluded due to missing model_info.id". + verbose_proxy_logger.warning( + "health_endpoint: scoped key %s has accessible models %s " + "but none of the matching deployments carry a model_info.id; " + "background health-check cache will return an empty result.", + user_api_key_dict.user_id, + list(user_api_key_dict.models), + ) + filtered["warnings"] = [ + "Some accessible deployments are missing model_info.id " + "and were excluded from this response. Ask a proxy admin " + "to populate model_info.id for these models." + ] + return _post_process(filtered) + if targeted_ids is not None: + # Admin caller targeting a specific model: filter the cache + # so the response (and the targeted-503 check) reflects only + # that deployment, not the global aggregate. + return _post_process( + _filter_health_check_results_by_model_ids( + health_check_results, targeted_ids + ) + ) + return _post_process(health_check_results) else: - return await _perform_health_check_and_save( + router_result = await _perform_health_check_and_save( model_list=_llm_model_list, target_model=target_model, cli_model=None, @@ -877,6 +1086,7 @@ async def health_endpoint( model_id=model_id, max_concurrency=health_check_concurrency, ) + return _post_process(router_result) except Exception as e: verbose_proxy_logger.error( "litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format( @@ -1242,7 +1452,7 @@ def callback_name(callback): tags=["health"], dependencies=[Depends(user_api_key_auth)], ) -async def health_readiness(): +async def health_readiness(response: Response): """ Unprotected endpoint for checking if worker can receive requests """ @@ -1275,8 +1485,8 @@ async def health_readiness(): try: index_info = await litellm.cache.cache._index_info() except Exception as e: - index_info = "index does not exist - error: " + str(e) - cache_type = {"type": cache_type, "index_info": index_info} + index_info = "index does not exist - error: " + str(e) # type: ignore[assignment] + cache_type = {"type": cache_type, "index_info": index_info} # type: ignore[assignment] # check log level log_level_name = logging.getLevelName(verbose_logger.getEffectiveLevel()) @@ -1285,6 +1495,12 @@ async def health_readiness(): # check DB if prisma_client is not None: # if db passed in, check if it's connected db_health_status = await _db_health_readiness_check() + # A configured DB that is not reachable means the worker cannot + # serve requests that depend on persisted state (keys, budgets, + # spend logs). Return 503 so orchestrators take this pod out of + # rotation; "Not connected" (no DB configured at all) stays 200. + if db_health_status["status"] != "connected": + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", "db": db_health_status["status"], diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 06b7d85789..2ee0588f19 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -164,13 +164,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage: BatchFileUsage, ) -> None: """ - Check rate limits and increment counters by the batch amounts. + Atomically check + increment rate-limit counters by the batch amounts. - Raises HTTPException if any limit would be exceeded. + Raises HTTPException if any descriptor would exceed its limit; in that + case no counter is modified. Backed by `atomic_check_and_increment_by_n` + which uses a Redis Lua script when available (multi-process atomic) and + falls back to a per-process asyncio.Lock + in-memory operation. """ - from litellm.types.caching import RedisPipelineIncrementOperation - - # Create descriptors and check if batch would exceed limits descriptors = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, @@ -179,73 +179,31 @@ class _PROXY_BatchRateLimiter(CustomLogger): model_has_failures=False, ) - # Check current usage without incrementing - rate_limit_response = await self.parallel_request_limiter.should_rate_limit( - descriptors=descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, - ) + increment: Dict[Literal["requests", "tokens"], int] = { + "requests": batch_usage.request_count, + "tokens": batch_usage.total_tokens, + } + increments: List[Dict[Literal["requests", "tokens"], int]] = [ + increment for _ in descriptors + ] - # Verify batch won't exceed any limits - for status in rate_limit_response["statuses"]: - rate_limit_type = status["rate_limit_type"] - limit_remaining = status["limit_remaining"] - - required_capacity = ( - batch_usage.request_count - if rate_limit_type == "requests" - else batch_usage.total_tokens if rate_limit_type == "tokens" else 0 - ) - - if required_capacity > limit_remaining: - self._raise_rate_limit_error( - status, descriptors, batch_usage, rate_limit_type - ) - - # Build pipeline operations for batch increments - # Reuse the same keys that descriptors check - pipeline_operations: List[RedisPipelineIncrementOperation] = [] - - for descriptor in descriptors: - key = descriptor["key"] - value = descriptor["value"] - rate_limit = descriptor.get("rate_limit") - - if rate_limit is None: - continue - - # Add RPM increment if limit is set - if rate_limit.get("requests_per_unit") is not None: - rpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="requests" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=rpm_key, - increment_value=batch_usage.request_count, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Add TPM increment if limit is set - if rate_limit.get("tokens_per_unit") is not None: - tpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="tokens" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=tpm_key, - increment_value=batch_usage.total_tokens, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Execute increments - if pipeline_operations: - await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation( - pipeline_operations=pipeline_operations, + rate_limit_response = ( + await self.parallel_request_limiter.atomic_check_and_increment_by_n( + descriptors=descriptors, + increments=increments, parent_otel_span=user_api_key_dict.parent_otel_span, ) + ) + + if rate_limit_response["overall_code"] == "OVER_LIMIT": + for status in rate_limit_response["statuses"]: + if status["code"] == "OVER_LIMIT": + self._raise_rate_limit_error( + status, + descriptors, + batch_usage, + status["rate_limit_type"], + ) async def count_input_file_usage( self, diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 72483d29cd..f7c0592992 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting import os from datetime import datetime -from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, Union from fastapi import HTTPException @@ -460,92 +460,128 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if priority_descriptors: descriptors_to_check.extend(priority_descriptors) - # PHASE 1: Read-only check of ALL limits (no increments) - check_response = await self.v3_limiter.should_rate_limit( - descriptors=descriptors_to_check, + # Atomic check-and-increment for the ENFORCED descriptor set: + # - model_saturation_check is always enforced + # - priority_model is enforced only when saturation crosses threshold + # + # Backed by a Redis Lua script (multi-process atomic) with an + # asyncio.Lock + in-memory fallback for single-process deployments. + # All-or-nothing: if any enforced descriptor would exceed its limit, + # no counter is modified and the response carries "OVER_LIMIT". + enforced_descriptors: List[RateLimitDescriptor] = [model_wide_descriptor] + if priority_descriptors and should_enforce_priority: + enforced_descriptors.extend(priority_descriptors) + + per_request_increment: Dict[Literal["requests", "tokens"], int] = { + "requests": 1, + "tokens": 0, + } + atomic_response = await self.v3_limiter.atomic_check_and_increment_by_n( + descriptors=enforced_descriptors, + increments=[per_request_increment for _ in enforced_descriptors], parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, # CRITICAL: Don't increment counters yet ) verbose_proxy_logger.debug( - f"Read-only check: {json.dumps(check_response, indent=2)}" + f"Atomic check+increment response: {json.dumps(atomic_response, indent=2)}" ) - # PHASE 2: Decide which limits to enforce - if check_response["overall_code"] == "OVER_LIMIT": - for status in check_response["statuses"]: - if status["code"] == "OVER_LIMIT": - descriptor_key = status["descriptor_key"] + if atomic_response["overall_code"] == "OVER_LIMIT": + for status in atomic_response["statuses"]: + if status["code"] != "OVER_LIMIT": + continue + descriptor_key = status["descriptor_key"] + if descriptor_key == "model_saturation_check": + raise HTTPException( + status_code=429, + detail={ + "error": f"Model capacity reached for {model}. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + }, + ) + if descriptor_key == "priority_model": + verbose_proxy_logger.debug( + f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " + f"priority: {priority}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": f"Priority-based rate limit exceeded. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + ) - # Model-wide limit exceeded (ALWAYS enforce) - if descriptor_key == "model_saturation_check": - raise HTTPException( - status_code=429, - detail={ - "error": f"Model capacity reached for {model}. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - }, - ) + # Fail-closed guard: overall_code says OVER_LIMIT but no status + # matched a descriptor key we know how to translate into a 429. + # Refuse the request rather than silently fall through and let an + # over-limit request proceed to the model. Without this, a future + # caller wiring an unfamiliar descriptor into enforced_descriptors + # would silently bypass the rate limit. + offending = next( + (s for s in atomic_response["statuses"] if s["code"] == "OVER_LIMIT"), + None, + ) + verbose_proxy_logger.error( + f"Dynamic rate limiter: OVER_LIMIT response with unknown " + f"descriptor_key(s) — refusing request. response={atomic_response}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": "Rate limit exceeded", + "descriptor_key": ( + offending["descriptor_key"] if offending else "unknown" + ), + "rate_limit_type": ( + str(offending["rate_limit_type"]) if offending else "unknown" + ), + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "x-litellm-priority": priority or "default", + }, + ) - # Priority limit exceeded (ONLY enforce when saturated) - elif descriptor_key == "priority_model" and should_enforce_priority: - verbose_proxy_logger.debug( - f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " - f"priority: {priority}" - ) - raise HTTPException( - status_code=429, - detail={ - "error": f"Priority-based rate limit exceeded. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}, " - f"Model saturation: {saturation:.1%}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - "x-litellm-saturation": f"{saturation:.2%}", - }, - ) - - # PHASE 3: Increment counters separately to avoid early-exit issues - # Model counter must ALWAYS increment, but priority counter might be over limit - # If we increment them together, v3_limiter's in-memory check will exit early - # and skip incrementing the model counter - - # Step 3a: Increment model-wide counter (always) - model_increment_response = await self.v3_limiter.should_rate_limit( - descriptors=[model_wide_descriptor], - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=False, - ) - - # Step 3b: Increment priority counter (may be over limit, but we still track it) - if priority_descriptors: - priority_increment_response = await self.v3_limiter.should_rate_limit( + # If priority is NOT enforced (saturation below threshold) but + # priority_descriptors exist, increment them for tracking only — no + # check, no rollback. This matches the prior tracking semantics. + # + # Using the non-atomic should_rate_limit (instead of + # atomic_check_and_increment_by_n) is intentional here: we don't want + # to enforce the limit, we only want to bump the counter so the + # priority allocation has accurate usage when it later becomes + # enforced. The increment-then-check semantics of should_rate_limit + # are fine because we ignore the OVER_LIMIT response. + if priority_descriptors and not should_enforce_priority: + priority_tracking_response = await self.v3_limiter.should_rate_limit( descriptors=priority_descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=False, ) - - # Combine responses for post-call hook - combined_response = { - "overall_code": model_increment_response["overall_code"], - "statuses": model_increment_response["statuses"] - + priority_increment_response["statuses"], + data["litellm_proxy_rate_limit_response"] = { + "overall_code": atomic_response["overall_code"], + "statuses": atomic_response["statuses"] + + priority_tracking_response["statuses"], } - data["litellm_proxy_rate_limit_response"] = combined_response else: - data["litellm_proxy_rate_limit_response"] = model_increment_response + data["litellm_proxy_rate_limit_response"] = atomic_response async def async_pre_call_hook( self, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index f29bbd2d9d..4497e64c17 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4,6 +4,7 @@ This is a rate limiter implementation based on a similar one by Envoy proxy. This is currently in development and not yet ready for production. """ +import asyncio import binascii import os from datetime import datetime @@ -80,6 +81,90 @@ end return results """ +CHECK_AND_INCREMENT_BY_N_SCRIPT = """ +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) +-- +-- Return on success: { 0, new_counter_1, new_counter_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local descriptor_count = #KEYS / 2 + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + if current_counter + increment > limit then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_expired = descriptor_state[i][1] + + if window_expired then + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end +end + +return results +""" + TOKEN_INCREMENT_SCRIPT = """ local results = {} @@ -162,15 +247,37 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TOKEN_INCREMENT_SCRIPT ) ) + self.check_and_increment_by_n_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + CHECK_AND_INCREMENT_BY_N_SCRIPT + ) + ) else: self.batch_rate_limiter_script = None self.token_increment_script = None + self.check_and_increment_by_n_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) # Batch rate limiter (lazy loaded) self._batch_rate_limiter: Optional[Any] = None + # Serializes multi-phase check+increment sequences (batch + dynamic + # limiters) within this process to close the TOCTOU window between + # read-only check and counter increment. Multi-replica deployments + # additionally rely on Redis Lua atomicity for cross-process safety. + # + # Coarse granularity: this single lock serializes ALL atomic check+ + # increment operations across batch and dynamic limiters on this + # instance. A slow batch input-file fetch (which happens upstream of + # the lock) does not block here, but Redis Lua latency does. If + # contention shows up under load (visible as p99 latency spikes + # correlated with batch traffic), shard to a per-descriptor-key lock + # via a `weakref.WeakValueDictionary[str, asyncio.Lock]`. Punted as a + # follow-up because Lua dominates wall-time and the lock is held for + # one round-trip. + self._check_and_increment_lock = asyncio.Lock() + def _get_batch_rate_limiter(self) -> Optional[Any]: """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: @@ -588,6 +695,281 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return rate_limit_response + async def atomic_check_and_increment_by_n( + self, + descriptors: List[RateLimitDescriptor], + increments: List[Dict[Literal["requests", "tokens"], int]], + parent_otel_span: Optional[Span] = None, + ) -> RateLimitResponse: + """ + Atomic check-and-increment-by-N across one or more descriptors. + + All-or-nothing: if any descriptor would exceed its limit, no counter is + modified and the response carries `overall_code = "OVER_LIMIT"` with + the offending descriptor's status. Closes the TOCTOU window between + read and increment in both single-process and multi-process (Redis) + deployments. + + Args: + descriptors: rate-limit descriptors to check + increments: per-descriptor increment amounts, indexed parallel to + `descriptors`. Each entry is `{"requests": int, "tokens": int}` + — values default to 0 when a descriptor has no matching limit. + + Returns: + RateLimitResponse with one status per (descriptor, rate_limit_type) + counter, mirroring `should_rate_limit`'s shape. + """ + if len(descriptors) != len(increments): + raise ValueError( + "atomic_check_and_increment_by_n: descriptors and increments " + "must have the same length" + ) + + keys: List[str] = [] + per_counter_meta: List[Dict[str, Any]] = [] + script_args: List[Any] = [] + + for descriptor, increment_amounts in zip(descriptors, increments): + descriptor_key = descriptor["key"] + descriptor_value = descriptor["value"] + rate_limit: RateLimitDescriptorRateLimitObject = ( + descriptor.get("rate_limit") or RateLimitDescriptorRateLimitObject() + ) + window_size = rate_limit.get("window_size") or self.window_size + window_key = f"{{{descriptor_key}:{descriptor_value}}}:window" + + for rate_limit_type in ("requests", "tokens"): + rlt: Literal["requests", "tokens"] = cast( + Literal["requests", "tokens"], rate_limit_type + ) + if rlt == "requests": + limit_value = rate_limit.get("requests_per_unit") + inc_amount = int(increment_amounts.get("requests", 0) or 0) + else: + limit_value = rate_limit.get("tokens_per_unit") + inc_amount = int(increment_amounts.get("tokens", 0) or 0) + if limit_value is None or inc_amount <= 0: + continue + counter_key = self.create_rate_limit_keys( + descriptor_key, descriptor_value, rlt + ) + # Counter-key TTL and window_size are conceptually distinct + # ("how long the counter Redis key lives" vs "how long the + # sliding window is"). They happen to be equal today because + # we have no descriptor type that needs them apart, but they + # are kept as separate variables here so a future custom-TTL + # descriptor doesn't reintroduce a silent expiry bug. Both + # the Lua script and the in-memory fallback read these from + # their respective ARGV / meta slots. + ttl_seconds = int(window_size) + window_size_seconds = int(window_size) + keys.extend([window_key, counter_key]) + # Per-counter 4-tuple matches the Lua ARGV layout exactly: + # [limit, increment, ttl_seconds, window_size_seconds]. + script_args.extend( + [ + int(limit_value), + inc_amount, + ttl_seconds, + window_size_seconds, + ] + ) + per_counter_meta.append( + { + "descriptor_key": descriptor_key, + "current_limit": int(limit_value), + "rate_limit_type": rlt, + "window_key": window_key, + "counter_key": counter_key, + "increment": inc_amount, + "ttl": ttl_seconds, + "window_size": window_size_seconds, + } + ) + + if not keys: + return RateLimitResponse(overall_code="OK", statuses=[]) + + # Multi-process atomicity via Redis Lua. Single-process atomicity + # falls back to the asyncio.Lock + in-memory sliding window below. + # Note: in-memory state diverges from Redis state — if Lua fails + # mid-write, retrying via in-memory may double-count. See fallback + # warning below. + if self.check_and_increment_by_n_script is not None: + try: + raw = await self.check_and_increment_by_n_script( + keys=keys, + args=script_args, + ) + return self._build_atomic_response(raw, per_counter_meta) + except Exception as e: + # Escalated from warning to error: Lua failures (script timeout, + # Redis OOM, network partition) leave counter state ambiguous. + # The fallback path below uses LOCAL DualCache, which is a + # different store from Redis — counters here will diverge from + # Redis until that key's window expires (TTL bounds divergence). + # Operators should alert on this log line; sustained occurrences + # indicate Redis health degradation that may erode rate-limit + # accuracy. + verbose_proxy_logger.error( + f"atomic_check_and_increment_by_n: Redis Lua execution " + f"failed ({type(e).__name__}: {e}). Falling back to " + f"in-memory enforcement — counters will diverge from Redis " + f"state until window expires (window_size={self.window_size}s)." + ) + + async with self._check_and_increment_lock: + return await self._atomic_check_and_increment_in_memory( + per_counter_meta=per_counter_meta, + parent_otel_span=parent_otel_span, + ) + + def _build_atomic_response( + self, + raw: List[Any], + per_counter_meta: List[Dict[str, Any]], + ) -> RateLimitResponse: + """Convert Lua script return value to RateLimitResponse. + + Indexing invariant: `per_counter_meta` and `KEYS` are parallel-indexed + at the COUNTER level, not the descriptor level. A descriptor with both + RPM and TPM limits emits two `(window_key, counter_key)` pairs and + two meta entries — one per counter. The Lua script's loop variable + `i` therefore enumerates counters, and the over-limit return tuple + `{1, i, ...}` carries a counter index that maps directly to + `per_counter_meta[i - 1]`. Keep these arrays parallel at the counter + level when modifying this code. + """ + if not raw: + return RateLimitResponse(overall_code="OK", statuses=[]) + + status_code = int(raw[0]) + if status_code == 1: + # Over limit: { 1, counter_index (1-based), current_counter, limit } + descriptor_index = int(raw[1]) - 1 + current_counter = int(raw[2]) + limit = int(raw[3]) + meta = per_counter_meta[descriptor_index] + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ + RateLimitStatus( + code="OVER_LIMIT", + current_limit=limit, + limit_remaining=max(0, limit - current_counter), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ], + ) + + statuses: List[RateLimitStatus] = [] + for meta, new_counter in zip(per_counter_meta, raw[1:]): + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - int(new_counter)), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + return RateLimitResponse(overall_code="OK", statuses=statuses) + + async def _atomic_check_and_increment_in_memory( + self, + per_counter_meta: List[Dict[str, Any]], + parent_otel_span: Optional[Span] = None, + ) -> RateLimitResponse: + """In-memory all-or-nothing check-and-increment. Caller holds lock. + + Reads/writes the LOCAL DualCache (`local_only=True`) — note this is + a different store from Redis. When this fallback fires after a Lua + failure, in-memory counters will diverge from Redis until each key's + window expires (TTL bounds divergence). + """ + # Use a single 'now' for the duration of this critical section so all + # descriptors evaluate window expiry consistently. + now_int = int(self._get_current_time().timestamp()) + + # Pass 1: read state, validate. + descriptor_state: List[Dict[str, Any]] = [] + for meta in per_counter_meta: + window_size = meta["window_size"] + window_start = await self.internal_usage_cache.async_get_cache( + key=meta["window_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + window_expired = ( + window_start is None or (now_int - int(window_start)) >= window_size + ) + current_counter = ( + 0 + if window_expired + else int( + await self.internal_usage_cache.async_get_cache( + key=meta["counter_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + or 0 + ) + ) + if current_counter + meta["increment"] > meta["current_limit"]: + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ + RateLimitStatus( + code="OVER_LIMIT", + current_limit=meta["current_limit"], + limit_remaining=max( + 0, meta["current_limit"] - current_counter + ), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ], + ) + descriptor_state.append( + {"window_expired": window_expired, "current": current_counter} + ) + + # Pass 2: apply increments. + statuses: List[RateLimitStatus] = [] + for meta, state in zip(per_counter_meta, descriptor_state): + new_counter = ( + meta["increment"] + if state["window_expired"] + else state["current"] + meta["increment"] + ) + if state["window_expired"]: + await self.internal_usage_cache.async_set_cache( + key=meta["window_key"], + value=str(now_int), + ttl=meta["window_size"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=meta["counter_key"], + value=new_counter, + ttl=meta["ttl"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - new_counter), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + return RateLimitResponse(overall_code="OK", statuses=statuses) + def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: Optional[str] = None ) -> List[RateLimitDescriptor]: diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index caaec12f7a..ceaef20a8d 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -236,13 +236,12 @@ async def _patch_key_caches_add_access_group( ) -> None: """Patch cached key objects to include access_group_id.""" for token in key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) + cached_key = await user_api_key_cache.async_get_cache( + key=token, + model_type=UserAPIKeyAuth, + ) if cached_key is None: continue - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if not isinstance(cached_key, UserAPIKeyAuth): - continue if cached_key.access_group_ids is None: cached_key.access_group_ids = [access_group_id] elif access_group_id not in cached_key.access_group_ids: @@ -267,12 +266,11 @@ async def _patch_key_caches_remove_access_group( ) -> None: """Patch cached key objects to remove access_group_id.""" for token in key_tokens: - cached_key = await user_api_key_cache.async_get_cache(key=token) - if cached_key is None: - continue - if isinstance(cached_key, dict): - cached_key = UserAPIKeyAuth(**cached_key) - if isinstance(cached_key, UserAPIKeyAuth) and cached_key.access_group_ids: + cached_key = await user_api_key_cache.async_get_cache( + key=token, + model_type=UserAPIKeyAuth, + ) + if cached_key is not None and cached_key.access_group_ids: cached_key.access_group_ids = [ ag for ag in cached_key.access_group_ids if ag != access_group_id ] diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d1f75e3706..a01f5e6321 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -27,7 +27,7 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._experimental.mcp_server.db import ( rotate_mcp_server_credentials_master_key, + rotate_mcp_user_credentials_master_key, ) from litellm.proxy._types import * from litellm.proxy._types import LiteLLM_VerificationToken @@ -1058,7 +1059,7 @@ async def _check_project_key_limits( project_id: str, data: Union[GenerateKeyRequest, UpdateKeyRequest], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> None: """ Validate that key's models and budget respect its project's limits. @@ -1833,7 +1834,7 @@ async def _process_single_key_update( user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str], prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: Any, llm_router: Optional[Router], user_custom_key_update: Optional[Callable] = None, @@ -3297,7 +3298,7 @@ async def _team_key_deletion_check( user_api_key_dict: UserAPIKeyAuth, key_info: LiteLLM_VerificationToken, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ): is_team_key = _is_team_key(data=key_info) @@ -3340,7 +3341,7 @@ async def _team_key_deletion_check( async def can_modify_verification_token( key_info: LiteLLM_VerificationToken, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, ) -> bool: @@ -3414,7 +3415,7 @@ async def can_modify_verification_token( async def delete_verification_tokens( tokens: List, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: @@ -3604,7 +3605,7 @@ async def _persist_deleted_verification_tokens( async def delete_key_aliases( key_aliases: List[str], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, @@ -3718,6 +3719,17 @@ async def _rotate_master_key( # noqa: PLR0915 "Failed to rotate MCP server credentials: %s", str(e) ) + # 4b. process MCP user-scoped credentials table (BYOK + OAuth2 tokens) + try: + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma_client, + new_master_key=new_master_key, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to rotate MCP user credentials: %s", str(e) + ) + # 5. process credentials table try: credentials = await prisma_client.db.litellm_credentialstable.find_many() @@ -3850,7 +3862,7 @@ async def _execute_virtual_key_regeneration( data: Optional[RegenerateKeyRequest], user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> GenerateKeyResponse: """Generate new token, update DB, invalidate cache, and return response.""" @@ -4140,7 +4152,7 @@ async def _check_proxy_or_team_admin_for_key( key_in_db: LiteLLM_VerificationToken, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> None: if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: return @@ -5161,7 +5173,7 @@ async def _check_key_admin_access( user_api_key_dict: UserAPIKeyAuth, hashed_token: str, prisma_client: Any, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, route: str, ) -> None: """ diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 46e7963da7..9dfc67370f 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -13,7 +13,9 @@ import base64 import hashlib import inspect import os +import re import secrets +from html import escape from copy import deepcopy from typing import ( TYPE_CHECKING, @@ -27,20 +29,23 @@ from typing import ( Union, cast, ) -from urllib.parse import urlencode, urlparse +from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: import httpx import jwt -from fastapi import APIRouter, Depends, HTTPException, Request, status +from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from fastapi.responses import RedirectResponse import litellm +from litellm.caching.dual_cache import DualCache from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching import DualCache from litellm.constants import ( + CLI_SSO_SESSION_CACHE_KEY_PREFIX, + CLI_SSO_SESSION_TTL_SECONDS, + LITELLM_CLI_SOURCE_IDENTIFIER, LITELLM_UI_SESSION_DURATION, MAX_SPENDLOG_ROWS_TO_QUERY, MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE, @@ -70,7 +75,11 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object -from litellm.proxy.auth.auth_utils import _has_user_setup_sso +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.auth.auth_utils import ( + _get_request_ip_address, + _has_user_setup_sso, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.admin_ui_utils import ( @@ -123,6 +132,250 @@ router = APIRouter() # Metadata fields (token_type, expires_in, scope) are intentionally kept so # response convertors see the same fields in the PKCE path as in the non-PKCE path. _OAUTH_TOKEN_FIELDS = frozenset({"access_token", "id_token", "refresh_token"}) +_CLI_SSO_FLOW_CACHE_KEY_PREFIX = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:flow" +_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX = ( + f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:start_rate_limit" +) +_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60 +_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30 +_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" +_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$") + + +def _hash_cli_sso_secret(secret: str) -> str: + return hashlib.sha256(secret.encode("utf-8")).hexdigest() + + +def _normalize_cli_sso_user_code(user_code: str) -> str: + return "".join(ch for ch in user_code.upper() if ch.isalnum()) + + +def _generate_cli_sso_user_code() -> str: + user_code = "".join(secrets.choice(_CLI_SSO_USER_CODE_ALPHABET) for _ in range(8)) + return f"{user_code[:4]}-{user_code[4:]}" + + +def _get_cli_sso_flow_cache_key(login_id: str) -> str: + return f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:{login_id}" + + +def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool: + return isinstance(login_id, str) and bool(_CLI_SSO_LOGIN_ID_RE.fullmatch(login_id)) + + +def _get_cli_sso_start_rate_limit_cache_key( + request: Request, use_x_forwarded_for: Optional[bool] = False +) -> str: + client_ip = ( + _get_request_ip_address( + request=request, use_x_forwarded_for=use_x_forwarded_for + ) + or "unknown" + ) + client_ip_hash = _hash_cli_sso_secret(client_ip) + return f"{_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX}:{client_ip_hash}" + + +def _check_cli_sso_start_rate_limit( + request: Request, + cache: DualCache, + use_x_forwarded_for: Optional[bool] = False, +) -> None: + rate_limit_cache_key = _get_cli_sso_start_rate_limit_cache_key( + request=request, use_x_forwarded_for=use_x_forwarded_for + ) + current_attempts = cache.increment_cache( + key=rate_limit_cache_key, + value=1, + ttl=_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS, + ) + if current_attempts > _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS: + raise HTTPException( + status_code=429, + detail="Too many CLI login attempts. Try again later.", + ) + + +def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict: + if not _is_valid_cli_sso_login_id(login_id): + raise HTTPException(status_code=400, detail="Invalid CLI login session") + + cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id)) + flow = cache.get_cache(key=cache_key) + if not isinstance(flow, dict) or "poll_secret_hash" not in flow: + raise HTTPException(status_code=400, detail="Invalid CLI login session") + return flow + + +def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None: + cache.set_cache( + key=_get_cli_sso_flow_cache_key(login_id), + value=flow, + ttl=CLI_SSO_SESSION_TTL_SECONDS, + ) + + +def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool: + expected_poll_secret_hash = flow.get("poll_secret_hash") + if not isinstance(expected_poll_secret_hash, str) or not isinstance( + poll_secret, str + ): + return False + supplied_poll_secret_hash = _hash_cli_sso_secret(poll_secret) + return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash) + + +def _render_cli_sso_verification_page( + verify_url: str, browser_complete_token: str +) -> str: + escaped_verify_url = escape(verify_url, quote=True) + escaped_browser_complete_token = escape(browser_complete_token, quote=True) + return f""" + + + + LiteLLM CLI Login + + + +
+

Complete CLI Login

+

Enter the verification code shown in your terminal to finish this login.

+
+ + + + +
+
+ + + """ + + +@router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False) +async def cli_sso_start(request: Request): + from litellm.proxy.proxy_server import general_settings, user_api_key_cache + + _check_cli_sso_start_rate_limit( + request=request, + cache=user_api_key_cache, + use_x_forwarded_for=bool( + (general_settings or {}).get("use_x_forwarded_for", False) + ), + ) + + login_id = f"cli-{secrets.token_urlsafe(24)}" + poll_secret = secrets.token_urlsafe(32) + user_code = _generate_cli_sso_user_code() + + flow = { + "poll_secret_hash": _hash_cli_sso_secret(poll_secret), + "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)), + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + + return { + "login_id": login_id, + "poll_secret": poll_secret, + "user_code": user_code, + "expires_in": CLI_SSO_SESSION_TTL_SECONDS, + } + + +@router.post( + "/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False +) +async def cli_sso_complete(request: Request, login_id: str): + from fastapi.responses import HTMLResponse + + from litellm.proxy.common_utils.html_forms.cli_sso_success import ( + render_cli_sso_success_page, + ) + from litellm.proxy.proxy_server import user_api_key_cache + + flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache) + if not flow.get("sso_complete") or not flow.get("session_data"): + raise HTTPException(status_code=400, detail="CLI login is not ready") + + body = (await request.body()).decode("utf-8") + form_values = parse_qs(body) + supplied_user_code = (form_values.get("user_code") or [""])[0] + supplied_browser_complete_token = ( + form_values.get("browser_complete_token") or [""] + )[0] + supplied_user_code_hash = _hash_cli_sso_secret( + _normalize_cli_sso_user_code(supplied_user_code) + ) + supplied_browser_complete_token_hash = _hash_cli_sso_secret( + supplied_browser_complete_token + ) + + expected_user_code_hash = flow.get("user_code_hash") + if not isinstance(expected_user_code_hash, str) or not secrets.compare_digest( + supplied_user_code_hash, expected_user_code_hash + ): + raise HTTPException(status_code=400, detail="Invalid verification code") + + expected_browser_complete_token_hash = flow.get("browser_complete_token_hash") + if not isinstance( + expected_browser_complete_token_hash, str + ) or not secrets.compare_digest( + supplied_browser_complete_token_hash, expected_browser_complete_token_hash + ): + raise HTTPException(status_code=400, detail="Invalid verification code") + + flow["user_code_verified"] = True + _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow) + + html_content = render_cli_sso_success_page() + return HTMLResponse(content=html_content, status_code=200) def normalize_email(email: Optional[str]) -> Optional[str]: @@ -333,6 +586,7 @@ async def google_login( from litellm.proxy.proxy_server import ( premium_user, prisma_client, + user_api_key_cache, user_custom_ui_sso_sign_in_handler, ) @@ -382,14 +636,15 @@ async def google_login( redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=request, sso_callback_route="sso/callback", - existing_key=existing_key, ) - # Store CLI key in state for OAuth flow + if source == LITELLM_CLI_SOURCE_IDENTIFIER: + _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) + + # Store CLI login handle in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( source=source, key=key, - existing_key=existing_key, ) # check if user defined a custom auth sso sign in handler, if yes, use it @@ -1050,7 +1305,7 @@ async def get_existing_user_info_from_db( user_id: Optional[str], user_email: Optional[str], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> Optional[LiteLLM_UserTable]: try: @@ -1074,7 +1329,7 @@ async def get_existing_user_info_from_db( async def get_user_info_from_db( result: Union[CustomOpenID, OpenID, dict], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, user_email: Optional[str], user_defined_values: Optional[SSOUserDefinedValues], @@ -1194,7 +1449,7 @@ async def _sync_user_role_from_jwt_role_map( received_response: Optional[dict], user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]], prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, user_defined_values: Optional[SSOUserDefinedValues], ) -> None: """ @@ -1233,11 +1488,8 @@ async def _sync_user_role_from_jwt_role_map( user_info.user_role = mapped_role.value await user_api_key_cache.async_set_cache( key=user_info.user_id, - value=( - user_info.model_dump() - if hasattr(user_info, "model_dump") - else dict(user_info) - ), + value=user_info, + model_type=LiteLLM_UserTable, ) @@ -1392,18 +1644,12 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): - # Extract the key ID and existing_key from the state - # State format: {PREFIX}:{key}:{existing_key} or {PREFIX}:{key} - state_parts = state.split(":", 2) # Split into max 3 parts + # State format: {PREFIX}:{login_id} + state_parts = state.split(":", 1) key_id = state_parts[1] if len(state_parts) > 1 else None - existing_key = state_parts[2] if len(state_parts) > 2 else None - verbose_proxy_logger.info( - f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}" - ) - return await cli_sso_callback( - request=request, key=key_id, existing_key=existing_key, result=result - ) + verbose_proxy_logger.info("CLI SSO callback detected") + return await cli_sso_callback(request=request, key=key_id, result=result) # Control-plane cross-origin: read return_to from cookie. # Starlette's cookie_parser already handles RFC 2109 unquoting. @@ -1424,13 +1670,10 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: async def cli_sso_callback( request: Request, key: Optional[str] = None, - existing_key: Optional[str] = None, result: Optional[Union[OpenID, dict]] = None, ): """CLI SSO callback - stores session info for JWT generation on polling""" - verbose_proxy_logger.info( - f"CLI SSO callback for key: {key}, existing_key: {existing_key}" - ) + verbose_proxy_logger.info("CLI SSO callback") from litellm.proxy.proxy_server import ( prisma_client, @@ -1438,11 +1681,7 @@ async def cli_sso_callback( user_api_key_cache, ) - if not key or not key.startswith("sk-"): - raise HTTPException( - status_code=400, - detail="Invalid key parameter. Must be a valid key ID starting with 'sk-'", - ) + flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache) if prisma_client is None: raise HTTPException( @@ -1480,9 +1719,6 @@ async def cli_sso_callback( status_code=500, detail="Failed to retrieve user information from SSO" ) - # Store session info in cache (10 min TTL) - from litellm.constants import CLI_SSO_SESSION_CACHE_KEY_PREFIX - # Get all teams from user_info - CLI will let user select which one teams: List[str] = [] if hasattr(user_info, "teams") and user_info.teams: @@ -1523,21 +1759,25 @@ async def cli_sso_callback( "team_details": team_details, } - cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key}" - user_api_key_cache.set_cache(key=cache_key, value=session_data, ttl=600) + flow["session_data"] = session_data + flow["sso_complete"] = True + browser_complete_token = secrets.token_urlsafe(32) + flow["browser_complete_token_hash"] = _hash_cli_sso_secret( + browser_complete_token + ) + _set_cli_sso_flow(login_id=cast(str, key), cache=user_api_key_cache, flow=flow) verbose_proxy_logger.info( f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}" ) - # Return success page from fastapi.responses import HTMLResponse - from litellm.proxy.common_utils.html_forms.cli_sso_success import ( - render_cli_sso_success_page, + verify_url = str(request.url_for("cli_sso_complete", login_id=key)) + html_content = _render_cli_sso_verification_page( + verify_url=verify_url, + browser_complete_token=browser_complete_token, ) - - html_content = render_cli_sso_success_page() return HTMLResponse(content=html_content, status_code=200) except Exception as e: @@ -1548,7 +1788,11 @@ async def cli_sso_callback( @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) -async def cli_poll_key(key_id: str, team_id: Optional[str] = None): +async def cli_poll_key( + key_id: str, + team_id: Optional[str] = None, + x_litellm_cli_poll_secret: Optional[str] = Header(default=None), +): """ CLI polling endpoint - retrieves session from cache and generates JWT. @@ -1557,22 +1801,25 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None): 2. Second poll (with team_id): Generates JWT with selected team and deletes session Args: - key_id: The session key ID + key_id: The CLI login session ID team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. """ - from litellm.constants import CLI_SSO_SESSION_CACHE_KEY_PREFIX from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken from litellm.proxy.proxy_server import user_api_key_cache - if not key_id.startswith("sk-"): - raise HTTPException(status_code=400, detail="Invalid key ID format") - try: - # Look up session in cache - cache_key = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:{key_id}" - session_data = user_api_key_cache.get_cache(key=cache_key) + flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache) + if not _verify_cli_sso_poll_secret( + flow=flow, poll_secret=x_litellm_cli_poll_secret + ): + raise HTTPException(status_code=403, detail="Invalid CLI polling secret") - if session_data: + if not flow.get("sso_complete") or not flow.get("user_code_verified"): + return {"status": "pending"} + + session_data = flow.get("session_data") + + if isinstance(session_data, dict): user_teams = session_data.get("teams", []) user_team_details = session_data.get("team_details") user_id = session_data["user_id"] @@ -1632,7 +1879,7 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None): ) # Delete cache entry (single-use) - user_api_key_cache.delete_cache(key=cache_key) + user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id)) verbose_proxy_logger.info( f"CLI JWT generated for user: {user_id}, team: {team_id}" @@ -1650,6 +1897,8 @@ async def cli_poll_key(key_id: str, team_id: Optional[str] = None): else: return {"status": "pending"} + except HTTPException: + raise except Exception as e: verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}") raise HTTPException( @@ -2393,20 +2642,15 @@ class SSOAuthenticationHandler: This is used to authenticate through the CLI login flow. - The state parameter format is: {PREFIX}:{key}:{existing_key} - - If existing_key is provided, it's included in the state + The state parameter format is: {PREFIX}:{login_id} - The state parameter is used to pass data through the OAuth flow without changing the callback URL """ from litellm.constants import ( LITELLM_CLI_SESSION_TOKEN_PREFIX, - LITELLM_CLI_SOURCE_IDENTIFIER, ) if source == LITELLM_CLI_SOURCE_IDENTIFIER and key: - if existing_key: - return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}:{existing_key}" - else: - return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" + return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" else: return None diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index e035168ca0..50339210a6 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -1,6 +1,5 @@ from typing import List, Optional -from litellm.caching import DualCache from litellm.proxy._types import ( KeyManagementRoutes, LiteLLM_TeamTableCachedObj, @@ -12,6 +11,7 @@ from litellm.proxy._types import ( ProxyException, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import PrismaClient @@ -65,7 +65,7 @@ class TeamMemberPermissionChecks: user_api_key_dict: UserAPIKeyAuth, route: KeyManagementRoutes, prisma_client: PrismaClient, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, existing_key_row: LiteLLM_VerificationToken, ): """ diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index 6bdff59da5..3b30fd3d63 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -3,6 +3,7 @@ Prometheus Auth Middleware - Pure ASGI implementation """ import json +from typing import Any, List, MutableMapping from fastapi import Request from starlette.types import ASGIApp, Receive, Scope, Send @@ -40,8 +41,17 @@ class PrometheusAuthMiddleware: # Only run auth if configured to do so if litellm.require_auth_for_metrics_endpoint is True: - # Construct Request only when auth is actually needed - request = Request(scope, receive) + # user_api_key_auth reads the request body, which consumes ASGI `receive`. + # Buffer those messages and replay them for the inner app; otherwise a + # successful auth would forward an exhausted receive and /metrics hangs. + buffered_messages: List[MutableMapping[str, Any]] = [] + + async def receive_for_auth() -> MutableMapping[str, Any]: + message = await receive() + buffered_messages.append(message) + return message + + request = Request(scope, receive_for_auth) api_key = request.headers.get(_AUTHORIZATION_HEADER) or "" try: @@ -70,5 +80,18 @@ class PrometheusAuthMiddleware: ) return + replay_idx = 0 + + async def receive_replay() -> MutableMapping[str, Any]: + nonlocal replay_idx + if replay_idx < len(buffered_messages): + msg = buffered_messages[replay_idx] + replay_idx += 1 + return msg + return await receive() + + await self.app(scope, receive_replay, send) + return + # Pass through to the inner application await self.app(scope, receive, send) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py index a8c5562d4d..6277f6b4a7 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/assembly_passthrough_logging_handler.py @@ -1,6 +1,7 @@ import asyncio import json import time +import urllib.parse from datetime import datetime from typing import Literal, Optional from urllib.parse import urlparse @@ -203,8 +204,16 @@ class AssemblyAIPassthroughLoggingHandler: ) if _api_key is None: raise ValueError("AssemblyAI API key not found") + if ( + any(c in transcript_id for c in ("/", "\\", "#", "?")) + or ".." in transcript_id + ): + raise ValueError( + f"Invalid transcript_id {transcript_id!r}: contains disallowed characters" + ) + safe_transcript_id = urllib.parse.quote(transcript_id, safe="") try: - url = f"{_base_url}/v2/transcript/{transcript_id}" + url = f"{_base_url}/v2/transcript/{safe_transcript_id}" headers = { "Authorization": f"Bearer {_api_key}", "Content-Type": "application/json", diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 77eb3a5ee0..714b5f3c7b 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -41,7 +41,6 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.passthrough import BasePassthroughUtils from litellm.proxy._types import ( - CommonProxyErrors, ConfigFieldInfo, ConfigFieldUpdate, LiteLLMRoutes, @@ -2325,12 +2324,14 @@ async def _register_pass_through_endpoint( dependencies = None if auth is not None and str(auth).lower() == "true": - if premium_user is not True: - raise ValueError( - "Error Setting Authentication on Pass Through Endpoint: {}".format( - CommonProxyErrors.not_premium_user.value - ) - ) + # Authentication on a pass-through endpoint used to be enterprise- + # only — which left the OSS tier with no safe configuration: the + # default was ``auth=False`` (unauthenticated forwarder) and the + # safe ``auth=True`` raised at startup unless the operator had a + # license. The default is now ``True`` (safe-by-default), and + # turning it on no longer requires a license: an unauthenticated + # forwarder is a deployment choice the operator should be allowed + # to make explicitly, but the safe option must always be free. dependencies = [Depends(user_api_key_auth)] if path not in LiteLLMRoutes.openai_routes.value: LiteLLMRoutes.openai_routes.value.append(path) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 11c7320efe..a2bdc3ab3c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -78,8 +78,11 @@ from litellm.proxy._types import ( InvitationNew, InvitationUpdate, Litellm_EntityType, + LiteLLM_EndUserTable, LiteLLM_JWTAuth, + LiteLLM_TagTable, LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, LitellmUserRoles, PassThroughGenericEndpoint, @@ -94,6 +97,7 @@ from litellm.proxy._types import ( UI_TEAM_ID, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.callback_utils import ( normalize_callback_names, process_callback, @@ -206,6 +210,7 @@ from litellm import Router from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, APSCHEDULER_COALESCE, @@ -1612,7 +1617,7 @@ prisma_client: Optional[PrismaClient] = None shared_aiohttp_session: Optional["ClientSession"] = ( None # Global shared session for connection reuse ) -user_api_key_cache = DualCache( +user_api_key_cache: UserApiKeyCache = UserApiKeyCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) spend_counter_cache = DualCache( @@ -1798,12 +1803,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: 3. Reseed from authoritative DB spend (counter expired, cross-pod stale) 4. Caller-supplied fallback (DB unavailable, cold start) """ - # 1. Try Redis first (cross-pod authoritative) + # 1. Redis first (cross-pod authoritative). On clean miss, skip + # in-memory: per-pod in-memory only has this pod's writes, so it + # would mask cross-pod increments. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: return float(val) + redis_clean_miss = True except Exception as e: verbose_proxy_logger.debug( "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", @@ -1811,10 +1820,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: e, ) - # 2. Fall back to in-memory counter (single-instance or Redis failure) - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + # 2. In-memory only when Redis is unreachable. + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. db_spend = await SpendCounterReseed.coalesced( @@ -1976,10 +1986,12 @@ async def _init_and_increment_spend_counter( base_spend = getattr(source, "spend", 0.0) or 0.0 if base_spend > 0: await spend_counter_cache.async_increment_cache( - key=counter_key, value=base_spend + key=counter_key, value=base_spend, refresh_ttl=True ) - await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, refresh_ttl=True + ) async def update_cache( # noqa: PLR0915 @@ -2007,14 +2019,16 @@ async def update_cache( # noqa: PLR0915 else: hashed_token = token verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token) - existing_spend_obj: LiteLLM_VerificationTokenView = await user_api_key_cache.async_get_cache(key=hashed_token) # type: ignore + existing_spend_obj = await user_api_key_cache.async_get_cache( + key=hashed_token, model_type=UserAPIKeyAuth + ) verbose_proxy_logger.debug( f"_update_key_cache: existing_spend_obj={existing_spend_obj}" ) if existing_spend_obj is None: return - else: - existing_spend = existing_spend_obj.spend + + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost @@ -2072,41 +2086,48 @@ async def update_cache( # noqa: PLR0915 existing_team_member_spend + response_cost ) - # Update the cost column for the given token + # Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns + # BaseModel values into dicts for Redis (same Codec path as async_set_cache). existing_spend_obj.spend = new_spend values_to_update_in_cache.append((hashed_token, existing_spend_obj)) ### UPDATE USER SPEND ### async def _update_user_cache(): ## UPDATE CACHE FOR USER ID + GLOBAL PROXY + if response_cost is None: + return user_ids = [user_id] try: for _id in user_ids: # Fetch the existing cost for the given user if _id is None: continue - existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id) - if existing_spend_obj is None: + cached_user = await user_api_key_cache.async_get_cache(key=_id) + if cached_user is None: # do nothing if there is no cache value return + existing_spend_obj = CacheCodec.deserialize( + cached_user, LiteLLM_UserTable + ) + if existing_spend_obj is None: + return verbose_proxy_logger.debug( f"_update_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj.json())) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_UserTable + ), + ) + ) ## UPDATE GLOBAL PROXY ## global_proxy_spend = await user_api_key_cache.async_get_cache( key="{}:spend".format(litellm_proxy_admin_name) @@ -2138,31 +2159,33 @@ async def update_cache( # noqa: PLR0915 _id = "end_user_id:{}".format(end_user_id) try: # Fetch the existing cost for the given user - existing_spend_obj = await user_api_key_cache.async_get_cache(key=_id) - if existing_spend_obj is None: + cached_end_user = await user_api_key_cache.async_get_cache(key=_id) + if cached_end_user is None: # if user does not exist in LiteLLM_UserTable, create a new user # do nothing if end-user not in api key cache return + existing_spend_obj = CacheCodec.deserialize( + cached_end_user, LiteLLM_EndUserTable + ) + if existing_spend_obj is None: + return verbose_proxy_logger.debug( f"_update_end_user_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if existing_spend_obj is None: - existing_spend = 0 - else: - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend + + existing_spend = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj.json())) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_EndUserTable + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update end user spend in cache. " @@ -2181,36 +2204,32 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: - # Fetch the existing cost for the given user - existing_spend_obj: Optional[LiteLLM_TeamTable] = ( - await user_api_key_cache.async_get_cache(key=_id) + cached_team = await user_api_key_cache.async_get_cache(key=_id) + if cached_team is None: + # do nothing if team not in api key cache + return + existing_spend_obj: Optional[LiteLLM_TeamTableCachedObj] = ( + CacheCodec.deserialize(cached_team, LiteLLM_TeamTableCachedObj) ) if existing_spend_obj is None: - # do nothing if team not in api key cache return verbose_proxy_logger.debug( f"_update_team_db: existing spend: {existing_spend_obj}; response_cost: {response_cost}" ) - if existing_spend_obj is None: - existing_spend: Optional[float] = 0.0 - else: - if isinstance(existing_spend_obj, dict): - existing_spend = existing_spend_obj["spend"] - else: - existing_spend = existing_spend_obj.spend - if existing_spend is None: - existing_spend = 0.0 + existing_spend: float = existing_spend_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the cost column for the given user - if isinstance(existing_spend_obj, dict): - existing_spend_obj["spend"] = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) - else: - existing_spend_obj.spend = new_spend - values_to_update_in_cache.append((_id, existing_spend_obj)) + existing_spend_obj.spend = new_spend + values_to_update_in_cache.append( + ( + _id, + CacheCodec.serialize( + existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update team spend in cache. " @@ -2237,32 +2256,32 @@ async def update_cache( # noqa: PLR0915 cache_key = f"tag:{tag_name}" # Fetch the existing tag object from cache - existing_tag_obj = await user_api_key_cache.async_get_cache( - key=cache_key - ) - if existing_tag_obj is None: + cached_tag = await user_api_key_cache.async_get_cache(key=cache_key) + if cached_tag is None: # do nothing if tag not in api key cache continue + existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable) + if existing_tag_obj is None: + continue + verbose_proxy_logger.debug( f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}" ) - if isinstance(existing_tag_obj, dict): - existing_spend = existing_tag_obj.get("spend", 0) or 0 - else: - existing_spend = getattr(existing_tag_obj, "spend", 0) or 0 - + existing_spend = existing_tag_obj.spend or 0.0 # Calculate the new cost by adding the existing cost and response_cost new_spend = existing_spend + response_cost - # Update the spend column for the given tag - if isinstance(existing_tag_obj, dict): - existing_tag_obj["spend"] = new_spend - values_to_update_in_cache.append((cache_key, existing_tag_obj)) - else: - existing_tag_obj.spend = new_spend - values_to_update_in_cache.append((cache_key, existing_tag_obj)) + existing_tag_obj.spend = new_spend + values_to_update_in_cache.append( + ( + cache_key, + CacheCodec.serialize( + existing_tag_obj, model_type=LiteLLM_TagTable + ), + ) + ) except Exception as e: verbose_proxy_logger.warning( "Spend tracking - failed to update tag spend in cache. " @@ -2930,8 +2949,9 @@ class ProxyConfig: def _init_cache( self, cache_params: dict, + enable_redis_auth_cache: bool = False, ): - global redis_usage_cache, llm_router + global redis_usage_cache, llm_router, general_settings from litellm import Cache if "default_in_memory_ttl" in cache_params: @@ -2947,7 +2967,29 @@ class ProxyConfig: ): ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache - spend_counter_cache.redis_cache = redis_usage_cache + spend_counter_cache.attach_redis_cache( + redis_usage_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + # Note: PKCE verifier storage uses redis_usage_cache directly (not + # user_api_key_cache) to avoid routing all API-key lookups through Redis. + if enable_redis_auth_cache is True: + user_api_key_cache.attach_redis_cache( + redis_usage_cache, + default_redis_ttl=litellm.default_redis_ttl, + ) + verbose_proxy_logger.info( + "enable_redis_auth_cache=True: attached Redis to " + "user_api_key_cache — virtual-key lookups are now " + "shared across all proxy workers." + ) + else: + verbose_proxy_logger.info( + "enable_redis_auth_cache is not set: user_api_key_cache " + "remains in-memory only (per-worker). Set " + "litellm_settings.enable_redis_auth_cache: true to share " + "the auth cache across workers and reduce DB load." + ) litellm_config_cache.redis_cache = redis_usage_cache # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. @@ -3273,7 +3315,13 @@ class ProxyConfig: cache_params[key] = get_secret(value) ## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = `, _redis.py checks for REDIS specific environment variables - self._init_cache(cache_params=cache_params) + self._init_cache( + cache_params=cache_params, + enable_redis_auth_cache=litellm_settings.get( + "enable_redis_auth_cache", False + ) + is True, + ) if litellm.cache is not None: verbose_proxy_logger.debug( f"{blue_color_code}Set Cache on LiteLLM Proxy{reset_color_code}" @@ -3544,21 +3592,23 @@ class ProxyConfig: verbose_proxy_logger.critical( "LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use." ) - ### USER API KEY CACHE IN-MEMORY TTL ### + ### USER API KEY CACHE TTL (in-memory + Redis when Redis auth sharing is enabled) ### user_api_key_cache_ttl = general_settings.get( "user_api_key_cache_ttl", None ) if user_api_key_cache_ttl is not None: + ttl = float(user_api_key_cache_ttl) + # Mirror TTL on Redis as well when ``litellm_settings.enable_redis_auth_cache`` + # attaches Redis to ``user_api_key_cache``; otherwise DualCache misses in + # memory fall back to a key that outlasts ``user_api_key_cache_ttl``. user_api_key_cache.update_cache_ttl( - default_in_memory_ttl=float(user_api_key_cache_ttl), - default_redis_ttl=None, # user_api_key_cache uses in-memory TTL only; Redis not configured for key lookups + default_in_memory_ttl=ttl, + default_redis_ttl=ttl, ) ### PKCE MULTI-INSTANCE PREREQUISITE CHECK ### # PKCE verifiers are stored in redis_usage_cache when available so they can # be read back by any instance (not just the one that started the auth flow). - # user_api_key_cache is intentionally left in-memory-only to avoid routing - # all API-key lookups through Redis. use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true" if use_pkce and redis_usage_cache is None: global _pkce_no_redis_warning_emitted @@ -6287,7 +6337,7 @@ class ProxyStartupEvent: cls, general_settings: dict, prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ): """Initialize JWT auth on startup""" if general_settings.get("litellm_jwtauth", None) is not None: @@ -6336,7 +6386,7 @@ class ProxyStartupEvent: async def _warm_global_spend_cache( cls, litellm_proxy_admin_name: str, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, prisma_client: PrismaClient, ) -> None: """Warm global spend cache once at startup to reduce impact of first wave of requests.""" @@ -6976,7 +7026,7 @@ class ProxyStartupEvent: cls, database_url: Optional[str], proxy_logging_obj: ProxyLogging, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> Optional[PrismaClient]: """ - Sets up prisma client @@ -10278,6 +10328,101 @@ def _paginate_models_response( } +def _team_models_resolve_to_names( + team_models: List[str], access_groups: Dict[str, Any] +) -> List[str]: + """Expand team model entries (including access group names) to concrete model names.""" + resolved: List[str] = [] + for name in team_models: + if name in access_groups: + resolved.extend(access_groups[name]) + else: + resolved.append(name) + return resolved + + +async def _load_team_object_for_model_filter( + team_id: str, prisma_client: PrismaClient +) -> Optional[LiteLLM_TeamTable]: + """Load team row from DB; returns None if missing or on error.""" + try: + team_db_object = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + if team_db_object is None: + verbose_proxy_logger.warning(f"Team {team_id} not found in database") + return None + return LiteLLM_TeamTable(**team_db_object.model_dump()) + except Exception as e: + verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}") + return None + + +async def _gather_team_accessible_model_ids( + team_object: LiteLLM_TeamTable, + team_id: str, + prisma_client: PrismaClient, + llm_router: Router, +) -> Set[str]: + """Collect model IDs the team can use from router config and DB.""" + team_accessible_model_ids: Set[str] = set() + access_groups = llm_router.get_model_access_groups() if llm_router else {} + + if ( + not team_object.models + or SpecialModelNames.all_proxy_models.value in team_object.models + ): + model_list = llm_router.get_model_list() if llm_router else [] + if model_list is not None: + for model in model_list: + model_id = model.get("model_info", {}).get("id", None) + if model_id is None: + continue + team_model_id = model.get("model_info", {}).get("team_id", None) + if team_model_id is None or team_model_id == team_id: + team_accessible_model_ids.add(model_id) + else: + resolved_model_names: Set[str] = set() + for model_name in team_object.models: + if model_name in access_groups: + resolved_model_names.update(access_groups[model_name]) + else: + resolved_model_names.add(model_name) + + for model_name in resolved_model_names: + _models = ( + llm_router.get_model_list(model_name=model_name, team_id=team_id) + if llm_router + else [] + ) + if _models is not None: + for model in _models: + model_id = model.get("model_info", {}).get("id", None) + if model_id is not None: + team_accessible_model_ids.add(model_id) + + try: + if ( + team_object.models + and SpecialModelNames.all_proxy_models.value not in team_object.models + ): + _resolved_names = _team_models_resolve_to_names( + team_object.models, access_groups + ) + db_models = await prisma_client.db.litellm_proxymodeltable.find_many( + where={"model_name": {"in": _resolved_names}} + ) + for db_model in db_models: + if db_model.model_id: + team_accessible_model_ids.add(db_model.model_id) + except Exception as e: + verbose_proxy_logger.debug( + f"Error querying database models for team {team_id}: {str(e)}" + ) + + return team_accessible_model_ids + + async def _filter_models_by_team_id( all_models: List[Dict[str, Any]], team_id: str, @@ -10300,78 +10445,13 @@ async def _filter_models_by_team_id( Returns: Filtered list of models """ - # Get team from database - try: - team_db_object = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - if team_db_object is None: - verbose_proxy_logger.warning(f"Team {team_id} not found in database") - # If team doesn't exist, return empty list - return [] - - team_object = LiteLLM_TeamTable(**team_db_object.model_dump()) - except Exception as e: - verbose_proxy_logger.exception(f"Error fetching team {team_id}: {str(e)}") + team_object = await _load_team_object_for_model_filter(team_id, prisma_client) + if team_object is None: return [] - # Get models accessible to this team (similar to _add_team_models_to_all_models) - team_accessible_model_ids: Set[str] = set() - - if ( - not team_object.models # empty list = all model access - or SpecialModelNames.all_proxy_models.value in team_object.models - ): - # Team has access to all models - model_list = llm_router.get_model_list() if llm_router else [] - if model_list is not None: - for model in model_list: - model_id = model.get("model_info", {}).get("id", None) - if model_id is None: - continue - # if team model id set, check if team id matches - team_model_id = model.get("model_info", {}).get("team_id", None) - can_add_model = False - if team_model_id is None: - can_add_model = True - elif team_model_id == team_id: - can_add_model = True - - if can_add_model: - team_accessible_model_ids.add(model_id) - else: - # Team has access to specific models - for model_name in team_object.models: - _models = ( - llm_router.get_model_list(model_name=model_name, team_id=team_id) - if llm_router - else [] - ) - if _models is not None: - for model in _models: - model_id = model.get("model_info", {}).get("id", None) - if model_id is not None: - team_accessible_model_ids.add(model_id) - - # Also search database for models accessible to this team - # This complements the config search done above - try: - if ( - team_object.models - and SpecialModelNames.all_proxy_models.value not in team_object.models - ): - # Team has specific models - check database for those model names - db_models = await prisma_client.db.litellm_proxymodeltable.find_many( - where={"model_name": {"in": team_object.models}} - ) - for db_model in db_models: - model_id = db_model.model_id - if model_id: - team_accessible_model_ids.add(model_id) - except Exception as e: - verbose_proxy_logger.debug( - f"Error querying database models for team {team_id}: {str(e)}" - ) + team_accessible_model_ids = await _gather_team_accessible_model_ids( + team_object, team_id, prisma_client, llm_router + ) # Filter models based on direct_access or access_via_team_ids # Models are already enriched with these fields before this function is called @@ -12404,9 +12484,20 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): @app.get("/get_logo_url", include_in_schema=False) def get_logo_url(): - """Get the current logo URL from environment""" + """Get the current logo URL from environment. + + Only HTTP(S) URLs are returned — those are intended to be loaded + directly by the browser from a public/internal CDN. Local file + paths set via ``UI_LOGO_PATH`` are NOT returned: they are admin- + only filesystem details, the dashboard falls back to ``/get_image`` + which serves the file only when it is a supported image. Without + this filter, the unauthenticated endpoint would disclose internal + hostnames or filesystem paths to any caller. + """ logo_path = os.getenv("UI_LOGO_PATH", "") - return {"logo_url": logo_path} + if logo_path.startswith(("http://", "https://")): + return {"logo_url": logo_path} + return {"logo_url": ""} @app.get("/get_image", include_in_schema=False) @@ -12445,61 +12536,44 @@ async def get_image(): if assets_dir != current_dir and not os.path.exists(default_logo): default_logo = default_site_logo - cache_dir = assets_dir if os.access(assets_dir, os.W_OK) else current_dir - cache_path = os.path.join(cache_dir, "cached_logo.jpg") - logo_path = os.getenv("UI_LOGO_PATH", default_logo) verbose_proxy_logger.debug("Reading logo from path: %s", logo_path) - # If UI_LOGO_PATH points to a local file, serve it directly (skip cache) + from litellm.proxy.common_utils.static_asset_utils import ( + resolve_validated_local_image_path, + ) + if logo_path != default_logo and not logo_path.startswith(("http://", "https://")): - if os.path.exists(logo_path): - return FileResponse(logo_path, media_type="image/jpeg") - # Custom path doesn't exist — fall back to default + safe_logo = resolve_validated_local_image_path(logo_path) + if safe_logo is not None: + safe_logo_path, media_type = safe_logo + return FileResponse(safe_logo_path, media_type=media_type) verbose_proxy_logger.warning( - f"UI_LOGO_PATH '{logo_path}' does not exist, falling back to default logo" + "UI_LOGO_PATH %r is not a supported image file or does not exist, " + "falling back to default logo", + logo_path, ) logo_path = default_logo - # [OPTIMIZATION] For HTTP URLs and default logo, check if the cached image exists - if os.path.exists(cache_path): - return FileResponse(cache_path, media_type="image/jpeg") - - # Check if the logo path is an HTTP/HTTPS URL + # Remote logo URLs are loaded by the browser. The proxy should not fetch + # arbitrary admin-configured URLs server-side. if logo_path.startswith(("http://", "https://")): - try: - # Download the image and cache it - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client - from litellm.types.llms.custom_http import httpxSpecialProvider + return RedirectResponse(url=logo_path) - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.UI, - params={"timeout": 5.0}, - ) - response = await async_client.get(logo_path) - if response.status_code == 200: - # Save the image to a local file - with open(cache_path, "wb") as f: - f.write(response.content) - - # Return the cached image as a FileResponse - return FileResponse(cache_path, media_type="image/jpeg") - else: - # Handle the case when the image cannot be downloaded - return FileResponse(default_logo, media_type="image/jpeg") - except Exception as e: - # Handle any exceptions during the download (e.g., timeout, connection error) - verbose_proxy_logger.debug(f"Error downloading logo from {logo_path}: {e}") - return FileResponse(default_logo, media_type="image/jpeg") - else: - # Return the local image file if the logo path is not an HTTP/HTTPS URL - return FileResponse(logo_path, media_type="image/jpeg") + # Default logo (resolved from the bundled asset, not user-controlled). + safe_logo = resolve_validated_local_image_path(logo_path) + if safe_logo is not None: + safe_logo_path, media_type = safe_logo + return FileResponse(safe_logo_path, media_type=media_type) + return FileResponse(default_site_logo, media_type="image/jpeg") @app.get("/get_favicon", include_in_schema=False) async def get_favicon(): """Get custom favicon for the admin UI.""" - from fastapi.responses import Response + from litellm.proxy.common_utils.static_asset_utils import ( + resolve_validated_local_image_path, + ) current_dir = os.path.dirname(os.path.abspath(__file__)) default_favicon = os.path.join(current_dir, "_experimental", "out", "favicon.ico") @@ -12512,42 +12586,17 @@ async def get_favicon(): raise HTTPException(status_code=404, detail="Default favicon not found") if favicon_url.startswith(("http://", "https://")): - try: - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client - from litellm.types.llms.custom_http import httpxSpecialProvider - - async_client = get_async_httpx_client( - llm_provider=httpxSpecialProvider.UI, - params={"timeout": 5.0}, - ) - response = await async_client.get(favicon_url) - if response.status_code == 200: - content_type = response.headers.get("content-type", "image/x-icon") - return Response( - content=response.content, - media_type=content_type, - ) - else: - verbose_proxy_logger.warning( - "Failed to fetch favicon from %s: status %s", - favicon_url, - response.status_code, - ) - if os.path.exists(default_favicon): - return FileResponse(default_favicon, media_type="image/x-icon") - raise HTTPException(status_code=404, detail="Favicon not found") - except HTTPException: - raise - except Exception as e: - verbose_proxy_logger.debug( - "Error downloading favicon from %s: %s", favicon_url, e - ) - if os.path.exists(default_favicon): - return FileResponse(default_favicon, media_type="image/x-icon") - raise HTTPException(status_code=404, detail="Favicon not found") + return RedirectResponse(url=favicon_url) else: - if os.path.exists(favicon_url): - return FileResponse(favicon_url, media_type="image/x-icon") + safe_favicon = resolve_validated_local_image_path(favicon_url) + if safe_favicon is not None: + safe_favicon_path, media_type = safe_favicon + return FileResponse(safe_favicon_path, media_type=media_type) + verbose_proxy_logger.warning( + "LITELLM_FAVICON_URL %r is not a supported image file or does not " + "exist, falling back to default favicon", + favicon_url, + ) if os.path.exists(default_favicon): return FileResponse(default_favicon, media_type="image/x-icon") raise HTTPException(status_code=404, detail="Favicon not found") @@ -12842,9 +12891,12 @@ async def update_config( # noqa: PLR0915 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - For Admin UI - allows admin to update config via UI + For Admin UI - allows admin to update config via UI. - Currently supports modifying General Settings + LiteLLM settings + Writes only the sections present in the request body to LiteLLM_Config rows + (one row per top-level section). Sections the caller did not send are left + untouched — this endpoint never persists pre-existing YAML values to DB as + a side effect of an unrelated update. """ global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client try: @@ -12852,109 +12904,96 @@ async def update_config( # noqa: PLR0915 raise HTTPException( status_code=403, detail="Only proxy admins can update config" ) - import base64 - """ - - Update the ConfigTable DB - - Run 'add_deployment' - """ if prisma_client is None: raise Exception("No DB Connected") - if store_model_in_db is not True: - raise HTTPException( - status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + async def _read_section(param_name: str) -> dict: + row = await prisma_client.db.litellm_config.find_first( + where={"param_name": param_name} + ) + if row is None or row.param_value is None: + return {} + return dict(row.param_value) + + async def _upsert_section(param_name: str, value: dict) -> None: + serialized = json.dumps(value) + await prisma_client.db.litellm_config.upsert( + where={"param_name": param_name}, + data={ + "create": {"param_name": param_name, "param_value": serialized}, + "update": {"param_value": serialized}, }, ) + # invalidate the DualCache entry so the next reader (this process + # or any other proxy in the cluster) goes to DB. + await invalidate_config_param(param_name) - updated_settings = config_info.json(exclude_none=True) - updated_settings = prisma_client.jsonify_object(updated_settings) - for k, v in updated_settings.items(): - if k == "router_settings": - await prisma_client.db.litellm_config.upsert( - where={"param_name": k}, - data={ - "create": {"param_name": k, "param_value": v}, - "update": {"param_value": v}, - }, - ) - await invalidate_config_param(k) - - ### OLD LOGIC [TODO] MOVE TO DB ### - - # Load existing config - config = await proxy_config.get_config() - verbose_proxy_logger.debug("Loaded config: %s", config) - - # update the general settings + # general_settings: merge per-key, with the alert_to_webhook_url side + # effect of auto-enabling slack alerting. if config_info.general_settings is not None: - config.setdefault("general_settings", {}) - updated_general_settings = config_info.general_settings.dict( - exclude_none=True - ) - - _existing_settings = config["general_settings"] - for k, v in updated_general_settings.items(): - # overwrite existing settings with updated values + existing = await _read_section("general_settings") + updates = config_info.general_settings.dict(exclude_none=True) + for k, v in updates.items(): if k == "alert_to_webhook_url": - # check if slack is already enabled. if not, enable it - if "alerting" not in _existing_settings: - _existing_settings = {"alerting": ["slack"]} - elif isinstance(_existing_settings["alerting"], list): - if "slack" not in _existing_settings["alerting"]: - _existing_settings["alerting"].append("slack") - _existing_settings[k] = v - config["general_settings"] = _existing_settings + if "alerting" not in existing: + existing["alerting"] = ["slack"] + elif ( + isinstance(existing["alerting"], list) + and "slack" not in existing["alerting"] + ): + existing["alerting"].append("slack") + existing[k] = v + await _upsert_section("general_settings", existing) + # environment_variables: encrypt request values, then merge into existing. if config_info.environment_variables is not None: - config.setdefault("environment_variables", {}) - _updated_environment_variables = config_info.environment_variables + existing = await _read_section("environment_variables") + for k, v in config_info.environment_variables.items(): + existing[k] = encrypt_value_helper(value=v) + await _upsert_section("environment_variables", existing) - # encrypt updated_environment_variables # - for k, v in _updated_environment_variables.items(): - encrypted_value = encrypt_value_helper(value=v) - _updated_environment_variables[k] = encrypted_value - - _existing_env_variables = config["environment_variables"] - - for k, v in _updated_environment_variables.items(): - # overwrite existing env variables with updated values - _existing_env_variables[k] = _updated_environment_variables[k] - - # update the litellm settings + # litellm_settings: merge existing + request, request wins (matching + # router_settings semantics — the caller's value for any given key is + # what gets persisted). success_callback is special-cased: it is + # always normalized + deduped, and unioned with any existing list, + # because callbacks are additive (callers send the new entry, not + # the full set). Normalizing on every write — not only when an + # existing entry is present — keeps the DB free of mixed-case + # entries that delete_callback (lowercase lookup) cannot find. if config_info.litellm_settings is not None: - config.setdefault("litellm_settings", {}) - updated_litellm_settings = config_info.litellm_settings - config["litellm_settings"] = { - **updated_litellm_settings, - **config["litellm_settings"], - } + existing = await _read_section("litellm_settings") + updated_litellm_settings = dict(config_info.litellm_settings) - # if litellm.success_callback in updated_litellm_settings and config["litellm_settings"] - if ( - "success_callback" in updated_litellm_settings - and "success_callback" in config["litellm_settings"] - ): - # check both success callback are lists - if isinstance( - config["litellm_settings"]["success_callback"], list - ) and isinstance(updated_litellm_settings["success_callback"], list): - updated_success_callbacks_normalized = normalize_callback_names( - updated_litellm_settings["success_callback"] - ) - combined_success_callback = ( - config["litellm_settings"]["success_callback"] - + updated_success_callbacks_normalized - ) - combined_success_callback = list(set(combined_success_callback)) - config["litellm_settings"][ - "success_callback" - ] = combined_success_callback + incoming_cb = updated_litellm_settings.get("success_callback") + if isinstance(incoming_cb, list): + updated_litellm_settings["success_callback"] = normalize_callback_names( + incoming_cb + ) - # Save the updated config - await proxy_config.save_config(new_config=config) + merged = {**existing, **updated_litellm_settings} + + incoming_cb = updated_litellm_settings.get("success_callback") + existing_cb = existing.get("success_callback") + if isinstance(incoming_cb, list): + if isinstance(existing_cb, list): + # Normalize the existing list too — a row written by a + # different code path may still hold mixed-case names, + # which would otherwise dedup-miss against the lowercase + # incoming entries. + merged["success_callback"] = list( + set(normalize_callback_names(existing_cb) + incoming_cb) + ) + else: + merged["success_callback"] = list(set(incoming_cb)) + + await _upsert_section("litellm_settings", merged) + + # router_settings: merge existing + request, request wins. + if config_info.router_settings is not None: + existing = await _read_section("router_settings") + updates = config_info.router_settings.dict(exclude_none=True) + await _upsert_section("router_settings", {**existing, **updates}) await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d2dfa17751..8c5fce8409 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -101,6 +101,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, should_create_missing_views, @@ -340,7 +341,7 @@ class ProxyLogging: def __init__( self, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, premium_user: bool = False, ): ## INITIALIZE LITELLM CALLBACKS ## @@ -5715,7 +5716,7 @@ async def get_available_models_for_user( include_model_access_groups: bool = False, only_model_access_groups: bool = False, return_wildcard_routes: bool = False, - user_api_key_cache: Optional["DualCache"] = None, + user_api_key_cache: Optional["UserApiKeyCache"] = None, ) -> List[str]: """ Get the list of models available to a user based on their API key and team permissions. diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 9868634362..e27585116c 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -163,19 +163,21 @@ def rerank( # noqa: PLR0915 model_response = RerankResponse() + rerank_litellm_params = { + "litellm_call_id": litellm_call_id, + "proxy_server_request": proxy_server_request, + "model_info": model_info, + "preset_cache_key": None, + "stream_response": {}, + **optional_params.model_dump(exclude_unset=True), + } + litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, user=user, optional_params=dict(optional_rerank_params), - litellm_params={ - "litellm_call_id": litellm_call_id, - "proxy_server_request": proxy_server_request, - "model_info": model_info, - "preset_cache_key": None, - "stream_response": {}, - **optional_params.model_dump(exclude_unset=True), - }, + litellm_params=dict(rerank_litellm_params), custom_llm_provider=_custom_llm_provider, ) @@ -214,6 +216,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.AZURE_AI: api_base = ( @@ -235,6 +238,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.INFINITY: # Implement Infinity rerank logic @@ -265,6 +269,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.TOGETHER_AI: # Implement Together AI rerank logic @@ -318,6 +323,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.NVIDIA_NIM: if dynamic_api_key is None: @@ -346,6 +352,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.BEDROCK: api_base = ( @@ -409,6 +416,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA: @@ -442,6 +450,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.FIREWORKS_AI: api_key = ( @@ -472,6 +481,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.VOYAGE: api_key = ( @@ -500,6 +510,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) elif _custom_llm_provider == litellm.LlmProviders.WATSONX: credentials = IBMWatsonXMixin.get_watsonx_credentials( @@ -527,6 +538,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) else: # Generic handler for all providers that use base_llm_http_handler @@ -559,6 +571,7 @@ def rerank( # noqa: PLR0915 headers=headers or litellm.headers or {}, client=client, model_response=model_response, + litellm_params=rerank_litellm_params, ) # Placeholder return diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 145ec3a641..da8da1b486 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,9 +1,12 @@ +from __future__ import annotations + import asyncio import json import time import traceback from datetime import datetime -from typing import Any, Dict, List, Optional +from functools import lru_cache +from typing import Any, Dict, List, Literal, Optional import httpx @@ -22,19 +25,26 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.types.llms.openai import ( - OutputTextDeltaEvent, - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIRequestParams, - ResponsesAPIResponse, - ResponsesAPIStreamEvents, - ResponsesAPIStreamingResponse, -) +from litellm.types.llms.openai import ResponsesAPIStreamEvents from litellm.types.utils import CallTypes from litellm.utils import CustomStreamWrapper, async_post_call_success_deployment_hook +@lru_cache(maxsize=1) +def _get_openai_response_types(): + from litellm.types.llms import openai as openai_types + + return openai_types + + +def _log_background_task_failure(task: "asyncio.Task[Any]", *, task_name: str) -> None: + if task.cancelled(): + return + exception = task.exception() + if exception is not None: + verbose_logger.error("%s failed: %s", task_name, exception) + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -46,7 +56,7 @@ class BaseResponsesAPIStreamingIterator: self, response: httpx.Response, model: str, - responses_api_provider_config: BaseResponsesAPIConfig, + responses_api_provider_config: Optional[BaseResponsesAPIConfig], logging_obj: LiteLLMLoggingObj, litellm_metadata: Optional[Dict[str, Any]] = None, custom_llm_provider: Optional[str] = None, @@ -58,9 +68,13 @@ class BaseResponsesAPIStreamingIterator: self.logging_obj = logging_obj self.finished = False self.responses_api_provider_config = responses_api_provider_config - self.completed_response: Optional[ResponsesAPIStreamingResponse] = None + self.completed_response: Optional[Any] = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called + self._completed_response_cached = False + self._completed_response_logged = False + self._completed_response_cache_hit: Optional[bool] = None + self._persist_completed_response_before_logging = True self._stream_created_time: float = time.time() # track request context for hooks @@ -101,7 +115,7 @@ class BaseResponsesAPIStreamingIterator: llm_provider=self.custom_llm_provider or "", ) - def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]: + def _process_chunk(self, chunk) -> Optional[Any]: """Process a single chunk of data from the stream""" if not chunk: return None @@ -122,6 +136,10 @@ class BaseResponsesAPIStreamingIterator: # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): + if self.responses_api_provider_config is None: + raise ValueError( + "responses_api_provider_config is required to process live streaming chunks" + ) openai_responses_api_chunk = ( self.responses_api_provider_config.transform_streaming_response( model=self.model, @@ -195,10 +213,11 @@ class BaseResponsesAPIStreamingIterator: if self.litellm_metadata and self.litellm_metadata.get( "encrypted_content_affinity_enabled" ): + openai_types = _get_openai_response_types() event_type = getattr(openai_responses_api_chunk, "type", None) if event_type in ( - ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, - ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): item = getattr(openai_responses_api_chunk, "item", None) if item: @@ -219,10 +238,11 @@ class BaseResponsesAPIStreamingIterator: # Store the completed response (also for incomplete/failed so logging still fires) _chunk_type = getattr(openai_responses_api_chunk, "type", None) + openai_types = _get_openai_response_types() if openai_responses_api_chunk and _chunk_type in ( - ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - ResponsesAPIStreamEvents.RESPONSE_FAILED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, ): self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True @@ -230,11 +250,11 @@ class BaseResponsesAPIStreamingIterator: litellm.include_cost_in_streaming_usage and self.logging_obj is not None ): - response_obj: Optional[ResponsesAPIResponse] = getattr( + response_obj: Optional[Any] = getattr( openai_responses_api_chunk, "response", None ) if response_obj: - usage_obj: Optional[ResponseAPIUsage] = getattr( + usage_obj: Optional[Any] = getattr( response_obj, "usage", None ) if usage_obj is not None: @@ -247,9 +267,13 @@ class BaseResponsesAPIStreamingIterator: if cost is not None: setattr(usage_obj, "cost", cost) except Exception: + # Best-effort usage cost annotation should not break stream replay. pass - if _chunk_type == ResponsesAPIStreamEvents.RESPONSE_FAILED: + if ( + _chunk_type + == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED + ): self._handle_logging_failed_response() else: self._handle_logging_completed_response() @@ -266,6 +290,59 @@ class BaseResponsesAPIStreamingIterator: self._handle_failure(e) raise + def _log_completed_response(self, *, is_async: bool) -> None: + if self._completed_response_logged: + return + self._completed_response_logged = True + + if self._persist_completed_response_before_logging: + self._persist_completed_response_to_cache(is_async=is_async) + + # Create a copy for logging to avoid modifying the response object that will be returned to the user + # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) + # to chat completion format (prompt_tokens/completion_tokens) for internal logging + # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with + # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) + logging_response = self.completed_response + if self.completed_response is not None and hasattr( + self.completed_response, "model_dump" + ): + try: + logging_response = type(self.completed_response).model_validate( + self.completed_response.model_dump() + ) + except Exception: + # Fallback to original if serialization fails + pass + + end_time = datetime.now() + if is_async: + asyncio.create_task( + self.logging_obj.async_success_handler( + result=logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + ) + ) + else: + run_async_function( + async_function=self.logging_obj.async_success_handler, + result=logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + ) + + executor.submit( + self.logging_obj.success_handler, + result=logging_response, + cache_hit=self._completed_response_cache_hit, + start_time=self.start_time, + end_time=end_time, + ) + self._run_post_success_hooks(end_time=end_time) + def _handle_logging_completed_response(self): """Base implementation - should be overridden by subclasses""" pass @@ -296,6 +373,88 @@ class BaseResponsesAPIStreamingIterator: ) self._handle_failure(exception) + def _get_completed_response_object(self) -> Optional[Any]: + openai_types = _get_openai_response_types() + completed_response = self.completed_response + if isinstance(completed_response, openai_types.ResponsesAPIResponse): + return completed_response + + response_obj = getattr(completed_response, "response", None) + if isinstance(response_obj, openai_types.ResponsesAPIResponse): + return response_obj + + return None + + def _persist_completed_response_to_cache(self, *, is_async: bool) -> None: + if self._completed_response_cached: + return + + completed_response = self.completed_response + openai_types = _get_openai_response_types() + if ( + getattr(completed_response, "type", None) + != openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + return + + response_obj = self._get_completed_response_object() + if response_obj is None: + return + + caching_handler = getattr(self.logging_obj, "_llm_caching_handler", None) + if caching_handler is None: + return + + request_kwargs = getattr(caching_handler, "request_kwargs", None) + if ( + not isinstance(request_kwargs, dict) + or request_kwargs.get("stream") is not True + ): + return + request_kwargs = request_kwargs.copy() + preset_cache_key = getattr(caching_handler, "preset_cache_key", None) + request_cache_key = request_kwargs.pop("cache_key", None) + if preset_cache_key is None: + preset_cache_key = request_cache_key + if request_kwargs.get("metadata") is None: + request_kwargs.pop("metadata", None) + request_kwargs.pop("custom_llm_provider", None) + if preset_cache_key is not None: + request_kwargs["cache_key"] = preset_cache_key + + if not caching_handler._should_store_result_in_cache( + original_function=caching_handler.original_function, + kwargs=request_kwargs, + ): + return + + if litellm.cache is None: + return + + cached_response = response_obj.model_dump_json() + if is_async: + cache_write_task = asyncio.create_task( + litellm.cache.async_add_cache( + cached_response, + dynamic_cache_object=getattr(caching_handler, "dual_cache", None), + **request_kwargs, + ) + ) + cache_write_task.add_done_callback( + lambda task: _log_background_task_failure( + task, + task_name="Responses stream cache write", + ) + ) + else: + litellm.cache.add_cache( + cached_response, + dynamic_cache_object=getattr(caching_handler, "dual_cache", None), + **request_kwargs, + ) + + self._completed_response_cached = True + async def _call_post_streaming_deployment_hook(self, chunk): """ Allow callbacks to modify streaming chunks before returning (parity with chat). @@ -480,7 +639,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __aiter__(self): return self - async def __anext__(self) -> ResponsesAPIStreamingResponse: + async def __anext__(self) -> Any: try: self._check_max_streaming_duration() while True: @@ -520,40 +679,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in async context""" - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr( - self.completed_response, "model_dump" - ): - try: - logging_response = type(self.completed_response).model_validate( - self.completed_response.model_dump() - ) - except Exception: - # Fallback to original if serialization fails - pass - - asyncio.create_task( - self.logging_obj.async_success_handler( - result=logging_response, - start_time=self.start_time, - end_time=datetime.now(), - cache_hit=None, - ) - ) - - executor.submit( - self.logging_obj.success_handler, - result=logging_response, - cache_hit=None, - start_time=self.start_time, - end_time=datetime.now(), - ) - self._run_post_success_hooks(end_time=datetime.now()) + self._log_completed_response(is_async=True) class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -627,39 +753,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _handle_logging_completed_response(self): """Handle logging for completed responses in sync context""" - # Create a copy for logging to avoid modifying the response object that will be returned to the user - # The logging handlers may transform usage from Responses API format (input_tokens/output_tokens) - # to chat completion format (prompt_tokens/completion_tokens) for internal logging - # Use model_dump + model_validate instead of deepcopy to avoid pickle errors with - # Pydantic ValidatorIterator when response contains tool_choice with allowed_tools (fixes #17192) - logging_response = self.completed_response - if self.completed_response is not None and hasattr( - self.completed_response, "model_dump" - ): - try: - logging_response = type(self.completed_response).model_validate( - self.completed_response.model_dump() - ) - except Exception: - # Fallback to original if serialization fails - pass - - run_async_function( - async_function=self.logging_obj.async_success_handler, - result=logging_response, - start_time=self.start_time, - end_time=datetime.now(), - cache_hit=None, - ) - - executor.submit( - self.logging_obj.success_handler, - result=logging_response, - cache_hit=None, - start_time=self.start_time, - end_time=datetime.now(), - ) - self._run_post_success_hooks(end_time=datetime.now()) + self._log_completed_response(is_async=False) class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): @@ -683,90 +777,441 @@ class MockResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): request_data: Optional[Dict[str, Any]] = None, call_type: Optional[str] = None, ): - super().__init__( - response=response, + transformed = responses_api_provider_config.transform_response_api_response( model=model, - responses_api_provider_config=responses_api_provider_config, + raw_response=response, + logging_obj=logging_obj, + ) + super().__init__( + response=httpx.Response(200), + model=model, + responses_api_provider_config=None, logging_obj=logging_obj, litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, request_data=request_data, call_type=call_type, ) + self._set_events_from_response(transformed=transformed, logging_obj=logging_obj) - # one-time transform - transformed = ( - self.responses_api_provider_config.transform_response_api_response( - model=self.model, - raw_response=response, - logging_obj=logging_obj, - ) + def _set_events_from_response( + self, + transformed: Any, + logging_obj: LiteLLMLoggingObj, + ) -> None: + self._events = _build_synthetic_response_events( + transformed=transformed, + logging_obj=logging_obj, + chunk_size=self.CHUNK_SIZE, ) - full_text = self._collect_text(transformed) - - # build a list of 5‑char delta events - deltas = [ - OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - delta=full_text[i : i + self.CHUNK_SIZE], - item_id=transformed.id, - output_index=0, - content_index=0, - ) - for i in range(0, len(full_text), self.CHUNK_SIZE) - ] - - # Add cost to usage object if include_cost_in_streaming_usage is True - if litellm.include_cost_in_streaming_usage and logging_obj is not None: - usage_obj: Optional[ResponseAPIUsage] = getattr(transformed, "usage", None) - if usage_obj is not None: - try: - cost: Optional[float] = logging_obj._response_cost_calculator( - result=transformed - ) - if cost is not None: - setattr(usage_obj, "cost", cost) - except Exception: - # If cost calculation fails, continue without cost - pass - - # append the completed event - self._events = deltas + [ - ResponseCompletedEvent( - type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - response=transformed, - ) - ] self._idx = 0 + self.completed_response = self._events[-1] def __aiter__(self): return self - async def __anext__(self) -> ResponsesAPIStreamingResponse: + async def __anext__(self) -> Any: if self._idx >= len(self._events): raise StopAsyncIteration evt = self._events[self._idx] self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=True) return evt def __iter__(self): return self - def __next__(self) -> ResponsesAPIStreamingResponse: + def __next__(self) -> Any: if self._idx >= len(self._events): raise StopIteration evt = self._events[self._idx] self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=False) return evt - def _collect_text(self, resp: ResponsesAPIResponse) -> str: - out = "" - for out_item in resp.output: - item_type = getattr(out_item, "type", None) - if item_type == "message": - for c in getattr(out_item, "content", []): - out += c.text - return out + +class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): + def __init__( + self, + response: Any, + logging_obj: LiteLLMLoggingObj, + request_data: Optional[Dict[str, Any]] = None, + call_type: Optional[str] = None, + ): + BaseResponsesAPIStreamingIterator.__init__( + self, + response=httpx.Response(200), + model=getattr(response, "model", ""), + responses_api_provider_config=None, + logging_obj=logging_obj, + litellm_metadata=None, + custom_llm_provider="cached_response", + request_data=request_data, + call_type=call_type, + ) + self._completed_response_cache_hit = True + self._persist_completed_response_before_logging = False + self._events: List[Any] = [] + self._idx = 0 + self._set_events_from_response(transformed=response, logging_obj=logging_obj) + + def _set_events_from_response( + self, + transformed: Any, + logging_obj: LiteLLMLoggingObj, + ) -> None: + self._events = _build_synthetic_response_events( + transformed=transformed, + logging_obj=logging_obj, + chunk_size=MockResponsesAPIStreamingIterator.CHUNK_SIZE, + ) + self._idx = 0 + self.completed_response = self._events[-1] + + def __aiter__(self): + return self + + async def __anext__(self) -> Any: + if self._idx >= len(self._events): + raise StopAsyncIteration + evt = self._events[self._idx] + self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=True) + return evt + + def __iter__(self): + return self + + def __next__(self) -> Any: + if self._idx >= len(self._events): + raise StopIteration + evt = self._events[self._idx] + self._idx += 1 + openai_types = _get_openai_response_types() + if ( + getattr(evt, "type", None) + == openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ): + self.completed_response = evt + self._log_completed_response(is_async=False) + return evt + + +def _dump_response_object(obj: Any) -> Dict[str, Any]: + if hasattr(obj, "model_dump"): + return obj.model_dump() + if isinstance(obj, dict): + return obj + return {} + + +def _build_response_status_event( + event_type: Literal[ + "response.created", + "response.in_progress", + ], + transformed: Any, +) -> Any: + openai_types = _get_openai_response_types() + in_progress_response = transformed.model_copy( + deep=True, + update={"status": "in_progress", "output": []}, + ) + if event_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED: + return openai_types.ResponseCreatedEvent( + type=event_type, response=in_progress_response + ) + return openai_types.ResponseInProgressEvent( + type=event_type, response=in_progress_response + ) + + +def _build_content_part_done_event( + *, + item_id: str, + output_index: int, + content_index: int, + part_payload: Dict[str, Any], +) -> Optional[Any]: + openai_types = _get_openai_response_types() + part_type = part_payload.get("type") + part: Any + if part_type == "output_text": + annotations = [ + openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) + for annotation in part_payload.get("annotations", []) or [] + ] + part = openai_types.ContentPartDonePartOutputText( + type="output_text", + text=str(part_payload.get("text") or ""), + annotations=annotations, + logprobs=part_payload.get("logprobs"), + ) + elif part_type == "refusal": + part = openai_types.ContentPartDonePartRefusal( + type="refusal", + refusal=str(part_payload.get("refusal") or ""), + ) + elif part_type == "reasoning_text": + part = openai_types.ContentPartDonePartReasoningText( + type="reasoning_text", + reasoning=str(part_payload.get("reasoning") or ""), + ) + else: + return None + + return openai_types.ContentPartDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_DONE, + item_id=item_id, + output_index=output_index, + content_index=content_index, + part=part, + ) + + +def _add_text_like_part_events( + *, + events: List[Any], + item_id: str, + output_index: int, + content_index: int, + part_payload: Dict[str, Any], + chunk_size: int, +) -> None: + openai_types = _get_openai_response_types() + part_type = part_payload.get("type") + if part_type == "output_text": + text = str(part_payload.get("text") or "") + for i in range(0, len(text), chunk_size): + events.append( + openai_types.OutputTextDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id=item_id, + output_index=output_index, + content_index=content_index, + delta=text[i : i + chunk_size], + ) + ) + for annotation_index, annotation in enumerate( + part_payload.get("annotations", []) or [] + ): + events.append( + openai_types.OutputTextAnnotationAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, + item_id=item_id, + output_index=output_index, + content_index=content_index, + annotation_index=annotation_index, + annotation=annotation, + ) + ) + events.append( + openai_types.OutputTextDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, + item_id=item_id, + output_index=output_index, + content_index=content_index, + text=text, + ) + ) + elif part_type == "refusal": + refusal = str(part_payload.get("refusal") or "") + for i in range(0, len(refusal), chunk_size): + events.append( + openai_types.RefusalDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DELTA, + item_id=item_id, + output_index=output_index, + content_index=content_index, + delta=refusal[i : i + chunk_size], + ) + ) + events.append( + openai_types.RefusalDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REFUSAL_DONE, + item_id=item_id, + output_index=output_index, + content_index=content_index, + refusal=refusal, + ) + ) + + +def _build_synthetic_response_events( + *, + transformed: Any, + logging_obj: LiteLLMLoggingObj, + chunk_size: int, +) -> List[Any]: + openai_types = _get_openai_response_types() + if litellm.include_cost_in_streaming_usage and logging_obj is not None: + usage_obj: Optional[Any] = getattr(transformed, "usage", None) + if usage_obj is not None: + try: + cost: Optional[float] = logging_obj._response_cost_calculator( + result=transformed + ) + if cost is not None: + setattr(usage_obj, "cost", cost) + except Exception: + pass + + events: List[Any] = [ + _build_response_status_event( + openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed + ), + _build_response_status_event( + openai_types.ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS, transformed + ), + ] + + sequence_number = 0 + for output_index, output_item in enumerate( + getattr(transformed, "output", []) or [] + ): + output_item_payload = _dump_response_object(output_item) + item_id = str(output_item_payload.get("id") or transformed.id) + item_type = output_item_payload.get("type") + + events.append( + openai_types.OutputItemAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=output_index, + item=openai_types.BaseLiteLLMOpenAIResponseObject( + **output_item_payload + ), + ) + ) + + if item_type == "message": + for content_index, part in enumerate( + output_item_payload.get("content", []) or [] + ): + part_payload = _dump_response_object(part) + events.append( + openai_types.ContentPartAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.CONTENT_PART_ADDED, + item_id=item_id, + output_index=output_index, + content_index=content_index, + part=openai_types.BaseLiteLLMOpenAIResponseObject( + **part_payload + ), + ) + ) + _add_text_like_part_events( + events=events, + item_id=item_id, + output_index=output_index, + content_index=content_index, + part_payload=part_payload, + chunk_size=chunk_size, + ) + done_event = _build_content_part_done_event( + item_id=item_id, + output_index=output_index, + content_index=content_index, + part_payload=part_payload, + ) + if done_event is not None: + events.append(done_event) + elif item_type == "function_call": + arguments = str(output_item_payload.get("arguments") or "") + for i in range(0, len(arguments), chunk_size): + events.append( + openai_types.FunctionCallArgumentsDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + item_id=item_id, + output_index=output_index, + delta=arguments[i : i + chunk_size], + ) + ) + events.append( + openai_types.FunctionCallArgumentsDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DONE, + item_id=item_id, + output_index=output_index, + arguments=arguments, + ) + ) + elif item_type == "reasoning": + for summary_index, summary in enumerate( + output_item_payload.get("summary", []) or [] + ): + summary_payload = _dump_response_object(summary) + summary_text = str(summary_payload.get("text") or "") + for i in range(0, len(summary_text), chunk_size): + events.append( + openai_types.ReasoningSummaryTextDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id=item_id, + output_index=output_index, + summary_index=summary_index, + delta=summary_text[i : i + chunk_size], + ) + ) + sequence_number += 1 + events.append( + openai_types.ReasoningSummaryTextDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, + item_id=item_id, + output_index=output_index, + sequence_number=sequence_number, + summary_index=summary_index, + text=summary_text, + ) + ) + sequence_number += 1 + events.append( + openai_types.ReasoningSummaryPartDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, + item_id=item_id, + output_index=output_index, + sequence_number=sequence_number, + summary_index=summary_index, + part=openai_types.BaseLiteLLMOpenAIResponseObject( + **summary_payload + ), + ) + ) + + sequence_number += 1 + events.append( + openai_types.OutputItemDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=output_index, + sequence_number=sequence_number, + item=openai_types.BaseLiteLLMOpenAIResponseObject( + **output_item_payload + ), + ) + ) + + events.append( + openai_types.ResponseCompletedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=transformed, + ) + ) + return events # --------------------------------------------------------------------------- @@ -951,8 +1396,8 @@ class ResponsesWebSocketStreaming: # --------------------------------------------------------------------------- _RESPONSE_CREATE_PARAMS: frozenset = ( - ResponsesAPIRequestParams.__required_keys__ - | ResponsesAPIRequestParams.__optional_keys__ + _get_openai_response_types().ResponsesAPIRequestParams.__required_keys__ + | _get_openai_response_types().ResponsesAPIRequestParams.__optional_keys__ ) _MANAGED_WS_SKIP_KWARGS: frozenset = frozenset( @@ -1085,7 +1530,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: Dict[str, Any] + completed_event: Dict[str, Any], ) -> List[Dict[str, Any]]: """ Convert the output items in a ``response.completed`` event into diff --git a/litellm/router.py b/litellm/router.py index 7448cdd1b4..50fd7eaed0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5261,11 +5261,34 @@ class Router: """ Initialize the Containers API endpoints on the router. - Container operations don't need model-based routing, so we call the - original function directly with the custom_llm_provider. + LiteLLM-managed container IDs (``cntr_...``) encode ``model_id`` and provider + metadata. When present, decode the ID, replace ``container_id`` with the + upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so + deployment credentials (e.g. regional ``api_base`` for Azure) match + :meth:`_init_responses_api_endpoints`. Otherwise call the handler directly. """ if custom_llm_provider and "custom_llm_provider" not in kwargs: kwargs["custom_llm_provider"] = custom_llm_provider + + from litellm.responses.utils import ResponsesAPIRequestUtils + + container_id = kwargs.get("container_id") + if isinstance(container_id, str): + decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) + original_id = decoded.get("response_id", container_id) + if original_id != container_id: + kwargs["container_id"] = original_id + decoded_provider = decoded.get("custom_llm_provider") + if decoded_provider and kwargs.get("custom_llm_provider") == "openai": + kwargs["custom_llm_provider"] = decoded_provider + model_id = decoded.get("model_id") + if model_id: + kwargs["model"] = model_id + return await self._ageneric_api_call_with_fallbacks( + original_function=original_function, + **kwargs, + ) + return await original_function(**kwargs) async def _init_responses_api_endpoints( diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index 0163f3bbd4..07143af38a 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -106,7 +106,8 @@ def _match_deployment( # check either didn't run (no request tags) or failed (step 1 returned # None). Block the regex path so it cannot circumvent the operator's # strict-tag policy. - strict_tag_check_failed = not match_any and bool(deployment_tags) + deployment_has_plain_tags = deployment_tags is not None and len(deployment_tags) > 0 + strict_tag_check_failed = not match_any and deployment_has_plain_tags if deployment_tag_regex and header_strings and not strict_tag_check_failed: regex_match = _is_valid_deployment_tag_regex( deployment_tag_regex, header_strings diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index bef42e2384..f6da26ccd7 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -23,21 +23,28 @@ def add_model_file_id_mappings( healthy_deployments: Union[List[Dict], Dict], responses: List["OpenAIFileObject"] ) -> dict: """ - Create a mapping of model name to file id + Create a mapping of model id to file id { "model_id": "file_id", "model_id": "file_id", } + + `healthy_deployments` may be either a list of deployment dicts (multiple + matched deployments) or a single deployment dict (when the router resolved + a specific deployment, e.g. because the requested model matched a + `model_info.id`). Both shapes must be handled by extracting + `model_info.id` from each deployment. """ - model_file_id_mapping = {} - if isinstance(healthy_deployments, list): - for deployment, response in zip(healthy_deployments, responses): - model_file_id_mapping[deployment.get("model_info", {}).get("id")] = ( - response.id - ) - elif isinstance(healthy_deployments, dict): - for model_id, file_id in healthy_deployments.items(): - model_file_id_mapping[model_id] = file_id + model_file_id_mapping: Dict[str, str] = {} + deployments_list: List[Dict] = ( + healthy_deployments + if isinstance(healthy_deployments, list) + else [healthy_deployments] + ) + for deployment, response in zip(deployments_list, responses): + model_id = deployment.get("model_info", {}).get("id") + if model_id is not None: + model_file_id_mapping[model_id] = response.id return model_file_id_mapping diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 2fd0c4ea97..986ec39f3b 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1482,6 +1482,7 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject): type: Literal[ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA] item_id: str output_index: int + summary_index: int = 0 delta: str @@ -1490,7 +1491,7 @@ class ReasoningSummaryTextDoneEvent(BaseLiteLLMOpenAIResponseObject): item_id: str output_index: int sequence_number: int - summary_index: int + summary_index: int = 0 text: str @@ -1499,7 +1500,7 @@ class ReasoningSummaryPartDoneEvent(BaseLiteLLMOpenAIResponseObject): item_id: str output_index: int sequence_number: int - summary_index: int + summary_index: int = 0 part: BaseLiteLLMOpenAIResponseObject diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 2e7d57cef2..87bf11a902 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -6,6 +6,13 @@ from typing_extensions import ( TypedDict, ) +from litellm.types.llms.openai import EmbeddingInput + +# Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit +# opt-in for combined embeddings — a provider-specific extension of the +# OpenAI-faithful EmbeddingInput shape. +GeminiEmbeddingInput = Union[EmbeddingInput, List[List[str]]] + class FunctionResponse(TypedDict): name: str diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bbe13442d6..8391fdb48f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -33391,6 +33391,72 @@ "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas", "supports_reasoning": true }, + "vertex_ai/xai/grok-4.1-fast-non-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "vertex_ai/xai/grok-4.1-fast-reasoning": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "vertex_ai/xai/grok-4.20-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, + "vertex_ai/xai/grok-4.20-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 2000000, + "max_output_tokens": 2000000, + "max_tokens": 2000000, + "mode": "chat", + "output_cost_per_token": 6e-06, + "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true + }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -34828,6 +34894,20 @@ "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "zai.glm-5": { + "input_cost_per_token": 1e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "source": "https://aws.amazon.com/bedrock/pricing/" + }, "zai.glm-4.7-flash": { "input_cost_per_token": 7e-08, "litellm_provider": "bedrock_converse", diff --git a/pyproject.toml b/pyproject.toml index 657632d69e..65b9bd2c98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ proxy = [ "azure-identity==1.25.2", "azure-storage-blob==12.28.0", "mcp==1.26.0", - "litellm-proxy-extras==0.4.69", + "litellm-proxy-extras==0.4.70", "litellm-enterprise==0.1.39", "RestrictedPython==8.1", "rich==13.9.4", @@ -149,6 +149,8 @@ dev = [ "parameterized==0.9.0", "openapi-core==0.22.0; python_version < '3.14'", "pytest-timeout==2.4.0", + "vcrpy==8.1.1", + "pytest-recording==0.13.4", ] proxy-dev = [ "prisma==0.11.0", diff --git a/tests/_flush_vcr_cache.py b/tests/_flush_vcr_cache.py new file mode 100644 index 0000000000..d236c88fa3 --- /dev/null +++ b/tests/_flush_vcr_cache.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import os +import sys + +import redis + +from tests._vcr_redis_persister import CASSETTE_REDIS_URL_ENV, _redis_url_from_env + +PREFIX = "litellm:vcr:cassette:" +SCAN_BATCH = 500 + + +def _client() -> redis.Redis: + url = _redis_url_from_env() + if not url: + sys.exit(f"Set {CASSETTE_REDIS_URL_ENV} to flush the VCR cache") + return redis.Redis.from_url( + url, + socket_timeout=5, + socket_connect_timeout=5, + decode_responses=False, + ) + + +def main() -> None: + client = _client() + deleted = 0 + pipeline = client.pipeline(transaction=False) + pending = 0 + for key in client.scan_iter(match=f"{PREFIX}*", count=SCAN_BATCH): + pipeline.delete(key) + pending += 1 + if pending >= SCAN_BATCH: + deleted += sum(pipeline.execute()) + pipeline = client.pipeline(transaction=False) + pending = 0 + if pending: + deleted += sum(pipeline.execute()) + print(f"Deleted {deleted} VCR cassette key(s) under {PREFIX!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/_vcr_redis_persister.py b/tests/_vcr_redis_persister.py new file mode 100644 index 0000000000..4d72a1142b --- /dev/null +++ b/tests/_vcr_redis_persister.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import logging +import os +from typing import Any, Optional + +from vcr.persisters.filesystem import CassetteNotFoundError +from vcr.serialize import deserialize, serialize + +CASSETTE_TTL_SECONDS = 24 * 60 * 60 +REDIS_KEY_PREFIX = "litellm:vcr:cassette:" +CASSETTE_REDIS_URL_ENV = "CASSETTE_REDIS_URL" +VCR_VERBOSE_ENV = "LITELLM_VCR_VERBOSE" +MAX_EPISODES_PER_CASSETTE = 50 + +_log = logging.getLogger(__name__) +_passed_by_cassette_key: dict[str, bool] = {} + + +def mark_test_outcome_for_cassette(cassette_path: str, passed: bool) -> None: + _passed_by_cassette_key[redis_key_for(cassette_path)] = passed + + +def redis_key_for(cassette_path: str) -> str: + rel = os.path.relpath(str(cassette_path)) + if rel.endswith(".yaml"): + rel = rel[: -len(".yaml")] + rel = rel.replace("/cassettes/", "/").lstrip("./") + return f"{REDIS_KEY_PREFIX}{rel}" + + +def _redis_url_from_env() -> Optional[str]: + return os.environ.get(CASSETTE_REDIS_URL_ENV) or None + + +def _build_default_client(): + import redis + from redis.backoff import ExponentialBackoff + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + from redis.retry import Retry + + url = _redis_url_from_env() + if not url: + raise RuntimeError( + f"Set {CASSETTE_REDIS_URL_ENV} to enable the VCR persister. " + "Cassette Redis is intentionally separate from the application " + "Redis (REDIS_URL/REDIS_HOST) to avoid being flushed by tests." + ) + return redis.Redis.from_url( + url, + socket_timeout=5, + socket_connect_timeout=5, + decode_responses=False, + retry=Retry(ExponentialBackoff(cap=2, base=0.1), retries=2), + retry_on_error=[RedisConnectionError, RedisTimeoutError], + ) + + +def make_redis_persister( + client: Optional[Any] = None, + ttl_seconds: int = CASSETTE_TTL_SECONDS, +): + redis_client = client if client is not None else _build_default_client() + + try: + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + _transient_errors: tuple = (RedisConnectionError, RedisTimeoutError) + except ImportError: # pragma: no cover - redis is a hard test dep + _transient_errors = () + + class _RedisPersister: + @staticmethod + def load_cassette(cassette_path, serializer): + try: + data = redis_client.get(redis_key_for(cassette_path)) + except _transient_errors as exc: + _log.warning( + "VCR redis load failed for %s; treating as cache miss: %s", + cassette_path, + exc, + ) + raise CassetteNotFoundError() from exc + if data is None: + raise CassetteNotFoundError() + if isinstance(data, bytes): + data = data.decode("utf-8") + return deserialize(data, serializer) + + @staticmethod + def save_cassette(cassette_path, cassette_dict, serializer): + key = redis_key_for(cassette_path) + passed = _passed_by_cassette_key.pop(key, True) + episode_count = len(cassette_dict.get("requests", []) or []) + if episode_count > MAX_EPISODES_PER_CASSETTE: + _log.warning( + "VCR redis save refused for %s; cassette has %d episodes " + "(> MAX_EPISODES_PER_CASSETTE=%d). The test likely produces " + "non-deterministic request bodies (e.g. uuid) and is " + "appending instead of replaying. Opt it out with the " + "no-vcr list in conftest, or stabilize its request body.", + cassette_path, + episode_count, + MAX_EPISODES_PER_CASSETTE, + ) + return + if not passed: + _log.info( + "VCR redis save skipped for %s; test did not pass — " + "leaving any prior cassette intact", + cassette_path, + ) + return + data = serialize(cassette_dict, serializer) + payload = data.encode("utf-8") if isinstance(data, str) else data + try: + redis_client.set(key, payload, ex=ttl_seconds) + except _transient_errors as exc: + _log.warning( + "VCR redis save failed for %s; cassette not persisted: %s", + cassette_path, + exc, + ) + + return _RedisPersister + + +def filter_non_2xx_response(response): + if not isinstance(response, dict): + return response + status = response.get("status") + code = status.get("code") if isinstance(status, dict) else status + if not isinstance(code, int): + return response + return response if 200 <= code < 300 else None + + +_PATCHED_AIOHTTP_RECORD = False + + +def patch_vcrpy_aiohttp_record_path() -> None: + """Re-feed the response body into aiohttp's StreamReader after vcrpy's + record_response drains it, so downstream consumers (e.g. + LiteLLMAiohttpTransport.AiohttpResponseStream) can still read it.""" + global _PATCHED_AIOHTTP_RECORD + if _PATCHED_AIOHTTP_RECORD: + return + import vcr.stubs.aiohttp_stubs as _aiohttp_stubs + + _orig_record_response = _aiohttp_stubs.record_response + + async def _record_response_preserving_body(cassette, vcr_request, response): + await _orig_record_response(cassette, vcr_request, response) + body = getattr(response, "_body", None) or b"" + if body: + response.content.unread_data(body) + + _aiohttp_stubs.record_response = _record_response_preserving_body + _PATCHED_AIOHTTP_RECORD = True + + +def vcr_verbose_enabled() -> bool: + return os.environ.get(VCR_VERBOSE_ENV) == "1" + + +def format_vcr_verdict(cassette: Any) -> str: + if cassette is None: + return "[VCR NOOP]" + played = getattr(cassette, "play_count", 0) or 0 + dirty = getattr(cassette, "dirty", False) + total = len(cassette) if hasattr(cassette, "__len__") else 0 + if played == 0 and not dirty: + return "[VCR NOOP] (no http traffic)" + if played > 0 and not dirty: + return f"[VCR HIT] {played} replayed, 0 new ({total} cassette entries)" + if played == 0 and dirty: + return f"[VCR MISS] 0 replayed, recorded new ({total} cassette entries)" + return ( + f"[VCR PARTIAL] {played} replayed + new recordings ({total} cassette entries)" + ) diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 4ecc0d0bc9..344e38da83 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -169,3 +169,4 @@ langchain-mcp-adapters: >=0.2.1 # MIT License langgraph: >=1.0.10 # MIT License langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE pytest-rerunfailures: >=15.1 # MPL 2.0 license +pytest-recording: >=0.13.4 # MIT license diff --git a/tests/llm_responses_api_testing/conftest.py b/tests/llm_responses_api_testing/conftest.py index 0b03348190..80f36e159a 100644 --- a/tests/llm_responses_api_testing/conftest.py +++ b/tests/llm_responses_api_testing/conftest.py @@ -1,5 +1,6 @@ # conftest.py +import asyncio import importlib import os import sys @@ -9,9 +10,161 @@ import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import litellm -import asyncio +import litellm # noqa: E402 + +from tests._vcr_redis_persister import ( # noqa: E402 + filter_non_2xx_response, + format_vcr_verdict, + make_redis_persister, + mark_test_outcome_for_cassette, + patch_vcrpy_aiohttp_record_path, + vcr_verbose_enabled, +) + + +_controller_pluginmanager = None +_controller_terminal_reporter = None + + +_FILTERED_REQUEST_HEADERS = ( + "authorization", + "x-api-key", + "anthropic-api-key", + "anthropic-version", + "openai-api-key", + "azure-api-key", + "api-key", + "cookie", + "x-amz-security-token", + "x-amz-date", + "x-amz-content-sha256", + "amz-sdk-invocation-id", + "amz-sdk-request", + "x-goog-api-key", + "x-goog-user-project", +) + +_FILTERED_RESPONSE_HEADERS = ( + "set-cookie", + "x-request-id", + "request-id", + "cf-ray", + "anthropic-organization-id", + "openai-organization", + "x-amzn-requestid", + "x-amzn-trace-id", + "date", +) + + +def _scrub_response(response): + if not isinstance(response, dict): + return response + headers = response.get("headers") or {} + if isinstance(headers, dict): + for header in list(headers): + if header.lower() in _FILTERED_RESPONSE_HEADERS: + headers.pop(header, None) + return response + + +def _before_record_response(response): + return filter_non_2xx_response(_scrub_response(response)) + + +@pytest.fixture(scope="module") +def vcr_config(): + return { + "filter_headers": list(_FILTERED_REQUEST_HEADERS), + "decode_compressed_response": True, + "record_mode": "new_episodes", + "allow_playback_repeats": True, + "match_on": ( + "method", + "scheme", + "host", + "port", + "path", + "query", + "body", + ), + "before_record_response": _before_record_response, + } + + +def _vcr_disabled() -> bool: + if os.environ.get("LITELLM_VCR_DISABLE") == "1": + return True + return not os.environ.get("CASSETTE_REDIS_URL") + + +def pytest_recording_configure(config, vcr): + if _vcr_disabled(): + return + vcr.register_persister(make_redis_persister()) + patch_vcrpy_aiohttp_record_path() + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + outcome = yield + rep = outcome.get_result() + setattr(item, f"rep_{rep.when}", rep) + + +@pytest.fixture(autouse=True) +def _vcr_outcome_gate(request, vcr): + yield + cassette = vcr + rep_call = getattr(request.node, "rep_call", None) + test_passed = bool(rep_call and rep_call.passed) + cassette_path = getattr(cassette, "_path", None) if cassette is not None else None + if cassette_path: + mark_test_outcome_for_cassette(cassette_path, test_passed) + + if not vcr_verbose_enabled(): + return + verdict = format_vcr_verdict(cassette) + request.node.user_properties.append(("vcr_verdict", verdict)) + + +def pytest_configure(config): + global _controller_pluginmanager + if os.environ.get("PYTEST_XDIST_WORKER"): + return + _controller_pluginmanager = config.pluginmanager + + +def _resolve_terminal_reporter(): + global _controller_terminal_reporter + if _controller_terminal_reporter is not None: + return _controller_terminal_reporter + if _controller_pluginmanager is None: + return None + _controller_terminal_reporter = _controller_pluginmanager.getplugin( + "terminalreporter" + ) + return _controller_terminal_reporter + + +def pytest_runtest_logreport(report): + if report.when != "teardown": + return + if os.environ.get("PYTEST_XDIST_WORKER"): + return + if not vcr_verbose_enabled(): + return + reporter = _resolve_terminal_reporter() + if reporter is None: + return + verdict = next( + (v for k, v in (report.user_properties or []) if k == "vcr_verdict"), + None, + ) + if not verdict: + return + reporter.write_line(f"{verdict} :: {report.nodeid}") @pytest.fixture(scope="session") @@ -61,15 +214,18 @@ def setup_and_teardown(): def pytest_collection_modifyitems(config, items): - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + if not _vcr_disabled(): + for item in items: + if item.get_closest_marker("vcr") is not None: + continue + item.add_marker(pytest.mark.vcr) + custom_logger_tests = [ item for item in items if "custom_logger" in item.parent.name ] other_tests = [item for item in items if "custom_logger" not in item.parent.name] - # Sort tests based on their names custom_logger_tests.sort(key=lambda x: x.name) other_tests.sort(key=lambda x: x.name) - # Reorder the items list items[:] = custom_logger_tests + other_tests diff --git a/tests/llm_responses_api_testing/test_responses_hooks.py b/tests/llm_responses_api_testing/test_responses_hooks.py index 3227fecdfb..3799a0b912 100644 --- a/tests/llm_responses_api_testing/test_responses_hooks.py +++ b/tests/llm_responses_api_testing/test_responses_hooks.py @@ -1,6 +1,9 @@ import asyncio +from contextlib import suppress from datetime import datetime +import json from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -8,8 +11,17 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.responses import streaming_iterator as streaming_module -from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator -from litellm.types.llms.openai import ResponsesAPIStreamEvents +from litellm.responses.streaming_iterator import ( + CachedResponsesAPIStreamingIterator, + MockResponsesAPIStreamingIterator, + ResponsesAPIStreamingIterator, + SyncResponsesAPIStreamingIterator, +) +from litellm.types.llms.openai import ( + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) from litellm.types.utils import CallTypes @@ -19,15 +31,19 @@ class _FakeLoggingObj: self.async_success_calls = 0 self.failure_calls = 0 self.async_failure_calls = 0 + self.last_success_kwargs = None + self.last_async_success_kwargs = None self.start_time = datetime.now() self.model_call_details = {"litellm_params": {}} # Signature alignment with Logging handlers def success_handler(self, *args, **kwargs): self.success_calls += 1 + self.last_success_kwargs = kwargs async def async_success_handler(self, *args, **kwargs): self.async_success_calls += 1 + self.last_async_success_kwargs = kwargs def failure_handler(self, *args, **kwargs): self.failure_calls += 1 @@ -36,6 +52,115 @@ class _FakeLoggingObj: self.async_failure_calls += 1 +def _make_completed_response(response_id: str = "resp_test") -> ResponseCompletedEvent: + return ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id=response_id, + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[ + { + "type": "message", + "id": f"msg_{response_id}", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "cached streamed response", + "annotations": [], + } + ], + } + ], + ), + ) + + +@pytest.mark.asyncio +async def test_log_background_task_failure_logs_task_exceptions(monkeypatch): + error_logger = MagicMock() + monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger) + + async def _boom(): + raise RuntimeError("boom") + + task = asyncio.create_task(_boom()) + with suppress(RuntimeError): + await task + + streaming_module._log_background_task_failure(task, task_name="cache write") + + error_logger.assert_called_once() + assert error_logger.call_args.args == ( + "%s failed: %s", + "cache write", + task.exception(), + ) + + +@pytest.mark.asyncio +async def test_log_background_task_failure_ignores_cancelled_tasks(monkeypatch): + error_logger = MagicMock() + monkeypatch.setattr(streaming_module.verbose_logger, "error", error_logger) + + task = asyncio.create_task(asyncio.sleep(1)) + task.cancel() + with suppress(asyncio.CancelledError): + await task + + streaming_module._log_background_task_failure(task, task_name="cache write") + + error_logger.assert_not_called() + + +def test_content_part_done_event_supports_refusal_and_reasoning_text(): + refusal_event = streaming_module._build_content_part_done_event( + item_id="msg_1", + output_index=0, + content_index=0, + part_payload={"type": "refusal", "refusal": "no"}, + ) + reasoning_event = streaming_module._build_content_part_done_event( + item_id="msg_1", + output_index=0, + content_index=1, + part_payload={"type": "reasoning_text", "reasoning": "because"}, + ) + unsupported_event = streaming_module._build_content_part_done_event( + item_id="msg_1", + output_index=0, + content_index=2, + part_payload={"type": "image"}, + ) + + assert refusal_event.part.type == "refusal" + assert refusal_event.part.refusal == "no" + assert reasoning_event.part.type == "reasoning_text" + assert reasoning_event.part.reasoning == "because" + assert unsupported_event is None + + +def test_dump_response_object_handles_model_and_unknown_values(): + response = ResponsesAPIResponse( + id="resp_dump", + created_at=int(datetime.now().timestamp()), + status="completed", + model="gpt-4.1-mini", + object="response", + output=[], + ) + + assert streaming_module._dump_response_object(response)["id"] == "resp_dump" + assert streaming_module._dump_response_object({"type": "message"}) == { + "type": "message" + } + assert streaming_module._dump_response_object(object()) == {} + + @pytest.mark.asyncio async def test_responses_streaming_triggers_hooks(monkeypatch): """ @@ -167,3 +292,768 @@ async def test_responses_streaming_failure_triggers_failure_handlers(): await asyncio.sleep(0.2) assert logging_obj.failure_calls >= 1 assert logging_obj.async_failure_calls >= 1 + + +def test_process_chunk_requires_provider_config(): + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=None, + logging_obj=_FakeLoggingObj(), + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + with pytest.raises(ValueError, match="responses_api_provider_config is required"): + iterator._process_chunk(json.dumps({"type": "response.completed"})) + + +def test_process_chunk_wraps_encrypted_content_with_model_id(): + openai_types = streaming_module._get_openai_response_types() + + class _EncryptedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.OutputItemAddedEvent( + type=openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + output_index=0, + item=openai_types.BaseLiteLLMOpenAIResponseObject( + id="rs_123", + type="reasoning", + encrypted_content="ciphertext", + ), + ) + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_EncryptedConfig(), + logging_obj=_FakeLoggingObj(), + litellm_metadata={ + "encrypted_content_affinity_enabled": True, + "model_info": {"id": "model-123"}, + }, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + event = iterator._process_chunk(json.dumps({"type": "response.output_item.added"})) + + assert event.item.encrypted_content.startswith("litellm_enc:") + assert event.item.encrypted_content.endswith(";ciphertext") + + +def test_process_chunk_completed_response_updates_id_and_usage_cost(monkeypatch): + original_include_cost = litellm.include_cost_in_streaming_usage + litellm.include_cost_in_streaming_usage = True + openai_types = streaming_module._get_openai_response_types() + + class _CompletedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.ResponseCompletedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_live", + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[], + usage=openai_types.ResponseAPIUsage( + input_tokens=1, + output_tokens=2, + total_tokens=3, + ), + ), + ) + + logging_obj = _FakeLoggingObj() + logging_obj._response_cost_calculator = MagicMock(return_value=1.23) + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_CompletedConfig(), + logging_obj=logging_obj, + litellm_metadata={"model_info": {"id": "model-123"}}, + custom_llm_provider="openai", + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + completion_handler = MagicMock() + monkeypatch.setattr( + iterator, "_handle_logging_completed_response", completion_handler + ) + + try: + # Chunk must include a top-level "response" key so BaseResponsesAPIStreamingIterator + # runs _update_responses_api_response_id_with_model_id (see streaming_iterator.py). + event = iterator._process_chunk( + json.dumps( + {"type": "response.completed", "response": {"id": "resp_live"}} + ) + ) + finally: + litellm.include_cost_in_streaming_usage = original_include_cost + + assert iterator.completed_response is event + assert event.response.id != "resp_live" + assert event.response.id.startswith("resp_") + assert event.response.usage.cost == 1.23 + completion_handler.assert_called_once() + + +def test_process_chunk_failed_response_triggers_failure_logging(monkeypatch): + openai_types = streaming_module._get_openai_response_types() + + class _FailedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.ResponseFailedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, + response=ResponsesAPIResponse( + id="resp_failed", + created_at=int(datetime.now().timestamp()), + status="failed", + model="test-model", + object="response", + output=[], + error={"message": "provider failed"}, + ), + ) + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_FailedConfig(), + logging_obj=_FakeLoggingObj(), + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + failure_handler = MagicMock() + monkeypatch.setattr(iterator, "_handle_logging_failed_response", failure_handler) + + event = iterator._process_chunk(json.dumps({"type": "response.failed"})) + + assert iterator.completed_response is event + failure_handler.assert_called_once() + + +@pytest.mark.asyncio +async def test_handle_logging_failed_response_uses_response_error_message(): + openai_types = streaming_module._get_openai_response_types() + logging_obj = _FakeLoggingObj() + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + iterator.completed_response = openai_types.ResponseFailedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, + response=ResponsesAPIResponse( + id="resp_failed_real", + created_at=int(datetime.now().timestamp()), + status="failed", + model="test-model", + object="response", + output=[], + error={"message": "provider failed"}, + ), + ) + + iterator._handle_logging_failed_response() + await asyncio.sleep(0.2) + + assert logging_obj.failure_calls == 1 + assert logging_obj.async_failure_calls == 1 + + +def test_process_chunk_returns_none_for_invalid_json_and_non_dict_payload(): + class _NoopConfig: + def transform_streaming_response(self, **kwargs): + raise AssertionError("should not be called") + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_NoopConfig(), + logging_obj=_FakeLoggingObj(), + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + + assert iterator._process_chunk("not-json") is None + assert iterator._process_chunk(json.dumps(["not", "a", "dict"])) is None + + +def test_process_chunk_cost_annotation_failure_is_nonfatal(monkeypatch): + original_include_cost = litellm.include_cost_in_streaming_usage + litellm.include_cost_in_streaming_usage = True + openai_types = streaming_module._get_openai_response_types() + + class _CompletedConfig: + def transform_streaming_response(self, **kwargs): + return openai_types.ResponseCompletedEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_cost_failure", + created_at=int(datetime.now().timestamp()), + status="completed", + model="test-model", + object="response", + output=[], + usage=openai_types.ResponseAPIUsage( + input_tokens=1, + output_tokens=2, + total_tokens=3, + ), + ), + ) + + logging_obj = _FakeLoggingObj() + logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom")) + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_CompletedConfig(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + completion_handler = MagicMock() + monkeypatch.setattr( + iterator, "_handle_logging_completed_response", completion_handler + ) + + try: + event = iterator._process_chunk(json.dumps({"type": "response.completed"})) + finally: + litellm.include_cost_in_streaming_usage = original_include_cost + + assert iterator.completed_response is event + assert event.response.usage.cost is None + completion_handler.assert_called_once() + + +def test_get_completed_response_object_accepts_direct_response(): + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + direct_response = _make_completed_response("resp_direct").response + iterator.completed_response = direct_response + + assert iterator._get_completed_response_object() is direct_response + + +@pytest.mark.asyncio +async def test_responses_streaming_completed_event_persists_async_cache(): + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": True, + "caching": True, + "cache_key": "stale-request-cache-key", + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.aresponses, + async_set_cache=AsyncMock(), + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + logging_obj._llm_caching_handler = caching_handler + + iterator = ResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data=caching_handler.request_kwargs, + call_type=CallTypes.aresponses.value, + ) + iterator.completed_response = _make_completed_response() + + iterator._handle_logging_completed_response() + await asyncio.sleep(0.2) + + litellm.cache.async_add_cache.assert_called_once() + assert litellm.cache.async_add_cache.call_args.kwargs["stream"] is True + assert ( + litellm.cache.async_add_cache.call_args.kwargs["cache_key"] + == "responses-stream-cache-key" + ) + assert "metadata" not in litellm.cache.async_add_cache.call_args.kwargs + assert "custom_llm_provider" not in litellm.cache.async_add_cache.call_args.kwargs + assert ( + json.loads(litellm.cache.async_add_cache.call_args.args[0])["id"] + == iterator.completed_response.response.id + ) + litellm.cache = original_cache + + +def test_responses_streaming_completed_event_persists_sync_cache(): + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": True, + "caching": True, + "cache_key": "stale-request-cache-key", + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.responses, + sync_set_cache=MagicMock(), + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + logging_obj._llm_caching_handler = caching_handler + + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data=caching_handler.request_kwargs, + call_type=CallTypes.responses.value, + ) + iterator.completed_response = _make_completed_response("resp_sync") + + iterator._handle_logging_completed_response() + + litellm.cache.add_cache.assert_called_once() + assert litellm.cache.add_cache.call_args.kwargs["stream"] is True + assert ( + litellm.cache.add_cache.call_args.kwargs["cache_key"] + == "responses-stream-cache-key" + ) + assert "metadata" not in litellm.cache.add_cache.call_args.kwargs + assert "custom_llm_provider" not in litellm.cache.add_cache.call_args.kwargs + assert ( + json.loads(litellm.cache.add_cache.call_args.args[0])["id"] + == iterator.completed_response.response.id + ) + litellm.cache = original_cache + + +def test_log_completed_response_sync_direct_path(monkeypatch): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + iterator._persist_completed_response_before_logging = False + iterator.completed_response = _make_completed_response("resp_log_sync") + + iterator._log_completed_response(is_async=False) + asyncio.run(asyncio.sleep(0.2)) + + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + + +def test_log_completed_response_falls_back_when_model_validate_fails(monkeypatch): + class _BadSerializableResponse: + @classmethod + def model_validate(cls, value): + raise RuntimeError("nope") + + def model_dump(self): + return {"id": "bad"} + + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + iterator._persist_completed_response_before_logging = False + iterator.completed_response = _BadSerializableResponse() + monkeypatch.setattr(iterator, "_run_post_success_hooks", MagicMock()) + + iterator._log_completed_response(is_async=False) + asyncio.run(asyncio.sleep(0.2)) + + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + + +@pytest.mark.parametrize( + "scenario", + [ + "already_cached", + "not_completed", + "missing_caching_handler", + "not_streaming", + "store_disabled", + "missing_cache_backend", + ], +) +def test_persist_completed_response_to_cache_guard_branches(monkeypatch, scenario): + logging_obj = _FakeLoggingObj() + iterator = SyncResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=SimpleNamespace(), + logging_obj=logging_obj, + request_data={"foo": "bar"}, + call_type=CallTypes.responses.value, + ) + openai_types = streaming_module._get_openai_response_types() + completed_event = _make_completed_response("resp_guard") + iterator.completed_response = completed_event + + if scenario == "already_cached": + iterator._completed_response_cached = True + elif scenario == "not_completed": + iterator.completed_response = openai_types.ResponseIncompleteEvent( + type=openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + response=completed_event.response, + ) + elif scenario == "missing_caching_handler": + logging_obj._llm_caching_handler = None + else: + logging_obj._llm_caching_handler = SimpleNamespace( + request_kwargs={ + "model": "test-model", + "input": "hello", + "stream": scenario != "not_streaming", + "cache_key": "request-cache-key", + "metadata": None, + "custom_llm_provider": "openai", + }, + preset_cache_key=None, + original_function=litellm.responses, + dual_cache=None, + _should_store_result_in_cache=lambda original_function, kwargs: ( + scenario != "store_disabled" + ), + ) + if scenario == "missing_cache_backend": + monkeypatch.setattr(streaming_module.litellm, "cache", None) + else: + monkeypatch.setattr( + streaming_module.litellm, + "cache", + SimpleNamespace(add_cache=MagicMock(), async_add_cache=AsyncMock()), + ) + + iterator._persist_completed_response_to_cache(is_async=False) + + expected_cached_flag = scenario == "already_cached" + assert iterator._completed_response_cached is expected_cached_flag + + +def test_build_synthetic_response_events_covers_annotations_function_calls_and_refusals(): + original_include_cost = litellm.include_cost_in_streaming_usage + litellm.include_cost_in_streaming_usage = True + logging_obj = _FakeLoggingObj() + logging_obj._response_cost_calculator = MagicMock(side_effect=RuntimeError("boom")) + transformed = ResponsesAPIResponse( + id="resp_events", + created_at=int(datetime.now().timestamp()), + status="completed", + model="gpt-4.1-mini", + object="response", + output=[ + { + "type": "message", + "id": "msg_events", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "hello world", + "annotations": [{"type": "file_citation", "file_id": "file_1"}], + }, + { + "type": "refusal", + "refusal": "no thanks", + }, + ], + }, + { + "type": "function_call", + "id": "fc_events", + "call_id": "call_123", + "name": "lookup", + "arguments": '{"id":1}', + }, + ], + ) + + try: + events = streaming_module._build_synthetic_response_events( + transformed=transformed, + logging_obj=logging_obj, + chunk_size=5, + ) + finally: + litellm.include_cost_in_streaming_usage = original_include_cost + + event_types = [ + event.type.value if hasattr(event.type, "value") else str(event.type) + for event in events + ] + + assert "response.output_text.annotation.added" in event_types + assert "response.refusal.delta" in event_types + assert "response.refusal.done" in event_types + assert "response.function_call_arguments.delta" in event_types + assert "response.function_call_arguments.done" in event_types + assert event_types[-1] == "response.completed" + + +@pytest.mark.asyncio +async def test_mock_responses_streaming_iterator_async_iteration_logs_completion( + monkeypatch, +): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + class _MockTransformConfig: + def transform_response_api_response(self, **kwargs): + return _make_completed_response("resp_mock").response + + logging_obj = _FakeLoggingObj() + + iterator = MockResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_MockTransformConfig(), + logging_obj=logging_obj, + request_data={"model": "test-model", "stream": True}, + call_type=CallTypes.responses.value, + ) + + streamed_events = [event async for event in iterator] + await asyncio.sleep(0.2) + + assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + + +def test_mock_responses_streaming_iterator_sync_iteration_logs_completion(monkeypatch): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + class _MockTransformConfig: + def transform_response_api_response(self, **kwargs): + return _make_completed_response("resp_mock_sync").response + + logging_obj = _FakeLoggingObj() + iterator = MockResponsesAPIStreamingIterator( + response=httpx.Response(200), + model="test-model", + responses_api_provider_config=_MockTransformConfig(), + logging_obj=logging_obj, + request_data={"model": "test-model", "stream": True}, + call_type=CallTypes.responses.value, + ) + + streamed_events = list(iterator) + asyncio.run(asyncio.sleep(0.2)) + + assert streamed_events[0].type == ResponsesAPIStreamEvents.RESPONSE_CREATED + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + + +@pytest.mark.asyncio +async def test_cached_responses_stream_async_hit_triggers_success_callbacks( + monkeypatch, +): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + logging_obj._llm_caching_handler = SimpleNamespace( + request_kwargs={"model": "test-model", "input": "hello", "stream": True}, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.aresponses, + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + + iterator = CachedResponsesAPIStreamingIterator( + response=_make_completed_response("resp_cached_async").response, + logging_obj=logging_obj, + request_data={"model": "test-model", "input": "hello", "stream": True}, + call_type=CallTypes.aresponses.value, + ) + + streamed_events = [event async for event in iterator] + await asyncio.sleep(0.2) + + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert logging_obj.last_success_kwargs["cache_hit"] is True + assert logging_obj.last_async_success_kwargs["cache_hit"] is True + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + litellm.cache.async_add_cache.assert_not_called() + litellm.cache.add_cache.assert_not_called() + litellm.cache = original_cache + + +def test_cached_responses_stream_sync_hit_triggers_success_callbacks(monkeypatch): + hook_calls = {"post_call": 0, "metadata": 0} + + async def fake_post_call(request_data, response, call_type): + hook_calls["post_call"] += 1 + + def fake_update_metadata(**kwargs): + hook_calls["metadata"] += 1 + + monkeypatch.setattr( + streaming_module, + "async_post_call_success_deployment_hook", + fake_post_call, + ) + monkeypatch.setattr( + streaming_module, + "update_response_metadata", + fake_update_metadata, + ) + + logging_obj = _FakeLoggingObj() + original_cache = litellm.cache + litellm.cache = SimpleNamespace( + async_add_cache=AsyncMock(), + add_cache=MagicMock(), + ) + logging_obj._llm_caching_handler = SimpleNamespace( + request_kwargs={"model": "test-model", "input": "hello", "stream": True}, + preset_cache_key="responses-stream-cache-key", + original_function=litellm.responses, + _should_store_result_in_cache=lambda original_function, kwargs: True, + ) + + iterator = CachedResponsesAPIStreamingIterator( + response=_make_completed_response("resp_cached_sync").response, + logging_obj=logging_obj, + request_data={"model": "test-model", "input": "hello", "stream": True}, + call_type=CallTypes.responses.value, + ) + + streamed_events = list(iterator) + asyncio.run(asyncio.sleep(0.2)) + + assert streamed_events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert logging_obj.success_calls == 1 + assert logging_obj.async_success_calls == 1 + assert logging_obj.last_success_kwargs["cache_hit"] is True + assert logging_obj.last_async_success_kwargs["cache_hit"] is True + assert hook_calls["post_call"] == 1 + assert hook_calls["metadata"] == 1 + litellm.cache.async_add_cache.assert_not_called() + litellm.cache.add_cache.assert_not_called() + litellm.cache = original_cache diff --git a/tests/llm_translation/Readme.md b/tests/llm_translation/Readme.md index db84e7c33c..958adbd975 100644 --- a/tests/llm_translation/Readme.md +++ b/tests/llm_translation/Readme.md @@ -1,3 +1,41 @@ -Unit tests for individual LLM providers. +Unit tests for individual LLM providers. -Name of the test file is the name of the LLM provider - e.g. `test_openai.py` is for OpenAI. \ No newline at end of file +Name of the test file is the name of the LLM provider - e.g. `test_openai.py` is for OpenAI. + +## Redis-backed VCR cache + +Every test in this directory is auto-decorated with `@pytest.mark.vcr` (via +`conftest.py`). The first time a test runs we hit the live provider and +record the HTTP exchange into Redis under +`litellm:vcr:cassette:`. Every subsequent run within 24h replays +from Redis without touching the network. The 24h TTL means each new day's +first run records again, so upstream API drift surfaces within a day. + +The persister, header scrubbing, and 2xx-only filtering are defined in +`tests/_vcr_redis_persister.py`. Files that already use `respx` (which +patches the same httpx transport vcrpy does) are excluded from the +auto-marker — see `_RESPX_CONFLICTING_FILES` in `conftest.py`. + +### Required environment + +`REDIS_HOST`, `REDIS_PORT`, `REDIS_PASSWORD` — same vars CircleCI uses for +its other Redis-backed jobs. Provider credentials +(`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `AWS_*`, etc.) are needed only on +cache-miss (the daily re-record), not on replay. + +### Flushing the cache + +When you want the next run to re-record immediately instead of waiting +for the 24h TTL: + +```bash +make test-llm-translation-flush-vcr-cache +``` + +### Disabling VCR + +Skip the cache entirely (every call goes live, no recording): + +```bash +LITELLM_VCR_DISABLE=1 uv run pytest tests/llm_translation/test_.py +``` diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index f3b1895323..aedc4f810c 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -401,7 +401,7 @@ class BaseLLMChatTest(ABC): { "type": "file", "file": { - "file_id": "https://upload.wikimedia.org/wikipedia/commons/2/20/Re_example.pdf" + "file_id": "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/tests/llm_translation/fixtures/dummy.pdf" }, }, ] diff --git a/tests/llm_translation/conftest.py b/tests/llm_translation/conftest.py index d315dc63bc..09da0520be 100644 --- a/tests/llm_translation/conftest.py +++ b/tests/llm_translation/conftest.py @@ -5,6 +5,7 @@ # - Function-scoped fixture resets litellm globals to true defaults # - Module-scoped reload only in single-process mode +import asyncio import importlib import os import sys @@ -14,9 +15,195 @@ import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import litellm -import asyncio +import litellm # noqa: E402 + +from tests._vcr_redis_persister import ( # noqa: E402 + filter_non_2xx_response, + format_vcr_verdict, + make_redis_persister, + mark_test_outcome_for_cassette, + patch_vcrpy_aiohttp_record_path, + vcr_verbose_enabled, +) + + +_controller_pluginmanager = None +_controller_terminal_reporter = None + + +# vcrpy and respx both patch the httpx transport — applying both makes one +# silently win, so respx-using files opt out of the auto-marker. +_RESPX_CONFLICTING_FILES = frozenset( + { + "test_azure_o_series.py", + "test_gpt4o_audio.py", + "test_nvidia_nim.py", + "test_openai.py", + "test_openai_o1.py", + "test_prompt_caching.py", + "test_text_completion_unit_tests.py", + "test_xai.py", + } +) +_VCR_AUTO_MARKER_SKIP_FILES = _RESPX_CONFLICTING_FILES | frozenset( + {"test_vcr_redis_persister.py"} +) + +# Tests that observe live cross-call provider state (e.g. prompt-cache +# warm-up between two consecutive calls); replay can't reproduce that state. +_VCR_INCOMPATIBLE_NODEID_SUFFIXES = frozenset( + { + "::test_prompt_caching", + "TestBedrockInvokeNovaJson::test_json_response_pydantic_obj", + "::test_bedrock_converse__streaming_passthrough", + } +) + + +def _is_vcr_incompatible(nodeid: str) -> bool: + return any(nodeid.endswith(suffix) for suffix in _VCR_INCOMPATIBLE_NODEID_SUFFIXES) + + +_FILTERED_REQUEST_HEADERS = ( + "authorization", + "x-api-key", + "anthropic-api-key", + "anthropic-version", + "openai-api-key", + "azure-api-key", + "api-key", + "cookie", + "x-amz-security-token", + "x-amz-date", + "x-amz-content-sha256", + "amz-sdk-invocation-id", + "amz-sdk-request", + "x-goog-api-key", + "x-goog-user-project", +) + +_FILTERED_RESPONSE_HEADERS = ( + "set-cookie", + "x-request-id", + "request-id", + "cf-ray", + "anthropic-organization-id", + "openai-organization", + "x-amzn-requestid", + "x-amzn-trace-id", + "date", +) + + +def _scrub_response(response): + if not isinstance(response, dict): + return response + headers = response.get("headers") or {} + if isinstance(headers, dict): + for header in list(headers): + if header.lower() in _FILTERED_RESPONSE_HEADERS: + headers.pop(header, None) + return response + + +def _before_record_response(response): + return filter_non_2xx_response(_scrub_response(response)) + + +@pytest.fixture(scope="module") +def vcr_config(): + return { + "filter_headers": list(_FILTERED_REQUEST_HEADERS), + "decode_compressed_response": True, + "record_mode": "new_episodes", + "allow_playback_repeats": True, + "match_on": ( + "method", + "scheme", + "host", + "port", + "path", + "query", + "body", + ), + "before_record_response": _before_record_response, + } + + +def _vcr_disabled() -> bool: + if os.environ.get("LITELLM_VCR_DISABLE") == "1": + return True + return not os.environ.get("CASSETTE_REDIS_URL") + + +def pytest_recording_configure(config, vcr): + if _vcr_disabled(): + return + vcr.register_persister(make_redis_persister()) + patch_vcrpy_aiohttp_record_path() + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + outcome = yield + rep = outcome.get_result() + setattr(item, f"rep_{rep.when}", rep) + + +@pytest.fixture(autouse=True) +def _vcr_outcome_gate(request, vcr): + yield + cassette = vcr + rep_call = getattr(request.node, "rep_call", None) + test_passed = bool(rep_call and rep_call.passed) + cassette_path = getattr(cassette, "_path", None) if cassette is not None else None + if cassette_path: + mark_test_outcome_for_cassette(cassette_path, test_passed) + + if not vcr_verbose_enabled(): + return + verdict = format_vcr_verdict(cassette) + request.node.user_properties.append(("vcr_verdict", verdict)) + + +def pytest_configure(config): + global _controller_pluginmanager + if os.environ.get("PYTEST_XDIST_WORKER"): + return + _controller_pluginmanager = config.pluginmanager + + +def _resolve_terminal_reporter(): + global _controller_terminal_reporter + if _controller_terminal_reporter is not None: + return _controller_terminal_reporter + if _controller_pluginmanager is None: + return None + _controller_terminal_reporter = _controller_pluginmanager.getplugin( + "terminalreporter" + ) + return _controller_terminal_reporter + + +def pytest_runtest_logreport(report): + if report.when != "teardown": + return + if os.environ.get("PYTEST_XDIST_WORKER"): + return + if not vcr_verbose_enabled(): + return + reporter = _resolve_terminal_reporter() + if reporter is None: + return + verdict = next( + (v for k, v in (report.user_properties or []) if k == "vcr_verdict"), + None, + ) + if not verdict: + return + reporter.write_line(f"{verdict} :: {report.nodeid}") + # --------------------------------------------------------------------------- # Capture TRUE defaults at conftest import time (before test modules pollute). @@ -48,7 +235,6 @@ def event_loop(): @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(event_loop): # Add event_loop as a dependency - curr_dir = os.getcwd() sys.path.insert(0, os.path.abspath("../..")) import litellm @@ -97,15 +283,23 @@ def setup_and_teardown(event_loop): # Add event_loop as a dependency def pytest_collection_modifyitems(config, items): - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + if not _vcr_disabled(): + for item in items: + filename = os.path.basename(str(item.fspath)) + if filename in _VCR_AUTO_MARKER_SKIP_FILES: + continue + if _is_vcr_incompatible(item.nodeid): + continue + if item.get_closest_marker("vcr") is not None: + continue + item.add_marker(pytest.mark.vcr) + custom_logger_tests = [ item for item in items if "custom_logger" in item.parent.name ] other_tests = [item for item in items if "custom_logger" not in item.parent.name] - # Sort tests based on their names custom_logger_tests.sort(key=lambda x: x.name) other_tests.sort(key=lambda x: x.name) - # Reorder the items list items[:] = custom_logger_tests + other_tests diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index fdf8c24ac9..371b27c5b2 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -870,7 +870,7 @@ from litellm.constants import RESPONSE_FORMAT_TOOL_NAME def test_anthropic_json_mode_and_tool_call_response( json_mode, tool_calls, expect_null_response ): - result = litellm.AnthropicConfig()._transform_response_for_json_mode( + result, _, _ = litellm.AnthropicConfig()._resolve_json_mode_non_streaming( json_mode=json_mode, tool_calls=tool_calls, ) @@ -1885,3 +1885,42 @@ def test_metadata_filter_applies_to_azure_anthropic(): headers={}, ) assert data.get("metadata") == {"user_id": "u2"} + + +def test_anthropic_basic_completion_replay(): + response = litellm.completion( + model="anthropic/claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello!"}], + ) + + assert response is not None + content = response.choices[0].message.content + assert isinstance(content, str) and content.strip(), content + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.choices[0].finish_reason in {"stop", "length"} + + +def test_anthropic_streaming_completion_replay(): + stream = litellm.completion( + model="anthropic/claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello!"}], + stream=True, + ) + + collected_text = "" + finish_reason = None + chunk_count = 0 + for chunk in stream: + chunk_count += 1 + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta and delta.content: + collected_text += delta.content + if chunk.choices[0].finish_reason: + finish_reason = chunk.choices[0].finish_reason + + assert chunk_count > 1, "expected multiple SSE chunks from streaming response" + assert collected_text.strip(), collected_text + assert finish_reason in {"stop", "length"} diff --git a/tests/llm_translation/test_vcr_redis_persister.py b/tests/llm_translation/test_vcr_redis_persister.py new file mode 100644 index 0000000000..853558150c --- /dev/null +++ b/tests/llm_translation/test_vcr_redis_persister.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +import os +import sys + +import fakeredis +import pytest +from redis.exceptions import ConnectionError as RedisConnectionError +from vcr.persisters.filesystem import CassetteNotFoundError +from vcr.request import Request +from vcr.serializers import yamlserializer + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) + +from tests._vcr_redis_persister import ( # noqa: E402 + CASSETTE_TTL_SECONDS, + MAX_EPISODES_PER_CASSETTE, + filter_non_2xx_response, + make_redis_persister, + mark_test_outcome_for_cassette, + redis_key_for, +) + + +def _sample_cassette_dict(): + request = Request( + method="POST", + uri="https://api.anthropic.com/v1/messages", + body=b'{"model":"claude","messages":[{"role":"user","content":"hi"}]}', + headers={"content-type": "application/json"}, + ) + response = { + "status": {"code": 200, "message": "OK"}, + "headers": {"content-type": ["application/json"]}, + "body": {"string": b'{"id":"msg_1","type":"message"}'}, + } + return {"requests": [request], "responses": [response]} + + +def _persister_with_fake_redis(): + fake = fakeredis.FakeStrictRedis() + return fake, make_redis_persister(client=fake) + + +def test_save_then_load_roundtrips_cassette_content(): + _, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_y" + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + requests, responses = persister.load_cassette(cassette_id, yamlserializer) + + assert len(requests) == 1 + assert len(responses) == 1 + assert requests[0].method == "POST" + assert requests[0].uri == "https://api.anthropic.com/v1/messages" + assert responses[0]["status"]["code"] == 200 + assert responses[0]["body"]["string"] == b'{"id":"msg_1","type":"message"}' + + +def test_saved_key_has_24h_ttl(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_ttl" + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + + ttl = fake.ttl(redis_key_for(cassette_id)) + assert CASSETTE_TTL_SECONDS - 5 <= ttl <= CASSETTE_TTL_SECONDS + + +def test_load_missing_key_raises_cassette_not_found(): + _, persister = _persister_with_fake_redis() + with pytest.raises(CassetteNotFoundError): + persister.load_cassette("never/recorded", yamlserializer) + + +def test_redis_key_normalizes_path_passed_by_pytest_recording(): + raw = "tests/llm_translation/cassettes/test_anthropic/test_streaming.yaml" + assert ( + redis_key_for(raw) + == "litellm:vcr:cassette:tests/llm_translation/test_anthropic/test_streaming" + ) + + +class _FlakyRedis: + def __init__(self, inner, fail_on: str): + self._inner = inner + self._fail_on = fail_on + + def get(self, *args, **kwargs): + if self._fail_on == "get": + raise RedisConnectionError("simulated outage") + return self._inner.get(*args, **kwargs) + + def set(self, *args, **kwargs): + if self._fail_on == "set": + raise RedisConnectionError("simulated outage") + return self._inner.set(*args, **kwargs) + + +def test_save_swallows_connection_errors_so_teardown_does_not_fail(): + flaky = _FlakyRedis(fakeredis.FakeStrictRedis(), fail_on="set") + persister = make_redis_persister(client=flaky) + + persister.save_cassette( + "tests/llm_translation/test_x/test_save_outage", + _sample_cassette_dict(), + yamlserializer, + ) + + +def test_save_skipped_when_test_marked_failed_and_prior_cassette_preserved(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_flaky" + key = redis_key_for(cassette_id) + + good = _sample_cassette_dict() + persister.save_cassette(cassette_id, good, yamlserializer) + good_payload = fake.get(key) + assert good_payload is not None + + mark_test_outcome_for_cassette(cassette_id, passed=False) + bad_response = { + "status": {"code": 200, "message": "OK"}, + "headers": {}, + "body": {"string": b'{"id":"BAD","type":"message"}'}, + } + bad = {"requests": good["requests"], "responses": [bad_response]} + persister.save_cassette(cassette_id, bad, yamlserializer) + + assert fake.get(key) == good_payload + + +def test_save_proceeds_when_test_marked_passed(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_passed" + key = redis_key_for(cassette_id) + + mark_test_outcome_for_cassette(cassette_id, passed=True) + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + + assert fake.get(key) is not None + + +def test_save_refused_when_cassette_exceeds_max_episodes(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_runaway" + key = redis_key_for(cassette_id) + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + seed_payload = fake.get(key) + + request = Request( + method="POST", + uri="https://api.anthropic.com/v1/messages", + body=b"x", + headers={"content-type": "application/json"}, + ) + response = { + "status": {"code": 200, "message": "OK"}, + "headers": {}, + "body": {"string": b"{}"}, + } + bloated = { + "requests": [request] * (MAX_EPISODES_PER_CASSETTE + 1), + "responses": [response] * (MAX_EPISODES_PER_CASSETTE + 1), + } + persister.save_cassette(cassette_id, bloated, yamlserializer) + + assert fake.get(key) == seed_payload + + +def test_save_proceeds_at_max_episodes_threshold(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_at_threshold" + key = redis_key_for(cassette_id) + + request = Request( + method="POST", + uri="https://api.anthropic.com/v1/messages", + body=b"x", + headers={"content-type": "application/json"}, + ) + response = { + "status": {"code": 200, "message": "OK"}, + "headers": {}, + "body": {"string": b"{}"}, + } + at_threshold = { + "requests": [request] * MAX_EPISODES_PER_CASSETTE, + "responses": [response] * MAX_EPISODES_PER_CASSETTE, + } + persister.save_cassette(cassette_id, at_threshold, yamlserializer) + + assert fake.get(key) is not None + + +def test_save_proceeds_when_outcome_unknown(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_no_marker" + key = redis_key_for(cassette_id) + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + + assert fake.get(key) is not None + + +def test_load_treats_connection_errors_as_cassette_miss(): + flaky = _FlakyRedis(fakeredis.FakeStrictRedis(), fail_on="get") + persister = make_redis_persister(client=flaky) + + with pytest.raises(CassetteNotFoundError): + persister.load_cassette( + "tests/llm_translation/test_x/test_load_outage", yamlserializer + ) + + +@pytest.mark.parametrize( + ("status_code", "expect_dropped"), + [ + (200, False), + (201, False), + (204, False), + (299, False), + (300, True), + (400, True), + (401, True), + (404, True), + (429, True), + (500, True), + (502, True), + (503, True), + ], +) +def test_only_2xx_responses_are_cached(status_code, expect_dropped): + response = { + "status": {"code": status_code, "message": "X"}, + "headers": {}, + "body": {"string": ""}, + } + result = filter_non_2xx_response(response) + assert (result is None) == expect_dropped + if not expect_dropped: + assert result is response diff --git a/tests/local_testing/test_caching_handler.py b/tests/local_testing/test_caching_handler.py index 806f72bfde..2b6712cbaa 100644 --- a/tests/local_testing/test_caching_handler.py +++ b/tests/local_testing/test_caching_handler.py @@ -19,9 +19,14 @@ import pytest import litellm from litellm import aembedding, completion, embedding, aresponses, responses from litellm.caching.caching import Cache +from litellm.responses.streaming_iterator import CachedResponsesAPIStreamingIterator from unittest.mock import AsyncMock, patch, MagicMock -from litellm.caching.caching_handler import LLMCachingHandler, CachingHandlerResponse +from litellm.caching.caching_handler import ( + LLMCachingHandler, + CachingHandlerResponse, + _should_defer_streaming_cache_hit_callbacks, +) from litellm.caching.caching import LiteLLMCacheType from litellm.types.utils import CallTypes from litellm.types.rerank import RerankResponse @@ -627,6 +632,55 @@ async def test_async_responses_api_caching(): assert cached_response.cached_result._hidden_params["cache_hit"] == True +@pytest.mark.asyncio +async def test_async_get_cache_updates_request_kwargs_for_streaming_responses(): + """ + Ensure streamed responses retain the normalized lookup kwargs so a later + cache write can reuse the exact cache key from the read path. + """ + setup_cache() + + caching_handler = LLMCachingHandler( + original_function=aresponses, + request_kwargs={"stale": True}, + start_time=datetime.now(), + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-4o", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + kwargs = { + "model": "gpt-4o", + "input": "hello", + "stream": True, + "caching": True, + } + + await caching_handler._async_get_cache( + model="gpt-4o", + original_function=aresponses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aresponses.value, + kwargs=kwargs, + ) + + assert "stale" not in caching_handler.request_kwargs + assert caching_handler.request_kwargs["model"] == "gpt-4o" + assert caching_handler.request_kwargs["input"] == "hello" + assert caching_handler.request_kwargs["stream"] is True + assert caching_handler.request_kwargs["cache_key"] == litellm.cache.get_cache_key( + **caching_handler.request_kwargs + ) + + def test_sync_responses_api_caching(): """ Test that synchronous responses API calls are properly cached and retrieved. @@ -769,6 +823,339 @@ def test_convert_cached_responses_api_result_to_model_response(): assert len(result.output) == 1 +def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): + litellm.set_verbose = True + setup_cache() + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + + original_model = "gpt-4o" + responses_api_response = ResponsesAPIResponse( + id="resp_stream_sync_hit", + created_at=int(time.time()), + status="completed", + model=original_model, + object="response", + output=[ + { + "type": "message", + "id": "msg_stream_sync_hit", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Sync streamed cache hit response.", + "annotations": [], + } + ], + } + ], + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.responses.value, + model=original_model, + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + + kwargs = { + "model": original_model, + "input": "Tell me a cached story", + "stream": True, + "caching": True, + } + + caching_handler.sync_set_cache(result=responses_api_response, kwargs=kwargs) + time.sleep(0.2) + + cached_response = caching_handler._sync_get_cache( + model=original_model, + original_function=responses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.responses.value, + kwargs=kwargs, + ) + + assert cached_response.cached_result is not None + assert isinstance( + cached_response.cached_result, CachedResponsesAPIStreamingIterator + ) + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + + +def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): + litellm.set_verbose = True + setup_cache() + caching_handler = LLMCachingHandler( + original_function=completion, request_kwargs={}, start_time=datetime.now() + ) + + original_model = "gpt-4o" + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.completion.value, + model=original_model, + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + + kwargs = { + "model": original_model, + "messages": [{"role": "user", "content": "Tell me a cached joke"}], + "stream": True, + "caching": True, + } + + caching_handler.sync_set_cache(result=chat_completion_response, kwargs=kwargs) + time.sleep(0.2) + + cached_response = caching_handler._sync_get_cache( + model=original_model, + original_function=completion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.completion.value, + kwargs=kwargs, + ) + + assert cached_response.cached_result is not None + logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + + +def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request(): + assert ( + _should_defer_streaming_cache_hit_callbacks( + kwargs={"stream": True}, + ) + is True + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + kwargs={"stream": False}, + ) + is False + ) + assert ( + _should_defer_streaming_cache_hit_callbacks( + kwargs={}, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_async_get_cache_defers_streaming_completion_hit_callbacks(): + litellm.set_verbose = True + setup_cache() + caching_handler = LLMCachingHandler( + original_function=completion, request_kwargs={}, start_time=datetime.now() + ) + + original_model = "gpt-4o" + kwargs = { + "model": original_model, + "messages": [{"role": "user", "content": "Tell me a cached joke"}], + "stream": True, + "caching": True, + } + + await caching_handler.async_set_cache( + result=chat_completion_response, + original_function=litellm.acompletion, + kwargs=kwargs, + ) + await asyncio.sleep(0.2) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.acompletion.value, + model=original_model, + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + caching_handler._async_log_cache_hit_on_callbacks = MagicMock() + + cached_response = await caching_handler._async_get_cache( + model=original_model, + original_function=litellm.acompletion, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.acompletion.value, + kwargs=kwargs, + ) + + assert cached_response is not None + assert cached_response.cached_result is not None + caching_handler._async_log_cache_hit_on_callbacks.assert_not_called() + + +def test_convert_cached_streaming_responses_result_to_iterator(): + """ + Test that cached streaming Responses results are replayed through a synthetic + streaming iterator instead of being returned as a full response object. + """ + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.responses.value, + model="gpt-4o", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + cached_result = { + "id": "resp_stream_cache_test", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "message", + "id": "msg_stream_cache_test", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Streaming cache replay test.", + "annotations": [], + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "test", "stream": True}, + logging_obj=logging_obj, + model="gpt-4o", + args=(), + ) + + assert isinstance(result, CachedResponsesAPIStreamingIterator) + assert result.completed_response is not None + assert result.completed_response.response.id == cached_result["id"] + + streamed_events = list(result) + assert streamed_events[0].type == "response.created" + assert streamed_events[1].type == "response.in_progress" + assert streamed_events[2].type == "response.output_item.added" + assert streamed_events[3].type == "response.content_part.added" + assert streamed_events[-4].type == "response.output_text.done" + assert streamed_events[-3].type == "response.content_part.done" + assert streamed_events[-2].type == "response.output_item.done" + assert streamed_events[-1].type == "response.completed" + assert streamed_events[-1].response.id == cached_result["id"] + assert streamed_events[-1].response.output[0].content[0].text == ( + "Streaming cache replay test." + ) + + +def test_convert_cached_streaming_reasoning_result_to_iterator(): + caching_handler = LLMCachingHandler( + original_function=responses, request_kwargs={}, start_time=datetime.now() + ) + + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.responses.value, + model="gpt-4o", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + cached_result = { + "id": "resp_stream_reasoning_cache_test", + "created_at": int(time.time()), + "status": "completed", + "model": "gpt-4o", + "object": "response", + "output": [ + { + "type": "reasoning", + "id": "rs_stream_cache_test", + "summary": [ + { + "type": "summary_text", + "text": "Cached reasoning summary.", + } + ], + } + ], + } + + result = caching_handler._convert_cached_result_to_model_response( + cached_result=cached_result, + call_type=CallTypes.responses.value, + kwargs={"model": "gpt-4o", "input": "test", "stream": True}, + logging_obj=logging_obj, + model="gpt-4o", + args=(), + ) + + assert isinstance(result, CachedResponsesAPIStreamingIterator) + + streamed_events = list(result) + streamed_event_types = [ + event.type.value if hasattr(event.type, "value") else str(event.type) + for event in streamed_events + ] + + assert streamed_event_types[:3] == [ + "response.created", + "response.in_progress", + "response.output_item.added", + ] + assert streamed_event_types[-4:] == [ + "response.reasoning_summary_text.done", + "response.reasoning_summary_part.done", + "response.output_item.done", + "response.completed", + ] + assert streamed_event_types.count("response.reasoning_summary_text.delta") >= 1 + + delta_events = [ + event + for event in streamed_events + if (event.type.value if hasattr(event.type, "value") else str(event.type)) + == "response.reasoning_summary_text.delta" + ] + text_done_event = streamed_events[-4] + part_done_event = streamed_events[-3] + output_item_done_event = streamed_events[-2] + + assert all(delta_event.summary_index == 0 for delta_event in delta_events) + assert text_done_event.text == "Cached reasoning summary." + assert text_done_event.summary_index == 0 + assert part_done_event.part.type == "summary_text" + assert part_done_event.part.text == "Cached reasoning summary." + assert output_item_done_event.item.type == "reasoning" + assert output_item_done_event.item.summary[0]["text"] == "Cached reasoning summary." + + @pytest.mark.asyncio async def test_responses_api_cache_with_different_inputs(): """ diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 010a071f73..14b9e8cd13 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -477,3 +477,4 @@ def test_get_llm_provider_use_proxy_arg_true_with_direct_args(): assert provider == "litellm_proxy" assert key == arg_api_key # Should use the argument key assert base == arg_api_base # Should use the argument base + diff --git a/tests/local_testing/test_responses_stream_cache_keys.py b/tests/local_testing/test_responses_stream_cache_keys.py new file mode 100644 index 0000000000..5637028f55 --- /dev/null +++ b/tests/local_testing/test_responses_stream_cache_keys.py @@ -0,0 +1,141 @@ +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm +from litellm import aresponses +from litellm._uuid import uuid +from litellm.caching.caching_handler import LLMCachingHandler +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.llms import openai as openai_types +from litellm.types.utils import CallTypes + + +@pytest.mark.asyncio +async def test_async_get_cache_reuses_preset_cache_key_for_responses(): + caching_handler = LLMCachingHandler( + original_function=aresponses, + request_kwargs={}, + start_time=datetime.now(), + ) + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-4.1-mini", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + original_cache = litellm.cache + mock_cache = MagicMock() + mock_cache.supported_call_types = [CallTypes.aresponses.value] + mock_cache._supports_async.return_value = True + mock_cache.get_cache_key.return_value = "responses-stream-cache-key" + mock_cache.async_get_cache = AsyncMock(return_value=None) + litellm.cache = mock_cache + + kwargs = { + "model": "gpt-4.1-mini", + "input": "hello", + "stream": True, + "litellm_params": {}, + } + await caching_handler._async_get_cache( + model="gpt-4.1-mini", + original_function=aresponses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aresponses.value, + kwargs=kwargs, + ) + + assert caching_handler.preset_cache_key == "responses-stream-cache-key" + mock_cache.async_get_cache.assert_awaited_once() + assert ( + mock_cache.async_get_cache.call_args.kwargs["cache_key"] + == "responses-stream-cache-key" + ) + + litellm.cache = original_cache + + +@pytest.mark.asyncio +async def test_async_get_cache_falls_back_to_sync_cache_for_responses(): + caching_handler = LLMCachingHandler( + original_function=aresponses, + request_kwargs={}, + start_time=datetime.now(), + ) + logging_obj = LiteLLMLogging( + litellm_call_id=str(datetime.now()), + call_type=CallTypes.aresponses.value, + model="gpt-4.1-mini", + messages=[], + function_id=str(uuid.uuid4()), + stream=True, + start_time=datetime.now(), + ) + + original_cache = litellm.cache + mock_cache = MagicMock() + mock_cache.supported_call_types = [CallTypes.aresponses.value] + mock_cache._supports_async.return_value = False + mock_cache.get_cache_key.return_value = "responses-stream-cache-key" + mock_cache.get_cache.return_value = None + litellm.cache = mock_cache + + kwargs = { + "model": "gpt-4.1-mini", + "input": "hello", + "stream": True, + "litellm_params": {}, + } + await caching_handler._async_get_cache( + model="gpt-4.1-mini", + original_function=aresponses, + logging_obj=logging_obj, + start_time=datetime.now(), + call_type=CallTypes.aresponses.value, + kwargs=kwargs, + ) + + assert caching_handler.preset_cache_key == "responses-stream-cache-key" + mock_cache.get_cache.assert_called_once() + assert mock_cache.get_cache.call_args.kwargs["cache_key"] == ( + "responses-stream-cache-key" + ) + + litellm.cache = original_cache + + +def test_reasoning_summary_events_default_summary_index(): + delta_event = openai_types.ReasoningSummaryTextDeltaEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DELTA, + item_id="rs_1", + output_index=0, + delta="abc", + ) + text_done_event = openai_types.ReasoningSummaryTextDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_TEXT_DONE, + item_id="rs_1", + output_index=0, + sequence_number=1, + text="abc", + ) + part_done_event = openai_types.ReasoningSummaryPartDoneEvent( + type=openai_types.ResponsesAPIStreamEvents.REASONING_SUMMARY_PART_DONE, + item_id="rs_1", + output_index=0, + sequence_number=2, + part=openai_types.BaseLiteLLMOpenAIResponseObject( + type="summary_text", + text="abc", + ), + ) + + assert delta_event.summary_index == 0 + assert text_done_event.summary_index == 0 + assert part_done_event.summary_index == 0 diff --git a/tests/otel_tests/test_e2e_model_access.py b/tests/otel_tests/test_e2e_model_access.py index ade319c2d4..7ea75a9d61 100644 --- a/tests/otel_tests/test_e2e_model_access.py +++ b/tests/otel_tests/test_e2e_model_access.py @@ -6,13 +6,19 @@ from httpx import AsyncClient from typing import Any, Optional, List, Literal +# The proxy strips client-supplied `mock_response` unless the calling key or +# team has this admin-metadata flag set. See `_UNTRUSTED_ROOT_CONTROL_FIELDS` +# in litellm/proxy/litellm_pre_call_utils.py. +_ALLOW_CLIENT_MOCK_METADATA = {"allow_client_mock_response": True} + + async def generate_key( session, models: Optional[List[str]] = None, team_id: Optional[str] = None ): """Helper function to generate a key with specific model access controls""" url = "http://0.0.0.0:4000/key/generate" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} - data = {} + data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)} if models is not None: data["models"] = models if team_id is not None: @@ -25,7 +31,7 @@ async def generate_team(session, models: Optional[List[str]] = None): """Helper function to generate a team with specific model access""" url = "http://0.0.0.0:4000/team/new" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} - data = {} + data: dict = {"metadata": dict(_ALLOW_CLIENT_MOCK_METADATA)} if models is not None: data["models"] = models async with session.post(url, headers=headers, json=data) as response: @@ -111,7 +117,12 @@ async def test_model_access_update(): # Create initial key with restricted access response = await client.post( - "/key/generate", json={"models": ["openai/gpt-4"]}, headers=headers + "/key/generate", + json={ + "models": ["openai/gpt-4"], + "metadata": dict(_ALLOW_CLIENT_MOCK_METADATA), + }, + headers=headers, ) assert response.status_code == 200 key_data = response.json() @@ -214,7 +225,11 @@ async def test_team_model_access_update(): # Create initial team with restricted access response = await client.post( "/team/new", - json={"models": ["openai/gpt-4"], "name": "test-team"}, + json={ + "models": ["openai/gpt-4"], + "name": "test-team", + "metadata": dict(_ALLOW_CLIENT_MOCK_METADATA), + }, headers=headers, ) assert response.status_code == 200 @@ -223,7 +238,12 @@ async def test_team_model_access_update(): # Generate a key for this team response = await client.post( - "/key/generate", json={"team_id": team_id}, headers=headers + "/key/generate", + json={ + "team_id": team_id, + "metadata": dict(_ALLOW_CLIENT_MOCK_METADATA), + }, + headers=headers, ) assert response.status_code == 200 key = response.json()["key"] diff --git a/tests/pass_through_unit_tests/test_assemblyai_unit_tests_passthrough.py b/tests/pass_through_unit_tests/test_assemblyai_unit_tests_passthrough.py index 963f1ad6ef..67bc4423d8 100644 --- a/tests/pass_through_unit_tests/test_assemblyai_unit_tests_passthrough.py +++ b/tests/pass_through_unit_tests/test_assemblyai_unit_tests_passthrough.py @@ -134,3 +134,62 @@ def test_is_assemblyai_route(): == False ) assert handler.is_assemblyai_route("") == False + + +# --- Security: SSRF via transcript_id path traversal --- + + +def test_get_assembly_transcript_rejects_slash_in_id(assembly_handler): + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="test-key", + ): + with pytest.raises(ValueError, match="disallowed characters"): + assembly_handler._get_assembly_transcript("../../admin/credentials") + + +def test_get_assembly_transcript_rejects_dotdot_in_id(assembly_handler): + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="test-key", + ): + with pytest.raises(ValueError, match="disallowed characters"): + assembly_handler._get_assembly_transcript("..evil") + + +def test_get_assembly_transcript_rejects_fragment_in_id(assembly_handler): + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="test-key", + ): + with pytest.raises(ValueError, match="disallowed characters"): + assembly_handler._get_assembly_transcript("abc#suffix") + + +def test_get_assembly_transcript_rejects_query_in_id(assembly_handler): + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="test-key", + ): + with pytest.raises(ValueError, match="disallowed characters"): + assembly_handler._get_assembly_transcript("abc?x=1") + + +def test_get_assembly_transcript_allows_valid_id( + assembly_handler, mock_transcript_response +): + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="test-key", + ): + with patch("httpx.get") as mock_get: + mock_get.return_value.json.return_value = mock_transcript_response + mock_get.return_value.raise_for_status.return_value = None + + transcript = assembly_handler._get_assembly_transcript( + "abc123-valid-id_xyz" + ) + assert transcript == mock_transcript_response + called_url = mock_get.call_args[0][0] + assert "abc123-valid-id_xyz" in called_url + assert ".." not in called_url diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/proxy_unit_tests/test_auth_checks.py index 86cd5c0c41..5636a55c95 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -16,6 +16,7 @@ import httpx from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import get_end_user_object from litellm.caching.caching import DualCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy._types import ( LiteLLM_EndUserTable, LiteLLM_BudgetTable, @@ -48,9 +49,15 @@ async def test_get_end_user_object(customer_spend, customer_budget): litellm_budget_table=_budget, blocked=False, ) - _cache = DualCache() + # UserApiKeyCache applies model_type on get/set; plain DualCache returns raw dicts + # and breaks get_end_user_object's typed async_get_cache path. + _cache = UserApiKeyCache() _key = "end_user_id:{}".format(end_user_id) - _cache.set_cache(key=_key, value=end_user_obj.model_dump()) + await _cache.async_set_cache( + key=_key, + value=end_user_obj, + model_type=LiteLLM_EndUserTable, + ) try: await get_end_user_object( end_user_id=end_user_id, diff --git a/tests/proxy_unit_tests/test_get_favicon.py b/tests/proxy_unit_tests/test_get_favicon.py index f17787e740..ddc8b1230a 100644 --- a/tests/proxy_unit_tests/test_get_favicon.py +++ b/tests/proxy_unit_tests/test_get_favicon.py @@ -1,6 +1,5 @@ import os import sys -from unittest import mock sys.path.insert(0, os.path.abspath("../..")) @@ -26,50 +25,30 @@ async def test_get_favicon_default(): @pytest.mark.asyncio -async def test_get_favicon_with_custom_url(): - """Test that get_favicon fetches from a custom URL.""" - os.environ["LITELLM_FAVICON_URL"] = "https://example.com/favicon.ico" +async def test_get_favicon_with_custom_url(monkeypatch): + """Test that get_favicon redirects browser-loaded custom URLs.""" + monkeypatch.setenv("LITELLM_FAVICON_URL", "https://example.com/favicon.ico") - mock_response = mock.Mock() - mock_response.status_code = 200 - mock_response.content = b"\x00\x00\x01\x00" - mock_response.headers = {"content-type": "image/x-icon"} + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://testserver", + ) as ac: + response = await ac.get("/get_favicon") - try: - with mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get: - mock_get.return_value = mock_response - - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=app), - base_url="http://testserver", - ) as ac: - response = await ac.get("/get_favicon") - - assert response.status_code == 200 - assert response.headers["content-type"] == "image/x-icon" - finally: - os.environ.pop("LITELLM_FAVICON_URL", None) + assert response.status_code == 307 + assert response.headers["location"] == "https://example.com/favicon.ico" @pytest.mark.asyncio -async def test_get_favicon_url_error_fallback(): - """Test that get_favicon falls back to default on error.""" - os.environ["LITELLM_FAVICON_URL"] = "https://invalid.com/favicon.ico" +async def test_get_favicon_remote_url_is_not_server_fetched(monkeypatch): + """Test that get_favicon does not validate remote URLs server-side.""" + monkeypatch.setenv("LITELLM_FAVICON_URL", "https://invalid.com/favicon.ico") - try: - with mock.patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" - ) as mock_get: - mock_get.side_effect = httpx.ConnectError("unreachable") + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url="http://testserver", + ) as ac: + response = await ac.get("/get_favicon") - async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=app), - base_url="http://testserver", - ) as ac: - response = await ac.get("/get_favicon") - - assert response.status_code in [200, 404] - finally: - os.environ.pop("LITELLM_FAVICON_URL", None) + assert response.status_code == 307 + assert response.headers["location"] == "https://invalid.com/favicon.ico" diff --git a/tests/proxy_unit_tests/test_get_image.py b/tests/proxy_unit_tests/test_get_image.py index ad8c267275..57e472f86c 100644 --- a/tests/proxy_unit_tests/test_get_image.py +++ b/tests/proxy_unit_tests/test_get_image.py @@ -5,85 +5,48 @@ from unittest import mock # Standard path insertion sys.path.insert(0, os.path.abspath("../..")) -import pytest import httpx +import pytest from litellm.proxy.proxy_server import app @pytest.mark.asyncio -async def test_get_image_error_handling(): +async def test_get_image_redirects_remote_logo_without_server_fetch(monkeypatch): """ - Test that get_image handles network errors gracefully and doesn't hang. + Remote logo URLs should be loaded by the browser, not fetched by the proxy. """ - # Set an unreachable URL - os.environ["UI_LOGO_PATH"] = "http://invalid-url-12345.com/logo.jpg" + monkeypatch.setenv("UI_LOGO_PATH", "http://invalid-url-12345.com/logo.jpg") - # Clear cache - parent_dir = os.path.dirname( - os.path.dirname( - app.__file__ - if hasattr(app, "__file__") - else "litellm/proxy/proxy_server.py" - ) - ) - cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg") - if os.path.exists(cache_path): - os.remove(cache_path) - - # Mock AsyncHTTPHandler to simulate a timeout or connection error with mock.patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" ) as mock_get: - mock_get.side_effect = httpx.ConnectError("Network is unreachable") - async with httpx.AsyncClient( transport=httpx.ASGITransport(app=app), base_url="http://testserver" ) as ac: response = await ac.get("/get_image") - assert response.status_code == 200 - assert response.headers["content-type"] == "image/jpeg" + assert response.status_code == 307 + assert response.headers["location"] == "http://invalid-url-12345.com/logo.jpg" + mock_get.assert_not_called() @pytest.mark.asyncio -async def test_get_image_cache_logic(): +async def test_get_image_remote_logo_does_not_use_stale_cache(monkeypatch, tmp_path): """ - Test that once cached, get_image doesn't hit the network. + A stale pre-fix cache file should not mask a configured remote logo URL. """ - os.environ["UI_LOGO_PATH"] = "http://example.com/logo.jpg" - - # Clear cache - parent_dir = os.path.dirname( - os.path.dirname( - app.__file__ - if hasattr(app, "__file__") - else "litellm/proxy/proxy_server.py" - ) - ) - cache_path = os.path.join(parent_dir, "proxy", "cached_logo.jpg") - if os.path.exists(cache_path): - os.remove(cache_path) - - # Mock response - mock_response = mock.Mock() - mock_response.status_code = 200 - mock_response.content = b"fake image data" + monkeypatch.setenv("UI_LOGO_PATH", "http://example.com/logo.jpg") + monkeypatch.setenv("LITELLM_ASSETS_PATH", str(tmp_path)) + (tmp_path / "cached_logo.jpg").write_bytes(b"\xff\xd8\xff cached logo") with mock.patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get" ) as mock_get: - mock_get.return_value = mock_response - async with httpx.AsyncClient( transport=httpx.ASGITransport(app=app), base_url="http://testserver" ) as ac: - # First call - should hit download logic - response1 = await ac.get("/get_image") - assert response1.status_code == 200 - assert mock_get.call_count == 1 + response = await ac.get("/get_image") - # Second call - should hit cache - response2 = await ac.get("/get_image") - assert response2.status_code == 200 - # If cache works, mock_get shouldn't be called again - assert mock_get.call_count == 1 + assert response.status_code == 307 + assert response.headers["location"] == "http://example.com/logo.jpg" + mock_get.assert_not_called() diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 86bbc5170e..cdcdc89e7f 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2768,40 +2768,40 @@ async def test_update_config_success_callback_normalization(): import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import ConfigYAML - # Ensure feature is enabled and prisma_client is set - setattr(proxy_server, "store_model_in_db", True) setattr(proxy_server, "proxy_logging_obj", MagicMock()) + existing_litellm_settings = {"success_callback": ["langfuse"]} + + class FakeRow: + def __init__(self, name, value): + self.param_name = name + self.param_value = value + + upserted = {} + + async def fake_find_first(where=None): + if where and where.get("param_name") == "litellm_settings": + return FakeRow("litellm_settings", existing_litellm_settings) + return None + + async def fake_upsert(where=None, data=None): + upserted[where["param_name"]] = json.loads(data["update"]["param_value"]) + class MockPrisma: def __init__(self): self.db = MagicMock() self.db.litellm_config = MagicMock() - self.db.litellm_config.upsert = AsyncMock() - - # proxy_server.update_config expects this to be sync returning a dict - def jsonify_object(self, obj): - return obj + self.db.litellm_config.find_first = AsyncMock(side_effect=fake_find_first) + self.db.litellm_config.upsert = AsyncMock(side_effect=fake_upsert) setattr(proxy_server, "prisma_client", MockPrisma()) class MockProxyConfig: - def __init__(self): - self.saved_config = None - - async def get_config(self): - # Existing config has one lowercase callback already - return {"litellm_settings": {"success_callback": ["langfuse"]}} - - async def save_config(self, new_config: dict): - self.saved_config = new_config - async def add_deployment(self, prisma_client=None, proxy_logging_obj=None): return None - mock_proxy_config = MockProxyConfig() - setattr(proxy_server, "proxy_config", mock_proxy_config) + setattr(proxy_server, "proxy_config", MockProxyConfig()) - # Update config with mixed-case callbacks - expect normalization to lowercase config_update = ConfigYAML(litellm_settings={"success_callback": ["SQS", "sQs"]}) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -2810,9 +2810,10 @@ async def test_update_config_success_callback_normalization(): ) await proxy_server.update_config(config_update, user_api_key_dict=admin_user) - saved = mock_proxy_config.saved_config - assert saved is not None, "save_config was not called" - callbacks = saved["litellm_settings"]["success_callback"] + assert ( + "litellm_settings" in upserted + ), "litellm_config.upsert was not called for litellm_settings" + callbacks = upserted["litellm_settings"]["success_callback"] # Deduped and normalized assert "sqs" in callbacks diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index e51f81561a..543cabb6b4 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -268,7 +268,12 @@ async def test_aaauser_personal_budgets(key_ownership): test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache") - assert test_user_cache.get_cache(key=hash_token(user_key)) == valid_token + assert ( + test_user_cache.get_cache( + key=hash_token(user_key), model_type=UserAPIKeyAuth + ) + == valid_token + ) try: await user_api_key_auth(request=request, api_key="Bearer " + user_key) diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index b93502e815..0ce2dec9b5 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1110,7 +1110,7 @@ def test_initialize_skills_endpoints(): async def test_init_containers_api_endpoints(): """ Test that _init_containers_api_endpoints calls the original function - directly without model-based routing. + directly when there is no managed container ID (no embedded model_id). """ router = Router(model_list=[]) @@ -1127,3 +1127,112 @@ async def test_init_containers_api_endpoints(): custom_llm_provider="openai", name="Test Container" ) assert result == mock_response + + +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks(): + """ + Managed ``cntr_`` IDs embed ``model_id``; router should decode and use + ``_ageneric_api_call_with_fallbacks`` so deployment credentials apply. + """ + from litellm.responses.utils import ResponsesAPIRequestUtils + + router = Router( + model_list=[ + { + "model_name": "azure-router-model", + "litellm_params": { + "model": "azure/gpt-4", + "api_key": "fake-key", + "api_base": "https://westus.api.cognitive.microsoft.com", + }, + } + ] + ) + router._ageneric_api_call_with_fallbacks = AsyncMock() + + managed_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="azure-router-model", + container_id="cfile_upstream_abc", + ) + + await router._init_containers_api_endpoints( + original_function=AsyncMock(), + custom_llm_provider="openai", + container_id=managed_id, + file_id="cfile_xyz", + ) + + router._ageneric_api_call_with_fallbacks.assert_called_once() + call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs + assert call_kw["model"] == "azure-router-model" + assert call_kw["container_id"] == "cfile_upstream_abc" + assert call_kw["file_id"] == "cfile_xyz" + assert call_kw["custom_llm_provider"] == "azure" + + +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps(): + """ + Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a + streaming response had no router metadata). The router must still unwrap the + managed ID before calling the upstream provider — otherwise the raw + ``cntr_...`` token leaks downstream and the provider rejects it. + """ + from litellm.responses.utils import ResponsesAPIRequestUtils + + router = Router(model_list=[]) + mock_original_function = AsyncMock(return_value={"ok": True}) + + managed_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="openai", + model_id=None, + container_id="cfile_upstream_abc", + ) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + container_id=managed_id, + file_id="cfile_xyz", + ) + + mock_original_function.assert_called_once() + call_kw = mock_original_function.call_args.kwargs + assert call_kw["container_id"] == "cfile_upstream_abc" + assert call_kw["file_id"] == "cfile_xyz" + assert call_kw["custom_llm_provider"] == "openai" + + +@pytest.mark.asyncio +async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider(): + """ + A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with + an empty ``model_id`` (streaming events without router ``model_info.id``). + The router must still apply the decoded provider so the request routes to + the correct upstream — not stay on the default ``openai``. + """ + from litellm.responses.utils import ResponsesAPIRequestUtils + + router = Router(model_list=[]) + mock_original_function = AsyncMock(return_value={"ok": True}) + + managed_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id=None, + container_id="cfile_upstream_abc", + ) + + await router._init_containers_api_endpoints( + original_function=mock_original_function, + custom_llm_provider="openai", + container_id=managed_id, + file_id="cfile_xyz", + ) + + mock_original_function.assert_called_once() + call_kw = mock_original_function.call_args.kwargs + assert call_kw["container_id"] == "cfile_upstream_abc" + assert call_kw["file_id"] == "cfile_xyz" + assert call_kw["custom_llm_provider"] == "azure" diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 8e50217576..6477472620 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -1,5 +1,6 @@ import asyncio import time +import uuid from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -260,3 +261,72 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re f"Expected None when in_memory_cache is absent and Redis fails, got {result!r}. " "Returning the delta (1.0) would silently miscalculate rate-limit counters." ) + + +def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): + """ + Typical lazy startup (sync): DualCache runs with in-memory only, then Redis + becomes available and is attached. New writes must reach Redis; keys written + before attach are not backfilled. Optional default_redis_ttl is applied on attach. + """ + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) + + mock_redis = MagicMock() + mock_redis.set_cache = MagicMock() + mock_redis.async_set_cache = AsyncMock() + + key_before = f"before_attach_{uuid.uuid4()}" + val_before = {"phase": "memory_only"} + dual_cache.set_cache(key_before, val_before) + + assert in_memory.get_cache(key_before) == val_before + + dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) + assert dual_cache.redis_cache is mock_redis + assert dual_cache.default_redis_ttl == 99.0 + + mock_redis.set_cache.assert_not_called() + + key_after = f"after_attach_{uuid.uuid4()}" + val_after = {"phase": "memory_and_redis"} + dual_cache.set_cache(key_after, val_after) + mock_redis.set_cache.assert_called_once() + assert mock_redis.set_cache.call_args[0][:2] == (key_after, val_after) + + assert in_memory.get_cache(key_after) == val_after + + +@pytest.mark.asyncio +async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async(): + """ + Typical lazy startup (async): DualCache runs with in-memory only, then Redis + becomes available and is attached. New writes must reach Redis; keys written + before attach are not backfilled. Optional default_redis_ttl is applied on attach. + """ + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory, redis_cache=None) + + mock_redis = MagicMock() + mock_redis.set_cache = MagicMock() + mock_redis.async_set_cache = AsyncMock() + + key_before = f"before_attach_{uuid.uuid4()}" + val_before = {"phase": "memory_only"} + await dual_cache.async_set_cache(key_before, val_before) + + assert in_memory.get_cache(key_before) == val_before + + dual_cache.attach_redis_cache(mock_redis, default_redis_ttl=99.0) + assert dual_cache.redis_cache is mock_redis + assert dual_cache.default_redis_ttl == 99.0 + + mock_redis.async_set_cache.assert_not_called() + + key_after = f"after_attach_{uuid.uuid4()}" + val_after = {"phase": "memory_and_redis"} + await dual_cache.async_set_cache(key_after, val_after) + mock_redis.async_set_cache.assert_called_once() + assert mock_redis.async_set_cache.call_args[0][:2] == (key_after, val_after) + + assert in_memory.get_cache(key_after) == val_after diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index b39eb42821..78192400fb 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping ) +@pytest.mark.asyncio +async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl( + monkeypatch, redis_no_ping +): + """With refresh_ttl=True, every increment should call expire() to bump + the TTL, even when the key already has a TTL (counter-style use).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment( + key="spend:team_member:u:t", value=0.05, refresh_ttl=True + ) + + mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60) + + +@pytest.mark.asyncio +async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( + monkeypatch, redis_no_ping +): + """Default (refresh_ttl=False) preserves window-style semantics: TTL is + set only on first creation, never refreshed (used by rate-limit windows).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment(key="rate_limit:window", value=1) + + mock_redis_instance.expire.assert_not_awaited() + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/test_litellm/containers/test_azure_container_transformation.py index a46046b318..32caaaffef 100644 --- a/tests/test_litellm/containers/test_azure_container_transformation.py +++ b/tests/test_litellm/containers/test_azure_container_transformation.py @@ -11,6 +11,7 @@ sys.path.insert(0, os.path.abspath("../../../")) import litellm from litellm.llms.azure.containers.transformation import AzureContainerConfig from litellm.llms.base_llm.containers.transformation import BaseContainerConfig +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.containers.main import ( ContainerFileListResponse, ContainerListResponse, @@ -323,6 +324,29 @@ class TestAzureContainerConfig: assert url_fc == expected_fc assert url_fc.index("/content") < url_fc.index("?") + def test_transform_requests_encode_path_ids_before_query_string(self): + from litellm.types.router import GenericLiteLLMParams + + api_base = ( + "https://my-resource.openai.azure.com/openai/v1/containers" + "?api-version=v1" + ) + + url, _ = self.config.transform_container_file_content_request( + container_id="../../other", + file_id="file?download=1#frag", + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + expected_url = ( + "https://my-resource.openai.azure.com/openai/v1/containers/" + "..%2F..%2Fother/files/file%3Fdownload%3D1%23frag/content" + "?api-version=v1" + ) + assert url == expected_url + def test_provider_config_manager_returns_azure_config(self): from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -518,3 +542,206 @@ class TestAzureContainerKnownFailureRegressions: c2 = _get_container_provider_config("azure_text") assert type(c1) is type(c2) assert isinstance(c1, AzureContainerConfig) + + @pytest.mark.asyncio + async def test_proxy_process_request_preserves_managed_container_id( + self, monkeypatch + ): + from starlette.requests import Request + + from litellm.proxy.container_endpoints import handler_factory + + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="model_abc123", + container_id="cntr_123", + ) + captured = {} + + async def _mock_base_process_llm_request( + self, + request, + fastapi_response, + user_api_key_dict, + route_type, + **kwargs, + ): + captured["data"] = self.data + captured["route_type"] = route_type + return {"id": "cfile_abc"} + + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + monkeypatch.setattr( + ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + _mock_base_process_llm_request, + ) + + request = Request( + { + "type": "http", + "method": "GET", + "path": "/v1/containers/id/files/id/content", + "headers": [], + "query_string": b"", + } + ) + fastapi_response = MagicMock() + + await handler_factory._process_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=MagicMock(), + route_type="alist_container_files", + path_params={"container_id": encoded_id}, + ) + + assert captured["route_type"] == "alist_container_files" + assert captured["data"]["container_id"] == encoded_id + assert captured["data"]["custom_llm_provider"] == "openai" + assert "model_id" not in captured["data"] + assert "api_base" not in captured["data"] + + @pytest.mark.asyncio + async def test_regression_binary_file_request_routes_through_proxy_processor( + self, monkeypatch + ): + from fastapi import Response + from starlette.requests import Request + + from litellm.proxy.container_endpoints import handler_factory + + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="model_abc123", + container_id="cntr_123", + ) + captured = {} + + async def _mock_base_process_llm_request( + self, + request, + fastapi_response, + user_api_key_dict, + route_type, + **kwargs, + ): + captured["data"] = self.data + captured["route_type"] = route_type + fastapi_response.headers["x-litellm-call-id"] = "call-123" + return b"csv-bytes" + + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + + monkeypatch.setattr( + ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + _mock_base_process_llm_request, + ) + + request = Request( + { + "type": "http", + "method": "GET", + "path": "/v1/containers/id/files/id/content", + "headers": [], + "query_string": b"", + } + ) + fastapi_response = Response() + + response = await handler_factory._process_binary_request( + request=request, + fastapi_response=fastapi_response, + container_id=encoded_id, + file_id="cfile_abc", + user_api_key_dict=MagicMock(), + ) + + assert captured["route_type"] == "aretrieve_container_file_content" + assert captured["data"]["container_id"] == encoded_id + assert captured["data"]["file_id"] == "cfile_abc" + assert captured["data"]["custom_llm_provider"] == "openai" + assert response.status_code == 200 + assert response.body == b"csv-bytes" + assert response.headers["x-litellm-call-id"] == "call-123" + + @pytest.mark.asyncio + async def test_regression_multipart_upload_request_uses_provider_from_managed_id( + self, monkeypatch + ): + from starlette.requests import Request + + from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + ) + from litellm.proxy.common_utils import http_parsing_utils + from litellm.proxy.container_endpoints import handler_factory + + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="model_abc123", + container_id="cntr_123", + ) + captured = {} + + async def _mock_get_form_data(request): + return {"file": "ignored"} + + async def _mock_convert_upload_files_to_file_data(form_data): + return {"file": [("data.csv", b"csv-bytes", "text/csv")]} + + async def _mock_base_process_llm_request( + self, + request, + fastapi_response, + user_api_key_dict, + route_type, + **kwargs, + ): + captured["data"] = self.data + captured["route_type"] = route_type + return {"id": "cfile_abc"} + + monkeypatch.setattr( + http_parsing_utils, + "get_form_data", + _mock_get_form_data, + ) + monkeypatch.setattr( + http_parsing_utils, + "convert_upload_files_to_file_data", + _mock_convert_upload_files_to_file_data, + ) + monkeypatch.setattr( + ProxyBaseLLMRequestProcessing, + "base_process_llm_request", + _mock_base_process_llm_request, + ) + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/containers/id/files", + "headers": [], + "query_string": b"", + } + ) + + await handler_factory._process_multipart_upload_request( + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=MagicMock(), + route_type="aupload_container_file", + container_id=encoded_id, + ) + + assert captured["route_type"] == "aupload_container_file" + assert captured["data"]["container_id"] == encoded_id + assert captured["data"]["custom_llm_provider"] == "openai" diff --git a/tests/test_litellm/containers/test_container_handler_url.py b/tests/test_litellm/containers/test_container_handler_url.py new file mode 100644 index 0000000000..1927580068 --- /dev/null +++ b/tests/test_litellm/containers/test_container_handler_url.py @@ -0,0 +1,28 @@ +import pytest + +from litellm.llms.custom_httpx.container_handler import _build_url + + +def test_build_url_encodes_path_params_and_preserves_query(): + url = _build_url( + api_base="https://example.com/v1/containers?api-version=v1", + path_template="/containers/{container_id}/files/{file_id}/content", + path_params={ + "container_id": "../../containers/other", + "file_id": "file?download=1#frag", + }, + ) + + assert ( + url + == "https://example.com/v1/containers/..%2F..%2Fcontainers%2Fother/files/file%3Fdownload%3D1%23frag/content?api-version=v1" + ) + + +def test_build_url_rejects_dot_segment_path_param(): + with pytest.raises(ValueError, match="container_id cannot be a dot path segment"): + _build_url( + api_base="https://example.com/v1/containers", + path_template="/containers/{container_id}", + path_params={"container_id": ".."}, + ) diff --git a/tests/test_litellm/containers/test_container_transformation.py b/tests/test_litellm/containers/test_container_transformation.py index 47b5b8dc56..555fe7773f 100644 --- a/tests/test_litellm/containers/test_container_transformation.py +++ b/tests/test_litellm/containers/test_container_transformation.py @@ -230,6 +230,23 @@ class TestOpenAIContainerTransformation: assert url == f"{api_base}/{container_id}" assert params == {} # No query params for retrieve + def test_transform_container_retrieve_request_encodes_path_traversal(self): + """Test container IDs are treated as a single upstream path segment.""" + api_base = "https://api.openai.com/v1/containers" + + url, params = self.config.transform_container_retrieve_request( + container_id="../../vector_stores?x=1#frag", + api_base=api_base, + litellm_params={}, + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/containers/..%2F..%2Fvector_stores%3Fx%3D1%23frag" + ) + assert params == {} + def test_transform_container_retrieve_response(self): """Test container retrieve response transformation.""" # Mock HTTP response diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/test_litellm/integrations/arize/test_arize_phoenix.py index 01f85af262..4a2eab29e8 100644 --- a/tests/test_litellm/integrations/arize/test_arize_phoenix.py +++ b/tests/test_litellm/integrations/arize/test_arize_phoenix.py @@ -280,3 +280,55 @@ class TestDynamicProjectNameOnSpan: if __name__ == "__main__": unittest.main() + + +# --- Security: SSRF via prompt_version_id path traversal --- + + +def test_arize_phoenix_client_sanitize_id_rejects_traversal(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + # dotdot without slashes + with pytest.raises(ValueError, match="path traversal"): + _sanitize_id("..something") + # full traversal (slash caught first) + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("../../projects") + + +def test_arize_phoenix_client_sanitize_id_rejects_slash(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("valid/extra") + + +def test_arize_phoenix_client_sanitize_id_rejects_fragment(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("abc#suffix") + + +def test_arize_phoenix_client_sanitize_id_rejects_query(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("abc?x=1") + + +def test_arize_phoenix_client_sanitize_id_allows_uuid(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + uid = "550e8400-e29b-41d4-a716-446655440000" + assert _sanitize_id(uid) == uid + + +def test_arize_phoenix_client_get_prompt_version_rejects_traversal(): + from litellm.integrations.arize.arize_phoenix_client import ArizePhoenixClient + + client = ArizePhoenixClient( + api_key="test-key", api_base="https://app.phoenix.arize.com" + ) + with pytest.raises(ValueError, match="disallowed characters"): + client.get_prompt_version("../../projects") diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py index a7b2d362ed..46cd1d6e76 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py @@ -11,6 +11,7 @@ sys.path.insert( import litellm from litellm.integrations.bitbucket import BitBucketPromptManager +from litellm.integrations.bitbucket.bitbucket_client import _sanitize_file_path @patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient") @@ -370,3 +371,45 @@ def test_bitbucket_prompt_manager_list_templates(mock_client_class): templates = manager.prompt_manager.list_templates() assert isinstance(templates, list) assert "test_prompt" in templates + + +# --- Security: path traversal / SSRF --- + + +def test_sanitize_file_path_rejects_traversal(): + with pytest.raises(ValueError, match="path traversal"): + _sanitize_file_path("../../etc/passwd") + + +def test_sanitize_file_path_rejects_fragment(): + with pytest.raises(ValueError, match="URL special characters"): + _sanitize_file_path("secret#.prompt") + + +def test_sanitize_file_path_rejects_query(): + with pytest.raises(ValueError, match="URL special characters"): + _sanitize_file_path("secret?.prompt") + + +def test_sanitize_file_path_encodes_special_chars(): + result = _sanitize_file_path("prompts/my prompt.prompt") + assert result == "prompts/my%20prompt.prompt" + + +def test_sanitize_file_path_allows_normal_paths(): + assert _sanitize_file_path("prompts/my-prompt") == "prompts/my-prompt" + assert _sanitize_file_path("simple") == "simple" + + +def test_bitbucket_client_rejects_traversal_in_get_file_content(): + from litellm.integrations.bitbucket.bitbucket_client import BitBucketClient + + client = BitBucketClient( + { + "workspace": "ws", + "repository": "repo", + "access_token": "tok", + } + ) + with pytest.raises(ValueError, match="path traversal"): + client.get_file_content("../../admin/credentials") diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py index 37dc491c26..758ff3ea38 100644 --- a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py +++ b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py @@ -147,6 +147,21 @@ class TestInteractionOperationUrls: assert "secret-key" not in url assert expected_suffix in url + def test_interaction_id_is_encoded_as_one_path_segment(self, config): + with patch(_PATCH_GET_API_KEY, return_value="secret-key"): + url, params = config.transform_cancel_interaction_request( + interaction_id="../../interactions/other?x=1#frag", + api_base="https://generativelanguage.googleapis.com", + litellm_params=GenericLiteLLMParams(api_key="secret-key"), + headers={}, + ) + + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/interactions/..%2F..%2Finteractions%2Fother%3Fx%3D1%23frag:cancel" + ) + assert params == {} + def test_get_interaction_raises_without_key(self, config): with patch(_PATCH_GET_API_KEY, return_value=None): with pytest.raises(ValueError, match="Google API key is required"): diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index d424cd8599..27a3ddb553 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -9,8 +9,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, BedrockConverseMessagesProcessor, BedrockImageProcessor, - anthropic_messages_pt, + _bedrock_converse_messages_pt, _convert_to_bedrock_tool_call_invoke, + _convert_to_bedrock_tool_call_result, + anthropic_messages_pt, convert_to_gemini_tool_call_result, ollama_pt, sanitize_messages_for_tool_calling, @@ -2485,10 +2487,6 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): inside the tool_result content. Reuses anthropic_process_openai_file_message, which already handles this for user messages. """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) - pdf_b64 = "JVBERi0xLjQKJeLjz9MK" message = { "tool_call_id": "toolu_pdf_1", @@ -2505,157 +2503,105 @@ def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): ], } - result = convert_to_anthropic_tool_result(message) + result = _convert_to_bedrock_tool_call_result(message) - assert result["type"] == "tool_result" - assert result["tool_use_id"] == "toolu_pdf_1" - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "document" - assert block["source"]["type"] == "base64" - assert block["source"]["media_type"] == "application/pdf" - assert block["source"]["data"] == pdf_b64 + tool_result = result["toolResult"] + assert len(tool_result["content"]) == 1 + assert "document" in tool_result["content"][0] + assert tool_result["content"][0]["document"]["format"] == "pdf" + assert tool_result["content"][0]["document"]["source"]["bytes"] == pdf_b64 -def test_convert_to_anthropic_tool_result_image_url_pdf_data_uri_becomes_document(): - """ - Regression: a PDF sent as an `image_url` data URI on the tool-result path - must translate to an Anthropic document block (not an image block — Anthropic - rejects image blocks whose media_type is a non-image like application/pdf). - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) +def test_bedrock_converse_messages_pt_document_various_formats(): + """Test that various document media types produce the correct format value.""" + test_cases = [ + ("application/pdf", "pdf"), + ("text/csv", "csv"), + ("text/html", "html"), + ("text/plain", "txt"), + ("text/markdown", "md"), + ( + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "docx", + ), + ] - pdf_b64 = "JVBERi0xLjQKJeLjz9MK" - message = { - "tool_call_id": "toolu_pdf_img_1", - "role": "tool", - "name": "fetch_document", - "content": [ + for media_type, expected_format in test_cases: + messages = [ { - "type": "image_url", - "image_url": { - "url": f"data:application/pdf;base64,{pdf_b64}", + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": media_type, + "data": "dGVzdA==", + }, + }, + ], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + doc_block = result[0]["content"][0] + assert doc_block["document"]["format"] == expected_format, ( + f"Expected format '{expected_format}' for media_type '{media_type}', " + f"got '{doc_block['document']['format']}'" + ) + + +def test_bedrock_converse_messages_pt_document_deterministic_name(): + """Test that the same document data always produces the same name.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, }, - }, - ], - } + ], + } + ] - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "document" - assert block["source"]["media_type"] == "application/pdf" - assert block["source"]["data"] == pdf_b64 - - -def test_convert_to_anthropic_tool_result_image_url_unsupported_mime_stays_image_path(): - """ - An `image_url` data URI whose mime is neither application/pdf nor text/plain - (e.g. application/json) must NOT be routed through the document path. Anthropic - only accepts application/pdf and text/plain as base64 document media_types — - anything else would produce a document block the API rejects. The old - (pre-fix) behavior was to wrap such data as an image block, which also - fails but stays on the image code path; preserve that failure mode rather - than switching to a document path that is equally broken. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, + result1 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + result2 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" ) - message = { - "tool_call_id": "toolu_json_1", - "role": "tool", - "name": "fetch_json", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:application/json;base64,eyJrIjoidiJ9", + name1 = result1[0]["content"][0]["document"]["name"] + name2 = result2[0]["content"][0]["document"]["name"] + assert name1 == name2 + + +def test_bedrock_converse_messages_pt_document_rejects_url_source(): + """Test that a URL-type document source raises a clear error instead of KeyError.""" + messages = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "url", + "url": "https://example.com/doc.pdf", + }, }, - }, - ], - } + ], + } + ] - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "image", ( - f"unsupported mime {block.get('source', {}).get('media_type')!r} " - f"should not be routed to document path; got {block}" - ) - - -def test_convert_to_anthropic_tool_result_image_url_text_plain_data_uri_becomes_document(): - """ - text/plain is one of the two mimes Anthropic accepts as a base64 document - media_type. Confirm it routes through the document path so tightening the - gate to {application/pdf, text/plain} (not "application/*") covers both. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) - - txt_b64 = "aGVsbG8=" # "hello" - message = { - "tool_call_id": "toolu_txt_1", - "role": "tool", - "name": "fetch_text", - "content": [ - { - "type": "image_url", - "image_url": { - "url": f"data:text/plain;base64,{txt_b64}", - }, - }, - ], - } - - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "document" - assert block["source"]["media_type"] == "text/plain" - assert block["source"]["data"] == txt_b64 - - -def test_convert_to_anthropic_tool_result_image_url_png_still_becomes_image(): - """ - Regression: image_url with a real image mime type must continue to translate - to an Anthropic image block. Locks in existing behavior after the - data-URI-mime-type branching for PDFs. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - convert_to_anthropic_tool_result, - ) - - png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg==" - message = { - "tool_call_id": "toolu_png_1", - "role": "tool", - "name": "fetch_image", - "content": [ - { - "type": "image_url", - "image_url": { - "url": f"data:image/png;base64,{png_b64}", - }, - }, - ], - } - - result = convert_to_anthropic_tool_result(message) - - content = result["content"] - assert isinstance(content, list) and len(content) == 1 - block = content[0] - assert block["type"] == "image" - assert block["source"]["media_type"] == "image/png" + with pytest.raises(ValueError, match="only supports base64-encoded"): + _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index d6281703a0..49d3c51e34 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -878,6 +878,39 @@ def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging): assert "invalid maxOutputTokens" in str(excinfo.value) +@pytest.mark.asyncio +async def test_async_streaming_read_timeout_triggers_midstream_fallback( + logging_obj: Logging, +): + """A mid-stream httpx.ReadTimeout must wrap into MidStreamFallbackError so + the Router's FallbackStreamWrapper can switch to a fallback model. + + Previously __anext__ caught httpx.TimeoutException and re-raised it raw, + which bypassed _handle_stream_fallback_error and prevented stream_timeout + from triggering fallbacks the way connection-phase timeout does. + """ + import httpx + + from litellm.exceptions import MidStreamFallbackError + + async def _raise_read_timeout(**kwargs): + raise httpx.ReadTimeout("Timeout on reading data from socket") + + response = CustomStreamWrapper( + completion_stream=None, + model="gpt-4", + logging_obj=logging_obj, + custom_llm_provider="openai", + make_call=_raise_read_timeout, + ) + + with pytest.raises(MidStreamFallbackError) as excinfo: + await response.__anext__() + + assert excinfo.value.is_pre_first_chunk is True + assert isinstance(excinfo.value.original_exception, Exception) + + def test_streaming_handler_with_created_time_propagation( initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging ): diff --git a/tests/test_litellm/litellm_core_utils/test_url_utils.py b/tests/test_litellm/litellm_core_utils/test_url_utils.py index efe3f2b67a..cef09f3f2b 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -4,7 +4,14 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils -from litellm.litellm_core_utils.url_utils import SSRFError, _is_blocked_ip, validate_url +from litellm.litellm_core_utils.url_utils import ( + SSRFError, + _is_blocked_ip, + assert_same_origin, + encode_url_path_segment, + encode_url_path_segments, + validate_url, +) @pytest.fixture @@ -80,6 +87,28 @@ class TestIsBlockedIp: assert _is_blocked_ip("::ffff:168.63.129.16") is True +class TestEncodeUrlPathSegment: + def test_encodes_path_delimiters_and_query_markers(self): + encoded = encode_url_path_segment("../../v1/files?limit=1#frag") + + assert encoded == "..%2F..%2Fv1%2Ffiles%3Flimit%3D1%23frag" + + def test_encodes_path_segments_without_collapsing_valid_model_paths(self): + encoded = encode_url_path_segments("@cf/meta/model?debug=1") + + assert encoded == "%40cf/meta/model%3Fdebug%3D1" + + @pytest.mark.parametrize("value", ["", ".", "..", None]) + def test_rejects_empty_and_dot_segments(self, value): + with pytest.raises(ValueError): + encode_url_path_segment(value, field_name="resource_id") + + @pytest.mark.parametrize("value", ["../model", "model/../other", "/model"]) + def test_rejects_dot_segments_in_multi_segment_paths(self, value): + with pytest.raises(ValueError): + encode_url_path_segments(value, field_name="model") + + class TestValidateUrl: def test_blocks_loopback(self): with pytest.raises(SSRFError): @@ -436,3 +465,73 @@ class TestProviderUrlDestinationAllowlist: "https://trusted.example:99999/v1/chat/completions", ["trusted.example"], ) + + +# ── assert_same_origin ──────────────────────────────────────────────────────── + + +def test_assert_same_origin_matches_scheme_host_port(): + """A polling URL on the same scheme + host + port as the api_base + passes — the upstream is trusted; the URL it returned points back at + the same upstream.""" + assert_same_origin( + "https://api.example.com/v1/operations/abc", + "https://api.example.com/v1/generate", + ) + + +def test_assert_same_origin_treats_default_ports_as_explicit(): + """``https://x/`` and ``https://x:443/`` are the same origin.""" + assert_same_origin("https://api.example.com/poll", "https://api.example.com:443/") + assert_same_origin("https://api.example.com:443/poll", "https://api.example.com/") + assert_same_origin("http://api.example.com/poll", "http://api.example.com:80/") + + +def test_assert_same_origin_rejects_different_host(): + with pytest.raises(SSRFError, match="host"): + assert_same_origin( + "https://attacker.example.com/poll", + "https://api.example.com/generate", + ) + + +def test_assert_same_origin_rejects_different_scheme(): + with pytest.raises(SSRFError, match="scheme"): + assert_same_origin( + "http://api.example.com/poll", "https://api.example.com/generate" + ) + + +def test_assert_same_origin_rejects_different_port(): + with pytest.raises(SSRFError, match="port"): + assert_same_origin( + "https://api.example.com:8443/poll", "https://api.example.com/generate" + ) + + +def test_assert_same_origin_rejects_non_http_scheme(): + """``file://`` polling URLs are rejected outright — the upstream + should never return a non-HTTP scheme.""" + with pytest.raises(SSRFError, match="scheme"): + assert_same_origin("file:///etc/passwd", "https://api.example.com/") + + +def test_assert_same_origin_case_insensitive_host(): + assert_same_origin( + "https://API.example.com/poll", "https://api.example.com/generate" + ) + + +def test_assert_same_origin_error_message_does_not_leak_hostnames(): + """Greptile P2: in the SSRF threat model the caller is the attacker. + The error message must not echo the operator's expected host or the + attacker-supplied candidate host back to the caller — only identify + *which* component mismatched.""" + with pytest.raises(SSRFError) as exc: + assert_same_origin( + "https://attacker.example.com:1234/poll", + "https://api.internal-corp.example/generate", + ) + detail = str(exc.value) + assert "attacker.example.com" not in detail + assert "api.internal-corp.example" not in detail diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index e1fe4befb5..e8da50f0ec 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -8,6 +8,7 @@ sys.path.insert( ) # Adds the parent directory to the system path from unittest.mock import MagicMock, patch +from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, @@ -38,6 +39,39 @@ def test_response_format_transformation_unit_test(): print(result) +def test_anthropic_json_mode_non_streaming_mixed_internal_and_user_tools(): + """Non-streaming + response_format: internal json tool must not require len(tool_calls)==1.""" + config = AnthropicConfig() + tool_calls = [ + { + "id": "toolu_json", + "type": "function", + "function": { + "name": RESPONSE_FORMAT_TOOL_NAME, + "arguments": '{"values": {"answer": 42}}', + }, + "index": 0, + }, + { + "id": "toolu_user", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NY"}', + }, + "index": 1, + }, + ] + replacement, filtered, extra = config._resolve_json_mode_non_streaming( + json_mode=True, + tool_calls=tool_calls, + ) + assert replacement is None + assert len(filtered) == 1 + assert filtered[0]["function"]["name"] == "get_weather" + assert extra == '{"answer": 42}' + + def test_calculate_usage(): """ Do not include cache_creation_input_tokens in the prompt_tokens diff --git a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py index e9509be9e1..9fc4981510 100644 --- a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py @@ -180,6 +180,19 @@ class TestAnthropicFilesConfig: assert url == "https://custom.api.com/v1/files/file-abc123" assert params == {} + def test_transform_retrieve_file_request_encodes_path_traversal(self): + url, params = self.config.transform_retrieve_file_request( + file_id="../../v1/messages/batches?limit=1#frag", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == f"{ANTHROPIC_FILES_API_BASE}/v1/files/..%2F..%2Fv1%2Fmessages%2Fbatches%3Flimit%3D1%23frag" + ) + assert params == {} + def test_transform_retrieve_file_response(self): mock_response = Mock(spec=httpx.Response) mock_response.json.return_value = { @@ -296,6 +309,14 @@ class TestAnthropicFilesConfig: assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/file-abc123/content" assert params == {} + def test_transform_file_content_request_rejects_dot_segment(self): + with pytest.raises(ValueError, match="file_id cannot be a dot path segment"): + self.config.transform_file_content_request( + file_content_request={"file_id": ".."}, + optional_params={}, + litellm_params={}, + ) + def test_transform_file_content_response(self): mock_response = Mock(spec=httpx.Response) result = self.config.transform_file_content_response( diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 9519b7c8a5..a4bd14d69f 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -96,6 +96,28 @@ def test_get_complete_url(): assert result == expected +@pytest.mark.serial +def test_response_id_path_requests_encode_response_id(): + config = AzureOpenAIResponsesAPIConfig() + api_base = ( + "https://litellm8397336933.openai.azure.com/openai/responses" + "?api-version=2024-05-01-preview" + ) + + url, params = config.transform_cancel_response_api_request( + response_id="../../responses/other?x=1#frag", + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://litellm8397336933.openai.azure.com/openai/responses/..%2F..%2Fresponses%2Fother%3Fx%3D1%23frag/cancel?api-version=2024-05-01-preview" + ) + assert params == {} + + @pytest.mark.serial def test_azure_o_series_responses_api_supported_params(): """Test that Azure OpenAI O-series responses API excludes temperature from supported parameters.""" diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py new file mode 100644 index 0000000000..f65573b7ae --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py @@ -0,0 +1,57 @@ +import pytest + +from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler + + +def test_should_encode_thread_id_in_azure_ai_agent_urls(): + handler = AzureAIAgentsHandler() + + assert ( + handler._build_messages_url( + "https://example.services.ai.azure.com/api/projects/proj", + "../../threads/other?x=1#frag", + "2024-05-01-preview", + ) + == "https://example.services.ai.azure.com/api/projects/proj/threads/..%2F..%2Fthreads%2Fother%3Fx%3D1%23frag/messages?api-version=2024-05-01-preview" + ) + assert ( + handler._build_runs_url( + "https://example.services.ai.azure.com/api/projects/proj", + "thread/abc", + "2024-05-01-preview", + ) + == "https://example.services.ai.azure.com/api/projects/proj/threads/thread%2Fabc/runs?api-version=2024-05-01-preview" + ) + + +def test_should_encode_thread_and_run_ids_in_azure_ai_agent_status_url(): + handler = AzureAIAgentsHandler() + + assert ( + handler._build_run_status_url( + "https://example.services.ai.azure.com/api/projects/proj", + "thread/abc", + "../runs/other#frag", + "2024-05-01-preview", + ) + == "https://example.services.ai.azure.com/api/projects/proj/threads/thread%2Fabc/runs/..%2Fruns%2Fother%23frag?api-version=2024-05-01-preview" + ) + + +def test_should_reject_dot_segments_in_azure_ai_agent_urls(): + handler = AzureAIAgentsHandler() + + with pytest.raises(ValueError, match="thread_id cannot be a dot path segment"): + handler._build_messages_url( + "https://example.services.ai.azure.com/api/projects/proj", + "..", + "2024-05-01-preview", + ) + + with pytest.raises(ValueError, match="run_id cannot be a dot path segment"): + handler._build_run_status_url( + "https://example.services.ai.azure.com/api/projects/proj", + "thread_123", + "..", + "2024-05-01-preview", + ) diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py new file mode 100644 index 0000000000..e638be68ec --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -0,0 +1,33 @@ +import pytest + +from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, +) + + +def test_should_encode_azure_document_intelligence_model_id(): + config = AzureDocumentIntelligenceOCRConfig() + + url = config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout?x=1#frag", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout%3Fx%3D1%23frag:analyze?api-version=2024-11-30" + ) + + +def test_should_reject_dot_segment_azure_document_intelligence_model_id(): + config = AzureDocumentIntelligenceOCRConfig() + + with pytest.raises(ValueError, match="model_id cannot be a dot path segment"): + config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="azure_ai/doc-intelligence/..", + optional_params={}, + litellm_params={}, + ) diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index eac022ec23..0de3f833a3 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -167,3 +167,24 @@ def test_tool_name_sanitization(): ] # Should be sanitized: only [a-zA-Z0-9_] assert tool_name == "my_tool_" + + +def test_count_tokens_endpoint_encodes_model_id(monkeypatch): + """Test model IDs are treated as a single Bedrock path segment.""" + config = BedrockCountTokensConfig() + + monkeypatch.setattr( + config, + "get_runtime_endpoint", + lambda **kwargs: ("https://bedrock-runtime.us-east-1.amazonaws.com", None), + ) + + endpoint = config.get_bedrock_count_tokens_endpoint( + model="bedrock/../../model/other?x=1#frag", + aws_region_name="us-east-1", + ) + + assert ( + endpoint + == "https://bedrock-runtime.us-east-1.amazonaws.com/model/..%2F..%2Fmodel%2Fother%3Fx%3D1%23frag/count-tokens" + ) diff --git a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py b/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py index ffd3526698..9e526e4778 100644 --- a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py +++ b/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py @@ -1,12 +1,9 @@ import base64 -import json import os import sys -from litellm._uuid import uuid -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -15,12 +12,7 @@ sys.path.insert( from litellm.llms.bedrock.chat.invoke_agent.transformation import ( AmazonInvokeAgentConfig, ) -from litellm.types.llms.bedrock_invoke_agents import ( - InvokeAgentEvent, - InvokeAgentEventHeaders, - InvokeAgentUsage, -) -from litellm.types.utils import Message, ModelResponse, Usage +from litellm.types.utils import ModelResponse class TestAmazonInvokeAgentConfig: @@ -270,3 +262,28 @@ class TestAmazonInvokeAgentConfig: "https://bedrock-runtime.us-east-1.amazonaws.com/agents/L1RT58GYRW/agentAliases/MFPSBCXYTW/sessions" in result ) + + @patch( + "litellm.llms.bedrock.chat.invoke_agent.transformation.convert_content_list_to_str" + ) + @patch.object(AmazonInvokeAgentConfig, "get_runtime_endpoint") + @patch.object(AmazonInvokeAgentConfig, "_get_aws_region_name") + def test_get_complete_url_encodes_session_id( + self, mock_region, mock_endpoint, mock_convert, config + ): + """Test get_complete_url encodes session ID path segment.""" + mock_endpoint.return_value = ( + "https://bedrock-runtime.us-east-1.amazonaws.com", + None, + ) + mock_region.return_value = "us-east-1" + + result = config.get_complete_url( + api_base=None, + api_key=None, + model="agent/L1RT58GYRW/MFPSBCXYTW", + optional_params={"sessionID": "../../sessions/other?x=1#frag"}, + litellm_params={}, + ) + + assert "sessions/..%2F..%2Fsessions%2Fother%3Fx%3D1%23frag/text" in result diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index d60d0487d0..7b04efa17d 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -28,6 +28,28 @@ def test_transform_search_request(): assert body["retrievalQuery"].get("text") == "hello" +def test_transform_search_request_encodes_vector_store_id(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + + url, body = config.transform_search_vector_store_request( + vector_store_id="../../knowledgebases/other?x=1#frag", + query="hello", + vector_store_search_optional_params={}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + extra_body=None, + ) + + assert ( + url + == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother%3Fx%3D1%23frag/retrieve" + ) + assert body["retrievalQuery"].get("text") == "hello" + + def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): config = BedrockVectorStoreConfig() mock_log = MagicMock() diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py index dd388fedc6..2f8cc5484b 100644 --- a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py @@ -87,6 +87,29 @@ class TestBytezChatConfig: assert response.choices[0].message.content == output_content # type: ignore + def test_get_complete_url_encodes_model_path_segment(self): + config = BytezChatConfig() + + assert ( + config.get_complete_url( + api_base=API_BASE, + api_key=TEST_API_KEY, + model="google/gemma?x=1#frag", + optional_params={}, + litellm_params={}, + ) + == f"{API_BASE}/google/gemma%3Fx%3D1%23frag" + ) + + with pytest.raises(ValueError, match="dot path segment"): + config.get_complete_url( + api_base=API_BASE, + api_key=TEST_API_KEY, + model="../../models/other", + optional_params={}, + litellm_params={}, + ) + def test_bytez_messages_adaptation(self): cases = [ dict( diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py new file mode 100644 index 0000000000..cecb6024de --- /dev/null +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py @@ -0,0 +1,27 @@ +import pytest + +from litellm.llms.cloudflare.chat.transformation import CloudflareChatConfig + + +def test_get_complete_url_encodes_model_path_segment(): + config = CloudflareChatConfig() + + assert ( + config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="@cf/meta/llama?x=1#frag", + optional_params={}, + litellm_params={}, + ) + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/%40cf/meta/llama%3Fx%3D1%23frag" + ) + + with pytest.raises(ValueError, match="dot path segment"): + config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="../../accounts/other", + optional_params={}, + litellm_params={}, + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 752b5ff090..b846cd600f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,3 +1,4 @@ +import asyncio import os import sys from unittest.mock import AsyncMock, Mock, patch @@ -8,6 +9,8 @@ import pytest sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, _google_genai_streaming_hidden_params, @@ -103,7 +106,9 @@ def test_fingerprint_agentic_tools_is_deterministic(): tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]} tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]} - assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b) + assert handler._fingerprint_agentic_tools( + tools_a + ) == handler._fingerprint_agentic_tools(tools_b) @pytest.mark.asyncio @@ -350,3 +355,70 @@ def test_google_genai_streaming_hidden_params_model_info_and_router_fallback(): response_headers=httpx.Headers({}), ) assert from_router["model_id"] == "router-model-id" + + +def _build_delete_response_mock(captured: dict): + """Returns a fake httpx delete that records its kwargs.""" + + def _response() -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{"id": "resp_x", "object": "response", "deleted": true}', + request=httpx.Request(method="DELETE", url="https://test.openai.azure.com"), + ) + + async def fake_async_delete(*args, **kwargs): + captured.update(kwargs) + return _response() + + def fake_sync_delete(*args, **kwargs): + captured.update(kwargs) + return _response() + + return fake_async_delete, fake_sync_delete + + +def test_async_delete_responses_omits_body_for_azure(): + """Azure responses DELETE rejects requests with any body. Verify the handler + does not pass `json=` to httpx when the transformer returns an empty dict.""" + captured: dict = {} + fake_async_delete, _ = _build_delete_response_mock(captured) + + async def run(): + with patch.object(AsyncHTTPHandler, "delete", new=fake_async_delete): + await litellm.adelete_responses( + response_id="resp_xyz", + custom_llm_provider="azure", + api_base="https://test.openai.azure.com", + api_key="test-key", + api_version="2025-03-01-preview", + ) + + asyncio.run(run()) + + assert "json" not in captured + assert "data" not in captured + assert captured["url"].endswith( + "/openai/responses/resp_xyz?api-version=2025-03-01-preview" + ) + + +def test_sync_delete_responses_omits_body_for_azure(): + captured: dict = {} + _, fake_sync_delete = _build_delete_response_mock(captured) + + with patch.object(HTTPHandler, "delete", new=fake_sync_delete): + litellm.delete_responses( + response_id="resp_xyz", + custom_llm_provider="azure", + api_base="https://test.openai.azure.com", + api_key="test-key", + api_version="2025-03-01-preview", + ) + + assert "json" not in captured + assert "data" not in captured + assert captured["url"].endswith( + "/openai/responses/resp_xyz?api-version=2025-03-01-preview" + ) diff --git a/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py new file mode 100644 index 0000000000..54e689dea6 --- /dev/null +++ b/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py @@ -0,0 +1,33 @@ +import pytest + +from litellm.llms.elevenlabs.text_to_speech.transformation import ( + ElevenLabsTextToSpeechConfig, +) + + +def test_should_encode_elevenlabs_voice_id_path_segment(): + config = ElevenLabsTextToSpeechConfig() + + url = config.get_complete_url( + model="elevenlabs/tts", + api_base="https://api.elevenlabs.io", + litellm_params={ + config.ELEVENLABS_VOICE_ID_KEY: "voice/../../models?x=1#frag", + }, + ) + + assert ( + url + == "https://api.elevenlabs.io/v1/text-to-speech/voice%2F..%2F..%2Fmodels%3Fx%3D1%23frag" + ) + + +def test_should_reject_dot_segment_elevenlabs_voice_id(): + config = ElevenLabsTextToSpeechConfig() + + with pytest.raises(ValueError, match="voice_id cannot be a dot path segment"): + config.get_complete_url( + model="elevenlabs/tts", + api_base="https://api.elevenlabs.io", + litellm_params={config.ELEVENLABS_VOICE_ID_KEY: ".."}, + ) diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py index 9cadba972e..33fe5d24d7 100644 --- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py +++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py @@ -141,44 +141,28 @@ class TestGoogleAIStudioFilesTransformation: assert url == "https://generativelanguage.googleapis.com/v1beta/files/test123" assert params == {} - def test_transform_retrieve_file_request_rejects_traversal_name(self): + def test_transform_retrieve_file_request_encodes_file_id_path_segment(self): + file_id = "files/../../models/gemini-pro?x=1#frag" litellm_params = {"api_key": "test-api-key"} - with pytest.raises(ValueError, match="Invalid Gemini file name"): + url, params = self.handler.transform_retrieve_file_request( + file_id=file_id, + optional_params={}, + litellm_params=litellm_params, + ) + + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" + ) + assert params == {} + + def test_transform_retrieve_file_request_rejects_dot_path_segment(self): + with pytest.raises(ValueError, match="file_id cannot be a dot path segment"): self.handler.transform_retrieve_file_request( - file_id="files/../secrets", + file_id="files/..", optional_params={}, - litellm_params=litellm_params, - ) - - def test_transform_retrieve_file_request_rejects_encoded_traversal_name(self): - litellm_params = {"api_key": "test-api-key"} - - with pytest.raises(ValueError, match="Invalid Gemini file name"): - self.handler.transform_retrieve_file_request( - file_id="files/..%2Fsecrets", - optional_params={}, - litellm_params=litellm_params, - ) - - def test_transform_retrieve_file_request_rejects_double_encoded_separator(self): - litellm_params = {"api_key": "test-api-key"} - - with pytest.raises(ValueError, match="Invalid Gemini file name"): - self.handler.transform_retrieve_file_request( - file_id="files/..%252Fsecrets", - optional_params={}, - litellm_params=litellm_params, - ) - - def test_transform_retrieve_file_request_rejects_deeply_encoded_separator(self): - litellm_params = {"api_key": "test-api-key"} - - with pytest.raises(ValueError, match="Invalid Gemini file name"): - self.handler.transform_retrieve_file_request( - file_id="files/..%2525252Fsecrets", - optional_params={}, - litellm_params=litellm_params, + litellm_params={"api_key": "test-api-key"}, ) @patch.dict("os.environ", {}, clear=True) @@ -385,9 +369,7 @@ class TestGoogleAIStudioFilesTransformation: litellm_params=litellm_params, ) - # Verify URL extraction - assert "files/test123" in url - assert "generativelanguage.googleapis.com" in url + assert url == "https://generativelanguage.googleapis.com/v1beta/files/test123" # Params should be empty (API key goes in header via validate_environment) assert params == {} @@ -439,28 +421,21 @@ class TestGoogleAIStudioFilesTransformation: assert url == "https://custom-gemini.example/v1beta/files/test123" assert params == {} - def test_transform_delete_file_request_rejects_encoded_traversal_url(self): + def test_transform_delete_file_request_encodes_file_id_path_segment(self): + file_id = "files/../../models/gemini-pro?x=1#frag" litellm_params = { "api_key": "test-api-key", "api_base": "https://generativelanguage.googleapis.com", } - with pytest.raises(ValueError, match="Invalid Gemini file name"): - self.handler.transform_delete_file_request( - file_id="https://generativelanguage.googleapis.com/v1beta/files/..%2Fsecrets", - optional_params={}, - litellm_params=litellm_params, - ) + url, params = self.handler.transform_delete_file_request( + file_id=file_id, + optional_params={}, + litellm_params=litellm_params, + ) - def test_transform_delete_file_request_rejects_deeply_encoded_traversal_url(self): - litellm_params = { - "api_key": "test-api-key", - "api_base": "https://generativelanguage.googleapis.com", - } - - with pytest.raises(ValueError, match="Invalid Gemini file name"): - self.handler.transform_delete_file_request( - file_id="https://generativelanguage.googleapis.com/v1beta/files/..%2525252Fsecrets", - optional_params={}, - litellm_params=litellm_params, - ) + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" + ) + assert params == {} diff --git a/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py b/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py index 10d66174c5..43ce030323 100644 --- a/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py +++ b/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py @@ -58,3 +58,18 @@ def test_transform_responses_api_request_adds_manus_params(): assert result["agent_profile"] == "manus-1.6" assert "input" in result assert "model" in result + + +def test_get_response_request_encodes_response_id(): + """Test response IDs are encoded before being appended to Manus URLs.""" + config = ManusResponsesAPIConfig() + + url, params = config.transform_get_response_api_request( + response_id="../../files?x=1#frag", + api_base="https://api.manus.im/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.manus.im/v1/responses/..%2F..%2Ffiles%3Fx%3D1%23frag" + assert params == {} diff --git a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py b/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py index 751e48ff7a..f39be511b9 100644 --- a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py +++ b/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py @@ -49,6 +49,17 @@ def test_get_complete_url_with_eval_id(config: OpenAIEvalsConfig): assert url == "https://api.openai.com/v1/evals/eval_123" +def test_get_complete_url_encodes_eval_id(config: OpenAIEvalsConfig): + """Test eval_id is treated as a single path segment.""" + url = config.get_complete_url( + api_base="https://api.openai.com", + endpoint="evals", + eval_id="../../files?x=1#frag", + ) + + assert url == "https://api.openai.com/v1/evals/..%2F..%2Ffiles%3Fx%3D1%23frag" + + def test_get_complete_url_without_eval_id(config: OpenAIEvalsConfig): """Test URL construction without eval_id""" url = config.get_complete_url( @@ -253,3 +264,20 @@ def test_transform_cancel_eval_response(config: OpenAIEvalsConfig): assert result.id == "eval_123" assert result.object == "eval" + + +def test_transform_run_requests_encode_eval_and_run_ids(config: OpenAIEvalsConfig): + """Test run path IDs are treated as single path segments.""" + url, _, request_body = config.transform_cancel_run_request( + eval_id="../../evals?x=1#frag", + run_id="../runs#other", + api_base="https://api.openai.com", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/evals/..%2F..%2Fevals%3Fx%3D1%23frag/runs/..%2Fruns%23other/cancel" + ) + assert request_body == {} 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 dae8784283..acb9fa9b64 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 @@ -265,6 +265,24 @@ class TestOpenAIResponsesAPIConfig: assert result == "https://custom-openai.example.com/v1/responses" + def test_response_id_path_requests_encode_response_id(self): + """Test response_id is treated as one upstream URL path segment.""" + api_base = "https://custom-openai.example.com/v1/responses" + response_id = "../../files?x=1#frag" + + url, data = self.config.transform_list_input_items_request( + response_id=response_id, + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://custom-openai.example.com/v1/responses/..%2F..%2Ffiles%3Fx%3D1%23frag/input_items" + ) + assert data["limit"] == 20 + def test_get_event_model_class_generic_event(self): """Test that get_event_model_class returns the correct event model class""" from litellm.types.llms.openai import GenericEvent @@ -547,7 +565,12 @@ class TestOpenAIResponsesAPIConfig: """Base helper strips ``namespace`` from custom_tool_call for every provider path.""" inp = [ {"type": "function_call", "call_id": "a", "name": "f", "namespace": "keep"}, - {"type": "custom_tool_call", "call_id": "b", "name": "c", "namespace": "drop"}, + { + "type": "custom_tool_call", + "call_id": "b", + "name": "c", + "namespace": "drop", + }, ] out = BaseResponsesAPIConfig.strip_custom_tool_call_namespace_from_responses_input( inp diff --git a/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py b/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py index a8e07cd364..aebd93d062 100644 --- a/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py +++ b/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py @@ -32,6 +32,21 @@ def test_get_complete_url(config: OpenAIVectorStoreFilesConfig): assert url == "https://api.example.com/v1/vector_stores/vs_123/files" +def test_get_complete_url_encodes_vector_store_id( + config: OpenAIVectorStoreFilesConfig, +): + url = config.get_complete_url( + api_base="https://api.example.com/v1", + vector_store_id="../vs_123?x=1#frag", + litellm_params={}, + ) + + assert ( + url + == "https://api.example.com/v1/vector_stores/..%2Fvs_123%3Fx%3D1%23frag/files" + ) + + def test_transform_create_request(config: OpenAIVectorStoreFilesConfig): api_base = "https://api.example.com/v1/vector_stores/vs_123/files" url, payload = config.transform_create_vector_store_file_request( @@ -60,6 +75,22 @@ def test_transform_list_request(config: OpenAIVectorStoreFilesConfig): assert params == {"limit": 2, "order": "asc"} +def test_transform_file_request_encodes_file_id(config: OpenAIVectorStoreFilesConfig): + api_base = "https://api.example.com/v1/vector_stores/vs_123/files" + + url, params = config.transform_retrieve_vector_store_file_content_request( + vector_store_id="vs_123", + file_id="../../files?x=1#frag", + api_base=api_base, + ) + + assert ( + url + == "https://api.example.com/v1/vector_stores/vs_123/files/..%2F..%2Ffiles%3Fx%3D1%23frag/content" + ) + assert params == {} + + def test_transform_create_response(config: OpenAIVectorStoreFilesConfig): response = httpx.Response( status_code=200, diff --git a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py b/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py index 053b107afb..e7b1aab45b 100644 --- a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py +++ b/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py @@ -64,3 +64,21 @@ class TestOpenAIVectorStoreAPIConfig: for i in range(16): assert f"key_{i}" in request_body["metadata"] assert request_body["metadata"][f"key_{i}"] == f"value_{i}" + + def test_transform_search_vector_store_request_encodes_vector_store_id(self): + config = OpenAIVectorStoreConfig() + + url, request_body = config.transform_search_vector_store_request( + vector_store_id="../../files?x=1#frag", + query="hello", + vector_store_search_optional_params={}, + api_base="https://api.openai.com/v1/vector_stores", + litellm_logging_obj=None, # type: ignore[arg-type] + litellm_params={}, + ) + + assert ( + url + == "https://api.openai.com/v1/vector_stores/..%2F..%2Ffiles%3Fx%3D1%23frag/search" + ) + assert request_body["query"] == "hello" diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py new file mode 100644 index 0000000000..c15554a46a --- /dev/null +++ b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py @@ -0,0 +1,70 @@ +from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.utils import encode_character_id_with_provider + + +def test_video_content_request_encodes_video_id_path_segment(): + config = OpenAIVideoConfig() + + url, params = config.transform_video_content_request( + video_id="../../responses?x=1#frag", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/videos/..%2F..%2Fresponses%3Fx%3D1%23frag/content" + ) + assert params == {} + + +def test_video_content_request_encodes_variant_query_param(): + """``variant`` is user-controlled and was previously interpolated raw + into the query string. A value like ``thumbnail&extra=1`` would + inject additional query parameters into the upstream request.""" + config = OpenAIVideoConfig() + + url, _ = config.transform_video_content_request( + video_id="vid_123", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + variant="thumbnail&extra=1", + ) + + # ``&`` and ``=`` must be percent-encoded so they cannot terminate + # the ``variant`` value or open a new query parameter. + assert "?variant=thumbnail%26extra%3D1" in url + # Sanity: the legitimate "thumbnail" value still round-trips cleanly. + url2, _ = config.transform_video_content_request( + video_id="vid_123", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + variant="thumbnail", + ) + assert url2.endswith("?variant=thumbnail") + + +def test_wrapped_character_id_is_decoded_then_encoded_as_path_segment(): + config = OpenAIVideoConfig() + character_id = encode_character_id_with_provider( + "../../characters?x=1#frag", + provider="openai", + model_id="sora", + ) + + url, params = config.transform_video_get_character_request( + character_id=character_id, + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/videos/characters/..%2F..%2Fcharacters%3Fx%3D1%23frag" + ) + assert params == {} diff --git a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py index e3343c3037..56953a574d 100644 --- a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py +++ b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py @@ -141,6 +141,24 @@ class TestPGVectorStoreConfig: assert headers["Authorization"] == "Bearer test_key" assert url == "https://example.com/v1/vector_stores" + def test_search_request_encodes_vector_store_id(self): + config = PGVectorStoreConfig() + + url, request_body = config.transform_search_vector_store_request( + vector_store_id="../../files?x=1#frag", + query="hello", + vector_store_search_optional_params={}, + api_base="https://example.com/v1/vector_stores", + litellm_logging_obj=Mock(), + litellm_params={}, + ) + + assert ( + url + == "https://example.com/v1/vector_stores/..%2F..%2Ffiles%3Fx%3D1%23frag/search" + ) + assert request_body["query"] == "hello" + def test_environment_variable_support(self): """ Test that environment variables are supported for configuration. diff --git a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py b/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py index ae43eac7ff..baf2ab3391 100644 --- a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py +++ b/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py @@ -117,6 +117,24 @@ class TestRAGFlowChatTransformation: == "http://localhost:9380/api/v1/agents_openai/my-agent-id/chat/completions" ) + def test_get_complete_url_encodes_entity_id(self): + """Test RAGFlow chat IDs are encoded as one upstream path segment.""" + config = RAGFlowConfig() + + url = config.get_complete_url( + api_base="http://localhost:9380", + api_key=None, + model="ragflow/chat/..%2F..%2Fagents_openai%2Fother/gpt-4o-mini", + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert ( + url + == "http://localhost:9380/api/v1/chats_openai/..%252F..%252Fagents_openai%252Fother/chat/completions" + ) + def test_get_complete_url_strips_v1(self): """Test URL construction when api_base ends with /v1.""" config = RAGFlowConfig() diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py index 755716c9da..24879ce83f 100644 --- a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py +++ b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py @@ -134,6 +134,21 @@ class TestRunwayMLVideoTransformation: with pytest.raises(ValueError, match="still processing"): self.config._extract_video_url_from_response(processing_response) + def test_transform_video_status_encodes_video_id_path_segment(self): + """Test task IDs are encoded before being appended to Runway URLs.""" + url, params = self.config.transform_video_status_retrieve_request( + video_id="../../tasks/other?x=1#frag", + api_base="https://api.dev.runwayml.com/v1", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.dev.runwayml.com/v1/tasks/..%2F..%2Ftasks%2Fother%3Fx%3D1%23frag" + ) + assert params == {} + def test_full_video_workflow(self): """Test complete video generation workflow from creation to status check.""" config = RunwayMLVideoConfig() diff --git a/tests/test_litellm/llms/test_polling_url_origin_match.py b/tests/test_litellm/llms/test_polling_url_origin_match.py new file mode 100644 index 0000000000..f1f910bc73 --- /dev/null +++ b/tests/test_litellm/llms/test_polling_url_origin_match.py @@ -0,0 +1,177 @@ +""" +VERIA-51: polling URLs returned by upstream APIs (Azure DALL-E, +Azure Document Intelligence, Black Forest Labs) used to be followed +without origin validation. The handlers attached the operator's API +key to the polling request, so an attacker who could influence the +upstream response (or a compromised upstream) could redirect the proxy +to send credentials anywhere. + +These tests assert each handler now rejects polling URLs that don't +share an origin with the original request URL. +""" + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + + +# Azure DALL-E sync + async paths route through ``assert_same_origin`` +# the same way as the cases below. The helper itself is unit-tested in +# ``tests/test_litellm/litellm_core_utils/test_url_utils.py``; the +# tests here exercise the wiring at sites with simpler signatures. + + +# ── Azure Document Intelligence polling ─────────────────────────────────────── + + +def test_azure_di_sync_rejects_cross_origin_polling(): + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, + ) + + config = AzureDocumentIntelligenceOCRConfig() + + raw_response = MagicMock() + raw_response.status_code = 202 + raw_response.headers = { + "Operation-Location": "https://attacker.example.com/results/xyz", + } + raw_response.request = MagicMock() + raw_response.request.url = ( + "https://eastus.cognitiveservices.azure.com/documentintelligence/.../analyze" + ) + raw_response.request.headers = {"Ocp-Apim-Subscription-Key": "leak-me"} + + with pytest.raises(ValueError, match="rejected polling URL"): + config.transform_ocr_response( + model="azure-doc-intel", + raw_response=raw_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + response={}, + ) + + +# ── Black Forest Labs polling ───────────────────────────────────────────────── + + +def test_bfl_image_generation_sync_rejects_cross_origin_polling(): + from litellm.llms.black_forest_labs.image_generation.handler import ( + BlackForestLabsImageGeneration, + ) + + handler = BlackForestLabsImageGeneration() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://attacker.example.com/get_result"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro" + + sync_client = MagicMock() + sync_client.get = MagicMock() + + with pytest.raises(Exception, match="Rejected polling URL"): + handler._poll_for_result_sync( + initial_response=initial_response, + headers={"x-key": "secret"}, + sync_client=sync_client, + ) + + sync_client.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_bfl_image_generation_async_rejects_cross_origin_polling(): + from litellm.llms.black_forest_labs.image_generation.handler import ( + BlackForestLabsImageGeneration, + ) + + handler = BlackForestLabsImageGeneration() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://attacker.example.com/get_result"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro" + + async_client = MagicMock() + async_client.get = MagicMock() + + with pytest.raises(Exception, match="Rejected polling URL"): + await handler._poll_for_result_async( + initial_response=initial_response, + headers={"x-key": "secret"}, + async_client=async_client, + ) + + async_client.get.assert_not_called() + + +def test_bfl_image_edit_sync_rejects_cross_origin_polling(): + from litellm.llms.black_forest_labs.image_edit.handler import ( + BlackForestLabsImageEdit, + ) + + handler = BlackForestLabsImageEdit() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://attacker.example.com/get_result"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro/edit" + + sync_client = MagicMock() + sync_client.get = MagicMock() + + with pytest.raises(Exception, match="Rejected polling URL"): + handler._poll_for_result_sync( + initial_response=initial_response, + headers={"x-key": "secret"}, + sync_client=sync_client, + ) + + sync_client.get.assert_not_called() + + +def test_bfl_image_generation_same_origin_polling_passes(): + """Sanity check: when the polling URL shares origin with the original + request, the origin check passes and polling proceeds.""" + from litellm.llms.black_forest_labs.image_generation.handler import ( + BlackForestLabsImageGeneration, + ) + + handler = BlackForestLabsImageGeneration() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://api.bfl.ai/v1/get_result?id=abc"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro" + + sync_client = MagicMock() + poll_response = MagicMock() + poll_response.status_code = 200 + poll_response.json = MagicMock(return_value={"status": "Ready"}) + sync_client.get = MagicMock(return_value=poll_response) + + result = handler._poll_for_result_sync( + initial_response=initial_response, + headers={"x-key": "secret"}, + sync_client=sync_client, + ) + + sync_client.get.assert_called_once() + assert result is poll_response diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py new file mode 100644 index 0000000000..bb4e6c67e9 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -0,0 +1,290 @@ +""" +Tests for Gemini batchEmbedContents transformation logic. + +Covers: +- Text-only inputs (single and batch) +- Multimodal inputs (data URIs, GCS URLs, file references) +- Mixed text + multimodal inputs +- Response processing with correct indices +""" + +import pytest + +from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + _build_part_for_input, + _is_multimodal_input, + process_response, + transform_openai_input_gemini_content, + transform_openai_input_gemini_embed_content, +) +from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject +from litellm.types.utils import EmbeddingResponse + + +IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" +GCS_URL = "gs://my-bucket/image.png" + + +class TestIsMultimodalInput: + def test_text_only_string(self): + assert _is_multimodal_input("hello world") is False + + def test_text_only_list(self): + assert _is_multimodal_input(["hello", "world"]) is False + + def test_data_uri(self): + assert _is_multimodal_input([IMAGE_DATA_URI]) is True + + def test_gcs_url(self): + assert _is_multimodal_input([GCS_URL]) is True + + def test_file_reference(self): + assert _is_multimodal_input(["files/abc123"]) is True + + def test_mixed_text_and_image(self): + assert _is_multimodal_input(["hello", IMAGE_DATA_URI]) is True + + def test_nested_text_is_not_multimodal(self): + """Nested list with text is not multimodal.""" + assert _is_multimodal_input([["text_a", "text_b"]]) is False + + def test_nested_list_with_image_is_multimodal(self): + assert _is_multimodal_input([["a red shoe", IMAGE_DATA_URI]]) is True + + +class TestBuildPartForInput: + def test_text_input(self): + part = _build_part_for_input("hello") + assert part["text"] == "hello" + assert part.get("inline_data") is None + + def test_data_uri_input(self): + part = _build_part_for_input(IMAGE_DATA_URI) + assert part.get("text") is None + assert part["inline_data"] is not None + assert part["inline_data"]["mime_type"] == "image/png" + + def test_gcs_url_input(self): + part = _build_part_for_input(GCS_URL) + assert part.get("text") is None + assert part["file_data"] is not None + assert part["file_data"]["mime_type"] == "image/png" + assert part["file_data"]["file_uri"] == GCS_URL + + def test_file_reference_resolved(self): + resolved = {"files/abc": {"mime_type": "image/jpeg", "uri": "https://example.com/abc"}} + part = _build_part_for_input("files/abc", resolved_files=resolved) + assert part["file_data"] is not None + assert part["file_data"]["mime_type"] == "image/jpeg" + + def test_file_reference_unresolved_raises(self): + with pytest.raises(ValueError, match="not resolved"): + _build_part_for_input("files/abc") + + +class TestTransformOpenaiInputGeminiContent: + """Test that transform_openai_input_gemini_content creates separate requests per input.""" + + def test_single_text(self): + result = transform_openai_input_gemini_content( + input="hello", model="gemini-embedding-2-preview", optional_params={} + ) + assert len(result["requests"]) == 1 + assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" + + def test_multiple_texts(self): + result = transform_openai_input_gemini_content( + input=["hello", "world"], model="gemini-embedding-2-preview", optional_params={} + ) + assert len(result["requests"]) == 2 + assert result["requests"][0]["content"]["parts"][0]["text"] == "hello" + assert result["requests"][1]["content"]["parts"][0]["text"] == "world" + + def test_multimodal_inputs_are_separate_requests(self): + """Key regression test for #24209: each input becomes its own request.""" + result = transform_openai_input_gemini_content( + input=["The food was delicious", IMAGE_DATA_URI], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert len(result["requests"]) == 2 + # First request is text + assert result["requests"][0]["content"]["parts"][0]["text"] == "The food was delicious" + # Second request is image + assert result["requests"][1]["content"]["parts"][0]["inline_data"] is not None + + def test_dimensions_mapped_to_output_dimensionality(self): + result = transform_openai_input_gemini_content( + input="hello", + model="gemini-embedding-2-preview", + optional_params={"dimensions": 256}, + ) + assert result["requests"][0]["outputDimensionality"] == 256 + + def test_model_name_prefixed(self): + result = transform_openai_input_gemini_content( + input="hello", model="gemini-embedding-2-preview", optional_params={} + ) + assert result["requests"][0]["model"] == "models/gemini-embedding-2-preview" + + def test_gcs_url_input(self): + result = transform_openai_input_gemini_content( + input=[GCS_URL], model="gemini-embedding-2-preview", optional_params={} + ) + assert len(result["requests"]) == 1 + assert result["requests"][0]["content"]["parts"][0]["file_data"] is not None + + def test_mixed_text_image_gcs(self): + result = transform_openai_input_gemini_content( + input=["hello", IMAGE_DATA_URI, GCS_URL], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert len(result["requests"]) == 3 + + def test_nested_input_combined_embedding(self): + """Nested list produces one request with multiple parts (combined embedding).""" + result = transform_openai_input_gemini_content( + input=[["a red shoe", IMAGE_DATA_URI]], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert len(result["requests"]) == 1 + parts = result["requests"][0]["content"]["parts"] + assert len(parts) == 2 + assert parts[0]["text"] == "a red shoe" + assert parts[1]["inline_data"] is not None + + def test_mixed_nested_and_flat(self): + """Mixed nested + flat produces correct number of requests.""" + result = transform_openai_input_gemini_content( + input=[["text", IMAGE_DATA_URI], "standalone"], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert len(result["requests"]) == 2 + # First: combined (2 parts) + assert len(result["requests"][0]["content"]["parts"]) == 2 + # Second: standalone (1 part) + assert len(result["requests"][1]["content"]["parts"]) == 1 + assert result["requests"][1]["content"]["parts"][0]["text"] == "standalone" + + +class TestTransformOpenaiInputGeminiEmbedContent: + """Test transform_openai_input_gemini_embed_content (vertex_ai / embedContent path).""" + + def test_text_and_image_combined(self): + result = transform_openai_input_gemini_embed_content( + input=["hello", IMAGE_DATA_URI], + model="gemini-embedding-2-preview", + optional_params={}, + ) + assert "content" in result + parts = result["content"]["parts"] + assert len(parts) == 2 + assert parts[0]["text"] == "hello" + assert parts[1]["inline_data"] is not None + + def test_gcs_url(self): + result = transform_openai_input_gemini_embed_content( + input=[GCS_URL], + model="gemini-embedding-2-preview", + optional_params={}, + ) + parts = result["content"]["parts"] + assert len(parts) == 1 + assert parts[0]["file_data"]["file_uri"] == GCS_URL + + def test_dimensions_mapped(self): + result = transform_openai_input_gemini_embed_content( + input="hello", + model="gemini-embedding-2-preview", + optional_params={"dimensions": 256}, + ) + assert result["outputDimensionality"] == 256 + + +class TestProcessResponse: + """Test that process_response sets correct indices.""" + + def test_single_embedding_index(self): + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [{"values": [0.1, 0.2]}] + } + model_response = EmbeddingResponse() + result = process_response( + input="hello", + model_response=model_response, + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 1 + assert result.data[0]["index"] == 0 + + def test_multiple_embeddings_have_correct_indices(self): + """Regression test: indices should be 0, 1, 2... not all 0.""" + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [ + {"values": [0.1, 0.2]}, + {"values": [0.3, 0.4]}, + {"values": [0.5, 0.6]}, + ] + } + model_response = EmbeddingResponse() + result = process_response( + input=["a", "b", "c"], + model_response=model_response, + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 3 + assert result.data[0]["index"] == 0 + assert result.data[1]["index"] == 1 + assert result.data[2]["index"] == 2 + + def test_multimodal_mixed_input(self): + """process_response works with mixed text + multimodal inputs.""" + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [{"values": [0.1, 0.2]}, {"values": [0.3, 0.4]}] + } + result = process_response( + input=["hello", IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 2 + assert result.data[0]["index"] == 0 + assert result.data[1]["index"] == 1 + # Should count tokens only for the text element, not the image + assert result.usage.prompt_tokens > 0 + + def test_nested_input_token_counting(self): + """Nested list: only plain-text sub-elements should be counted.""" + predictions: VertexAIBatchEmbeddingsResponseObject = { + "embeddings": [{"values": [0.1, 0.2]}] + } + result = process_response( + input=[["a red shoe", IMAGE_DATA_URI]], + model_response=EmbeddingResponse(), + model="gemini-embedding-2-preview", + _predictions=predictions, + ) + assert len(result.data) == 1 + assert result.usage.prompt_tokens > 0 + + def test_nested_empty_list_raises(self): + with pytest.raises(ValueError, match="must not be empty"): + transform_openai_input_gemini_content( + input=[[]], + model="gemini-embedding-2-preview", + optional_params={}, + ) + + def test_nested_non_string_element_raises(self): + with pytest.raises(ValueError, match="must be strings"): + transform_openai_input_gemini_content( + input=[[["doubly", "nested"]]], + model="gemini-embedding-2-preview", + optional_params={}, + ) diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 6905cda076..fe5b5a69c9 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -373,6 +373,20 @@ class TestVertexAIImagenImageGenerationConfig: assert request["parameters"]["sampleCount"] == 2 assert request["parameters"]["aspectRatio"] == "16:9" + def test_transform_image_generation_request_labels_from_metadata(self): + """Billing labels from litellm_params.metadata.requester_metadata on predict body.""" + request = self.config.transform_image_generation_request( + model="imagegeneration@006", + prompt="A cat", + optional_params={}, + litellm_params={ + "metadata": {"requester_metadata": {"team": "platform", "env": "prod"}} + }, + headers={}, + ) + assert request["labels"] == {"team": "platform", "env": "prod"} + assert "labels" not in request["parameters"] + def test_transform_image_generation_response(self): """Test response transformation""" mock_response = MagicMock(spec=httpx.Response) diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index d451fb2487..c2ea6f6fab 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -216,6 +216,22 @@ class TestVertexAIRerankTransform: ) assert request_data_default["ignoreRecordDetailsInResponse"] == False + def test_transform_rerank_request_user_labels_from_metadata(self): + """Discovery Engine Rank API uses userLabels (string map) for billing.""" + optional_params = { + "query": "q", + "documents": ["a", "b"], + } + request_data = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params=optional_params, + headers={}, + litellm_params={ + "metadata": {"requester_metadata": {"app": "litellm", "tier": "1"}} + }, + ) + assert request_data["userLabels"] == {"app": "litellm", "tier": "1"} + def test_transform_rerank_request_missing_required_params(self): """Test that transform_rerank_request handles missing required parameters.""" # Test missing query diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py new file mode 100644 index 0000000000..0b28fa9abc --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_userlabels_e2e.py @@ -0,0 +1,182 @@ +""" +End-to-end tests for Vertex AI rerank `userLabels` propagation. + +These tests go through the full `litellm.rerank()` call path with the HTTP +layer mocked, so they catch plumbing bugs (e.g. `litellm_params` losing +`metadata` between the rerank entrypoint and the Vertex transform) that +unit tests on `VertexAIRerankConfig.transform_rerank_request` miss. +""" + +import asyncio +import json +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +import litellm +import litellm.llms.vertex_ai.rerank.transformation + + +def _extract_body(call_kwargs): + """The rerank handler sends `data=json.dumps(...)`, not `json=...`.""" + if "json" in call_kwargs and call_kwargs["json"] is not None: + return call_kwargs["json"] + raw = call_kwargs.get("data") + if isinstance(raw, (bytes, bytearray)): + raw = raw.decode("utf-8") + return json.loads(raw) if raw else None + + +def _make_mock_rank_response(): + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = { + "records": [ + {"id": "0", "score": 0.9, "title": "doc 0", "content": "hello"}, + {"id": "1", "score": 0.1, "title": "doc 1", "content": "world"}, + ] + } + mock_response.text = '{"records": []}' + return mock_response + + +def _make_async_mock_rank_response(): + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json = MagicMock( + return_value={ + "records": [ + {"id": "0", "score": 0.9, "title": "doc 0", "content": "hello"}, + ] + } + ) + mock_response.text = '{"records": []}' + return mock_response + + +@pytest.fixture +def clean_vertex_env(): + saved = {} + for var in ( + "GOOGLE_APPLICATION_CREDENTIALS", + "GOOGLE_CLOUD_PROJECT", + "VERTEXAI_PROJECT", + "VERTEXAI_CREDENTIALS", + "VERTEX_AI_CREDENTIALS", + "VERTEX_PROJECT", + "VERTEX_LOCATION", + "VERTEX_AI_PROJECT", + ): + if var in os.environ: + saved[var] = os.environ.pop(var) + yield + for var, value in saved.items(): + os.environ[var] = value + + +def _patch_vertex_auth(): + return patch.object( + litellm.llms.vertex_ai.rerank.transformation.VertexAIRerankConfig, + "_ensure_access_token", + return_value=("test-access-token", "test-project-2049"), + ) + + +def test_rerank_userlabels_propagates_from_metadata_sync(clean_vertex_env): + """ + `litellm.rerank(metadata={"requester_metadata": {...}})` must end up as + `userLabels` on the Discovery Engine `:rank` request body. + """ + captured = {} + + def fake_post(*args, **kwargs): + captured["body"] = _extract_body(kwargs) + return _make_mock_rank_response() + + with ( + _patch_vertex_auth(), + patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + side_effect=fake_post, + ), + ): + litellm.rerank( + model="vertex_ai/semantic-ranker-default@latest", + query="what is gemini?", + documents=["hello", "world"], + vertex_project="test-project-2049", + vertex_credentials='{"type": "service_account"}', + metadata={"requester_metadata": {"team": "platform", "env": "prod"}}, + ) + + body = captured["body"] + assert body is not None, "expected POST body to be captured" + assert "userLabels" in body, ( + "Vertex rerank request body is missing `userLabels` — metadata was " + "lost between litellm.rerank() and transform_rerank_request. " + f"body keys: {sorted(body.keys())}" + ) + assert body["userLabels"] == {"team": "platform", "env": "prod"} + + +def test_rerank_userlabels_propagates_from_metadata_async(clean_vertex_env): + """Same as the sync test, but through `litellm.arerank`.""" + captured = {} + + async def fake_post(*args, **kwargs): + captured["body"] = _extract_body(kwargs) + return _make_async_mock_rank_response() + + with ( + _patch_vertex_auth(), + patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + side_effect=fake_post, + ), + ): + asyncio.run( + litellm.arerank( + model="vertex_ai/semantic-ranker-default@latest", + query="what is gemini?", + documents=["hello", "world"], + vertex_project="test-project-2049", + vertex_credentials='{"type": "service_account"}', + metadata={"requester_metadata": {"team": "platform"}}, + ) + ) + + body = captured["body"] + assert body is not None + assert body.get("userLabels") == {"team": "platform"} + + +def test_rerank_userlabels_absent_when_no_metadata(clean_vertex_env): + """No metadata → no `userLabels` key (don't send empty maps).""" + captured = {} + + def fake_post(*args, **kwargs): + captured["body"] = _extract_body(kwargs) + return _make_mock_rank_response() + + with ( + _patch_vertex_auth(), + patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + side_effect=fake_post, + ), + ): + litellm.rerank( + model="vertex_ai/semantic-ranker-default@latest", + query="what is gemini?", + documents=["hello", "world"], + vertex_project="test-project-2049", + vertex_credentials='{"type": "service_account"}', + ) + + body = captured["body"] + assert body is not None + assert "userLabels" not in body diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py index 33b3bfce44..8a135dac3b 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -93,6 +93,51 @@ def test_vertex_ai_cancel_batch(): assert ":cancel" in call_args.kwargs["url"] +def test_vertex_ai_cancel_batch_encodes_batch_id(): + """Test that vertex_ai cancel_batch encodes user-controlled batch IDs.""" + handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket") + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "name": "projects/test-project/locations/us-central1/batchPredictionJobs/123456", + "state": "JOB_STATE_CANCELLING", + "createTime": "2024-03-17T10:00:00.000000Z", + "inputConfig": {"gcsSource": {"uris": ["gs://test-bucket/input.jsonl"]}}, + "outputConfig": { + "gcsDestination": {"outputUriPrefix": "gs://test-bucket/output"} + }, + } + + with patch( + "litellm.llms.vertex_ai.batches.handler._get_httpx_client" + ) as mock_client: + mock_client.return_value.post.return_value = mock_response + mock_client.return_value.get.return_value = mock_response + + with patch.object(handler, "_ensure_access_token") as mock_auth: + mock_auth.return_value = ("fake-token", "test-project") + + handler.cancel_batch( + _is_async=False, + batch_id="../../batchPredictionJobs/other?x=1#frag", + api_base=None, + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=600.0, + max_retries=None, + ) + + post_url = mock_client.return_value.post.call_args.kwargs["url"] + get_url = mock_client.return_value.get.call_args.kwargs["url"] + assert ( + "/..%2F..%2FbatchPredictionJobs%2Fother%3Fx%3D1%23frag:cancel" + in post_url + ) + assert "/..%2F..%2FbatchPredictionJobs%2Fother%3Fx%3D1%23frag" in get_url + + def test_vertex_ai_cancel_batch_forwards_timeout(): """Test that timeout is forwarded to the POST (cancel) HTTP call. 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 95507390df..7cb3faf617 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 @@ -15,7 +15,9 @@ from litellm.llms.vertex_ai.common_utils import ( convert_anyof_null_to_nullable, get_vertex_location_from_url, get_vertex_project_id_from_url, + pop_vertex_request_labels, set_schema_property_ordering, + vertex_request_labels_from_litellm_params, ) @@ -1444,3 +1446,65 @@ def test_add_object_type_does_not_add_type_when_anyof_present(): # Verify type was not added (anyOf handles the type) assert "type" not in input_schema, "type should not be added when anyOf is present" + + +def test_vertex_request_labels_from_litellm_params_extracts_requester_metadata(): + assert vertex_request_labels_from_litellm_params(None) is None + assert vertex_request_labels_from_litellm_params({}) is None + assert vertex_request_labels_from_litellm_params({"metadata": None}) is None + lp = {"metadata": {"requester_metadata": {"team": "analytics", "count": 3}}} + assert vertex_request_labels_from_litellm_params(lp) == {"team": "analytics"} + + +def test_vertex_request_labels_from_litellm_params_accepts_litellm_metadata(): + lp = { + "litellm_metadata": { + "requester_metadata": {"team": "platform", "count": 3} + } + } + assert vertex_request_labels_from_litellm_params(lp) == {"team": "platform"} + + +def test_vertex_request_labels_prefers_metadata_over_litellm_metadata(): + lp = { + "metadata": {"requester_metadata": {"source": "metadata"}}, + "litellm_metadata": {"requester_metadata": {"source": "litellm_metadata"}}, + } + assert vertex_request_labels_from_litellm_params(lp) == {"source": "metadata"} + + +def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata(): + optional = {"labels": {"env": "prod"}} + litellm_params = {"metadata": {"requester_metadata": {"team": "x"}}} + assert pop_vertex_request_labels(optional, litellm_params) == {"env": "prod"} + assert "labels" not in optional + + optional2: dict = {} + assert pop_vertex_request_labels(optional2, litellm_params) == {"team": "x"} + + optional3 = {"labels": {"team": 123}} + assert pop_vertex_request_labels(optional3, litellm_params) == {"team": "x"} + + +def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent(): + optional: dict = {} + litellm_params = { + "litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}} + } + assert pop_vertex_request_labels(optional, litellm_params) == { + "team": "from_litellm_meta" + } + + +def test_vertex_text_embedding_request_includes_labels_from_metadata(): + import litellm + + req = litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request( + input="hi", + optional_params={}, + model="text-embedding-004", + litellm_params={ + "metadata": {"requester_metadata": {"project_id": "cost-center-1"}} + }, + ) + assert req.get("labels") == {"project_id": "cost-center-1"} diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py new file mode 100644 index 0000000000..b6329f33ae --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py @@ -0,0 +1,40 @@ +import pytest + +from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( + VertexSearchAPIVectorStoreConfig, +) + + +def test_should_encode_vertex_search_vector_store_id_in_complete_url(): + config = VertexSearchAPIVectorStoreConfig() + + url = config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vertex_collection_id": "default/collection", + "vector_store_id": "../../dataStores/other?x=1#frag", + }, + ) + + assert ( + url + == "https://discoveryengine.googleapis.com/v1/projects/test-project/locations/global/collections/default%2Fcollection/dataStores/..%2F..%2FdataStores%2Fother%3Fx%3D1%23frag/servingConfigs/default_config" + ) + + +def test_should_reject_dot_segment_vertex_search_vector_store_id(): + config = VertexSearchAPIVectorStoreConfig() + + with pytest.raises( + ValueError, match="vector_store_id cannot be a dot path segment" + ): + config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vector_store_id": "..", + }, + ) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py new file mode 100644 index 0000000000..91261b6325 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py @@ -0,0 +1,41 @@ +"""Vertex Model Garden: OpenAPI base URL for publisher/model ids vs per-endpoint path.""" + +import pytest + +from litellm.llms.vertex_ai.vertex_model_garden.main import ( + _vertex_model_garden_model_id_in_json_body, + create_vertex_url, +) + + +@pytest.mark.parametrize( + "model,expect_openapi_base", + [ + ("xai/grok-4.1-fast-reasoning", True), + ("openai/foo/bar", True), + ("5464397967697903616", False), + ("gpt-oss-20b-maas", False), + ], +) +def test_create_vertex_url_openapi_vs_deployed_endpoint( + model: str, expect_openapi_base: bool +) -> None: + url = create_vertex_url( + vertex_location="us-central1", + vertex_project="my-project", + stream=False, + model=model, + ) + if expect_openapi_base: + assert "/v1/projects/my-project/locations/us-central1/endpoints/openapi" in url + else: + assert ( + "/v1beta1/projects/my-project/locations/us-central1/endpoints/" + f"{model}" in url + ) + assert "openapi" not in url + + +def test_model_id_in_json_body_heuristic() -> None: + assert _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") is True + assert _vertex_model_garden_model_id_in_json_body("5464397967697903616") is False diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py index 623d162ccf..13571e63c7 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -101,6 +101,23 @@ class TestVolcengineResponsesAPITransformation: ) assert api_base_full == "https://custom.volc.com/api/v3/responses" + def test_response_id_path_requests_encode_response_id(self): + """response_id should be encoded before building Volcengine URLs.""" + config = VolcEngineResponsesAPIConfig() + + url, params = config.transform_cancel_response_api_request( + response_id="../../responses/other?x=1#frag", + api_base="https://custom.volc.com/api/v3/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://custom.volc.com/api/v3/responses/..%2F..%2Fresponses%2Fother%3Fx%3D1%23frag/cancel" + ) + assert params == {} + @pytest.mark.parametrize( "litellm_params, expected_key", [ diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py new file mode 100644 index 0000000000..5a236de900 --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -0,0 +1,39 @@ +import os +import sys + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.xai.chat.transformation import XAIChatConfig + + +class TestXAIParallelToolCalls: + """Test suite for XAI parallel tool calls functionality.""" + + def test_get_supported_openai_params_includes_parallel_tool_calls(self): + """Test that parallel_tool_calls is in supported parameters.""" + config = XAIChatConfig() + supported_params = config.get_supported_openai_params( + "xai/grok-4.20" + ) + assert "parallel_tool_calls" in supported_params + + def test_transform_request_preserves_parallel_tool_calls(self): + """Test that transform_request preserves parallel_tool_calls parameter.""" + config = XAIChatConfig() + + messages = [{"role": "user", "content": "What's the weather like?"}] + optional_params = {"parallel_tool_calls": True} + + result = config.transform_request( + model="xai/grok-4.20", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result.get("parallel_tool_calls") is True + assert len(result["messages"]) == 1 + assert result["messages"][0]["role"] == "user" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py new file mode 100644 index 0000000000..078adf72d4 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -0,0 +1,402 @@ +""" +Tests for the encrypted-at-rest persistence of MCP user credentials. + +The ``LiteLLM_MCPUserCredentials.credential_b64`` column previously stored +both BYOK API keys and OAuth2 access tokens as plain ``urlsafe_b64encode`` +of the raw value, leaving credentials readable from any DB read. The fix +runs every write through ``encrypt_value_helper`` (nacl SecretBox) and +keeps a plain-base64 fallback on read so existing rows continue to work. +""" + +import base64 +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy._experimental.mcp_server.db import ( + _decode_user_credential, + get_user_credential, + get_user_oauth_credential, + list_user_oauth_credentials, + rotate_mcp_user_credentials_master_key, + store_user_credential, + store_user_oauth_credential, +) +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + +SALT_KEY = "test-salt-key-for-byok-credential-tests-1234" + + +@pytest.fixture(autouse=True) +def _set_salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) + + +def _make_prisma_with_existing(row): + """Build a MagicMock prisma_client whose user-credentials table returns ``row`` + for find_unique and behaves async-correctly for upsert/find_many.""" + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_mcpusercredentials.upsert = AsyncMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[]) + return prisma + + +def _legacy_row(payload: str): + """A row exactly as the pre-fix code would have written it: plain + ``urlsafe_b64encode`` of the raw payload, no encryption.""" + row = MagicMock() + row.credential_b64 = base64.urlsafe_b64encode(payload.encode()).decode() + row.user_id = "alice" + row.server_id = "srv-1" + return row + + +def _stored_value(prisma) -> str: + """Pull the credential_b64 value passed to the most recent upsert call.""" + call = prisma.db.litellm_mcpusercredentials.upsert.call_args + data = call.kwargs["data"] + create_value = data["create"]["credential_b64"] + update_value = data["update"]["credential_b64"] + assert create_value == update_value, "create/update must agree" + return create_value + + +# ── BYOK round-trip ─────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_store_user_credential_does_not_persist_plaintext(): + # Stored bytes must not just be base64 of the secret — that's the regression. + secret = "sk-proj-very-secret-byok-key" + prisma = _make_prisma_with_existing(row=None) + + await store_user_credential(prisma, "alice", "srv-1", secret) + + stored = _stored_value(prisma) + plain_b64 = base64.urlsafe_b64encode(secret.encode()).decode() + assert stored != plain_b64 + # And the secret must not appear anywhere in a plain-b64 decode of the column. + try: + decoded_bytes = base64.urlsafe_b64decode(stored) + except Exception: + decoded_bytes = b"" + assert secret.encode() not in decoded_bytes + + +@pytest.mark.asyncio +async def test_byok_round_trip_returns_plaintext(): + secret = "sk-proj-very-secret-byok-key" + prisma = _make_prisma_with_existing(row=None) + await store_user_credential(prisma, "alice", "srv-1", secret) + + stored = _stored_value(prisma) + row = MagicMock() + row.credential_b64 = stored + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + + result = await get_user_credential(prisma, "alice", "srv-1") + assert result == secret + + +@pytest.mark.asyncio +async def test_byok_get_returns_plaintext_for_legacy_row(): + # Backward-compat: rows persisted by the pre-fix code (plain base64) must + # still decrypt-or-decode cleanly. + legacy_secret = "legacy-byok-key" + prisma = _make_prisma_with_existing(row=_legacy_row(legacy_secret)) + + result = await get_user_credential(prisma, "alice", "srv-1") + assert result == legacy_secret + + +@pytest.mark.asyncio +async def test_byok_get_returns_none_for_missing_row(): + prisma = _make_prisma_with_existing(row=None) + result = await get_user_credential(prisma, "alice", "srv-1") + assert result is None + + +# ── OAuth2 round-trip ───────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_store_user_oauth_credential_does_not_persist_plaintext(): + access_token = "ya29.a0AfH6SMBverysecretaccesstoken" + prisma = _make_prisma_with_existing(row=None) + + await store_user_oauth_credential( + prisma, "alice", "srv-1", access_token, refresh_token="rfr-xyz" + ) + + stored = _stored_value(prisma) + try: + decoded_bytes = base64.urlsafe_b64decode(stored) + except Exception: + decoded_bytes = b"" + assert access_token.encode() not in decoded_bytes + assert b"rfr-xyz" not in decoded_bytes + + +@pytest.mark.asyncio +async def test_oauth_round_trip_returns_payload(): + access_token = "ya29.a0AfH6SMBverysecretaccesstoken" + prisma = _make_prisma_with_existing(row=None) + await store_user_oauth_credential( + prisma, + "alice", + "srv-1", + access_token, + refresh_token="rfr-xyz", + scopes=["a", "b"], + ) + + stored = _stored_value(prisma) + row = MagicMock() + row.credential_b64 = stored + row.server_id = "srv-1" + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=row) + + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is not None + assert result["type"] == "oauth2" + assert result["access_token"] == access_token + assert result["refresh_token"] == "rfr-xyz" + assert result["scopes"] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_oauth_get_returns_payload_for_legacy_row(): + payload = { + "type": "oauth2", + "access_token": "legacy-token", + "connected_at": "2024-01-01T00:00:00Z", + } + legacy_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() + row = MagicMock() + row.credential_b64 = legacy_b64 + row.server_id = "srv-1" + prisma = _make_prisma_with_existing(row=row) + + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is not None + assert result["access_token"] == "legacy-token" + + +@pytest.mark.asyncio +async def test_oauth_get_returns_none_for_byok_row(): + # A row that holds a BYOK string must not leak as an OAuth payload. + prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-not-json")) + result = await get_user_oauth_credential(prisma, "alice", "srv-1") + assert result is None + + +# ── BYOK guard inside store_user_oauth_credential ───────────────────────────── + + +@pytest.mark.asyncio +async def test_byok_guard_rejects_overwriting_legacy_byok(): + prisma = _make_prisma_with_existing(row=_legacy_row("plain-byok-key")) + with pytest.raises(ValueError, match="could not be verified as an OAuth2"): + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") + + +@pytest.mark.asyncio +async def test_byok_guard_rejects_overwriting_encrypted_byok(): + # Simulate a row written by the new (encrypted) code path: write a BYOK, + # then attempt to overwrite with an OAuth token. + prisma = _make_prisma_with_existing(row=None) + await store_user_credential(prisma, "alice", "srv-1", "sk-secret-byok") + + encrypted_row = MagicMock() + encrypted_row.credential_b64 = _stored_value(prisma) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock( + return_value=encrypted_row + ) + + with pytest.raises(ValueError, match="could not be verified as an OAuth2"): + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok") + + +@pytest.mark.asyncio +async def test_byok_guard_allows_overwriting_existing_oauth(): + # Refresh path: row already holds an OAuth payload, write must succeed. + prisma = _make_prisma_with_existing(row=None) + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok-1") + + oauth_row = MagicMock() + oauth_row.credential_b64 = _stored_value(prisma) + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=oauth_row) + + await store_user_oauth_credential(prisma, "alice", "srv-1", "tok-2") + # Final upsert wrote a new payload (different from the first) + assert _stored_value(prisma) != oauth_row.credential_b64 + + +# ── list_user_oauth_credentials ─────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_list_oauth_credentials_filters_byok_and_returns_payloads(): + # Three rows: one encrypted OAuth, one legacy-plaintext OAuth, one BYOK. + # Only the two OAuth rows should come back. + prisma = _make_prisma_with_existing(row=None) + + await store_user_oauth_credential(prisma, "alice", "srv-encrypted", "tok-enc") + encrypted_b64 = _stored_value(prisma) + encrypted_row = MagicMock() + encrypted_row.credential_b64 = encrypted_b64 + encrypted_row.server_id = "srv-encrypted" + + legacy_payload = { + "type": "oauth2", + "access_token": "tok-legacy", + "connected_at": "2024-01-01T00:00:00Z", + } + legacy_row = MagicMock() + legacy_row.credential_b64 = base64.urlsafe_b64encode( + json.dumps(legacy_payload).encode() + ).decode() + legacy_row.server_id = "srv-legacy" + + byok_row = MagicMock() + byok_row.credential_b64 = base64.urlsafe_b64encode(b"plain-byok-key").decode() + byok_row.server_id = "srv-byok" + + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[encrypted_row, legacy_row, byok_row] + ) + + results = await list_user_oauth_credentials(prisma, "alice") + + server_ids = {r["server_id"] for r in results} + assert server_ids == {"srv-encrypted", "srv-legacy"} + tokens = {r["access_token"] for r in results} + assert tokens == {"tok-enc", "tok-legacy"} + + +# ── _decode_user_credential helper ──────────────────────────────────────────── + + +def test_decode_user_credential_handles_garbage(): + # Malformed input must return None, not raise. + assert _decode_user_credential("not-base64-and-not-encrypted!!!") is None + + +def test_decode_user_credential_handles_none(): + # Defensive: a null DB value must return None, not propagate TypeError. + assert _decode_user_credential(None) is None + + +def test_decode_user_credential_legacy_path(): + plain = "legacy-secret" + stored = base64.urlsafe_b64encode(plain.encode()).decode() + assert _decode_user_credential(stored) == plain + + +# ── master-key rotation ─────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_rotate_re_encrypts_byok_with_new_key(monkeypatch): + # Encrypt a row under the current salt, then rotate to a new key, then + # confirm the stored ciphertext decrypts under the NEW key — and not under + # the old one. + prisma = _make_prisma_with_existing(row=None) + secret = "sk-original-byok-key" + await store_user_credential(prisma, "alice", "srv-1", secret) + encrypted_old = _stored_value(prisma) + + row = MagicMock() + row.user_id = "alice" + row.server_id = "srv-1" + row.credential_b64 = encrypted_old + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[row]) + prisma.db.litellm_mcpusercredentials.update = AsyncMock() + + new_master_key = "rotated-salt-key-9999-9999-9999-9999" + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma, new_master_key=new_master_key + ) + + update_call = prisma.db.litellm_mcpusercredentials.update.call_args + new_stored = update_call.kwargs["data"]["credential_b64"] + assert new_stored != encrypted_old, "rotation must produce different ciphertext" + + # Decrypt the rotated value under the NEW salt key — round-trips to plaintext. + monkeypatch.setenv("LITELLM_SALT_KEY", new_master_key) + assert ( + decrypt_value_helper( + value=new_stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + == secret + ) + + +@pytest.mark.asyncio +async def test_rotate_migrates_legacy_plaintext_rows(monkeypatch): + # A legacy plain-base64 row must also get re-encrypted under the new key. + prisma = _make_prisma_with_existing(row=None) + + legacy_row = MagicMock() + legacy_row.user_id = "alice" + legacy_row.server_id = "srv-legacy" + legacy_row.credential_b64 = base64.urlsafe_b64encode(b"legacy-plain").decode() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[legacy_row] + ) + prisma.db.litellm_mcpusercredentials.update = AsyncMock() + + new_key = "another-rotation-key-aaaa-bbbb-cccc-dddd" + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma, new_master_key=new_key + ) + + new_stored = prisma.db.litellm_mcpusercredentials.update.call_args.kwargs["data"][ + "credential_b64" + ] + monkeypatch.setenv("LITELLM_SALT_KEY", new_key) + assert ( + decrypt_value_helper( + value=new_stored, + key="mcp_user_credential", + exception_type="debug", + return_original_value=False, + ) + == "legacy-plain" + ) + + +@pytest.mark.asyncio +async def test_rotate_skips_undecodable_rows(): + # One bad row must not abort the rotation for the rest. + prisma = _make_prisma_with_existing(row=None) + + bad_row = MagicMock() + bad_row.user_id = "alice" + bad_row.server_id = "srv-corrupt" + bad_row.credential_b64 = "!!! not base64 and not encrypted !!!" + + good_row = MagicMock() + good_row.user_id = "bob" + good_row.server_id = "srv-ok" + good_row.credential_b64 = base64.urlsafe_b64encode(b"good-byok").decode() + + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[bad_row, good_row] + ) + prisma.db.litellm_mcpusercredentials.update = AsyncMock() + + await rotate_mcp_user_credentials_master_key( + prisma_client=prisma, new_master_key="new-key-xxxx" + ) + + # Only one update call — the good row. + assert prisma.db.litellm_mcpusercredentials.update.call_count == 1 + where = prisma.db.litellm_mcpusercredentials.update.call_args.kwargs["where"] + assert where["user_id_server_id"]["server_id"] == "srv-ok" 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 558c677d2d..85d5d6ba46 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 @@ -23,6 +23,21 @@ def mock_mcp_client_ip(): yield +@pytest.fixture +def trust_xff(): + """Force ``IPAddressUtils.is_request_from_trusted_proxy`` to True. + + Tests that exercise X-Forwarded-* parsing logic opt into this fixture. + The trust gate's own behaviour is covered by + ``test_get_request_base_url_xff_trust_gate``. + """ + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ): + yield + + @pytest.mark.asyncio async def test_authorize_endpoint_includes_response_type(): """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" @@ -505,6 +520,7 @@ async def test_register_client_remote_registration_success(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_authorize_endpoint_respects_x_forwarded_proto(): """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" try: @@ -572,6 +588,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_token_endpoint_respects_x_forwarded_proto(): """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" try: @@ -650,6 +667,7 @@ async def test_token_endpoint_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_oauth_protected_resource_respects_x_forwarded_proto(): """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" try: @@ -704,6 +722,7 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_oauth_authorization_server_respects_x_forwarded_proto(): """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" try: @@ -759,6 +778,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_register_client_respects_x_forwarded_proto(): """Test that register_client uses X-Forwarded-Proto for redirect_uris""" try: @@ -796,6 +816,7 @@ async def test_register_client_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_authorize_endpoint_respects_x_forwarded_host(): """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" try: @@ -869,6 +890,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_token_endpoint_respects_x_forwarded_host(): """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" try: @@ -1071,7 +1093,12 @@ async def test_token_endpoint_respects_x_forwarded_host(): def test_get_request_base_url_comprehensive( base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url ): - """Comprehensive test for get_request_base_url with various header combinations""" + """Comprehensive test for get_request_base_url with various header combinations. + + These cases exercise the X-Forwarded-* parsing logic, so the trust gate + is patched True; the gate's own behaviour is covered by the + ``test_get_request_base_url_xff_trust_gate`` matrix below. + """ try: from fastapi import Request @@ -1081,11 +1108,9 @@ def test_get_request_base_url_comprehensive( except ImportError: pytest.skip("MCP discoverable endpoints not available") - # Create mock request mock_request = MagicMock(spec=Request) mock_request.base_url = base_url - # Build headers dict headers = {} if x_forwarded_proto: headers["X-Forwarded-Proto"] = x_forwarded_proto @@ -1094,16 +1119,17 @@ def test_get_request_base_url_comprehensive( if x_forwarded_port: headers["X-Forwarded-Port"] = x_forwarded_port - # Mock headers.get() to return our test values def mock_get(header_name, default=None): return headers.get(header_name, default) mock_request.headers.get = mock_get - # Test the function - result = get_request_base_url(mock_request) + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ): + result = get_request_base_url(mock_request) - # Verify result assert result == expected_url, ( f"Expected '{expected_url}' but got '{result}'\n" f"Input: base_url={base_url}, " @@ -1113,6 +1139,131 @@ def test_get_request_base_url_comprehensive( ) +@pytest.mark.parametrize( + "general_settings,direct_ip,expect_xff_honoured", + [ + # Default: use_x_forwarded_for not set -> ignore X-Forwarded-* entirely. + ({}, "127.0.0.1", False), + # XFF enabled, no trusted ranges -> still ignored (no way to tell a trusted + # reverse proxy from a direct attacker). + ({"use_x_forwarded_for": True}, "127.0.0.1", False), + # XFF enabled, ranges set, but caller IP outside any range -> ignored. + ( + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + "203.0.113.5", + False, + ), + # XFF enabled, caller in trusted range -> headers honoured. + ( + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + "10.0.0.7", + True, + ), + # Loopback example (common dev / single-host deploy). + ( + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["127.0.0.0/8"], + }, + "127.0.0.1", + True, + ), + ], +) +def test_get_request_base_url_xff_trust_gate( + general_settings, direct_ip, expect_xff_honoured +): + """Verify the X-Forwarded-* trust gate. + + With XFF poisoning attempted, the helper must return either the literal + base_url (gate denies) or the forwarded URL (gate allows), never the + forwarded URL when the gate denies. + """ + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = direct_ip + + headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "attacker.example.com", + } + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + mock_request.headers.__contains__ = lambda self_, name: name in headers + + with patch( + "litellm.proxy.proxy_server.general_settings", + general_settings, + create=True, + ): + result = get_request_base_url(mock_request) + + if expect_xff_honoured: + assert result == "https://attacker.example.com" + else: + assert result == "http://localhost:4000" + + +def test_xff_misconfig_warning_emitted_once(caplog): + """Operators upgrading from the old "always trust X-Forwarded-*" behaviour + get a one-shot warning when they have ``use_x_forwarded_for`` enabled + but no ``mcp_trusted_proxy_ranges`` configured. The warning must NOT + spam every request.""" + try: + from fastapi import Request + + from litellm.proxy import auth as proxy_auth_pkg # noqa: F401 + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, + ) + from litellm.proxy.auth import ip_address_utils + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Reset the module-level one-shot flag so the test is deterministic. + ip_address_utils._warned_xff_without_trusted_ranges = False + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.5" + headers = {"X-Forwarded-Host": "attacker.example.com"} + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + + misconfig = {"use_x_forwarded_for": True} + + import logging + + with ( + caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"), + patch("litellm.proxy.proxy_server.general_settings", misconfig, create=True), + ): + for _ in range(3): + get_request_base_url(mock_request) + + matching = [ + rec for rec in caplog.records if "mcp_trusted_proxy_ranges" in rec.getMessage() + ] + assert ( + len(matching) == 1 + ), f"expected exactly one warning, got {len(matching)}: {[r.getMessage() for r in matching]}" + + # ------------------------------------------------------------------- # Tests for scopes_supported when mcp_server.scopes is None # ------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 84c556b8dd..649a08e874 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -673,6 +673,106 @@ class TestHookHeaderMergePriority: assert headers["X-OAuth"] == "yes" assert headers["X-Trace-Id"] == "trace-123" + @pytest.mark.asyncio + async def test_m2m_oauth2_does_not_forward_litellm_caller_authorization(self): + """M2M must not put caller Bearer (LiteLLM API key) into extra_headers (#23652).""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="Test Server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://auth.example.com/token", + ) + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer sk-1234"}, + raw_headers={"authorization": "Bearer sk-1234"}, + proxy_logging_obj=None, + hook_extra_headers=None, + ) + except Exception: + pass + + assert captured_extra_headers.get("value") is None + + @pytest.mark.asyncio + async def test_m2m_oauth2_skips_authorization_in_configured_extra_headers(self): + """M2M must not take Authorization from raw_headers even if extra_headers lists it.""" + manager = MCPServerManager() + server = MCPServer( + server_id="test-id", + name="Test Server", + server_name="test_server", + url="https://example.com", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://auth.example.com/token", + extra_headers=["Authorization", "X-Custom"], + ) + + captured_extra_headers: Dict[str, Any] = {} + + async def fake_create_mcp_client( + server, mcp_auth_header=None, extra_headers=None, stdio_env=None + ): + captured_extra_headers["value"] = extra_headers + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + return mock_client + + with patch.object( + manager, "_create_mcp_client", side_effect=fake_create_mcp_client + ): + with patch.object(manager, "_build_stdio_env", return_value=None): + try: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="test_tool", + arguments={"key": "val"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer sk-1234"}, + raw_headers={ + "authorization": "Bearer sk-1234", + "x-custom": "from-client", + }, + proxy_logging_obj=None, + hook_extra_headers=None, + ) + except Exception: + pass + + headers = captured_extra_headers.get("value") or {} + assert "Authorization" not in headers + assert headers.get("X-Custom") == "from-client" + class TestUserAPIKeyAuthJwtClaims: """Tests that UserAPIKeyAuth correctly carries jwt_claims.""" 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 9df6408b0d..06f95159c0 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 @@ -17,6 +17,7 @@ from litellm.proxy._types import ( MCPTransport, UserAPIKeyAuth, ) +from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -135,6 +136,152 @@ def test_prepare_mcp_server_headers_case_insensitive_extra_headers(): assert extra_headers == {"Authorization": "Bearer token"} +def test_prepare_mcp_server_headers_oauth2_m2m_omits_litellm_caller_authorization(): + """M2M OAuth must not put caller Bearer (LiteLLM API key) into extra_headers (#23652).""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + except ImportError: + pytest.skip("MCP server not available") + + server = MCPServer( + server_id="m2m-server", + name="m2m", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://auth.example.com/token", + ) + caller_key = {"Authorization": "Bearer sk-litellm-caller"} + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=None, + mcp_auth_header=None, + oauth2_headers=caller_key, + raw_headers=None, + ) + + assert server_auth_header is None + assert extra_headers is None + + +def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers(): + """Interactive OAuth still forwards the user's OAuth token in extra_headers.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + except ImportError: + pytest.skip("MCP server not available") + + user_oauth = {"Authorization": "Bearer upstream-user-token"} + + server = MCPServer( + server_id="3lo-server", + name="3lo", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow=None, + ) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=None, + mcp_auth_header=None, + oauth2_headers=user_oauth, + raw_headers=None, + ) + + assert server_auth_header is None + assert extra_headers == user_oauth + + +def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers(): + """M2M must not merge caller Authorization from raw_headers when extra_headers lists it.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + _prepare_mcp_server_headers, + ) + except ImportError: + pytest.skip("MCP server not available") + + server = MCPServer( + server_id="m2m-raw", + name="m2m", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://auth.example.com/token", + extra_headers=["Authorization", "X-Custom"], + ) + + server_auth_header, extra_headers = _prepare_mcp_server_headers( + server=server, + mcp_server_auth_headers=None, + mcp_auth_header=None, + oauth2_headers={"Authorization": "Bearer sk-1234"}, + raw_headers={ + "authorization": "Bearer sk-1234", + "x-custom": "trace", + }, + ) + + assert server_auth_header is None + assert extra_headers is not None + assert "Authorization" not in extra_headers + assert extra_headers.get("X-Custom") == "trace" + + +@pytest.mark.asyncio +async def test_call_tool_m2m_skips_authorization_headers(): + """M2M call_tool must not forward caller Authorization in oauth2/raw headers.""" + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + server = MCPServer( + server_id="m2m-call-tool", + name="m2m-call-tool", + server_name="m2m-call-tool", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://auth.example.com/token", + client_id="cid", + client_secret="csecret", + extra_headers=["Authorization", "X-Custom"], + ) + + mock_client = MagicMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock()) + + with patch.object( + manager, "_create_mcp_client", new=AsyncMock(return_value=mock_client) + ) as create_client_mock: + await manager._call_regular_mcp_tool( + mcp_server=server, + original_tool_name="echo", + arguments={"message": "hello"}, + tasks=[], + mcp_auth_header=None, + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer sk-1234"}, + raw_headers={"authorization": "Bearer sk-1234", "x-custom": "trace"}, + proxy_logging_obj=None, + ) + + create_kwargs = create_client_mock.await_args.kwargs + extra_headers = create_kwargs["extra_headers"] or {} + assert "Authorization" not in extra_headers + assert extra_headers.get("X-Custom") == "trace" + + @pytest.mark.asyncio async def test_get_prompts_from_mcp_servers_success(): try: @@ -2288,6 +2435,79 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab assert spend_meta["per_server_tool_counts"]["server_a"] == 1 +@pytest.mark.asyncio +async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fails(): + """ + Regression test: list_tools should still return fetched tools even if + async_success_handler raises (e.g. serialization errors in logging path). + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + except ImportError: + pytest.skip("MCP server not available") + + user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + server_a = MagicMock(name="server_a_obj") + server_a.name = "server_a" + server_a.alias = "server_a" + server_a.server_name = "server_a" + server_a.server_id = "a" + server_a.auth_type = None + server_a.extra_headers = None + + tool_1 = MagicMock() + tool_1.name = "server_a-tool_1" + + dummy_logging_obj = MagicMock() + dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}} + dummy_logging_obj.async_success_handler = AsyncMock( + side_effect=TypeError("Object of type Tool is not JSON serializable") + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(return_value=[server_a]), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers", + return_value=(None, None), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + ) as mock_manager, + patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + side_effect=lambda tools, _server: tools, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + new=AsyncMock(side_effect=lambda tools, **_: tools), + ), + patch( + "litellm.proxy._experimental.mcp_server.server.function_setup", + return_value=(dummy_logging_obj, None), + ), + ): + mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1]) + + tools = await _get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["server_a"], + mcp_server_auth_headers=None, + log_list_tools_to_spendlogs=True, + list_tools_log_source="mcp_protocol", + ) + + assert tools == [tool_1] + dummy_logging_obj.async_success_handler.assert_awaited_once() + + def test_tool_name_matches_case_insensitive(): """Test that _tool_name_matches performs case-insensitive comparison. @@ -2719,3 +2939,177 @@ class TestGatewayCreateInitializationOptions: _mcp_gateway_initialize_instructions.reset(tok) opts = server.create_initialization_options() assert getattr(opts, "instructions", None) is None + + +@pytest.mark.asyncio +async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow(): + """ + P1 Regression: list_tools path must apply _resolve_oauth2_flow to legacy DB + rows where oauth2_flow is NULL but M2M credentials are present. + + Without this fix, has_client_credentials returns False and the caller's + Authorization header is forwarded upstream instead of being blocked. + """ + try: + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp import MCPAuth + except ImportError: + pytest.skip("MCP server not available") + + user_auth = UserAPIKeyAuth(api_key="sk-1234", user_id="test-user") + + # Simulate a legacy DB row: OAuth2 with M2M credentials but oauth2_flow=None + legacy_server = MagicMock(name="legacy_m2m_server") + legacy_server.name = "legacy_m2m" + legacy_server.alias = "legacy_m2m" + legacy_server.server_name = "legacy_m2m" + legacy_server.server_id = "legacy-m2m-id" + legacy_server.auth_type = MCPAuth.oauth2 + legacy_server.oauth2_flow = None # Legacy: field not set in DB + legacy_server.token_url = "https://oauth.example.com/token" + legacy_server.authorization_url = None + legacy_server.client_id = "client-id" + legacy_server.client_secret = "client-secret" + legacy_server.extra_headers = None + legacy_server.has_client_credentials = False # This is the bug: should be True + legacy_server.model_copy = MagicMock( + side_effect=lambda update: MCPServer( + server_id=legacy_server.server_id, + name=legacy_server.name, + transport=MCPTransport.http, + auth_type=legacy_server.auth_type, + oauth2_flow=update.get("oauth2_flow", legacy_server.oauth2_flow), + token_url=legacy_server.token_url, + authorization_url=legacy_server.authorization_url, + client_id=legacy_server.client_id, + client_secret=legacy_server.client_secret, + ) + ) + + tool_1 = MagicMock() + tool_1.name = "legacy_m2m-tool" + + captured_extra_headers = None + + async def capture_extra_headers(*args, **kwargs): + nonlocal captured_extra_headers + captured_extra_headers = kwargs.get("extra_headers") + return [tool_1] + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + ) as mock_manager, + patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools", + side_effect=lambda tools, _server: tools, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions", + new=AsyncMock(side_effect=lambda tools, **_: tools), + ), + ): + mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["legacy-m2m-id"]) + mock_manager.get_mcp_server_by_id = MagicMock(return_value=legacy_server) + mock_manager.filter_server_ids_by_ip_with_info = MagicMock( + return_value=(["legacy-m2m-id"], 0) + ) + mock_manager._get_tools_from_server = AsyncMock( + side_effect=capture_extra_headers + ) + + tools = await _get_tools_from_mcp_servers( + user_api_key_auth=user_auth, + mcp_auth_header=None, + mcp_servers=["legacy_m2m"], + mcp_server_auth_headers=None, + oauth2_headers={"Authorization": "Bearer sk-1234"}, # Caller's token + ) + + # With P1 fix: _get_allowed_mcp_servers applies _resolve_oauth2_flow, + # so has_client_credentials becomes True and extra_headers should be None + # (caller's Authorization blocked) + assert captured_extra_headers is None, ( + "P1 security issue: caller's Authorization header was forwarded to M2M server. " + "Expected None, got: " + str(captured_extra_headers) + ) + assert tools == [tool_1] + + +@pytest.mark.asyncio +async def test_call_tool_empty_extra_headers_returns_none(): + """ + P2 Regression: When all configured extra_headers are filtered out (e.g. + Authorization for M2M), the resulting extra_headers should be None, not {}. + + Downstream code that checks `if extra_headers is None` will behave + differently if an empty dict is passed instead. + """ + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + from litellm.types.mcp import MCPAuth + except ImportError: + pytest.skip("MCP server not available") + + manager = MCPServerManager() + + # M2M server with only Authorization in extra_headers + m2m_server = MCPServer( + server_id="m2m-srv", + name="m2m_test", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://oauth.example.com/token", + client_id="client-id", + client_secret="client-secret", + extra_headers=["Authorization"], # Will be filtered out for M2M + ) + + raw_headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + + captured_extra_headers = None + + async def capture_create_mcp_client(*args, **kwargs): + nonlocal captured_extra_headers + captured_extra_headers = kwargs.get("extra_headers") + # Return a mock client + mock_client = AsyncMock() + mock_client.call_tool = AsyncMock(return_value=MagicMock(content=[])) + return mock_client + + with ( + patch.object( + manager, + "_create_mcp_client", + side_effect=capture_create_mcp_client, + ), + patch.object( + manager, + "get_mcp_server_by_id", + return_value=m2m_server, + ), + ): + try: + await manager._call_regular_mcp_tool( + mcp_server=m2m_server, + original_tool_name="test_tool", + arguments={}, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=raw_headers, + ) + except Exception: + pass # We only care about the captured headers + + # With P2 fix: extra_headers should be None (not {}) when all headers filtered + assert captured_extra_headers is None, ( + "P2 API consistency issue: expected None for empty extra_headers, got: " + + str(captured_extra_headers) + ) + 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 8614cb3ba7..6cbdcd2820 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 @@ -4,6 +4,7 @@ import logging import os import sys from datetime import datetime +from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -719,7 +720,8 @@ class TestMCPServerManager: return_value=mock_client, ): servers, scopes = await manager._fetch_oauth_metadata_from_resource( - "https://protected.example.com/.well-known/oauth" + "https://protected.example.com/.well-known/oauth", + "https://protected.example.com/mcp", ) assert servers == [ @@ -772,7 +774,8 @@ class TestMCPServerManager: mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp") mock_fetch_auth.assert_awaited_once_with( - ["https://login.microsoftonline.com/test-tenant-id/v2.0"] + ["https://login.microsoftonline.com/test-tenant-id/v2.0"], + "http://localhost:8001/mcp", ) assert result is mock_metadata assert result.scopes == ["api://some-scope/.default"] @@ -784,7 +787,7 @@ class TestMCPServerManager: manager = MCPServerManager() issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0" - def build_response(url: str): + def build_response(url: str, **kwargs): mock_response = MagicMock() if url == f"{issuer}/.well-known/openid-configuration": mock_response.json.return_value = { @@ -810,7 +813,12 @@ class TestMCPServerManager: "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", return_value=mock_client, ): - result = await manager._fetch_single_authorization_server_metadata(issuer) + # The Azure issuer is cross-origin against the server_url — use + # the issuer itself as server_url so the test exercises the + # well-known fetch logic without needing real DNS. + result = await manager._fetch_single_authorization_server_metadata( + issuer, issuer + ) assert result is not None assert ( @@ -846,7 +854,9 @@ class TestMCPServerManager: "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", return_value=mock_client, ): - result = await manager._fetch_single_authorization_server_metadata(issuer) + result = await manager._fetch_single_authorization_server_metadata( + issuer, issuer + ) assert result is not None assert ( @@ -910,7 +920,7 @@ class TestMCPServerManager: ): result = await manager._descovery_metadata(server_url) - mock_fetch_auth.assert_awaited_once_with(["https://example.com"]) + mock_fetch_auth.assert_awaited_once_with(["https://example.com"], server_url) assert result is mock_metadata assert result.scopes == ["read"] @@ -2947,5 +2957,359 @@ class TestMCPServerManagerExpandToolPermissions: assert sorted(result["uuid-a"]) == ["read_file", "write_file"] +class TestOAuthDiscoverySSRFGuard: + """SSRF guard for the OAuth metadata discovery follow-up fetches. + + The vulnerability: a malicious MCP server returns a ``WWW-Authenticate`` + header pointing at an attacker-chosen ``resource_metadata`` URL, then a + PRM JSON whose ``authorization_servers[0]`` points at internal/loopback + addresses, coercing the proxy into making blind GETs to cloud-metadata + services, internal admin panels, or loopback debug endpoints. + """ + + @staticmethod + def _patch_resolves(monkeypatch, mapping): + """Patch ``socket.getaddrinfo`` for a deterministic SSRF-guard test. + + ``mapping`` is ``{hostname: [ip-string, ...]}``; unknown hosts raise + ``gaierror`` (treated as "unresolvable" -> blocked by async_safe_get). + """ + import socket as _socket + + def fake_getaddrinfo(host, port, *args, **kwargs): + if host not in mapping: + raise _socket.gaierror(f"unknown host {host}") + family = _socket.AF_INET + return [ + (family, _socket.SOCK_STREAM, 0, "", (ip, port)) for ip in mapping[host] + ] + + monkeypatch.setattr( + "litellm.litellm_core_utils.url_utils.socket.getaddrinfo", + fake_getaddrinfo, + ) + + def test_same_authority_url_is_direct_fetch_eligible(self): + # Same scheme + host + port skips DNS entirely — the well-known + # endpoint construction in _attempt_well_known_discovery always + # produces same-authority URLs against the admin's server_url. + assert MCPServerManager._is_same_authority_metadata_url( + "https://example.com/.well-known/oauth-protected-resource", + "https://example.com/mcp", + ) + + def test_same_host_different_port_uses_safe_fetch_path(self): + assert not MCPServerManager._is_same_authority_metadata_url( + "https://example.com:9999/.well-known/oauth-protected-resource", + "https://example.com/mcp", + ) + + @pytest.mark.parametrize( + "ip", + [ + "127.0.0.1", # loopback + "10.0.0.5", # RFC1918 + "172.16.0.1", # RFC1918 + "192.168.1.1", # RFC1918 + "169.254.169.254", # AWS / Azure / GCP IMDS + "100.100.100.200", # Alibaba Cloud metadata + "0.0.0.0", # unspecified + "::1", # IPv6 loopback + "fe80::1", # IPv6 link-local + "fc00::1", # IPv6 ULA + ], + ) + @pytest.mark.asyncio + async def test_cross_origin_blocked_when_resolves_to_unsafe_ip( + self, monkeypatch, ip + ): + self._patch_resolves(monkeypatch, {"attacker.example.com": [ip]}) + manager = MCPServerManager() + + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata( + "https://attacker.example.com", + "https://legit-mcp.example.com/mcp", + ) + + assert result is None + mock_client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_cross_origin_allowed_when_resolves_to_public_ip(self, monkeypatch): + self._patch_resolves( + monkeypatch, {"login.microsoftonline.com": ["20.190.151.7"]} + ) + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.is_redirect = False + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = { + "authorization_servers": ["https://login.microsoftonline.com/tenant/v2.0"], + "scopes_supported": ["mcp.read"], + } + + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + servers, scopes = await manager._fetch_oauth_metadata_from_resource( + "https://login.microsoftonline.com/tenant/v2.0/.well-known/openid-configuration", + "https://atlassian-mcp.example.com/mcp", + ) + + assert servers == ["https://login.microsoftonline.com/tenant/v2.0"] + assert scopes == ["mcp.read"] + mock_client.get.assert_awaited_once() + assert mock_client.get.await_args.kwargs["follow_redirects"] is False + assert ( + mock_client.get.await_args.kwargs["headers"]["Host"] + == "login.microsoftonline.com" + ) + + @pytest.mark.asyncio + async def test_cross_origin_blocked_when_unresolvable(self, monkeypatch): + self._patch_resolves(monkeypatch, {}) + manager = MCPServerManager() + + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + servers, scopes = await manager._fetch_oauth_metadata_from_resource( + "https://nope.example.invalid/.well-known/oauth-authorization-server", + "https://legit-mcp.example.com/mcp", + ) + + assert servers == [] + assert scopes is None + mock_client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_non_http_scheme_is_not_safe(self): + manager = MCPServerManager() + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + servers, scopes = await manager._fetch_oauth_metadata_from_resource( + "file:///etc/passwd", + "https://example.com/mcp", + ) + result = await manager._fetch_single_authorization_server_metadata( + "gopher://example.com/", + "https://example.com/mcp", + ) + + assert servers == [] + assert scopes is None + assert result is None + mock_client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_dual_resolution_blocked_if_any_ip_unsafe(self, monkeypatch): + # If the attacker controls a DNS record returning multiple A records, + # one of which is private, async_safe_get rejects before any network call. + self._patch_resolves( + monkeypatch, {"dual-stack.example.com": ["8.8.8.8", "127.0.0.1"]} + ) + manager = MCPServerManager() + + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + servers, scopes = await manager._fetch_oauth_metadata_from_resource( + "https://dual-stack.example.com/.well-known/oauth-authorization-server", + "https://legit-mcp.example.com/mcp", + ) + + assert servers == [] + assert scopes is None + mock_client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_fetch_oauth_metadata_refuses_unsafe_url(self, monkeypatch): + # End-to-end: a malicious WWW-Authenticate redirecting to a loopback + # resource_metadata URL must not produce a network call. + self._patch_resolves(monkeypatch, {"attacker.example.com": ["127.0.0.1"]}) + manager = MCPServerManager() + + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + servers, scopes = await manager._fetch_oauth_metadata_from_resource( + "https://attacker.example.com/meta", + "https://legit-mcp.example.com/mcp", + ) + + assert servers == [] + assert scopes is None + mock_client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_empty_getaddrinfo_result_blocks_url(self, monkeypatch): + # POSIX doesn't strictly forbid an empty success-list from getaddrinfo. + # async_safe_get must fail closed rather than making a network call. + monkeypatch.setattr( + "litellm.litellm_core_utils.url_utils.socket.getaddrinfo", + lambda *a, **k: [], + ) + manager = MCPServerManager() + + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + servers, scopes = await manager._fetch_oauth_metadata_from_resource( + "https://no-records.example.com/.well-known/oauth-authorization-server", + "https://legit-mcp.example.com/mcp", + ) + + assert servers == [] + assert scopes is None + mock_client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_cross_origin_redirect_is_revalidated(self, monkeypatch): + self._patch_resolves( + monkeypatch, + { + "provider.example.com": ["8.8.8.8"], + "127.0.0.1": ["127.0.0.1"], + }, + ) + manager = MCPServerManager() + + redirect_response = MagicMock() + redirect_response.is_redirect = True + redirect_response.headers = {"location": "http://127.0.0.1/admin"} + + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=redirect_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata( + "https://provider.example.com", + "https://legit-mcp.example.com/mcp", + ) + + assert result is None + assert mock_client.get.await_count == 3 + + @pytest.mark.asyncio + async def test_same_authority_fetch_does_not_follow_redirects(self): + # Same-authority URLs may be internal admin-configured MCP servers, so + # they are fetched directly. Redirects are still disabled because a + # Location target would not inherit the same-authority guarantee. + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.json.return_value = { + "authorization_servers": ["https://auth.example.com"], + } + mock_response.raise_for_status = MagicMock() + + captured_kwargs: Dict[str, Any] = {} + + async def fake_get(url, **kwargs): + captured_kwargs.update(kwargs) + return mock_response + + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=fake_get) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + await manager._fetch_oauth_metadata_from_resource( + "https://protected.example.com/.well-known/oauth", + "https://protected.example.com/mcp", + ) + assert captured_kwargs.get("follow_redirects") is False + + @pytest.mark.asyncio + async def test_same_authority_auth_server_fetch_does_not_follow_redirects(self): + # Same redirect-bypass concern for the authorization-server fetch path. + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.json.return_value = { + "authorization_endpoint": "https://provider.example.com/authorize", + "token_endpoint": "https://provider.example.com/token", + } + mock_response.raise_for_status = MagicMock() + + captured_kwargs: Dict[str, Any] = {} + + async def fake_get(url, **kwargs): + captured_kwargs.update(kwargs) + return mock_response + + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=fake_get) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + await manager._fetch_single_authorization_server_metadata( + "https://provider.example.com", + "https://provider.example.com", + ) + assert captured_kwargs.get("follow_redirects") is False + + @pytest.mark.asyncio + async def test_fetch_authorization_server_refuses_unsafe_issuer(self, monkeypatch): + # Mirrors the GHSA-mrfv repro: PRM lists a loopback issuer URL. + self._patch_resolves(monkeypatch, {"attacker.example.com": ["127.0.0.1"]}) + manager = MCPServerManager() + + mock_client = MagicMock() + mock_client.get = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata( + "http://attacker.example.com:19999", + "https://legit-mcp.example.com/mcp", + ) + + assert result is None + mock_client.get.assert_not_called() + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a848db27fc..4c21d0ec64 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -922,19 +922,21 @@ async def test_get_tag_objects_batch(): # Simulate 5 tags: 2 cached, 3 uncached tag_names = ["cached-1", "uncached-1", "cached-2", "uncached-2", "uncached-3"] - # Mock cached tags - cached_tag_1 = { - "tag_name": "cached-1", - "spend": 10.0, - "models": [], - "litellm_budget_table": None, - } - cached_tag_2 = { - "tag_name": "cached-2", - "spend": 20.0, - "models": [], - "litellm_budget_table": None, - } + # Mock cached tags — must be LiteLLM_TagTable instances: the mocked async_get_cache + # bypasses UserApiKeyCache deserialization, so returning plain dicts would flow through + # as dict (production returns models after Codec.deserialize inside the cache). + cached_tag_1 = LiteLLM_TagTable( + tag_name="cached-1", + spend=10.0, + models=[], + litellm_budget_table=None, + ) + cached_tag_2 = LiteLLM_TagTable( + tag_name="cached-2", + spend=20.0, + models=[], + litellm_budget_table=None, + ) # Mock DB response for uncached tags uncached_tag_1 = MagicMock() @@ -980,13 +982,13 @@ async def test_get_tag_objects_batch(): ) # Mock cache behavior - return cached tags, None for uncached - async def mock_get_cache(key): + async def mock_get_cache(*args, **kwargs): + key = kwargs.get("key") if key == "tag:cached-1": return cached_tag_1 - elif key == "tag:cached-2": + if key == "tag:cached-2": return cached_tag_2 - else: - return None + return None mock_cache.async_get_cache = AsyncMock(side_effect=mock_get_cache) mock_cache.async_set_cache = AsyncMock() diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 91f300b88c..b82cb35519 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -964,3 +964,129 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ) is True ) + + +# ── is_request_body_safe nested-config recursion (VERIA-6) ──────────────────── + + +class TestIsRequestBodySafeNestedConfig: + """The Milvus vector store transformer unpacks + ``litellm_embedding_config`` as ``**kwargs`` into ``litellm.embedding(...)`` + — same SSRF / credential-exfil surface as a top-level ``api_base`` in + the request body. ``is_request_body_safe`` must recurse into this + nested dict so a banned param can't be smuggled in via nesting.""" + + def test_root_level_api_base_blocked_when_no_opt_in(self): + """Sanity check: pre-existing root-level enforcement still works.""" + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={"api_base": "https://attacker.example.com"}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_nested_api_base_in_embedding_config_blocked(self): + """Smuggling ``api_base`` inside ``litellm_embedding_config`` is + the VERIA-6 bypass — must be blocked by the recursive check.""" + with pytest.raises(ValueError, match="api_base"): + is_request_body_safe( + request_body={ + "litellm_embedding_config": { + "api_base": "https://attacker.example.com", + "api_key": "leaked-key", + } + }, + general_settings={}, + llm_router=None, + model="milvus-store", + ) + + def test_nested_langfuse_host_in_embedding_config_blocked(self): + """The recursion uses the *full* banned-param list, not a special + subset — so any flag that's banned at the root is also banned + when nested.""" + with pytest.raises(ValueError, match="langfuse_host"): + is_request_body_safe( + request_body={ + "litellm_embedding_config": { + "langfuse_host": "https://attacker.example.com" + } + }, + general_settings={}, + llm_router=None, + model="milvus-store", + ) + + def test_nested_api_base_allowed_when_admin_opts_in(self): + """Admins who explicitly enable client-side credential passthrough + keep the existing escape hatch — same UX as for root-level.""" + assert ( + is_request_body_safe( + request_body={ + "litellm_embedding_config": { + "api_base": "https://my-azure.example.com" + } + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="milvus-store", + ) + is True + ) + + def test_safe_nested_config_accepted(self): + """A nested config without any banned params passes — there's no + false-positive on legitimate ``api_version`` / model params.""" + assert ( + is_request_body_safe( + request_body={ + "litellm_embedding_config": { + "api_version": "2024-02-15-preview", + } + }, + general_settings={}, + llm_router=None, + model="milvus-store", + ) + is True + ) + + def test_non_dict_nested_config_does_not_break_check(self): + """A bogus type for ``litellm_embedding_config`` (string, list, + None) must not crash the validator — it should just fall through.""" + assert ( + is_request_body_safe( + request_body={"litellm_embedding_config": "not-a-dict"}, + general_settings={}, + llm_router=None, + model="x", + ) + is True + ) + + def test_deeply_nested_config_does_not_recurse(self): + """Greptile P1: ``is_request_body_safe`` is iterative single-level — + a deeply-nested ``litellm_embedding_config`` cannot exhaust the + Python call stack to trigger a 500 ``RecursionError``. Build a + body 1000 levels deep; the validator must complete in O(1) + descent.""" + body = {"litellm_embedding_config": {}} + cur = body["litellm_embedding_config"] + for _ in range(1000): + cur["litellm_embedding_config"] = {} + cur = cur["litellm_embedding_config"] + # Banned param at the deepest level shouldn't be reached — single + # level only. + cur["api_base"] = "https://attacker.example.com" + + # No exception raised: deeper levels aren't checked. + assert ( + is_request_body_safe( + request_body=body, + general_settings={}, + llm_router=None, + model="x", + ) + is True + ) diff --git a/tests/test_litellm/proxy/auth/test_cli_auth.py b/tests/test_litellm/proxy/auth/test_cli_auth.py index 82497fcadf..a4f72ef90e 100644 --- a/tests/test_litellm/proxy/auth/test_cli_auth.py +++ b/tests/test_litellm/proxy/auth/test_cli_auth.py @@ -6,11 +6,12 @@ This module tests the auth commands and their associated functionality. import pytest import requests -from unittest.mock import AsyncMock, patch, Mock, call +from unittest.mock import patch, Mock, call from litellm.proxy.client.cli.commands.auth import ( _normalize_teams, _poll_for_ready_data, _poll_for_authentication, + _start_cli_sso_flow, ) @@ -57,6 +58,18 @@ async def test_normalize_teams_with_details_with_aliases(): ] +@patch("litellm.proxy.client.cli.commands.auth.requests.post") +def test_start_cli_sso_flow_rejects_invalid_response(request_mock): + """Test CLI SSO start rejects malformed server responses""" + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = {"login_id": "cli-session", "user_code": "ABCD-EFGH"} + request_mock.return_value = response + + with pytest.raises(ValueError, match="Invalid CLI SSO start response"): + _start_cli_sso_flow("https://litellm.com") + + @pytest.mark.asyncio @patch( "litellm.proxy.client.cli.commands.auth.requests.get", @@ -195,10 +208,11 @@ async def test_poll_for_ready_connection_failure(sleep_mock, click_mock, request @patch("litellm.proxy.client.cli.commands.auth.click.echo") async def test_poll_for_authentication_no_data(click_mock, poll_mock, handle_mock): """Test poll_for_authentication function""" - actual = _poll_for_authentication("https://litellm.com", "key-123") + actual = _poll_for_authentication("https://litellm.com", "key-123", "poll-secret") assert actual is None poll_mock.assert_called_once_with( "https://litellm.com/sso/cli/poll/key-123", + headers={"x-litellm-cli-poll-secret": "poll-secret"}, pending_message="Still waiting for authentication...", ) handle_mock.assert_not_called() @@ -214,10 +228,11 @@ async def test_poll_for_authentication_no_data(click_mock, poll_mock, handle_moc @patch("litellm.proxy.client.cli.commands.auth.click.echo") async def test_poll_for_authentication_no_teams(click_mock, poll_mock, handle_mock): """Test poll_for_authentication function""" - actual = _poll_for_authentication("https://litellm.com", "key-123") + actual = _poll_for_authentication("https://litellm.com", "key-123", "poll-secret") assert actual is None poll_mock.assert_called_once_with( "https://litellm.com/sso/cli/poll/key-123", + headers={"x-litellm-cli-poll-secret": "poll-secret"}, pending_message="Still waiting for authentication...", ) handle_mock.assert_not_called() @@ -243,7 +258,7 @@ async def test_poll_for_authentication_team_selection_success( click_mock, poll_mock, handle_mock ): """Test poll_for_authentication function""" - actual = _poll_for_authentication("https://litellm.com", "key-123") + actual = _poll_for_authentication("https://litellm.com", "key-123", "poll-secret") assert actual == { "api_key": "jwt-123", "user_id": "user-123", @@ -252,11 +267,13 @@ async def test_poll_for_authentication_team_selection_success( } poll_mock.assert_called_once_with( "https://litellm.com/sso/cli/poll/key-123", + headers={"x-litellm-cli-poll-secret": "poll-secret"}, pending_message="Still waiting for authentication...", ) handle_mock.assert_called_once_with( base_url="https://litellm.com", key_id="key-123", + poll_secret="poll-secret", teams=[ {"team_id": "1", "team_alias": None}, {"team_id": "2", "team_alias": None}, @@ -283,15 +300,17 @@ async def test_poll_for_authentication_team_selection_cancelled( click_mock, poll_mock, handle_mock ): """Test poll_for_authentication function""" - actual = _poll_for_authentication("https://litellm.com", "key-123") + actual = _poll_for_authentication("https://litellm.com", "key-123", "poll-secret") assert actual is None poll_mock.assert_called_once_with( "https://litellm.com/sso/cli/poll/key-123", + headers={"x-litellm-cli-poll-secret": "poll-secret"}, pending_message="Still waiting for authentication...", ) handle_mock.assert_called_once_with( base_url="https://litellm.com", key_id="key-123", + poll_secret="poll-secret", teams=[{"team_id": "team-1", "team_alias": None}], ) click_mock.assert_called_once() @@ -314,7 +333,7 @@ async def test_poll_for_authentication_auto_assigned_team( click_mock, poll_mock, handle_mock ): """Test poll_for_authentication function""" - actual = _poll_for_authentication("https://litellm.com", "key-123") + actual = _poll_for_authentication("https://litellm.com", "key-123", "poll-secret") assert actual == { "api_key": "jwt-456", "user_id": "user-456", @@ -323,6 +342,7 @@ async def test_poll_for_authentication_auto_assigned_team( } poll_mock.assert_called_once_with( "https://litellm.com/sso/cli/poll/key-123", + headers={"x-litellm-cli-poll-secret": "poll-secret"}, pending_message="Still waiting for authentication...", ) handle_mock.assert_not_called() diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 9085469268..47e513dc59 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -405,9 +405,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_role_change(): mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args assert call_kwargs.kwargs["key"] == "u1" - assert ( - call_kwargs.kwargs["value"]["user_role"] == LitellmUserRoles.PROXY_ADMIN.value - ) + assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable) + assert call_kwargs.kwargs["value"].user_role == LitellmUserRoles.PROXY_ADMIN.value + assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable @pytest.mark.asyncio @@ -452,7 +452,9 @@ async def test_sync_user_role_and_teams_cache_invalidation_on_team_change(): mock_cache.async_set_cache.assert_called_once() call_kwargs = mock_cache.async_set_cache.call_args assert call_kwargs.kwargs["key"] == "u1" - assert set(call_kwargs.kwargs["value"]["teams"]) == {"team1", "team2"} + assert isinstance(call_kwargs.kwargs["value"], LiteLLM_UserTable) + assert set(call_kwargs.kwargs["value"].teams) == {"team1", "team2"} + assert call_kwargs.kwargs["model_type"] == LiteLLM_UserTable @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 45d55a8d06..2e738ff900 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1,17 +1,15 @@ import json import os import sys -import tempfile import time from pathlib import Path -from unittest.mock import MagicMock, Mock, mock_open, patch +from unittest.mock import Mock, mock_open, patch sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -import pytest from click.testing import CliRunner from litellm.proxy.client.cli.commands.auth import ( @@ -26,6 +24,22 @@ from litellm.proxy.client.cli.commands.auth import ( ) +def _mock_cli_sso_start_response( + login_id: str = "cli-session-uuid-456", + poll_secret: str = "poll-secret", + user_code: str = "ABCD-EFGH", +) -> Mock: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "login_id": login_id, + "poll_secret": poll_secret, + "user_code": user_code, + } + mock_response.raise_for_status = Mock() + return mock_response + + class TestTokenUtilities: """Test token file utility functions""" @@ -217,6 +231,50 @@ class TestTokenUtilities: result = get_stored_api_key() assert result is None + def test_get_stored_api_key_base_url_match(self): + """Stored key is returned when expected_base_url matches stored origin""" + token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"} + with patch( + "litellm.litellm_core_utils.cli_token_utils.load_cli_token", + return_value=token_data, + ): + assert ( + get_stored_api_key(expected_base_url="https://real-proxy.com") + == "sk-prod" + ) + + def test_get_stored_api_key_base_url_match_trailing_slash(self): + """Trailing slash on expected_base_url is normalised before comparison""" + token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"} + with patch( + "litellm.litellm_core_utils.cli_token_utils.load_cli_token", + return_value=token_data, + ): + assert ( + get_stored_api_key(expected_base_url="https://real-proxy.com/") + == "sk-prod" + ) + + def test_get_stored_api_key_base_url_mismatch(self): + """Stored key is NOT returned when expected_base_url differs from stored origin""" + token_data = {"key": "sk-prod", "base_url": "https://real-proxy.com"} + with patch( + "litellm.litellm_core_utils.cli_token_utils.load_cli_token", + return_value=token_data, + ): + assert get_stored_api_key(expected_base_url="https://evil.com") is None + + def test_get_stored_api_key_old_token_no_base_url(self): + """Old tokens without a base_url field are rejected when origin check is requested""" + token_data = {"key": "sk-old-token"} + with patch( + "litellm.litellm_core_utils.cli_token_utils.load_cli_token", + return_value=token_data, + ): + assert ( + get_stored_api_key(expected_base_url="https://real-proxy.com") is None + ) + class TestLoginCommand: """Test login CLI command""" @@ -243,12 +301,15 @@ class TestLoginCommand: with ( patch("webbrowser.open") as mock_browser, + patch( + "requests.post", + return_value=_mock_cli_sso_start_response(login_id="cli-test-uuid-123"), + ) as mock_post, patch("requests.get", return_value=mock_response) as mock_get, patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save, patch( "litellm.proxy.client.cli.interface.show_commands" ) as mock_show_commands, - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -261,7 +322,13 @@ class TestLoginCommand: mock_browser.assert_called_once() call_args = mock_browser.call_args[0][0] assert "https://test.example.com/sso/key/generate" in call_args - assert "sk-test-uuid-123" in call_args + assert "cli-test-uuid-123" in call_args + assert "Verification code: ABCD-EFGH" in result.output + mock_post.assert_called_once() + mock_get.assert_called() + assert mock_get.call_args.kwargs["headers"] == { + "x-litellm-cli-poll-secret": "poll-secret" + } # Verify JWT was saved mock_save.assert_called_once() @@ -284,9 +351,9 @@ class TestLoginCommand: with ( patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", return_value=mock_response), - patch("time.sleep") as mock_sleep, - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), + patch("time.sleep"), ): # Mock time.sleep to avoid actual delays in tests @@ -306,9 +373,9 @@ class TestLoginCommand: with ( patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", return_value=mock_response), patch("time.sleep"), - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -325,12 +392,12 @@ class TestLoginCommand: with ( patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), patch( "requests.get", side_effect=requests.RequestException("Connection failed"), ), patch("time.sleep"), - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -345,8 +412,8 @@ class TestLoginCommand: with ( patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", side_effect=KeyboardInterrupt), - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -369,9 +436,9 @@ class TestLoginCommand: with ( patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", return_value=mock_response), patch("time.sleep"), - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -386,8 +453,8 @@ class TestLoginCommand: with ( patch("webbrowser.open"), + patch("requests.post", return_value=_mock_cli_sso_start_response()), patch("requests.get", side_effect=ValueError("Invalid value")), - patch("litellm._uuid.uuid.uuid4", return_value="test-uuid-123"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -556,6 +623,12 @@ class TestCLIKeyRegenerationFlow: # Simulate user selecting team #2 (team-beta) with ( patch("webbrowser.open") as mock_browser, + patch( + "requests.post", + return_value=_mock_cli_sso_start_response( + login_id="cli-session-uuid-456" + ), + ), patch( "requests.get", side_effect=[mock_first_response, mock_second_response] ) as mock_get, @@ -563,7 +636,6 @@ class TestCLIKeyRegenerationFlow: patch( "litellm.proxy.client.cli.interface.show_commands" ) as mock_show_commands, - patch("litellm._uuid.uuid.uuid4", return_value="session-uuid-456"), patch("click.prompt", return_value="2"), ): # User selects index 2 @@ -585,8 +657,11 @@ class TestCLIKeyRegenerationFlow: # First poll should be without team_id first_poll_url = mock_get.call_args_list[0][0][0] - assert "sk-session-uuid-456" in first_poll_url + assert "cli-session-uuid-456" in first_poll_url assert "team_id=" not in first_poll_url + assert mock_get.call_args_list[0].kwargs["headers"] == { + "x-litellm-cli-poll-secret": "poll-secret" + } # Second poll should include team_id=team-beta second_poll_url = mock_get.call_args_list[1][0][0] @@ -621,10 +696,15 @@ class TestCLIKeyRegenerationFlow: with ( patch("webbrowser.open") as mock_browser, + patch( + "requests.post", + return_value=_mock_cli_sso_start_response( + login_id="cli-session-uuid-solo" + ), + ), patch("requests.get", return_value=mock_response), patch("litellm.proxy.client.cli.commands.auth.save_token") as mock_save, patch("litellm.proxy.client.cli.interface.show_commands"), - patch("litellm._uuid.uuid.uuid4", return_value="session-uuid-solo"), ): result = self.runner.invoke(login, obj=mock_context.obj) @@ -637,7 +717,7 @@ class TestCLIKeyRegenerationFlow: call_args = mock_browser.call_args[0][0] assert "https://test.example.com/sso/key/generate" in call_args assert "source=litellm-cli" in call_args - assert "key=sk-session-uuid-solo" in call_args + assert "key=cli-session-uuid-solo" in call_args # Verify JWT was saved mock_save.assert_called_once() diff --git a/tests/test_litellm/proxy/common_utils/test_cache_codec.py b/tests/test_litellm/proxy/common_utils/test_cache_codec.py new file mode 100644 index 0000000000..044d4c2d1a --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_cache_codec.py @@ -0,0 +1,126 @@ +import logging +from typing import Optional +from unittest.mock import patch + +import pytest +from pydantic import BaseModel, ValidationError + +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + + +class _SampleModel(BaseModel): + name: str + count: Optional[int] = None + + +class _SampleSubModel(_SampleModel): + pass + + +class TestCacheCodecSerialize: + def test_without_model_type_base_model_dumped_json_safe(self): + m = _SampleModel(name="a", count=1) + out = CacheCodec.serialize(m) + assert out == {"name": "a", "count": 1} + + def test_without_model_type_dict_unchanged(self): + d = {"name": "x"} + assert CacheCodec.serialize(d) is d + + def test_without_model_type_primitive_unchanged(self): + assert CacheCodec.serialize(42) == 42 + + def test_with_model_type_dict_validated_and_dumped(self): + out = CacheCodec.serialize({"name": "b", "count": 2}, model_type=_SampleModel) + assert out == {"name": "b", "count": 2} + + def test_with_model_type_base_model_validated_and_dumped(self): + m = _SampleModel(name="c", count=None) + out = CacheCodec.serialize(m, model_type=_SampleModel) + assert out == {"name": "c"} + + def test_with_model_type_exclude_none_on_dump(self): + out = CacheCodec.serialize({"name": "d"}, model_type=_SampleModel) + assert out == {"name": "d"} + assert "count" not in out + + def test_with_model_type_non_dict_non_model_passthrough(self): + assert CacheCodec.serialize("raw", model_type=_SampleModel) == "raw" + + def test_with_model_type_invalid_dict_raises(self): + with pytest.raises(ValidationError): + CacheCodec.serialize({"count": 1}, model_type=_SampleModel) + + def test_with_model_type_already_correct_instance_skips_revalidation(self): + """Fast-path: value is already model_type — model_validate must NOT be called.""" + m = _SampleModel(name="fast", count=7) + with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate: + out = CacheCodec.serialize(m, model_type=_SampleModel) + assert out == {"name": "fast", "count": 7} + mock_validate.assert_not_called() + + def test_with_model_type_subclass_instance_skips_revalidation(self): + """Subclass is isinstance of base → should also take the fast path.""" + sub = _SampleSubModel(name="sub", count=2) + with patch.object(_SampleModel, "model_validate", wraps=_SampleModel.model_validate) as mock_validate: + out = CacheCodec.serialize(sub, model_type=_SampleModel) + assert out == {"name": "sub", "count": 2} + mock_validate.assert_not_called() + + def test_with_model_type_dict_input_goes_through_model_validate(self): + """A dict value (not yet an instance) must still go through model_validate.""" + raw = {"name": "via-dict", "count": 5} + with patch.object( + _SampleModel, "model_validate", wraps=_SampleModel.model_validate + ) as mock_validate: + out = CacheCodec.serialize(raw, model_type=_SampleModel) + assert out == {"name": "via-dict", "count": 5} + mock_validate.assert_called_once() + + def test_with_model_type_incompatible_model_raises_validation_error(self): + """Passing a BaseModel whose fields don't satisfy model_type's required fields raises. + + _IncompatibleModel only has `foo: int`, so when Pydantic v2 extracts its + data and validates it against _SampleModel (which requires `name: str`), + a ValidationError is raised. + """ + + class _IncompatibleModel(BaseModel): + foo: int # missing required 'name' field of _SampleModel + + with pytest.raises(ValidationError): + CacheCodec.serialize(_IncompatibleModel(foo=1), model_type=_SampleModel) + + +class TestCacheCodecDeserialize: + def test_none_returns_none(self): + assert CacheCodec.deserialize(None, _SampleModel) is None + + def test_dict_validates_to_model(self): + m = CacheCodec.deserialize({"name": "e", "count": 3}, _SampleModel) + assert isinstance(m, _SampleModel) + assert m.name == "e" + assert m.count == 3 + + def test_instance_same_type_returned_as_is(self): + original = _SampleModel(name="f") + m = CacheCodec.deserialize(original, _SampleModel) + assert m is original + + def test_subclass_instance_accepted(self): + sub = _SampleSubModel(name="g") + m = CacheCodec.deserialize(sub, _SampleModel) + assert m is sub + + def test_wrong_type_returns_none(self): + assert CacheCodec.deserialize("not-a-dict", _SampleModel) is None + + def test_invalid_dict_returns_none_and_logs_warning(self, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = CacheCodec.deserialize({"count": 1}, _SampleModel) + assert out is None + assert any( + "CacheCodec.deserialize" in r.message and "_SampleModel" in r.message + for r in caplog.records + if r.levelno >= logging.WARNING + ), f"Expected deserialize validation warning. Records: {[r.message for r in caplog.records]}" diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 379ccf4d9a..5c86f9057a 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1049,3 +1049,159 @@ def test_reset_budget_windows_query_error_does_not_break_team_path(monkeypatch): asyncio.run(job.reset_budget_windows()) # must not raise prisma_client.db.litellm_teamtable.update.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Counter invalidation on budget reset +# --------------------------------------------------------------------------- + + +def _make_counter_invalidation_job(monkeypatch): + """Stub spend_counter_cache so we can observe invalidation calls.""" + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache.set_cache = MagicMock() + spend_counter_cache.redis_cache = MagicMock() + spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + + return spend_counter_cache + + +def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): + """Team-member budget reset clears the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + expired_budget = type("B", (), {"budget_id": "budget-1"}) + membership = type( + "Membership", + (), + {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_teammembership.find_many = AsyncMock( + return_value=[membership] + ) + prisma_client.db.litellm_teammembership.update_many = AsyncMock( + return_value={"count": 1} + ) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team_member:alice:team-x", value=0.0, ttl=60 + ) + counter_cache.redis_cache.async_set_cache.assert_any_await( + key="spend:team_member:alice:team-x", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_keys_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Key budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["key"] = [ + type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": now, + "id": "key-1", + "token": "sk-abc", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-abc", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_users_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """User budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + { + "spend": 50.0, + "budget_duration": "7d", + "budget_reset_at": now, + "id": "user-1", + "user_id": "alice", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:user:alice", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_teams_invalidates_redis_counter( + reset_budget_job, mock_prisma_client, monkeypatch +): + """Team budget reset must clear the Redis spend counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + now = datetime.now(timezone.utc) + mock_prisma_client.data["team"] = [ + type( + "Team", + (), + { + "spend": 200.0, + "budget_duration": "1mo", + "budget_reset_at": now, + "id": "team-1", + "team_id": "team-x", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:team:team-x", value=0.0, ttl=60 + ) + + +def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch): + """Resetting keys via budget tier must clear each linked key's counter.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + + expired_budget = type("B", (), {"budget_id": "budget-1"}) + linked_key = type("Key", (), {"token": "sk-linked"}) + + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[linked_key] + ) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 1} + ) + + job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) + asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) + + counter_cache.in_memory_cache.set_cache.assert_any_call( + key="spend:key:sk-linked", value=0.0, ttl=60 + ) diff --git a/tests/test_litellm/proxy/common_utils/test_static_asset_utils.py b/tests/test_litellm/proxy/common_utils/test_static_asset_utils.py new file mode 100644 index 0000000000..93f7ccc92c --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_static_asset_utils.py @@ -0,0 +1,97 @@ +""" +Unit tests for unauthenticated logo / favicon endpoint helpers. + +Local image paths are an existing deployment workflow, so the helper keeps +arbitrary local image paths working while refusing non-image files like +``/etc/passwd`` or ``/proc/self/environ``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.common_utils.static_asset_utils import ( + detect_local_image_media_type, + resolve_validated_local_image_path, +) + + +@pytest.mark.parametrize( + ("body", "media_type"), + [ + (b"\x89PNG\r\n\x1a\nfake png body", "image/png"), + (b"GIF89a fake gif body", "image/gif"), + (b"\xff\xd8\xff fake jpeg body", "image/jpeg"), + (b"RIFF\x00\x00\x00\x00WEBP fake webp body", "image/webp"), + (b"\x00\x00\x01\x00 fake ico body", "image/x-icon"), + ], +) +def test_detect_local_image_media_type_accepts_supported_images(body, media_type): + assert detect_local_image_media_type(body) == media_type + + +def test_detect_local_image_media_type_rejects_non_images(): + assert detect_local_image_media_type(b"root:x:0:0:root:/root:/bin/bash") is None + + +class TestResolveValidatedLocalImagePath: + def test_returns_resolved_path_for_arbitrary_local_image(self, tmp_path): + logo = tmp_path / "logo.png" + logo.write_bytes(b"\x89PNG\r\n\x1a\nfake png body") + + result = resolve_validated_local_image_path(str(logo)) + + assert result == (str(logo.resolve()), "image/png") + + def test_rejects_etc_passwd(self): + result = resolve_validated_local_image_path("/etc/passwd") + assert result is None + + def test_rejects_proc_self_environ(self): + result = resolve_validated_local_image_path("/proc/self/environ") + assert result is None + + def test_rejects_symlink_pointing_to_non_image(self, tmp_path): + secret = tmp_path / "secret.txt" + secret.write_text("password=hunter2") + symlink = tmp_path / "logo.png" + os.symlink(str(secret), str(symlink)) + + result = resolve_validated_local_image_path(str(symlink)) + + assert result is None + + def test_accepts_symlink_pointing_to_image(self, tmp_path): + logo = tmp_path / "real_logo.png" + logo.write_bytes(b"\x89PNG\r\n\x1a\nfake png body") + symlink = tmp_path / "logo.png" + os.symlink(str(logo), str(symlink)) + + result = resolve_validated_local_image_path(str(symlink)) + + assert result == (str(logo.resolve()), "image/png") + + def test_rejects_path_traversal_to_non_image(self, tmp_path): + assets_dir = tmp_path / "assets" + assets_dir.mkdir() + secret = tmp_path / "secret.txt" + secret.write_text("nope") + traversal = str(assets_dir / ".." / "secret.txt") + + result = resolve_validated_local_image_path(traversal) + + assert result is None + + def test_rejects_directory(self, tmp_path): + result = resolve_validated_local_image_path(str(tmp_path)) + assert result is None + + def test_rejects_nonexistent_file(self, tmp_path): + result = resolve_validated_local_image_path(str(tmp_path / "missing.jpg")) + assert result is None + + def test_rejects_empty_path(self): + assert resolve_validated_local_image_path("") is None diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py new file mode 100644 index 0000000000..8667348d22 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -0,0 +1,219 @@ +import json +from typing import Any + +import pytest + +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.proxy_server import UserAPIKeyCacheTTLEnum + + +class CapturingInMemoryCache(InMemoryCache): + """Records ``ttl`` passed into ``set_cache`` (what DualCache injects).""" + + def __init__(self) -> None: + super().__init__() + self.last_ttl: Any = None + + def set_cache(self, key, value, **kwargs): # type: ignore[override] + self.last_ttl = kwargs.get("ttl") + super().set_cache(key, value, **kwargs) + + +class FakeRedisCache(RedisCache): + """ + In-memory fake that enforces the UserApiKeyCache Redis payload contract. + + For user_api_key_cache entries we expect Redis to store a JSON object (dict) + produced by `CacheCodec.serialize(..., model_type=...)`. + + This fake: + - raises TypeError if the value is not a dict + - raises TypeError if the dict is not JSON-serializable + + Records the ``ttl`` kwarg DualCache forwards on each Redis write for tests. + """ + + def __init__(self): # noqa: super().__init__ skipped intentionally + self._store: dict[str, str] = {} + self.last_ttl: Any = None + + def set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] + if not isinstance(value, dict): + raise TypeError("FakeRedisCache only accepts dict payloads") + self.last_ttl = kwargs.get("ttl") + self._store[key] = json.dumps(value) + return True + + def get_cache(self, key: str, **kwargs): # type: ignore[override] + raw = self._store.get(key) + if raw is None: + return None + return json.loads(raw) + + async def async_set_cache(self, key: str, value: Any, **kwargs): # type: ignore[override] + if not isinstance(value, dict): + raise TypeError("FakeRedisCache only accepts dict payloads") + self.last_ttl = kwargs.get("ttl") + self._store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs): # type: ignore[override] + raw = self._store.get(key) + if raw is None: + return None + return json.loads(raw) + + def delete_cache(self, key: str): # type: ignore[override] + self._store.pop(key, None) + + async def async_delete_cache(self, key: str): # type: ignore[override] + self._store.pop(key, None) + + +def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: + # Minimal object (UserAPIKeyAuth inherits token from base view). + return UserAPIKeyAuth(token=token) + + +class TestUserApiKeyCache: + @pytest.mark.asyncio + async def test_async_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + self, + ): + """ + If ``general_settings.user_api_key_cache_ttl`` is absent, the proxy never + calls ``update_cache_ttl``; ``user_api_key_cache`` keeps + ``default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl``. + DualCache must forward that as the in-memory ``ttl`` kwarg on each set. + """ + mem = CapturingInMemoryCache() + cache = UserApiKeyCache( + in_memory_cache=mem, + redis_cache=FakeRedisCache(), + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + ) + await cache.async_set_cache( + "k", + _make_key_obj("t"), + model_type=UserAPIKeyAuth, + ) + expected = UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value + assert mem.last_ttl == expected + + def test_sync_set_in_memory_gets_enum_default_when_user_api_key_cache_ttl_omitted( + self, + ): + mem = CapturingInMemoryCache() + cache = UserApiKeyCache( + in_memory_cache=mem, + redis_cache=FakeRedisCache(), + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value, + ) + cache.set_cache("sk", _make_key_obj("s"), model_type=UserAPIKeyAuth) + assert mem.last_ttl == UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value + + @pytest.mark.asyncio + async def test_async_set_forwards_default_in_memory_ttl_to_redis_layer(self): + """ + DualCache injects missing ``ttl`` from ``default_in_memory_ttl`` into kwargs + before calling ``redis_cache.async_set_cache`` — Redis should receive the same + TTL as memory (matches proxy defaults: enum 60s). + """ + fake = FakeRedisCache() + cache = UserApiKeyCache( + redis_cache=fake, + default_in_memory_ttl=60, + ) + + await cache.async_set_cache( + key="ttl-key", + value=_make_key_obj("ttl-tok"), + model_type=UserAPIKeyAuth, + ) + + assert fake.last_ttl == 60 + + @pytest.mark.asyncio + async def test_async_set_explicit_ttl_override_reaches_redis(self): + fake = FakeRedisCache() + cache = UserApiKeyCache( + redis_cache=fake, + default_in_memory_ttl=60, + ) + + await cache.async_set_cache( + key="k", + value=_make_key_obj("x"), + model_type=UserAPIKeyAuth, + ttl=900, + ) + + assert fake.last_ttl == 900 + + def test_sync_set_forwards_default_in_memory_ttl_to_redis_layer(self): + fake = FakeRedisCache() + cache = UserApiKeyCache( + redis_cache=fake, + default_in_memory_ttl=45, + ) + cache.set_cache( + "sk", + _make_key_obj("sync"), + model_type=UserAPIKeyAuth, + ) + assert fake.last_ttl == 45 + + @pytest.mark.asyncio + async def test_async_set_typed_stores_serialized_payload_in_memory_and_redis(self): + cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + obj = _make_key_obj("abc") + + await cache.async_set_cache("k", obj, model_type=UserAPIKeyAuth) + + # In-memory hit should still be raw dict (not BaseModel) because wrapper + # stores the serialized payload into both layers. + raw = await cache.in_memory_cache.async_get_cache("k") # type: ignore[union-attr] + assert isinstance(raw, dict) + assert raw["token"] == "abc" + + # Redis should also hold the same serialized dict + redis_raw = await cache.redis_cache.async_get_cache("k") # type: ignore[union-attr] + assert redis_raw == raw + + @pytest.mark.asyncio + async def test_async_get_typed_returns_model_on_valid_hit(self): + cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + await cache.async_set_cache("k", {"token": "abc"}, model_type=UserAPIKeyAuth) + + value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) + assert value is not None + assert isinstance(value, UserAPIKeyAuth) + assert value.token == "abc" + + @pytest.mark.asyncio + async def test_async_get_typed_returns_none_on_validation_failure_after_hit(self): + cache = UserApiKeyCache(redis_cache=FakeRedisCache()) + + # Bypass UserApiKeyCache.serialize: CacheCodec rejects non-dict cached values + # for dict-based models (deserialize returns None). + await cache.in_memory_cache.async_set_cache( + key="k", value="invalid-payload-not-a-dict" + ) + + value = await cache.async_get_cache("k", model_type=UserAPIKeyAuth) + assert value is None + + def test_fake_redis_cache_rejects_non_json_serializable_values(self): + fake = FakeRedisCache() + + class NotSerializable: + pass + + with pytest.raises(TypeError): + fake.set_cache("k", NotSerializable()) + + with pytest.raises(TypeError): + fake.set_cache("k2", {"ok": NotSerializable()}) diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index a35f358f36..434f7953c2 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -4,7 +4,7 @@ Test to verify the Google GenAI proxy API endpoints """ import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -13,520 +13,171 @@ sys.path.insert( ) # Adds the parent directory to the system path -def test_google_generate_content_endpoint(): - """Test that the google_generate_content endpoint correctly routes requests""" - # Skip this test if we can't import the required modules due to missing dependencies - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient +def _build_test_client(): + from fastapi import FastAPI + from fastapi.testclient import TestClient - from litellm.proxy.google_endpoints.endpoints import router as google_router + from litellm.proxy.google_endpoints.endpoints import router as google_router + + app = FastAPI() + app.include_router(google_router) + return TestClient(app) + + +def _patch_base_process(return_value=None): + """Patch ProxyBaseLLMRequestProcessing.base_process_llm_request so endpoint + tests don't run the full pipeline. Returns the AsyncMock so callers can + inspect call args.""" + if return_value is None: + return_value = {"test": "response"} + return patch( + "litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new_callable=AsyncMock, + return_value=return_value, + ) + + +def test_google_generate_content_endpoint(): + """generateContent routes through ProxyBaseLLMRequestProcessing with the + agenerate_content route_type — that pipeline runs pre_call_hook + + during_call_hook + post_call_success_hook for every guardrail callback.""" + try: + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock the router's agenerate_content method - with patch("litellm.proxy.proxy_server.llm_router") as mock_router: - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Send a request to the endpoint + with _patch_base_process() as mock_base: response = client.post( "/v1beta/models/test-model:generateContent", json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - # Verify the response assert response.status_code == 200 - assert response.json() == {"test": "response"} - - # Verify that agenerate_content was called - mock_router.agenerate_content.assert_called_once() + mock_base.assert_called_once() + kwargs = mock_base.call_args.kwargs + assert kwargs["route_type"] == "agenerate_content" + assert kwargs["model"] == "test-model" def test_google_stream_generate_content_endpoint(): - """Test that the google_stream_generate_content endpoint correctly routes streaming requests""" - # Skip this test if we can't import the required modules due to missing dependencies + """streamGenerateContent must route through the same processor with the + streaming route_type so the guardrail pipeline runs.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock the router's agenerate_content_stream method to return a stream - async def mock_stream_generator(): - yield 'data: {"test": "stream_chunk_1"}\n\n' - yield 'data: {"test": "stream_chunk_2"}\n\n' - yield "data: [DONE]\n\n" - - with patch("litellm.proxy.proxy_server.llm_router") as mock_router: - mock_router.agenerate_content_stream = AsyncMock( - return_value=mock_stream_generator() - ) - - # Send a request to the endpoint + with ( + _patch_base_process() as mock_base, + patch( + "litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__", + return_value=None, + ) as mock_init, + ): response = client.post( "/v1beta/models/test-model:streamGenerateContent", json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) - # Verify the response assert response.status_code == 200 + mock_base.assert_called_once() + kwargs = mock_base.call_args.kwargs + assert kwargs["route_type"] == "agenerate_content_stream" + assert kwargs["model"] == "test-model" - # Verify that agenerate_content_stream was called with correct parameters - mock_router.agenerate_content_stream.assert_called_once() - call_args = mock_router.agenerate_content_stream.call_args - assert call_args[1]["stream"] is True - assert call_args[1]["model"] == "test-model" - assert call_args[1]["contents"] == [ + # stream=True must be forced into the data the processor receives. + init_kwargs = mock_init.call_args.kwargs + assert init_kwargs["data"]["stream"] is True + assert init_kwargs["data"]["model"] == "test-model" + assert init_kwargs["data"]["contents"] == [ {"role": "user", "parts": [{"text": "Hello"}]} ] -def test_google_generate_content_with_cost_tracking_metadata(): - """Test that the google_generate_content endpoint includes user metadata for cost tracking""" +def test_google_generate_content_data_flows_through_processor(): + """The body the client sends must reach ProxyBaseLLMRequestProcessing + intact so the pipeline can apply guardrails to it.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), + _patch_base_process(), patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, + "litellm.proxy.google_endpoints.endpoints.ProxyBaseLLMRequestProcessing.__init__", + return_value=None, + ) as mock_init, ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - # Simulate adding user metadata - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - "user_api_key_team_id": "test-team-id", - "user_api_key": "hashed-key", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request to the endpoint - response = client.post( + client.post( "/v1beta/models/test-model:generateContent", - json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, - headers={"Authorization": "Bearer sk-test-key"}, - ) - - # Verify the response - assert response.status_code == 200 - - # Verify that add_litellm_data_to_request was called - mock_add_data.assert_called_once() - - # Verify that agenerate_content was called with metadata - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that litellm_metadata exists and contains user information - assert "litellm_metadata" in called_data - assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" - assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id" - - -def test_google_stream_generate_content_with_cost_tracking_metadata(): - """Test that the google_stream_generate_content endpoint includes user metadata for cost tracking""" - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router - except ImportError as e: - pytest.skip(f"Skipping test due to missing dependency: {e}") - - # Create a FastAPI app and include the router (required for FastAPI 0.120+) - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock the router's agenerate_content_stream method to return a stream - mock_stream = AsyncMock() - mock_stream.__aiter__ = lambda self: mock_stream - mock_stream.__anext__.side_effect = StopAsyncIteration - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) - - # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - # Simulate adding user metadata - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - "user_api_key_team_id": "test-team-id", - "user_api_key": "hashed-key", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request to the endpoint - response = client.post( - "/v1beta/models/test-model:streamGenerateContent", - json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, - headers={"Authorization": "Bearer sk-test-key"}, - ) - - # Verify the response - assert response.status_code == 200 - - # Verify that add_litellm_data_to_request was called - mock_add_data.assert_called_once() - - # Verify that agenerate_content_stream was called with metadata - mock_router.agenerate_content_stream.assert_called_once() - call_args = mock_router.agenerate_content_stream.call_args - called_data = call_args[1] - - # Verify that litellm_metadata exists and contains user information - assert "litellm_metadata" in called_data - assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" - assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id" - # Verify stream is set to True - assert called_data["stream"] is True - - -def test_google_generate_content_with_system_instruction(): - """ - Test that systemInstruction is correctly passed through from the endpoint to the router. - - This test verifies the fix for systemInstruction being dropped when forwarding - requests to Vertex AI through the Google GenAI endpoint. - """ - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router - except ImportError as e: - pytest.skip(f"Skipping test due to missing dependency: {e}") - - # Create a FastAPI app and include the router - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Define the systemInstruction to test - system_instruction = {"parts": [{"text": "Your name is Doodle."}]} - - # Send a request with systemInstruction - response = client.post( - "/v1beta/models/gemini-2.5-pro:generateContent", json={ - "systemInstruction": system_instruction, - "contents": [ - {"parts": [{"text": "What is your name?"}], "role": "user"} - ], - }, - headers={"Authorization": "Bearer sk-test-key"}, - ) - - # Verify the response - assert response.status_code == 200 - - # Verify that agenerate_content was called - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that systemInstruction is present in the call arguments - assert "systemInstruction" in called_data - assert called_data["systemInstruction"] == system_instruction - assert ( - called_data["systemInstruction"]["parts"][0]["text"] - == "Your name is Doodle." - ) - - # Verify contents are also present - assert "contents" in called_data - assert len(called_data["contents"]) == 1 - assert called_data["contents"][0]["role"] == "user" - - -def test_google_generate_content_with_image_config(): - """ - Test that imageConfig is correctly passed through from generationConfig to the router. - - This test verifies that imageConfig parameters (aspectRatio, imageSize) are preserved - when forwarding requests to Google GenAI through the endpoint. - """ - try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router - except ImportError as e: - pytest.skip(f"Skipping test due to missing dependency: {e}") - - # Create a FastAPI app and include the router - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to pass through data unchanged - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request with generationConfig containing imageConfig - response = client.post( - "/v1beta/models/gemini-3-pro-image-preview:generateContent", - json={ - "contents": [ - { - "role": "user", - "parts": [ - { - "text": "Create a vibrant infographic about photosynthesis" - } - ], - } - ], + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "systemInstruction": {"parts": [{"text": "Your name is Doodle."}]}, "generationConfig": { "responseModalities": ["TEXT", "IMAGE"], "imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"}, }, }, - headers={"Authorization": "Bearer sk-test-key"}, ) - # Verify the response - assert response.status_code == 200 - - # Verify that agenerate_content was called - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that config is present in the call arguments - assert "config" in called_data - - # Verify that imageConfig is preserved in the config - assert "imageConfig" in called_data["config"] - assert called_data["config"]["imageConfig"]["aspectRatio"] == "9:16" - assert called_data["config"]["imageConfig"]["imageSize"] == "4K" - - # Verify that responseModalities is also preserved - assert "responseModalities" in called_data["config"] - assert called_data["config"]["responseModalities"] == ["TEXT", "IMAGE"] - - # Verify contents are also present - assert "contents" in called_data - assert len(called_data["contents"]) == 1 - assert called_data["contents"][0]["role"] == "user" + data = mock_init.call_args.kwargs["data"] + assert data["model"] == "test-model" + assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}] + assert data["systemInstruction"] == { + "parts": [{"text": "Your name is Doodle."}] + } + # generationConfig arrives intact here; the rename to `config` is + # done downstream in route_request (see test_route_llm_request). + assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"] + assert data["generationConfig"]["imageConfig"]["aspectRatio"] == "9:16" -def test_google_generate_content_metadata_and_trace_id_callbacks(): - """Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)""" +def test_google_generate_content_forwards_call_id_header(): + """The endpoint must forward the x-litellm-call-id header to the processor + so the helper can stamp it on the logging object. Trace continuity from + client → callbacks (S3, Langfuse, etc.) depends on this header surviving + the hop through these endpoints.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - # Create a FastAPI app and include the router - app = FastAPI() - app.include_router(google_router) - - # Create a test client - client = TestClient(app) - - # Mock all required proxy server dependencies - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, - ): - mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) - - # Mock add_litellm_data_to_request to return data with metadata - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - # Simulate adding user metadata - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - # Send a request to the endpoint with x-litellm-call-id header - test_call_id = "test-custom-call-id" - response = client.post( + with _patch_base_process() as mock_base: + client.post( "/v1beta/models/test-model:generateContent", json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, - headers={ - "Authorization": "Bearer sk-test-key", - "x-litellm-call-id": test_call_id, - }, + headers={"x-litellm-call-id": "trace-abc-123"}, ) - assert response.status_code == 200 - - mock_router.agenerate_content.assert_called_once() - call_args = mock_router.agenerate_content.call_args - called_data = call_args[1] - - # Verify that the litellm_logging_obj got assigned in the final called_data to router - assert "litellm_logging_obj" in called_data - assert "litellm_call_id" in called_data - assert called_data["litellm_call_id"] == test_call_id + forwarded_request = mock_base.call_args.kwargs["request"] + assert forwarded_request.headers.get("x-litellm-call-id") == "trace-abc-123" -def test_google_stream_generate_content_metadata_and_trace_id_callbacks(): - """Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks""" +def test_google_count_tokens_unchanged(): + """countTokens has its own path and isn't affected by the pipeline change.""" try: - from fastapi import FastAPI - from fastapi.testclient import TestClient - - from litellm.proxy.google_endpoints.endpoints import router as google_router + client = _build_test_client() except ImportError as e: pytest.skip(f"Skipping test due to missing dependency: {e}") - app = FastAPI() - app.include_router(google_router) - client = TestClient(app) + fake_response = MagicMock() + fake_response.original_response = { + "totalTokens": 7, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}], + } + fake_response.total_tokens = 7 - mock_stream = AsyncMock() - mock_stream.__aiter__ = lambda self: mock_stream - mock_stream.__anext__.side_effect = StopAsyncIteration - - with ( - patch("litellm.proxy.proxy_server.llm_router") as mock_router, - patch("litellm.proxy.proxy_server.general_settings", {}), - patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, - patch("litellm.proxy.proxy_server.version", "1.0.0"), - patch( - "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" - ) as mock_add_data, + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + return_value=fake_response, ): - mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) - - async def mock_add_litellm_data( - data, request, user_api_key_dict, proxy_config, general_settings, version - ): - data["litellm_metadata"] = { - "user_api_key_user_id": "test-user-id", - } - return data - - mock_add_data.side_effect = mock_add_litellm_data - - test_call_id = "test-custom-stream-call-id" response = client.post( - "/v1beta/models/test-model:streamGenerateContent", - json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]}, - headers={ - "Authorization": "Bearer sk-test-key", - "x-litellm-call-id": test_call_id, - }, + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, ) assert response.status_code == 200 - - mock_router.agenerate_content_stream.assert_called_once() - call_args = mock_router.agenerate_content_stream.call_args - called_data = call_args[1] - - assert "litellm_logging_obj" in called_data - assert "litellm_call_id" in called_data - assert called_data["litellm_call_id"] == test_call_id + body = response.json() + assert body["totalTokens"] == 7 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 55d92e9141..716b4470d2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -220,6 +220,27 @@ class TestToolPermissionGuardrail: assert tool_calls[0].id == "call_123" assert tool_calls[0].function.name == "Read" + def test_extract_tool_calls_legacy_function_call_format(self): + response = ModelResponse( + choices=[ + Choices( + message={ + "function_call": { + "name": "Read", + "arguments": '{"file_path": "/test/file.txt"}', + }, + } + ) + ] + ) + + tool_calls = self.guardrail._extract_tool_calls_from_response(response) + assert len(tool_calls) == 1 + assert isinstance(tool_calls[0], ChatCompletionMessageToolCall) + assert tool_calls[0].id == "legacy_function_call_0" + assert tool_calls[0].function.name == "Read" + assert tool_calls[0].function.arguments == '{"file_path": "/test/file.txt"}' + def test_extract_tool_calls_empty_response(self): response = ModelResponse(choices=[]) tool_calls = self.guardrail._extract_tool_calls_from_response(response) @@ -271,6 +292,31 @@ class TestToolPermissionGuardrail: data=data, user_api_key_dict=user_api_key_dict, response=response ) + @pytest.mark.asyncio + async def test_async_post_call_success_hook_with_denied_legacy_function_call_raises( + self, + ): + response = ModelResponse( + choices=[ + Choices( + message={ + "function_call": { + "name": "Read", + "arguments": "{}", + }, + } + ) + ] + ) + user_api_key_dict = UserAPIKeyAuth() + data = {"guardrails": ["test-tool-permission"]} + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + @pytest.mark.asyncio async def test_async_post_call_success_hook_param_patterns_allow(self): guardrail = ToolPermissionGuardrail( @@ -379,7 +425,9 @@ class TestToolPermissionGuardrail: assert "berri" in choice.message.content @pytest.mark.asyncio - async def test_async_post_call_success_hook_missing_arguments_default_allows(self): + async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule( + self, + ): guardrail = ToolPermissionGuardrail( guardrail_name="mail-guardrail", rules=[ @@ -405,9 +453,52 @@ class TestToolPermissionGuardrail: data = {"guardrails": ["mail-guardrail"]} with patch.object(guardrail, "should_run_guardrail", return_value=True): - await guardrail.async_post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response - ) + with pytest.raises(GuardrailRaisedException): + await guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "arguments", + [ + "{not-json", + '["owner@berri.ai"]', + ], + ) + async def test_async_post_call_success_hook_malformed_arguments_blocks_param_rule( + self, arguments + ): + guardrail = ToolPermissionGuardrail( + guardrail_name="mail-guardrail", + rules=[ + { + "id": "deny_gmail", + "tool_name": r"^mail_mcp-send_email$", + "decision": "deny", + "allowed_param_patterns": {"to[]": r"^.+@gmail\.com$"}, + } + ], + default_action="allow", + on_disallowed_action="block", + ) + + tool_call = { + "function": { + "name": "mail_mcp-send_email", + "arguments": arguments, + }, + "type": "function", + } + response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})]) + user_api_key_dict = UserAPIKeyAuth() + data = {"guardrails": ["mail-guardrail"]} + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) @pytest.mark.asyncio async def test_async_pre_call_hook_block_mode(self): @@ -430,6 +521,65 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_legacy_functions(self): + data = { + "functions": [ + {"name": "Bash", "description": "allowed"}, + {"name": "Read", "description": "denied"}, + ] + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await self.guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + assert excinfo.value.status_code == 400 + + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_named_legacy_function_call(self): + data = { + "functions": [{"name": "Bash"}], + "function_call": {"name": "Read"}, + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await self.guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + assert excinfo.value.status_code == 400 + + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_named_tool_choice(self): + data = { + "tools": [{"type": "function", "function": {"name": "Bash"}}], + "tool_choice": {"type": "function", "function": {"name": "Read"}}, + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await self.guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + assert excinfo.value.status_code == 400 + @pytest.mark.asyncio async def test_async_pre_call_hook_uses_custom_template(self): guardrail = ToolPermissionGuardrail( @@ -491,6 +641,41 @@ class TestToolPermissionGuardrail: assert "Bash" in tool_names assert "Read" not in tool_names + @pytest.mark.asyncio + async def test_async_pre_call_hook_rewrite_mode_filters_legacy_functions(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="test-tool-permission", + rules=self.test_rules, + default_action="deny", + on_disallowed_action="rewrite", + ) + data = { + "functions": [ + {"name": "Bash", "description": "allowed"}, + {"name": "Read", "description": "denied"}, + ], + "function_call": {"name": "Read"}, + "tools": [ + {"type": "function", "function": {"name": "Bash"}}, + ], + "tool_choice": {"type": "function", "function": {"name": "Read"}}, + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + new_data = await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + + assert isinstance(new_data, dict) + assert [function["name"] for function in new_data["functions"]] == ["Bash"] + assert new_data["function_call"] == "none" + assert new_data["tool_choice"] == "none" + def test_modify_response_with_permission_errors(self): # Setup a response with one tool_call tool_call = ChatCompletionMessageToolCall( @@ -522,6 +707,40 @@ class TestToolPermissionGuardrail: assert isinstance(choice.message.content, str) assert "Permission denied" in choice.message.content + def test_modify_response_with_permission_errors_filters_legacy_function_call(self): + response = ModelResponse( + choices=[ + Choices( + message={ + "function_call": { + "name": "Read", + "arguments": "{}", + }, + "content": "", + } + ) + ] + ) + tool_call = self.guardrail._extract_tool_calls_from_response(response)[0] + denied_tools = [ + ( + tool_call, + PermissionError( + tool_name="Read", + rule_id="deny_read", + message="Tool 'Read' denied by rule 'deny_read'", + ), + ) + ] + + self.guardrail._modify_response_with_permission_errors(response, denied_tools) + + choice = response.choices[0] + assert isinstance(choice, Choices) + assert choice.message.function_call is None + assert isinstance(choice.message.content, str) + assert "Permission denied" in choice.message.content + class TestToolPermissionGuardrailIntegration: """Integration tests for Tool Permission Guardrail""" diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index ba26014235..353e67c9f7 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -778,3 +778,797 @@ def test_get_callback_identifier_custom_logger_registry_and_fallback(): result = get_callback_identifier(my_callback_function) # Should fall back to callback_name() which returns __name__ assert result == "my_callback_function" + + +# --------------------------------------------------------------------------- +# /health response shape: model-access scoping and display-field allowlist +# --------------------------------------------------------------------------- +# These tests pin the contract that the /health response (a) only includes +# deployments the calling key is allowed to see, and (b) does not return +# provider routing fields like api_base / api_version. They guard against +# regressions that would widen the response shape. + + +@pytest.mark.asyncio +async def test_health_endpoint_filters_model_list_by_user_access(): + """ + health_endpoint() should restrict _llm_model_list to deployments whose + model_name appears in user_api_key_dict.models before running the health + check. A key scoped to ["model-a"] should only see model-a in the result, + not other deployments configured on the proxy. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-b.test", + "api_version": "2024-10-21", + }, + "model_info": {"id": "id-b"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["model-a"], + ) + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [], + "healthy_count": 0, + "unhealthy_count": 0, + } + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + from fastapi import Response + + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + + assert ( + "model_list" in captured + ), "health_endpoint did not call _perform_health_check_and_save" + returned_names = {m["model_name"] for m in captured["model_list"]} + assert returned_names == { + "model-a" + }, f"health_endpoint did not scope model_list to caller access: {returned_names}" + + +@pytest.mark.asyncio +async def test_health_endpoint_filters_background_cache_by_user_access(): + """ + When background_health_checks is enabled, health_endpoint() should also + scope the cached result to the caller's allowed models rather than + returning the cache verbatim. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-b.test", + }, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "api_base": "https://example-a.test", + }, + { + "model": "openai/gpt-4o", + "model_id": "id-b", + "api_base": "https://example-b.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["model-a"], + ) + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + from fastapi import Response + + # Pass model=None, model_id=None explicitly: direct calls to the + # handler skip FastAPI's Query() resolution, so unspecified params + # would otherwise carry the Query() sentinel (which is truthy). + result = await health_endpoint( + response=Response(), + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + # Sanity: the source cache had two entries before scoping; the scoping + # step is what reduces it to one. (This guards against the test passing + # vacuously when the cache filter drops everything because cached + # entries lack the model_id key — both entries carry model_id above.) + assert len(cached_results["healthy_endpoints"]) == 2 + assert all( + ep.get("model_id") for ep in cached_results["healthy_endpoints"] + ), "test fixture invariant: every cached entry must carry a model_id" + + # The non-admin caller must not see api_base on the returned cache entries. + returned = result.get("healthy_endpoints", []) + assert ( + len(returned) == 1 + ), f"expected exactly one cached entry after scoping, got {len(returned)}" + assert returned[0]["model_id"] == "id-a" + assert "api_base" not in returned[0] + assert result["healthy_count"] == 1 + assert result["unhealthy_count"] == 0 + + +@pytest.mark.asyncio +async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not(): + """ + A proxy admin should still see ``api_base`` and ``api_version`` in the + /health response so they can tell which Vertex region / Azure resource + + API version is healthy. A non-admin caller must not — both fields + should be stripped, and the response should carry a notice header so + non-admin clients can detect the change programmatically. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + ] + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "api_base": "https://us-central1-aiplatform.googleapis.com/v1/projects/p", + "api_version": "2024-10-21", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + admin_key = UserAPIKeyAuth( + api_key="hashed-admin-key", + models=["model-a"], + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + non_admin_key = UserAPIKeyAuth( + api_key="hashed-user-key", + models=["model-a"], + ) + + common_patches = [ + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ] + + for p in common_patches: + p.start() + try: + admin_response = Response() + non_admin_response = Response() + admin_result = await health_endpoint( + response=admin_response, + user_api_key_dict=admin_key, + model=None, + model_id=None, + ) + non_admin_result = await health_endpoint( + response=non_admin_response, + user_api_key_dict=non_admin_key, + model=None, + model_id=None, + ) + finally: + for p in common_patches: + p.stop() + + admin_eps = admin_result.get("healthy_endpoints", []) + non_admin_eps = non_admin_result.get("healthy_endpoints", []) + + assert len(admin_eps) == 1 + assert ( + admin_eps[0]["api_base"] + == "https://us-central1-aiplatform.googleapis.com/v1/projects/p" + ), "admin must see the full api_base so they can identify the region" + assert ( + admin_eps[0]["api_version"] == "2024-10-21" + ), "admin must see api_version so they can distinguish provider deployments" + + assert len(non_admin_eps) == 1 + assert "api_base" not in non_admin_eps[0] + assert "api_version" not in non_admin_eps[0] + + # Non-admin response must advertise that api_base/api_version were + # withheld so clients that previously parsed them can detect the change. + assert ( + non_admin_response.headers.get("Litellm-Health-Field-Notice") + == "api_base and api_version are admin-only on this endpoint" + ) + assert "Litellm-Health-Field-Notice" not in admin_response.headers + + # Stripping must produce a copy — the shared cache must still carry the + # routing fields so the next admin caller can read them. + cached_first = cached_results["healthy_endpoints"][0] + assert ( + cached_first["api_base"] + == "https://us-central1-aiplatform.googleapis.com/v1/projects/p" + ) + assert cached_first["api_version"] == "2024-10-21" + + +@pytest.mark.asyncio +async def test_health_endpoint_warns_when_scoped_models_lack_model_id(): + """ + When a scoped key's accessible models exist on the proxy but none of the + matching deployments expose a ``model_info.id``, the cache filter drops + everything. The response should include a structured ``warnings`` field + so the caller can distinguish "no deployments configured" from + "deployments excluded due to missing model_info.id". + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + # Intentionally no model_info.id — this is the misconfiguration + # the warnings field is meant to flag. + "model_info": {}, + }, + ] + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "api_base": "https://example-a.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-user-key", + models=["model-a"], + ) + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 0 + assert "warnings" in result, ( + "empty cache result must surface a warnings field so the caller " + "can distinguish 'no deployments' from 'deployments excluded'" + ) + assert any("model_info.id" in w for w in result["warnings"]) + + +@pytest.mark.asyncio +async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cache(): + """ + A non-admin scoped to model-a must not be able to read model-b's cached + health entry by guessing its model_id. Before the fix, + _resolve_targeted_model_ids returned {model_id} unconditionally, so the + cache filter was driven by an unvalidated ID and the global cache + leaked id-b's entry to the caller. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", # caller has no access + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-b", + "api_base": "https://leaky-internal.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-scoped", + models=["model-a"], + ) + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + # llm_router None here means the model_id 404 lookup short-circuits; + # we patch _llm_model_list directly instead to drive the cache path. + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + # Calling with model="model-b" rather than model_id="id-b" because + # the model_id branch raises 404 when llm_router is None. The bug + # being verified is the same: targeted resolver must drop entries + # not in the caller's scoped model_list. With the fix, the result + # has no leaked endpoints and the targeted-503 path fires. + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-b", + model_id=None, + ) + + leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])} + leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])} + assert ( + "id-b" not in leaked_ids + ), "background cache leaked an out-of-scope deployment to a scoped caller" + assert result["healthy_count"] == 0 + assert response.status_code == 503 + + +@pytest.mark.asyncio +async def test_health_endpoint_503_for_targeted_unhealthy_model_under_background_cache_admin(): + """ + With background_health_checks enabled, an admin calling /health?model=foo + must get 503 when foo specifically has zero healthy endpoints — even if + other unrelated models in the cache are healthy. Without the cache-path + filter, the global healthy_count would mask the targeted failure. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", # the unhealthy target + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", # an unrelated healthy model + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-b"}, + ], + "unhealthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"}, + ], + "healthy_count": 1, + "unhealthy_count": 1, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + model_id=None, + ) + + assert response.status_code == 503 + # Body must be scoped to the targeted model — not the global cache. + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 1 + returned_ids = {ep["model_id"] for ep in result.get("unhealthy_endpoints", [])} + assert returned_ids == {"id-a"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_endpoints(): + """ + /health?model=foo must return 503 when the targeted model resolves but + has zero healthy endpoints. Body shape stays the same so existing + parsers still work; only the HTTP status changes. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "error": "boom", + } + ], + "healthy_count": 0, + "unhealthy_count": 1, + } + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + ) + + assert response.status_code == 503 + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 1 + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endpoints(): + """ + /health?model=foo with a healthy endpoint must keep returning the + default 200. Verifies the 503 path doesn't fire when healthy_count > 0. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [{"model": "openai/gpt-4o", "model_id": "id-a"}], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + response = Response() + # Default Response() exposes status_code as None; the endpoint should + # leave it alone for the healthy path. + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + ) + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_health_endpoint_no_model_param_returns_200_even_when_zero_healthy(): + """ + The non-targeted /health (no model / model_id query) preserves the + legacy 200 behavior even when healthy_count == 0. Existing K8s probes + and dashboards depend on this; only the targeted call became 5xx-aware. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"} + ], + "healthy_count": 0, + "unhealthy_count": 1, + } + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + # Pass model=None, model_id=None explicitly: when invoked through + # FastAPI, the Query(None) defaults resolve to None, but direct + # function calls in unit tests receive Query() sentinel objects + # (which are truthy). The explicit None mirrors production routing. + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_health_readiness_returns_503_when_db_disconnected(): + """ + When a Prisma client is configured but its health_check fails, the + readiness probe should mark the worker as unhealthy via the HTTP + status — not just a body field — so K8s removes the pod from the + Service endpoints. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope")) + mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope")) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result["db"] == "disconnected" + assert result["status"] == "healthy" # body shape unchanged for back-compat + + +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_db_connected(): + """Happy path: connected DB keeps the legacy 200.""" + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result["db"] == "connected" + + +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_no_db_configured(): + """ + `prisma_client is None` means the operator chose not to use a DB. That + is a valid configuration — the worker should still report ready. We + only flip to 503 when a DB *was* configured but is unreachable. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", None): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result["db"] == "Not connected" + + +def test_clean_endpoint_data_strips_credentials_keeps_routing_fields(): + """ + _clean_endpoint_data() drops credentials but leaves api_base / + api_version intact — the per-caller hide/show happens in the endpoint + layer based on user role, not in the cleaning helper. This guarantees + proxy admins continue to see those fields in the /health response. + """ + from litellm.proxy.health_check import _clean_endpoint_data + + raw = { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "api_base": "https://example.test/v1", + "api_version": "2024-10-21", + "aws_access_key_id": "AKIAEXAMPLE", + } + + cleaned = _clean_endpoint_data(raw, details=True) + + assert "api_key" not in cleaned + assert "aws_access_key_id" not in cleaned + assert cleaned.get("api_base") == "https://example.test/v1" + assert cleaned.get("api_version") == "2024-10-21" diff --git a/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py new file mode 100644 index 0000000000..ceea5de799 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -0,0 +1,489 @@ +""" +Tests validating TOCTOU race condition in batch + dynamic rate limiters. + +Issue: rate-limit check (read_only=True) and counter increment happen as two +separate awaits. Concurrent requests all observe the same pre-increment state, +all pass validation, then all increment — bypassing the limit. + +Vulnerable code paths: +- litellm/proxy/hooks/batch_rate_limiter.py:181-248 + (_check_and_increment_batch_counters: should_rate_limit(read_only=True) + → validate → async_increment_tokens_with_ttl_preservation) +- litellm/proxy/hooks/dynamic_rate_limiter_v3.py:463-548 + (_check_rate_limits PHASE 1 read_only check → PHASE 3 increment) + +These tests EXPECTED to fail against current (vulnerable) code and pass once +check-and-increment becomes atomic. +""" + +import asyncio +import os +import sys +from typing import List + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm import DualCache, Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler, +) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.proxy.utils import InternalUsageCache, hash_token + + +def _make_phase1_barrier(num_concurrent: int, timeout: float = 0.1): + """ + Sync primitive that, pre-fix, forces all N concurrent coroutines to finish + their read-only Phase 1 check before any proceeds to Phase 3 increment — + mimicking asyncio I/O scheduling under load on the vulnerable code. + + Wraps `should_rate_limit` so on `read_only=True` calls it waits until N + callers arrive (TOCTOU window opened) OR `timeout` elapses (post-fix path: + the limiter's serialization lock prevents N from ever reaching the + barrier; the timeout lets the holder proceed so the lock can do its job). + + Pre-fix: barrier fills before timeout → all see same state → bypass observed. + Post-fix: only lock-holder reaches barrier → times out → serial execution + enforces limit. + """ + arrived = 0 + all_arrived = asyncio.Event() + + def wrap(original): + async def patched(*args, **kwargs): + result = await original(*args, **kwargs) + if kwargs.get("read_only"): + nonlocal arrived + arrived += 1 + if arrived >= num_concurrent: + all_arrived.set() + try: + await asyncio.wait_for(all_arrived.wait(), timeout=timeout) + except asyncio.TimeoutError: + pass + return result + + return patched + + return wrap + + +@pytest.mark.asyncio +async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou(): + """ + 5 concurrent batch submissions of 40 tokens each against TPM=100 limit. + + Sequential semantics: only 2 batches fit (2 * 40 = 80 ≤ 100, 3rd at 120 fails). + With TOCTOU: all 5 succeed → 200 tokens consumed, 100% over limit. + + Demonstrates batch_rate_limiter.py:183-248 multi-phase flaw. + """ + NUM_CONCURRENT = 5 + BATCH_TOKENS = 40 + TPM_LIMIT = 100 + + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("toctou-batch-key"), + tpm_limit=TPM_LIMIT, + rpm_limit=1000, + ) + batch_usage = BatchFileUsage(total_tokens=BATCH_TOKENS, request_count=1) + + barrier = _make_phase1_barrier(NUM_CONCURRENT) + rate_limiter.should_rate_limit = barrier(rate_limiter.should_rate_limit) + + results = await asyncio.gather( + *[ + batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=batch_usage, + ) + for _ in range(NUM_CONCURRENT) + ], + return_exceptions=True, + ) + + successes = [r for r in results if not isinstance(r, Exception)] + rejections = [r for r in results if isinstance(r, Exception)] + total_consumed = len(successes) * BATCH_TOKENS + max_allowed_successes = TPM_LIMIT // BATCH_TOKENS # 2 + + assert len(successes) <= max_allowed_successes, ( + f"TOCTOU bypass: {len(successes)}/{NUM_CONCURRENT} concurrent batches " + f"passed despite TPM={TPM_LIMIT}. Consumed {total_consumed} tokens " + f"({total_consumed - TPM_LIMIT} over limit). " + f"Atomic check-and-increment would allow ≤{max_allowed_successes}. " + f"Rejections: {len(rejections)}" + ) + + +@pytest.mark.asyncio +async def test_batch_limiter_uses_atomic_check_and_increment(): + """ + Regression test: batch limiter routes through + `atomic_check_and_increment_by_n` rather than the legacy two-phase + pattern (read_only=True check + separate async_increment_tokens_with_ttl_preservation). + + Ensures future refactors don't reintroduce the TOCTOU window. + """ + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + call_log: List[str] = [] + original_atomic = rate_limiter.atomic_check_and_increment_by_n + original_should = rate_limiter.should_rate_limit + + async def logging_atomic(*args, **kwargs): + call_log.append("atomic_check_and_increment_by_n") + return await original_atomic(*args, **kwargs) + + async def logging_should(*args, **kwargs): + call_log.append(f"should_rate_limit(read_only={kwargs.get('read_only')})") + return await original_should(*args, **kwargs) + + rate_limiter.atomic_check_and_increment_by_n = logging_atomic + rate_limiter.should_rate_limit = logging_should + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("atomic-test-key"), + tpm_limit=10000, + rpm_limit=1000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=50, request_count=1), + ) + + assert "atomic_check_and_increment_by_n" in call_log, ( + f"Batch limiter must route through atomic_check_and_increment_by_n. " + f"Calls observed: {call_log}" + ) + legacy_calls = [c for c in call_log if c.startswith("should_rate_limit(")] + assert not legacy_calls, ( + f"Batch limiter must not call should_rate_limit directly (legacy " + f"two-phase pattern). Observed: {legacy_calls}" + ) + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): + """ + DynamicRateLimitHandler PHASE 1 (read_only check) → PHASE 3 (increment) + is non-atomic: dynamic_rate_limiter_v3.py:463-548. + + With TPM=100 model capacity and 5 concurrent priority="high" requests + each consuming the full model_saturation_check counter, all observe the + same Phase 1 state (counter=0), all pass, all proceed to Phase 3. + + Sequential atomic semantics would block requests once the model counter + reaches its limit. TOCTOU lets all pass Phase 1 simultaneously. + """ + NUM_CONCURRENT = 10 + MODEL_RPM = 2 + # Sequential bound: dynamic limiter rejects when `counter > current_limit` + # (strict `>`), so a request whose Phase 1 sees counter=RPM still passes + # (RPM is not strictly greater). Atomic execution therefore admits up to + # RPM + 1 successes before the next sees counter > RPM. + MAX_SEQUENTIAL_SUCCESSES = MODEL_RPM + 1 + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "toctou-dyn-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": MODEL_RPM, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + barrier = _make_phase1_barrier(NUM_CONCURRENT) + handler.v3_limiter.should_rate_limit = barrier(handler.v3_limiter.should_rate_limit) + + from litellm.types.router import ModelGroupInfo + + model_group_info = ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=MODEL_RPM, + tpm=None, + ) + + async def one_request(idx: int): + user = UserAPIKeyAuth(api_key=hash_token(f"dyn-key-{idx}")) + user.metadata = {"priority": "high"} + try: + await handler._check_rate_limits( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + return "OK" + except Exception as e: + return e + + results = await asyncio.gather( + *[one_request(i) for i in range(NUM_CONCURRENT)], + return_exceptions=True, + ) + successes = [r for r in results if r == "OK"] + + assert len(successes) <= MAX_SEQUENTIAL_SUCCESSES, ( + f"TOCTOU bypass in DynamicRateLimitHandler: {len(successes)}/{NUM_CONCURRENT} " + f"concurrent requests passed Phase 1 + Phase 3 despite model RPM={MODEL_RPM}. " + f"Atomic check-and-increment would block once counter > RPM " + f"(at most {MAX_SEQUENTIAL_SUCCESSES} sequential successes)." + ) + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): + """ + Regression test: dynamic limiter's enforced descriptors flow through + `atomic_check_and_increment_by_n`, not the legacy + read_only=True check followed by a separate read_only=False increment. + + When priority is enforced (saturation >= threshold), priority_model is + bundled into the atomic call alongside model_saturation_check. When not + enforced, priority counter is incremented for tracking only. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "atomic-dyn-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + atomic_descriptors_observed: List[List[str]] = [] + original_atomic = handler.v3_limiter.atomic_check_and_increment_by_n + + async def logging_atomic(*args, **kwargs): + ds = kwargs.get("descriptors") or (args[0] if args else []) + atomic_descriptors_observed.append([d["key"] for d in ds]) + return await original_atomic(*args, **kwargs) + + handler.v3_limiter.atomic_check_and_increment_by_n = logging_atomic + + from litellm.types.router import ModelGroupInfo + + user = UserAPIKeyAuth(api_key=hash_token("dyn-atomic-key")) + user.metadata = {"priority": "high"} + + await handler._check_rate_limits( + model=model, + model_group_info=ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=None, + tpm=1000, + ), + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + + assert atomic_descriptors_observed, ( + "Dynamic limiter must route enforced descriptors through " + "atomic_check_and_increment_by_n (no legacy read_only=True / " + "separate-increment pattern)." + ) + assert "model_saturation_check" in atomic_descriptors_observed[0], ( + f"Expected model_saturation_check in atomic descriptor set. " + f"Got: {atomic_descriptors_observed}" + ) + + +@pytest.mark.asyncio +async def test_batch_zero_token_consumes_rpm_only(): + """ + Zero-token batch (e.g. metadata-only call) should still increment RPM + counter but NOT TPM counter. + + Edge case from review: `if inc_amount <= 0: continue` in + `atomic_check_and_increment_by_n` skips descriptor counters whose + increment is zero. Verifies asymmetric quota consumption is intentional + and observable: an RPM-bounded but TPM-free request path stays bounded + by RPM alone. + """ + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("zero-token-key"), + tpm_limit=100, + rpm_limit=3, + ) + zero_batch = BatchFileUsage(total_tokens=0, request_count=1) + + # 3 zero-token batches must succeed (RPM=3 allows). 4th must hit RPM cap, + # NOT TPM (because token counter never increments past 0). + for i in range(3): + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=zero_batch, + ) + + # Inspect counters: RPM key incremented to 3, TPM key absent (or 0). + rpm_key = rate_limiter.create_rate_limit_keys( + "api_key", user_api_key_dict.api_key or "", "requests" + ) + tpm_key = rate_limiter.create_rate_limit_keys( + "api_key", user_api_key_dict.api_key or "", "tokens" + ) + rpm_val = await internal_usage_cache.async_get_cache( + key=rpm_key, litellm_parent_otel_span=None, local_only=True + ) + tpm_val = await internal_usage_cache.async_get_cache( + key=tpm_key, litellm_parent_otel_span=None, local_only=True + ) + assert int(rpm_val or 0) == 3, f"RPM counter must reach 3, got {rpm_val}" + assert tpm_val in ( + None, + 0, + "0", + ), f"TPM counter must remain unset/0 for zero-token batches, got {tpm_val}" + + # 4th attempt: RPM exhausted -> 429. + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=zero_batch, + ) + assert exc.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): + """ + Fail-closed guard: when atomic_check_and_increment_by_n returns + overall_code=OVER_LIMIT but with a descriptor_key the dispatcher does + not recognize, the dynamic limiter must raise 429 rather than silently + fall through. + + Reproduces by patching atomic_check_and_increment_by_n to return an + OVER_LIMIT response carrying an unknown descriptor_key. + """ + from fastapi import HTTPException + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fail-closed-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + async def fake_atomic(*args, **kwargs): + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "tokens", + "descriptor_key": "future_unrecognized_descriptor", + } + ], + } + + handler.v3_limiter.atomic_check_and_increment_by_n = fake_atomic + + from litellm.types.router import ModelGroupInfo + + user = UserAPIKeyAuth(api_key=hash_token("fail-closed-key")) + user.metadata = {"priority": "high"} + + with pytest.raises(HTTPException) as exc: + await handler._check_rate_limits( + model=model, + model_group_info=ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=None, + tpm=1000, + ), + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + assert ( + exc.value.status_code == 429 + ), f"Expected 429 fail-closed on unknown descriptor; got {exc.value.status_code}" diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index cd2eb78958..016e10859b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -738,15 +738,32 @@ def test_delete_access_group_patches_cached_team_and_key( return_value=None ) - # Build cached key object (returned from user_api_key_cache) - if key_cache_group_ids is not None: - cached_key = UserAPIKeyAuth( - token="hashed-key-1", - access_group_ids=list(key_cache_group_ids), + # user_api_key_cache is queried both for teams (fallback after dual_cache) and + # hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key) + # would wrongly serve the key blob for ``team_id:team-1`` and trigger team patching. + # Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects + # inconsistently across Python/unittest versions; sync returns are awaited as immediate results. + def user_cache_get_side_effect(*args, **kwargs): + cache_key = ( + kwargs.get("key") if "key" in kwargs else (args[0] if args else None) ) - mock_cache.async_get_cache = AsyncMock(return_value=cached_key) - else: - mock_cache.async_get_cache = AsyncMock(return_value=None) + if cache_key == "team_id:team-1": + if team_cache_group_ids is None: + return None + return LiteLLM_TeamTableCachedObj( + team_id="team-1", + access_group_ids=list(team_cache_group_ids), + ) + if cache_key == "hashed-key-1": + if key_cache_group_ids is None: + return None + return UserAPIKeyAuth( + token="hashed-key-1", + access_group_ids=list(key_cache_group_ids), + ) + return None + + mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_side_effect) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 @@ -803,7 +820,7 @@ def test_delete_access_group_patches_cached_team_and_key( def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): - """Delete correctly patches a key cached as a raw dict (not UserAPIKeyAuth).""" + """Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize).""" client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( client_and_mocks ) @@ -826,12 +843,24 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): return_value=None ) - # Key cached as a plain dict (as can happen with Redis serialization) + # Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that. + cached_key_payload = { + "token": "hashed-key-dict", + "access_group_ids": ["ag-to-delete", "ag-other"], + } + + def user_cache_get_dict_when_key_matches(*args, **kwargs): + cache_key = ( + kwargs.get("key") if "key" in kwargs else (args[0] if args else None) + ) + if cache_key == "team_id:team-1": + return None + if cache_key == "hashed-key-dict": + return UserAPIKeyAuth.model_validate(cached_key_payload) + return None + mock_cache.async_get_cache = AsyncMock( - return_value={ - "token": "hashed-key-dict", - "access_group_ids": ["ag-to-delete", "ag-other"], - } + side_effect=user_cache_get_dict_when_key_matches ) resp = client.delete("/v1/access_group/ag-to-delete") 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 0362d6f97d..e668672dd2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -5512,6 +5512,9 @@ async def test_update_team_guardrails_with_org_id(): return_value=mock_updated_team ) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + # async_get_cache must be an AsyncMock so `await` in get_org_object works + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() # Mock llm_router mock_router = MagicMock() diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index eecfcaa035..a0ae95df58 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2,9 +2,9 @@ import asyncio import json import os import sys +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch -import httpx import pytest from fastapi import HTTPException, Request @@ -25,7 +25,6 @@ from litellm.proxy.management_endpoints.ui_sso import ( SSOAuthenticationHandler, _setup_team_mappings, _sync_user_role_from_jwt_role_map, - determine_role_from_groups, normalize_email, process_sso_jwt_access_token, ) @@ -1471,13 +1470,13 @@ class TestAuthCallbackRouting: from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX # Test CLI state detection logic - cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-test123" + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-test1234567890" # This mimics the logic in auth_callback if cli_state and cli_state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): - # Extract the key ID from the state + # Extract the login ID from the state key_id = cli_state.split(":", 1)[1] - assert key_id == "sk-test123" + assert key_id == "cli-test1234567890" else: assert False, "CLI state should have been detected" @@ -1510,13 +1509,13 @@ class TestGoogleLoginCLIIntegration: # Test the CLI state generation logic used in google_login source = "litellm-cli" - key = "sk-test123" + key = "cli-test1234567890" cli_state = SSOAuthenticationHandler._get_cli_state(source=source, key=key) assert cli_state is not None assert cli_state.startswith("litellm-session-token:") - assert "sk-test123" in cli_state + assert "cli-test1234567890" in cli_state def test_google_login_no_cli_state_when_missing_params(self): """Test that google_login doesn't generate CLI state when CLI parameters are missing""" @@ -1526,8 +1525,8 @@ class TestGoogleLoginCLIIntegration: test_cases = [ (None, None), ("litellm-cli", None), - (None, "sk-test123"), - ("wrong-source", "sk-test123"), + (None, "cli-test1234567890"), + ("wrong-source", "cli-test1234567890"), ] for source, key in test_cases: @@ -1634,19 +1633,19 @@ class TestSSOStateHandling: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler state = SSOAuthenticationHandler._get_cli_state( - source="litellm-cli", key="sk-test123" + source="litellm-cli", key="cli-test1234567890" ) assert state is not None assert state.startswith("litellm-session-token:") - assert "sk-test123" in state + assert "cli-test1234567890" in state def test_get_cli_state_invalid_source(self): """Test generating CLI state with invalid source""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler state = SSOAuthenticationHandler._get_cli_state( - source="invalid_source", key="sk-test123" + source="invalid_source", key="cli-test1234567890" ) assert state is None @@ -1663,40 +1662,40 @@ class TestSSOStateHandling: """Test generating CLI state without source""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - state = SSOAuthenticationHandler._get_cli_state(source=None, key="sk-test123") + state = SSOAuthenticationHandler._get_cli_state( + source=None, key="cli-test1234567890" + ) assert state is None - def test_get_cli_state_with_existing_key(self): - """Test generating CLI state with existing_key embedded in state parameter""" + def test_get_cli_state_ignores_existing_key(self): + """Test CLI state does not embed an existing key""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler state = SSOAuthenticationHandler._get_cli_state( source="litellm-cli", - key="sk-new-key-123", + key="cli-new-key-1234567890", existing_key="sk-existing-key-456", ) assert state is not None assert state.startswith("litellm-session-token:") - assert "sk-new-key-123" in state - assert "sk-existing-key-456" in state - # Verify the format: {PREFIX}:{key}:{existing_key} - assert state == "litellm-session-token:sk-new-key-123:sk-existing-key-456" + assert "cli-new-key-1234567890" in state + assert "sk-existing-key-456" not in state + assert state == "litellm-session-token:cli-new-key-1234567890" def test_get_cli_state_without_existing_key(self): """Test generating CLI state without existing_key""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler state = SSOAuthenticationHandler._get_cli_state( - source="litellm-cli", key="sk-new-key-789", existing_key=None + source="litellm-cli", key="cli-new-key-789123456", existing_key=None ) assert state is not None assert state.startswith("litellm-session-token:") - assert "sk-new-key-789" in state - # Verify the format: {PREFIX}:{key} (no third part) - assert state == "litellm-session-token:sk-new-key-789" + assert "cli-new-key-789123456" in state + assert state == "litellm-session-token:cli-new-key-789123456" assert state.count(":") == 1 # Only one colon separator @@ -1708,44 +1707,37 @@ class TestStateRouting: from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX # Test CLI state format - cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-test123" + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-test1234567890" assert cli_state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:") # Test extraction of key from state key_id = cli_state.split(":", 1)[1] - assert key_id == "sk-test123" + assert key_id == "cli-test1234567890" - def test_cli_state_parsing_with_existing_key(self): - """Test parsing CLI state with existing_key embedded""" + def test_cli_state_parsing_uses_single_login_id(self): + """Test parsing CLI state with a single login ID""" from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX - # State format: {PREFIX}:{key}:{existing_key} - cli_state = ( - f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-key-456:sk-existing-key-789" - ) + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-key-456123" # Parse as done in auth_callback - state_parts = cli_state.split(":", 2) # Split into max 3 parts + state_parts = cli_state.split(":", 1) key_id = state_parts[1] if len(state_parts) > 1 else None - existing_key = state_parts[2] if len(state_parts) > 2 else None - assert key_id == "sk-new-key-456" - assert existing_key == "sk-existing-key-789" + assert key_id == "cli-new-key-456123" - def test_cli_state_parsing_without_existing_key(self): - """Test parsing CLI state without existing_key""" + def test_cli_state_parsing_without_extra_segments(self): + """Test parsing CLI state uses a single login ID""" from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX # State format: {PREFIX}:{key} - cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-key-999" + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-key-999123" # Parse as done in auth_callback - state_parts = cli_state.split(":", 2) # Split into max 3 parts + state_parts = cli_state.split(":", 1) key_id = state_parts[1] if len(state_parts) > 1 else None - existing_key = state_parts[2] if len(state_parts) > 2 else None - assert key_id == "sk-new-key-999" - assert existing_key is None + assert key_id == "cli-new-key-999123" def test_non_cli_state_detection(self): """Test detection of non-CLI state parameters""" @@ -2007,6 +1999,178 @@ class TestCustomUISSO: class TestCLIKeyRegenerationFlow: """Test the end-to-end CLI key regeneration flow""" + def test_cli_sso_login_id_validation_restricts_charset(self): + """Test CLI SSO login IDs only allow the generated character set""" + from litellm.proxy.management_endpoints.ui_sso import ( + _is_valid_cli_sso_login_id, + ) + + assert _is_valid_cli_sso_login_id("cli-test_1234567890") + assert not _is_valid_cli_sso_login_id("cli-session") + assert not _is_valid_cli_sso_login_id("cli-test\n1234567890") + assert not _is_valid_cli_sso_login_id("cli-test\x001234567890") + assert not _is_valid_cli_sso_login_id("sk-test1234567890") + + @pytest.mark.asyncio + async def test_cli_sso_start_creates_bound_flow(self): + """Test CLI SSO start creates a polling secret bound flow""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + _normalize_cli_sso_user_code, + cli_sso_start, + ) + + mock_request = MagicMock(spec=Request) + mock_request.client = SimpleNamespace(host="127.0.0.1") + mock_request.headers = {} + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 1 + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + result = await cli_sso_start(request=mock_request) + + assert result["login_id"].startswith("cli-") + assert result["poll_secret"] + assert result["user_code"] + + mock_cache.increment_cache.assert_called_once() + assert mock_cache.increment_cache.call_args.kwargs["ttl"] == 60 + mock_cache.set_cache.assert_called_once() + flow_data = mock_cache.set_cache.call_args.kwargs["value"] + assert flow_data["poll_secret_hash"] == _hash_cli_sso_secret( + result["poll_secret"] + ) + assert flow_data["user_code_hash"] == _hash_cli_sso_secret( + _normalize_cli_sso_user_code(result["user_code"]) + ) + assert flow_data["poll_secret_hash"] != result["poll_secret"] + assert flow_data["user_code_hash"] != result["user_code"] + + @pytest.mark.asyncio + async def test_cli_sso_start_rate_limits_by_client_ip(self): + """Test CLI SSO start enforces a coarse per-client rate limit""" + from litellm.proxy.management_endpoints.ui_sso import cli_sso_start + + mock_request = MagicMock(spec=Request) + mock_request.client = SimpleNamespace(host="127.0.0.1") + mock_request.headers = {} + mock_cache = MagicMock() + mock_cache.increment_cache.return_value = 31 + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with pytest.raises(HTTPException) as exc_info: + await cli_sso_start(request=mock_request) + + assert exc_info.value.status_code == 429 + mock_cache.set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_cli_sso_complete_verifies_user_code(self): + """Test CLI SSO complete marks a session as verified""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + _normalize_cli_sso_user_code, + cli_sso_complete, + ) + + mock_request = MagicMock(spec=Request) + mock_request.body = AsyncMock( + return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" + ) + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "user_code_hash": _hash_cli_sso_secret( + _normalize_cli_sso_user_code("ABCD-EFGH") + ), + "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"), + "sso_complete": True, + "user_code_verified": False, + "session_data": {"user_id": "test-user-123"}, + } + + with ( + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache), + patch( + "litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", + return_value="Success", + ), + ): + result = await cli_sso_complete( + request=mock_request, login_id="cli-session-4567890" + ) + + assert result.status_code == 200 + flow_data = mock_cache.set_cache.call_args.kwargs["value"] + assert flow_data["user_code_verified"] is True + + @pytest.mark.asyncio + async def test_cli_sso_complete_requires_callback_token(self): + """Test CLI SSO complete requires the callback-delivered token""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + _normalize_cli_sso_user_code, + cli_sso_complete, + ) + + mock_request = MagicMock(spec=Request) + mock_request.body = AsyncMock(return_value=b"user_code=ABCD-EFGH") + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "user_code_hash": _hash_cli_sso_secret( + _normalize_cli_sso_user_code("ABCD-EFGH") + ), + "browser_complete_token_hash": _hash_cli_sso_secret("browser-token"), + "sso_complete": True, + "user_code_verified": False, + "session_data": {"user_id": "test-user-123"}, + } + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with pytest.raises(HTTPException) as exc_info: + await cli_sso_complete( + request=mock_request, login_id="cli-session-4567890" + ) + + assert exc_info.value.status_code == 400 + mock_cache.set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_cli_sso_complete_waits_for_callback_before_token_checks(self): + """Test CLI SSO complete returns not-ready before verification checks""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + _normalize_cli_sso_user_code, + cli_sso_complete, + ) + + mock_request = MagicMock(spec=Request) + mock_request.body = AsyncMock( + return_value=b"user_code=ABCD-EFGH&browser_complete_token=browser-token" + ) + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "user_code_hash": _hash_cli_sso_secret( + _normalize_cli_sso_user_code("ABCD-EFGH") + ), + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with pytest.raises(HTTPException) as exc_info: + await cli_sso_complete( + request=mock_request, login_id="cli-session-4567890" + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "CLI login is not ready" + mock_request.body.assert_not_awaited() + mock_cache.set_cache.assert_not_called() + @pytest.mark.asyncio async def test_cli_sso_callback_stores_session(self): """Test CLI SSO callback stores session data in cache for JWT generation""" @@ -2017,7 +2181,7 @@ class TestCLIKeyRegenerationFlow: mock_request = MagicMock(spec=Request) # Test data - session_key = "sk-session-456" + session_key = "cli-session-4567890" # Mock user info mock_user_info = LiteLLM_UserTable( @@ -2032,6 +2196,16 @@ class TestCLIKeyRegenerationFlow: # Mock cache mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": "poll-secret-hash", + "user_code_hash": "user-code-hash", + "sso_complete": False, + "user_code_verified": False, + "session_data": None, + } + mock_request.url_for.return_value = ( + "https://test.example.com/sso/cli/complete/cli-session-4567890" + ) with ( patch( @@ -2049,7 +2223,6 @@ class TestCLIKeyRegenerationFlow: result = await cli_sso_callback( request=mock_request, key=session_key, - existing_key=None, result=mock_sso_result, ) @@ -2062,14 +2235,18 @@ class TestCLIKeyRegenerationFlow: assert session_key in call_args.kwargs["key"] # Verify session data structure - session_data = call_args.kwargs["value"] + flow_data = call_args.kwargs["value"] + session_data = flow_data["session_data"] + assert flow_data["sso_complete"] is True + assert flow_data["user_code_verified"] is False + assert isinstance(flow_data["browser_complete_token_hash"], str) assert session_data["user_id"] == "test-user-123" assert session_data["user_role"] == "internal_user" assert session_data["teams"] == ["team1", "team2"] assert session_data["models"] == ["gpt-4"] # Verify TTL - assert call_args.kwargs["ttl"] == 600 # 10 minutes + assert call_args.kwargs["ttl"] == 600 assert result.status_code == 200 # Verify response contains success message (response is HTML) @@ -2078,10 +2255,13 @@ class TestCLIKeyRegenerationFlow: @pytest.mark.asyncio async def test_cli_poll_key_returns_teams_for_selection(self): """Test CLI poll endpoint returns teams for user selection when multiple teams exist""" - from litellm.proxy.management_endpoints.ui_sso import cli_poll_key + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) # Test data - session_key = "sk-session-789" + session_key = "cli-session-789123" session_data = { "user_id": "test-user-456", "user_role": "internal_user", @@ -2091,11 +2271,20 @@ class TestCLIKeyRegenerationFlow: # Mock cache mock_cache = MagicMock() - mock_cache.get_cache.return_value = session_data + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): # Act - First poll without team_id - result = await cli_poll_key(key_id=session_key, team_id=None) + result = await cli_poll_key( + key_id=session_key, + team_id=None, + x_litellm_cli_poll_secret="poll-secret", + ) # Assert - should return teams list for selection assert result["status"] == "ready" @@ -2108,16 +2297,72 @@ class TestCLIKeyRegenerationFlow: mock_cache.delete_cache.assert_not_called() @pytest.mark.asyncio - async def test_auth_callback_routes_to_cli_with_existing_key(self): - """Test that auth_callback properly routes CLI requests and extracts existing_key from state parameter""" + async def test_cli_poll_key_requires_poll_secret(self): + """Test CLI poll endpoint rejects callers without the polling secret""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": { + "user_id": "test-user-456", + "user_role": "internal_user", + "teams": [], + "models": ["gpt-4"], + }, + } + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + with pytest.raises(HTTPException) as exc_info: + await cli_poll_key(key_id="cli-session-789123", team_id=None) + + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_cli_poll_key_waits_for_user_code_verification(self): + """Test CLI poll endpoint stays pending until user code verification""" + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) + + mock_cache = MagicMock() + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": False, + "session_data": { + "user_id": "test-user-456", + "user_role": "internal_user", + "teams": [], + "models": ["gpt-4"], + }, + } + + with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_cache): + result = await cli_poll_key( + key_id="cli-session-789123", + team_id=None, + x_litellm_cli_poll_secret="poll-secret", + ) + + assert result == {"status": "pending"} + + @pytest.mark.asyncio + async def test_auth_callback_routes_to_cli(self): + """Test that auth_callback properly routes CLI requests""" from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX from litellm.proxy.management_endpoints.ui_sso import auth_callback - # Mock request (no query params needed - existing_key is in state) + # Mock request mock_request = MagicMock(spec=Request) - # CLI state with existing_key embedded: {PREFIX}:{key}:{existing_key} - cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-session-key-456:sk-existing-cli-key-123" + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:cli-new-session-key-456" # Mock the CLI callback and required proxy server components mock_result = {"user_id": "test-user", "email": "test@example.com"} @@ -2142,16 +2387,14 @@ class TestCLIKeyRegenerationFlow: # Act await auth_callback(request=mock_request, state=cli_state) - # Assert - existing_key should be extracted from state parameter mock_cli_callback.assert_called_once_with( request=mock_request, - key="sk-new-session-key-456", - existing_key="sk-existing-cli-key-123", + key="cli-new-session-key-456", result=mock_result, ) def test_get_redirect_url_does_not_include_existing_key_in_url(self): - """Test that redirect URL generation does NOT include existing_key in URL (uses state parameter instead)""" + """Test that redirect URL generation does NOT include existing_key in URL""" from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Mock request @@ -2194,10 +2437,13 @@ class TestCLIKeyRegenerationFlow: async def test_cli_poll_key_generates_jwt_with_team(self): """Test CLI poll endpoint generates JWT when team_id is provided""" from litellm.proxy._types import LiteLLM_UserTable - from litellm.proxy.management_endpoints.ui_sso import cli_poll_key + from litellm.proxy.management_endpoints.ui_sso import ( + _hash_cli_sso_secret, + cli_poll_key, + ) # Test data - session_key = "sk-session-999" + session_key = "cli-session-999123" selected_team = "team-b" session_data = { "user_id": "test-user-789", @@ -2217,7 +2463,12 @@ class TestCLIKeyRegenerationFlow: # Mock cache mock_cache = MagicMock() - mock_cache.get_cache.return_value = session_data + mock_cache.get_cache.return_value = { + "poll_secret_hash": _hash_cli_sso_secret("poll-secret"), + "sso_complete": True, + "user_code_verified": True, + "session_data": session_data, + } mock_jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.test.token" @@ -2235,7 +2486,11 @@ class TestCLIKeyRegenerationFlow: ) # Act - Second poll with team_id - result = await cli_poll_key(key_id=session_key, team_id=selected_team) + result = await cli_poll_key( + key_id=session_key, + team_id=selected_team, + x_litellm_cli_poll_secret="poll-secret", + ) # Assert - should return JWT assert result["status"] == "ready" @@ -2901,7 +3156,7 @@ class TestGetGenericSSORedirectParams: from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler # Arrange - cli_state = "litellm-session-token:sk-test123" + cli_state = "litellm-session-token:cli-test1234567890" with patch.dict(os.environ, {"GENERIC_CLIENT_STATE": "env_state_value"}): # Act diff --git a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py index 9fd244d9c3..310ee11573 100644 --- a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py +++ b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py @@ -26,6 +26,15 @@ async def fake_valid_auth(request, api_key): return +async def fake_valid_auth_reads_body(request, api_key, **kwargs): + """ + Like real user_api_key_auth, consumes the ASGI body stream. Regression test + for successful auth passing a drained receive to the inner app (hang). + """ + await request.body() + return + + async def fake_invalid_auth(request, api_key): print("running fake invalid auth", request, api_key) # Simulate invalid auth by raising an exception. @@ -62,6 +71,28 @@ def app_with_middleware(): return app +def test_valid_auth_metrics_after_body_consumed(app_with_middleware, monkeypatch): + """ + Auth that reads the request body must not cause /metrics to hang on success. + """ + litellm.require_auth_for_metrics_endpoint = True + monkeypatch.setattr( + "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth", + fake_valid_auth_reads_body, + ) + + client = TestClient(app_with_middleware) + headers = {SpecialHeaders.openai_authorization.value: "valid"} + + response = client.get("/metrics", headers=headers) + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"} + + response = client.get("/metrics/", headers=headers) + assert response.status_code == 200, response.text + assert response.json() == {"msg": "metrics OK"} + + def test_valid_auth_metrics(app_with_middleware, monkeypatch): """ Test that a request to /metrics (and /metrics/) with valid auth headers passes. diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_auth_default.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_auth_default.py new file mode 100644 index 0000000000..4cac1cb4d3 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_auth_default.py @@ -0,0 +1,136 @@ +""" +Regression tests for the pass-through endpoint auth-default fix +(GHSA-7h34-mmrh-6g58). + +Two failures the fix closes: + +1. ``PassThroughGenericEndpoint.auth`` defaulted to ``False`` — an + admin who added a pass-through to ``general_settings`` without + explicitly setting ``auth: true`` shipped an unauthenticated + forwarder. +2. Setting ``auth: true`` was rejected at startup unless the operator + had a LiteLLM Enterprise license, leaving OSS deployments with no + safe configuration. + +The fix flips the default to ``True`` (safe-by-default) and removes +the enterprise gate so OSS operators can register an authenticated +pass-through. The runtime check in ``user_api_key_auth.py`` also now +defaults to ``True`` so a config dict (raw, not Pydantic) without an +``auth`` key still requires authentication. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import PassThroughGenericEndpoint +from litellm.proxy.auth.user_api_key_auth import ( + check_api_key_for_custom_headers_or_pass_through_endpoints, +) +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + _register_pass_through_endpoint, +) + + +def test_passthrough_auth_defaults_to_true(): + # Regression: an admin who configures a pass-through without setting + # auth explicitly used to ship an unauthenticated forwarder. The + # default is now safe. + endpoint = PassThroughGenericEndpoint( + path="/canary-forwarder", + target="https://postman-echo.com/get", + ) + assert endpoint.auth is True + + +def test_passthrough_auth_can_still_be_explicitly_disabled(): + # Operators who genuinely need an unauthenticated forwarder (e.g. + # public webhook receiver) can opt in explicitly. + endpoint = PassThroughGenericEndpoint( + path="/public-webhook", + target="https://example.com/webhook", + auth=False, + ) + assert endpoint.auth is False + + +@pytest.mark.asyncio +async def test_register_passthrough_with_auth_true_works_for_oss(monkeypatch): + # Regression: setting ``auth: true`` used to raise at startup + # unless ``premium_user`` was True, leaving OSS with no safe + # configuration. + app = MagicMock(spec=FastAPI) + visited: set = set() + + endpoint = PassThroughGenericEndpoint( + path="/forwarder", + target="https://example.com", + auth=True, + ) + + # Should not raise; OSS premium_user=False is allowed to use auth=True. + await _register_pass_through_endpoint( + endpoint=endpoint, + app=app, + premium_user=False, + visited_endpoints=visited, + ) + + +@pytest.mark.asyncio +async def test_runtime_check_treats_missing_auth_key_as_authenticated(): + # The runtime dispatch in user_api_key_auth pulls + # pass_through_endpoints from general_settings as raw dicts (the + # Pydantic default never applies). A dict without an ``auth`` key + # must default to "authenticated" — without this, the previous + # behaviour (``endpoint.get("auth") is not True`` -> True -> empty + # auth) ships an unauthenticated forwarder. + request = MagicMock() + request.headers = {} + raw_endpoint_no_auth_key = { + "path": "/forwarder", + "target": "https://example.com", + # ``auth`` deliberately omitted + } + + result = await check_api_key_for_custom_headers_or_pass_through_endpoints( + request=request, + route="/forwarder", + pass_through_endpoints=[raw_endpoint_no_auth_key], + api_key="sk-1234", + ) + + # Result is the api_key string (auth is REQUIRED for this endpoint + # — flow continues to normal key validation), NOT an empty + # ``UserAPIKeyAuth()`` (which was the unauthenticated-forwarder + # bug). + assert result == "sk-1234" + + +@pytest.mark.asyncio +async def test_runtime_check_explicit_auth_false_still_skips_validation(): + # Operators who explicitly set ``auth: False`` get the legacy + # behaviour — an empty UserAPIKeyAuth, no key required. + from litellm.proxy._types import UserAPIKeyAuth + + request = MagicMock() + request.headers = {} + raw_endpoint_auth_false = { + "path": "/public-webhook", + "target": "https://example.com", + "auth": False, + } + + result = await check_api_key_for_custom_headers_or_pass_through_endpoints( + request=request, + route="/public-webhook", + pass_through_endpoints=[raw_endpoint_auth_false], + api_key="", + ) + + assert isinstance(result, UserAPIKeyAuth) diff --git a/tests/test_litellm/proxy/test_filter_models_by_team_access_group.py b/tests/test_litellm/proxy/test_filter_models_by_team_access_group.py new file mode 100644 index 0000000000..2d8a9f30c1 --- /dev/null +++ b/tests/test_litellm/proxy/test_filter_models_by_team_access_group.py @@ -0,0 +1,236 @@ +""" +Tests for _filter_models_by_team_id resolving access group names. + +Verifies that when a team's `models` field contains an access group name +(e.g., "Group-A"), the filter resolves it to the member model names before +looking up deployments — matching the behavior of the auth path in +auth_checks.py:model_in_access_group(). +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.proxy_server import _filter_models_by_team_id + + +def _make_model(model_name: str, model_id: str, access_groups: list[str] = None): + """Helper to build a model dict matching the router's format.""" + return { + "model_name": model_name, + "litellm_params": {"model": model_name}, + "model_info": { + "id": model_id, + "access_groups": access_groups or [], + }, + } + + +def _make_team(models: list[str], team_id: str = "team_alpha"): + """Helper to build a mock team DB object.""" + mock = MagicMock() + mock.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Team Alpha", + "models": models, + "max_budget": None, + "spend": 0.0, + "blocked": False, + "members_with_roles": [], + "metadata": {}, + } + return mock + + +@pytest.mark.asyncio +async def test_filter_resolves_access_group_names(): + """ + When team.models contains an access group name, _filter_models_by_team_id + should resolve it to the member models and return only those deployments. + """ + # Models on the proxy + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + gpt5 = _make_model("gpt-5", "id-2", ["Group-A"]) + claude = _make_model("claude-3", "id-3", ["Group-B"]) + + all_models = [gpt4o, gpt5, claude] + + # Router mock + mock_router = MagicMock() + # get_model_access_groups returns {group_name: [model_names]} + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o", "gpt-5"], + "Group-B": ["claude-3"], + } + + # get_model_list returns deployments matching a model_name + def fake_get_model_list(model_name=None, team_id=None): + return [m for m in all_models if m["model_name"] == model_name] + + mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list) + + # Team has models: ["Group-A"] — an access group name, not a literal model + team_db = _make_team(models=["Group-A"]) + + # Prisma mock + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + result_ids = {m["model_info"]["id"] for m in result} + # Should include gpt-4o and gpt-5 (Group-A), but NOT claude-3 (Group-B) + assert result_ids == { + "id-1", + "id-2", + }, f"Expected Group-A models only, got {result_ids}" + + # Verify DB fallback query received resolved model names, not access group name + call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1] + assert set(call_kwargs["where"]["model_name"]["in"]) == { + "gpt-4o", + "gpt-5", + }, "find_many should receive resolved model names, not the access group name" + + +@pytest.mark.asyncio +async def test_filter_resolves_mix_of_access_groups_and_literal_names(): + """ + When team.models contains both an access group name and a literal model name, + both should be resolved correctly. + """ + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + gpt5 = _make_model("gpt-5", "id-2", ["Group-A"]) + claude = _make_model("claude-3", "id-3", ["Group-B"]) + mistral = _make_model("mistral-large", "id-4", []) # no access group + + all_models = [gpt4o, gpt5, claude, mistral] + + mock_router = MagicMock() + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o", "gpt-5"], + "Group-B": ["claude-3"], + } + + def fake_get_model_list(model_name=None, team_id=None): + return [m for m in all_models if m["model_name"] == model_name] + + mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list) + + # Team has access to Group-A (access group) + mistral-large (literal name) + team_db = _make_team(models=["Group-A", "mistral-large"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + result_ids = {m["model_info"]["id"] for m in result} + # Group-A models + mistral-large, but NOT claude-3 + assert result_ids == { + "id-1", + "id-2", + "id-4", + }, f"Expected Group-A + mistral-large, got {result_ids}" + + +@pytest.mark.asyncio +async def test_filter_excludes_models_from_other_access_group(): + """ + Models belonging only to a different access group must not appear in results. + """ + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + claude = _make_model("claude-3", "id-3", ["Group-B"]) + llama = _make_model("llama-4", "id-4", ["Group-B"]) + + all_models = [gpt4o, claude, llama] + + mock_router = MagicMock() + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o"], + "Group-B": ["claude-3", "llama-4"], + } + + def fake_get_model_list(model_name=None, team_id=None): + return [m for m in all_models if m["model_name"] == model_name] + + mock_router.get_model_list = MagicMock(side_effect=fake_get_model_list) + + team_db = _make_team(models=["Group-A"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + result_names = {m["model_name"] for m in result} + assert "claude-3" not in result_names, "Group-B model should not be accessible" + assert "llama-4" not in result_names, "Group-B model should not be accessible" + assert "gpt-4o" in result_names, "Group-A model should be accessible" + + +@pytest.mark.asyncio +async def test_filter_db_fallback_receives_resolved_model_names(): + """ + When get_model_list returns no results (forcing the DB fallback path), + the DB query should receive resolved model names, not the raw access group name. + """ + gpt4o = _make_model("gpt-4o", "id-1", ["Group-A"]) + all_models = [gpt4o] + + mock_router = MagicMock() + mock_router.get_model_access_groups.return_value = { + "Group-A": ["gpt-4o", "gpt-5"], + } + # get_model_list returns nothing — forces reliance on the DB fallback + mock_router.get_model_list = MagicMock(return_value=[]) + + team_db = _make_team(models=["Group-A"]) + + # DB returns a model that the router didn't find + mock_db_model = MagicMock() + mock_db_model.model_id = "id-db-1" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_db) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( + return_value=[mock_db_model] + ) + + result = await _filter_models_by_team_id( + all_models=all_models, + team_id="team_alpha", + prisma_client=mock_prisma, + llm_router=mock_router, + ) + + # Verify DB query received resolved names, not "Group-A" + call_kwargs = mock_prisma.db.litellm_proxymodeltable.find_many.call_args[1] + queried_names = set(call_kwargs["where"]["model_name"]["in"]) + assert queried_names == { + "gpt-4o", + "gpt-5", + }, f"DB query should receive resolved model names, got {queried_names}" + assert "Group-A" not in queried_names, "Raw access group name should not be in DB query" diff --git a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py index a5605a281f..8bc39c93ee 100644 --- a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py +++ b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py @@ -1,39 +1,35 @@ -from fastapi import FastAPI -from fastapi.openapi.utils import get_openapi -from fastapi.routing import APIRoute - -from litellm.proxy._lazy_openapi_snapshot import _routes_with_stable_unique_ids +from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids -def test_routes_with_stable_unique_ids_splits_multi_method_routes() -> None: - app = FastAPI() - - async def proxy_route() -> dict: - return {} - - app.add_api_route( - "/proxy/{endpoint:path}", - proxy_route, - methods=["GET", "POST", "DELETE"], - ) - - routes = [route for route in app.routes if isinstance(route, APIRoute)] - stable_routes = _routes_with_stable_unique_ids(routes) - - assert [route.methods for route in stable_routes] == [ - {"DELETE"}, - {"GET"}, - {"POST"}, - ] - - openapi = get_openapi(title="test", version="1", routes=stable_routes) - path_ops = openapi["paths"]["/proxy/{endpoint}"] - - operation_ids = { - method: operation["operationId"] for method, operation in path_ops.items() +def test_normalize_operation_ids_uses_each_http_method(): + paths = { + "/proxy/{endpoint}": { + "delete": {"operationId": "proxy_route_proxy__endpoint__put"}, + "get": {"operationId": "proxy_route_proxy__endpoint__put"}, + "post": {"operationId": "proxy_route_proxy__endpoint__put"}, + "put": {"operationId": "proxy_route_proxy__endpoint__put"}, + } } - assert operation_ids == { - "delete": "proxy_route_proxy__endpoint__delete", - "get": "proxy_route_proxy__endpoint__get", - "post": "proxy_route_proxy__endpoint__post", + + _normalize_operation_ids(paths) + + operations = paths["/proxy/{endpoint}"] + assert operations["delete"]["operationId"] == "proxy_route_proxy__endpoint__delete" + assert operations["get"]["operationId"] == "proxy_route_proxy__endpoint__get" + assert operations["post"]["operationId"] == "proxy_route_proxy__endpoint__post" + assert operations["put"]["operationId"] == "proxy_route_proxy__endpoint__put" + + +def test_normalize_operation_ids_preserves_custom_ids(): + paths = { + "/proxy/{endpoint}": { + "get": {"operationId": "custom_operation"}, + "post": {"operationId": "custom_operation"}, + } } + + _normalize_operation_ids(paths) + + operations = paths["/proxy/{endpoint}"] + assert operations["get"]["operationId"] == "custom_operation" + assert operations["post"]["operationId"] == "custom_operation" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7a96f6cbd1..465ce579e0 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -457,6 +457,59 @@ def test_fallback_login_has_no_deprecation_banner(client_no_auth): assert " import` that would defeat lazy loading. + # Importing proxy_server in a subprocess and diffing sys.modules + # would also work, but takes 60-120 s and flakes on slow CI runners. + import re + from pathlib import Path - check = ( - "import sys; " - "from litellm.proxy.proxy_server import app; " # noqa: F401 - "heavy = [" - "'litellm.proxy._experimental.mcp_server.rest_endpoints'," - "'litellm.proxy._experimental.mcp_server.server'," - "'litellm.proxy.management_endpoints.config_override_endpoints'," - "'litellm.proxy.guardrails.guardrail_endpoints'," - "'litellm.proxy.openai_evals_endpoints.endpoints'," - "]; " - "still_present = [m for m in heavy if m in sys.modules]; " - "print('PRESENT_AT_STARTUP:', still_present)" + from litellm.proxy._lazy_features import LAZY_FEATURES + + proxy_server_src = ( + Path(__file__).resolve().parents[3] / "litellm/proxy/proxy_server.py" + ).read_text() + + leaks = [] + for feat in LAZY_FEATURES: + # Anchor at column 0 — indented imports inside function bodies + # are fine (deferred until the function runs). + pattern = ( + rf"^(from\s+{re.escape(feat.module_path)}\s+import|" + rf"import\s+{re.escape(feat.module_path)})" + ) + if re.search(pattern, proxy_server_src, re.MULTILINE): + leaks.append(feat.module_path) + + assert not leaks, ( + "proxy_server.py top-level imports a lazy feature module — these " + f"should be loaded via LazyFeatureMiddleware: {leaks}" ) - result = subprocess.run( - [sys.executable, "-c", check], - capture_output=True, - text=True, - timeout=120, - ) - # Last non-empty line of stdout (skip warnings printed before) - out_lines = [ - line for line in result.stdout.strip().splitlines() if line.strip() - ] - report = next((line for line in out_lines if "PRESENT_AT_STARTUP" in line), "") - assert report, f"no report emitted (stderr: {result.stderr[-500:]})" - assert ( - "PRESENT_AT_STARTUP: []" in report - ), f"expected no heavy modules at startup, got: {report}" class TestLazyFeatureMiddleware: @@ -5723,3 +5943,90 @@ class TestLazyFeatureMiddleware: assert attempts == [ "called" ], f"failing register_fn should be invoked once, not on every request; got {attempts}" + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory(): + """When Redis is reachable and cleanly returns None (TTL expired, + counter genuinely absent), the read must reseed from DB - NOT fall + through to per-pod in-memory which only contains this pod's writes. + + Pre-fix in multi-pod deployments, in-memory contained a stale local + subset (e.g. $30) while DB had the true cross-pod total ($500). The + fall-through returned $30, enforcement passed, bypass. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + + # Per-pod stale in-memory: only this pod's writes, not cross-pod truth. + counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0) + + # Redis cleanly returns None (key expired or never written on this pod). + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_increment = AsyncMock(return_value=500.0) + counter_cache.redis_cache = fake_redis + + # DB has the authoritative cross-pod spend. + db_row = MagicMock() + db_row.spend = 500.0 + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend == 500.0, ( + f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} " + f"(stale per-pod in-memory $30 would have caused multi-pod bypass)" + ) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_error_falls_back_to_in_memory(): + """When Redis raises, the read should still degrade to in-memory rather + than going straight to DB - in-memory is at least same-pod-fresh and + cheaper than a DB query during a Redis outage.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + + counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + return_value=MagicMock(spend=999.0) + ) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend == 42.0, ( + f"expected in-memory fallback 42.0 on Redis error, got {spend} " + f"(should not have hit DB when Redis errored)" + ) + # DB query should NOT have fired - in-memory short-circuits. + fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py new file mode 100644 index 0000000000..d0cb5ec546 --- /dev/null +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -0,0 +1,145 @@ +""" +Tests for the enable_redis_auth_cache litellm_settings flag. + +Verifies that _init_cache attaches Redis to user_api_key_cache only when +the flag is explicitly set to True, and leaves it in-memory-only otherwise. +""" + +from contextlib import contextmanager +import json +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +import litellm.proxy.proxy_server as ps +from litellm.caching.caching import RedisCache +from litellm.caching.dual_cache import DualCache + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeRedisCache(RedisCache): + """ + Minimal RedisCache subclass that passes isinstance checks without + requiring a real Redis connection. __init__ is bypassed so no + network calls are made. + """ + + def __init__(self): # noqa: super().__init__ skipped intentionally + self._store = {} + + def set_cache(self, key, value, **kwargs): # type: ignore[override] + # Enforce Redis JSON-serializable payload contract. + self._store[key] = json.dumps(value) + return True + + def get_cache(self, key, **kwargs): # type: ignore[override] + raw = self._store.get(key) + if raw is None: + return None + return json.loads(raw) + + +@contextmanager +def _patched_init_cache(litellm_settings: dict, cache_params: dict): + """ + Context manager that: + 1. Replaces the module-level globals with fresh DualCache instances. + 2. Patches ``litellm.Cache`` (locally imported inside _init_cache) so + it returns a fake cache whose ``.cache`` attribute is a + _FakeRedisCache (passes the isinstance guard in _init_cache). + 3. Extracts enable_redis_auth_cache from litellm_settings and passes it + as the second argument to _init_cache (matching production behaviour). + 4. Yields (user_api_key_cache, spend_counter_cache) after calling + _init_cache, then restores everything. + """ + fake_redis = _FakeRedisCache() + + mock_litellm_cache = MagicMock() + mock_litellm_cache.cache = fake_redis + + fresh_user_cache = DualCache() + fresh_spend_cache = DualCache() + + enable_redis_auth_cache = litellm_settings.get("enable_redis_auth_cache", False) + + with ( + patch.object(ps, "user_api_key_cache", fresh_user_cache), + patch.object(ps, "spend_counter_cache", fresh_spend_cache), + patch.object(ps, "llm_router", None), + # Cache is locally imported inside _init_cache: patch it at source. + patch("litellm.Cache", return_value=mock_litellm_cache), + ): + litellm.cache = None + ps.ProxyConfig()._init_cache(cache_params, enable_redis_auth_cache) + yield fresh_user_cache, fresh_spend_cache + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestRedisAuthCacheFlag: + def test_flag_true_attaches_redis_to_user_api_key_cache(self): + """When enable_redis_auth_cache=True, user_api_key_cache.redis_cache must be set.""" + with _patched_init_cache( + litellm_settings={"enable_redis_auth_cache": True}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is not None, ( + "Redis should be attached to user_api_key_cache when " + "enable_redis_auth_cache=True" + ) + + def test_flag_false_leaves_user_api_key_cache_in_memory_only(self): + """When enable_redis_auth_cache=False, user_api_key_cache must stay in-memory.""" + with _patched_init_cache( + litellm_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is None, ( + "user_api_key_cache must remain in-memory-only when " + "enable_redis_auth_cache=False" + ) + + def test_flag_absent_leaves_user_api_key_cache_in_memory_only(self): + """When enable_redis_auth_cache is not set at all, default is in-memory-only.""" + with _patched_init_cache( + litellm_settings={}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, _): + assert user_cache.redis_cache is None, ( + "user_api_key_cache must remain in-memory-only when " + "enable_redis_auth_cache is absent from litellm_settings" + ) + + def test_spend_counter_cache_always_gets_redis_regardless_of_flag(self): + """spend_counter_cache must receive Redis regardless of the auth-cache flag.""" + for flag_value in (True, False, None): + ls = ( + {"enable_redis_auth_cache": flag_value} + if flag_value is not None + else {} + ) + with _patched_init_cache( + litellm_settings=ls, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (_, spend_cache): + assert spend_cache.redis_cache is not None, ( + f"spend_counter_cache must always get Redis " + f"(enable_redis_auth_cache={flag_value!r})" + ) + + def test_flag_false_spend_gets_redis_but_user_cache_does_not(self): + """Explicit False: spend cache wired, auth cache left in-memory.""" + with _patched_init_cache( + litellm_settings={"enable_redis_auth_cache": False}, + cache_params={"type": "redis", "host": "localhost", "port": 6379}, + ) as (user_cache, spend_cache): + assert spend_cache.redis_cache is not None + assert user_cache.redis_cache is None diff --git a/tests/test_litellm/proxy/test_route_llm_request.py b/tests/test_litellm/proxy/test_route_llm_request.py index 96870b6cc7..bfea21e705 100644 --- a/tests/test_litellm/proxy/test_route_llm_request.py +++ b/tests/test_litellm/proxy/test_route_llm_request.py @@ -239,3 +239,55 @@ async def test_route_request_with_router_settings_override_preserves_existing(): assert call_kwargs["num_retries"] == 10 # Key/team timeout should be applied since not in request assert call_kwargs["timeout"] == 30 + + +@pytest.mark.parametrize( + "route_type", ["agenerate_content", "agenerate_content_stream"] +) +@pytest.mark.asyncio +async def test_route_request_maps_generation_config_for_google_routes(route_type): + """For Google generate_content routes, route_request must rename + `generationConfig` (Google's wire format) to `config` (the kwarg the + router method expects). Without this mapping the request reaches the + LLM with the field under the wrong name and the config is dropped.""" + data = { + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "generationConfig": { + "responseModalities": ["TEXT", "IMAGE"], + "imageConfig": {"aspectRatio": "9:16", "imageSize": "4K"}, + }, + } + llm_router = MagicMock() + getattr(llm_router, route_type).return_value = "ok" + + await route_request(data, llm_router, None, route_type) + + call_kwargs = getattr(llm_router, route_type).call_args[1] + assert "generationConfig" not in call_kwargs + assert "config" in call_kwargs + assert call_kwargs["config"]["responseModalities"] == ["TEXT", "IMAGE"] + assert call_kwargs["config"]["imageConfig"]["aspectRatio"] == "9:16" + assert call_kwargs["config"]["imageConfig"]["imageSize"] == "4K" + + +@pytest.mark.parametrize( + "route_type", ["agenerate_content", "agenerate_content_stream"] +) +@pytest.mark.asyncio +async def test_route_request_preserves_existing_config_for_google_routes(route_type): + """If the caller already supplies `config`, route_request must not + overwrite it with `generationConfig`.""" + data = { + "model": "gemini-2.5-flash", + "contents": [{"role": "user", "parts": [{"text": "Hello"}]}], + "config": {"existing": True}, + "generationConfig": {"shouldNotWin": True}, + } + llm_router = MagicMock() + getattr(llm_router, route_type).return_value = "ok" + + await route_request(data, llm_router, None, route_type) + + call_kwargs = getattr(llm_router, route_type).call_args[1] + assert call_kwargs["config"] == {"existing": True} diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 4424c68f1d..a6e39ec3c0 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -346,6 +346,34 @@ def test_tag_routing_with_list_of_tags_match_all(): assert not is_valid_deployment_tag(["default"], ["teamA"], match_any=False) +def test_strict_tag_routing_without_request_tags_blocks_header_regex_fallback(): + """ + When tag_filtering_match_any=False, deployments with plain tags must require + those request tags before header regex can match. A spoofed User-Agent must + not route to a tagged deployment when the request has no tags. + """ + from litellm.router_strategy.tag_based_routing import _match_deployment + + deployment = { + "model_name": "restricted-model", + "litellm_params": { + "model": "gpt-4o", + "tags": ["internal"], + "tag_regex": ["^User-Agent: internal-tool"], + }, + } + + assert ( + _match_deployment( + deployment=deployment, + request_tags=None, + header_strings=["User-Agent: internal-tool"], + match_any=False, + ) + is None + ) + + @pytest.mark.asyncio() async def test_router_free_paid_tier_with_responses_api(): """ diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 02241d4bc9..465c6669ce 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -6,6 +6,7 @@ import pytest from litellm import Router from litellm.router_utils.common_utils import ( _deployment_supports_web_search, + add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, ) @@ -362,3 +363,112 @@ def test_invalidate_model_group_info_cache(): # Invalidate and verify cache is cleared router._invalidate_model_group_info_cache() assert router._cached_get_model_group_info.cache_info().currsize == 0 + + +class TestAddModelFileIdMappings: + """Test cases for add_model_file_id_mappings. + + The router may pass either a list of deployment dicts (multiple matched + deployments) or a single deployment dict (when a specific deployment was + resolved, e.g. because the requested model matched a `model_info.id`). + Both shapes must produce a `{model_id: file_id}` mapping by extracting + `model_info.id` from each deployment. + """ + + @staticmethod + def _make_response(file_id: str): + response = Mock() + response.id = file_id + return response + + def test_should_map_each_deployment_id_when_given_list(self): + deployments = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "deployment-2"}, + }, + ] + responses = [self._make_response("file-1"), self._make_response("file-2")] + + result = add_model_file_id_mappings(deployments, responses) + + assert result == {"deployment-1": "file-1", "deployment-2": "file-2"} + + def test_should_extract_model_info_id_when_given_single_deployment_dict(self): + """Regression test: when `_common_checks_available_deployment` resolves + a specific deployment (returned as a dict, not a list), the function + must still extract `model_info.id` rather than iterate over the + deployment's own keys (`model_name`, `litellm_params`, `model_info`). + """ + deployment = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "sk-test"}, + "model_info": {"id": "deployment-1", "mode": "chat"}, + } + responses = [self._make_response("file-1")] + + result = add_model_file_id_mappings(deployment, responses) + + assert result == {"deployment-1": "file-1"} + assert all(isinstance(v, str) for v in result.values()) + + def test_should_handle_batch_model_when_id_matches_model_name(self): + """Regression test for the batch-model case: when `model_info.id` is + intentionally set equal to `model_name`, the router resolves a single + deployment via `has_model_id` and returns it as a dict. The mapping + must contain only `{id: file_id}` with string values so the resulting + `LiteLLM_ManagedFileTable` Pydantic validation passes. + """ + deployment = { + "model_name": "openai/openai/gpt-5.5-batch", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "tpm": 40000000, + "rpm": 15000, + }, + "model_info": { + "id": "openai/openai/gpt-5.5-batch", + "mode": "batch", + "base_model": "gpt-5.5", + "access_groups": ["default-models"], + }, + } + responses = [self._make_response("file-batch-1")] + + result = add_model_file_id_mappings(deployment, responses) + + # Bug case would have produced keys ["model_name", "litellm_params", + # "model_info"] with non-string values. + assert result == {"openai/openai/gpt-5.5-batch": "file-batch-1"} + assert "litellm_params" not in result + assert "model_info" not in result + + def test_should_skip_deployment_when_model_info_id_missing(self): + deployments = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "deployment-2"}, + }, + ] + responses = [self._make_response("file-1"), self._make_response("file-2")] + + result = add_model_file_id_mappings(deployments, responses) + + assert result == {"deployment-2": "file-2"} + + def test_should_return_empty_mapping_when_given_empty_list(self): + result = add_model_file_id_mappings([], []) + assert result == {} diff --git a/tests/test_litellm/test_anthropic_skills_transformation.py b/tests/test_litellm/test_anthropic_skills_transformation.py index 6761d67180..1b917f08ca 100644 --- a/tests/test_litellm/test_anthropic_skills_transformation.py +++ b/tests/test_litellm/test_anthropic_skills_transformation.py @@ -70,6 +70,15 @@ class TestAnthropicSkillsConfigURLConstruction: ) assert url == f"{FAKE_API_BASE}/v1/skills/skill_abc123" + def test_url_with_skill_id_encodes_path_segment(self): + url = self.config.get_complete_url( + api_base=FAKE_API_BASE, + endpoint="skills", + skill_id="../../files?x=1#frag", + ) + + assert url == f"{FAKE_API_BASE}/v1/skills/..%2F..%2Ffiles%3Fx%3D1%23frag" + def test_url_falls_back_to_anthropic_default(self): with patch( "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_base", diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts index 8834724f76..c3bd848902 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -141,7 +141,9 @@ test.describe("Add Model", () => { await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") - await expect(page.getByText(/Showing \d+ - \d+ of \d+ results/)).toBeVisible({ timeout: 15_000 }); + await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + timeout: 15_000, + }); // Verify the model name appears in the table body const tableBody = page.locator("table tbody"); @@ -181,7 +183,9 @@ test.describe("Add Model", () => { await page.waitForTimeout(1000); // Verify the model appears in the results count (not "Showing 0 results") - await expect(page.getByText(/Showing \d+ - \d+ of \d+ results/)).toBeVisible({ timeout: 15_000 }); + await expect(page.getByTestId("models-results-count")).toHaveText(/Showing \d+ - \d+ of \d+ results/, { + timeout: 15_000, + }); // Verify the wildcard model appears in the table body (wildcard models show as "cohere/*") const tableBody = page.locator("table tbody"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts index 85c8b25645..79976f5462 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts @@ -6,8 +6,8 @@ import { deriveErrorMessage, handleError, } from "@/components/networking"; -import { all_admin_roles } from "@/utils/roles"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { all_admin_roles } from "@/utils/roles"; // ── Types ──────────────────────────────────────────────────────────────────── @@ -81,7 +81,6 @@ export const useProjects = () => { return useQuery({ queryKey: projectKeys.list({}), queryFn: async () => fetchProjects(accessToken!), - enabled: - Boolean(accessToken) && all_admin_roles.includes(userRole || ""), + enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 514ae673d0..7c162e2056 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -46,6 +46,8 @@ interface GlobalRetryPolicyObject { [retryPolicyKey: string]: number; } +const HEALTH_PAGE_SIZE = 50; + const ModelsAndEndpointsView: React.FC = ({ premiumUser, teams }) => { const { accessToken, token, userRole, userId: userID } = useAuthorized(); const [addModelForm] = Form.useForm(); @@ -62,6 +64,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te const [selectedModelId, setSelectedModelId] = useState(null); const [selectedTeamId, setSelectedTeamId] = useState(null); const [selectedTabIndex, setSelectedTabIndex] = useState(0); + const [healthCurrentPage, setHealthCurrentPage] = useState(1); const [showMissingProviderBanner, setShowMissingProviderBanner] = useState(() => { if (typeof window !== "undefined") { return localStorage.getItem("hideMissingProviderBanner") !== "true"; @@ -71,6 +74,10 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te const queryClient = useQueryClient(); const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo(); + const { data: healthModelDataResponse, isLoading: isLoadingHealthModels } = useModelsInfo( + healthCurrentPage, + HEALTH_PAGE_SIZE, + ); const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap(); const { data: credentialsResponse, isLoading: isLoadingCredentials } = useCredentials(); const credentialsList = credentialsResponse?.credentials || []; @@ -104,12 +111,12 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te return modelDataResponse.data.map((model: any) => model.model_name); }, [modelDataResponse?.data]); - const allModelIdsOnProxy = useMemo(() => { - if (!modelDataResponse?.data) return []; - return modelDataResponse.data + const healthModelIdsOnProxy = useMemo(() => { + if (!healthModelDataResponse?.data) return []; + return healthModelDataResponse.data .map((model: any) => model.model_info?.id) .filter((id: string | undefined): id is string => Boolean(id)); - }, [modelDataResponse?.data]); + }, [healthModelDataResponse?.data]); const getProviderFromModel = (model: string) => { if (modelCostMapData !== null && modelCostMapData !== undefined) { @@ -125,6 +132,20 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te return transformModelData(modelDataResponse, getProviderFromModel); }, [modelDataResponse?.data, getProviderFromModel]); + const processedHealthModelData = useMemo(() => { + if (!healthModelDataResponse?.data) return { data: [] }; + return transformModelData(healthModelDataResponse, getProviderFromModel); + }, [healthModelDataResponse?.data, getProviderFromModel]); + + const healthPaginationMeta = useMemo(() => { + return { + total_count: healthModelDataResponse?.total_count ?? 0, + current_page: healthModelDataResponse?.current_page ?? healthCurrentPage, + total_pages: healthModelDataResponse?.total_pages ?? 1, + size: healthModelDataResponse?.size ?? HEALTH_PAGE_SIZE, + }; + }, [healthModelDataResponse, healthCurrentPage]); + const isProxyAdmin = userRole && isProxyAdminRole(userRole); const isInternalUser = userRole && internalUserRoles.includes(userRole); const isUserTeamAdmin = userID && isUserTeamAdminForAnyTeam(teams, userID); @@ -166,7 +187,8 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te const handleRefreshClick = () => { const currentDate = new Date(); - setLastRefreshed(currentDate.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })); + setLastRefreshed(currentDate.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })); + setHealthCurrentPage(1); queryClient.invalidateQueries({ queryKey: ["models", "list"] }); refetchModels(); }; @@ -441,11 +463,16 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te = ({ teams, organizations }) => { } }, [isAdmin, userID]); - // For non-admins, always pass their own user_id - const effectiveUserId = isAdmin ? selectedUserId : userID || null; + // For non-admins or "my-usage" view, always pass their own user_id + const effectiveUserId = usageView === "my-usage" || !isAdmin ? userID || null : selectedUserId; const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]); const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]); @@ -477,10 +477,10 @@ const UsagePage: React.FC = ({ teams, organizations }) => { } /> )} - {/* Your Usage Panel */} - {usageView === "global" && ( + {/* Your Usage / Global Usage Panel */} + {(usageView === "global" || usageView === "my-usage") && ( <> - {isAdmin && ( + {isAdmin && usageView === "global" && (
Filter by user