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/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/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e1240b436c..a815442c2f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4725,7 +4725,7 @@ class StandardLoggingPayloadSetup: ): for key, value in litellm_params["metadata"].items(): # Skip non-serializable objects like UserAPIKeyAuth - if key == "user_api_key_auth": + if key in {"user_api_key_auth", "user_api_key_budget_reservation"}: continue merged_metadata[key] = value 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 a65d0892aa..f295b4a299 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -199,6 +199,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/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/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_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 76c247aea8..bb5ebaa6a5 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -17,6 +17,7 @@ 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, @@ -599,6 +600,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 +722,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/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/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/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a34b73b531..dc625918b9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1007,6 +1007,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 +1027,7 @@ class BaseLLMHTTPHandler: model=model, optional_rerank_params=optional_rerank_params, headers=headers, + litellm_params=litellm_params, ) ## LOGGING @@ -2535,10 +2537,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 +2627,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( 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/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/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/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/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/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/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/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f96350500d..9923c3ce4b 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -169,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] = {} @@ -342,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, @@ -679,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), @@ -2426,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, @@ -2489,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: @@ -2501,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 @@ -2536,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.json b/litellm/proxy/_lazy_openapi_snapshot.json index b8e9eb6c26..eb35dd6cb3 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -14008,7 +14008,7 @@ "/mcp-rest/test/connection": { "post": { "description": "Test if we can connect to the provided MCP server before adding it", - "operationId": "test_connection_mcp_rest_test_connection_post", + "operationId": "test_connection_mcp_rest_test_connection_post_2", "requestBody": { "content": { "application/json": { @@ -14053,7 +14053,7 @@ "/mcp-rest/test/tools/list": { "post": { "description": "Preview tools available from MCP server before adding it", - "operationId": "test_tools_list_mcp_rest_test_tools_list_post", + "operationId": "test_tools_list_mcp_rest_test_tools_list_post_2", "requestBody": { "content": { "application/json": { @@ -14098,7 +14098,7 @@ "/mcp-rest/tools/call": { "post": { "description": "REST API to call a specific MCP tool with the provided arguments", - "operationId": "call_tool_rest_api_mcp_rest_tools_call_post", + "operationId": "call_tool_rest_api_mcp_rest_tools_call_post_2", "responses": { "200": { "content": { @@ -14123,7 +14123,7 @@ "/mcp-rest/tools/list": { "get": { "description": "List all available tools with information about the server they belong to.\n\nExample response:\n{\n \"tools\": [\n {\n \"name\": \"create_zap\",\n \"description\": \"Create a new zap\",\n \"inputSchema\": \"tool_input_schema\",\n \"mcp_info\": {\n \"server_name\": \"zapier\",\n \"logo_url\": \"https://www.zapier.com/logo.png\",\n }\n }\n ],\n \"error\": null,\n \"message\": \"Successfully retrieved tools\"\n}", - "operationId": "list_tool_rest_api_mcp_rest_tools_list_get", + "operationId": "list_tool_rest_api_mcp_rest_tools_list_get_2", "parameters": [ { "description": "The server id to list tools for", @@ -21896,7 +21896,7 @@ "/policies/usage/overview": { "get": { "description": "Return policy performance overview for the dashboard.", - "operationId": "policies_usage_overview_policies_usage_overview_get", + "operationId": "policies_usage_overview_policies_usage_overview_get_2", "parameters": [ { "description": "YYYY-MM-DD", @@ -22521,7 +22521,7 @@ "/policies/attachments/estimate-impact": { "post": { "description": "Estimate how many keys and teams would be affected by a policy attachment.\n\nUse this before creating an attachment to preview the blast radius.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/attachments/estimate-impact\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"policy_name\": \"hipaa-compliance\",\n \"tags\": [\"healthcare\", \"health-*\"]\n }'\n```", - "operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post", + "operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post_2", "requestBody": { "content": { "application/json": { @@ -22568,7 +22568,7 @@ "/policies/resolve": { "post": { "description": "Resolve which policies and guardrails apply for a given context.\n\nUse this endpoint to debug \"what guardrails would apply to a request\nwith this team/key/model/tags combination?\"\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/resolve\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"tags\": [\"healthcare\"],\n \"model\": \"gpt-4\"\n }'\n```", - "operationId": "resolve_policies_for_context_policies_resolve_post", + "operationId": "resolve_policies_for_context_policies_resolve_post_2", "parameters": [ { "description": "Force a DB sync before resolving. Default uses in-memory cache.", @@ -28329,7 +28329,7 @@ "/v1/vector_stores": { "get": { "description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list", - "operationId": "vector_store_list_v1_vector_stores_get", + "operationId": "vector_store_list_v1_vector_stores_get_2", "parameters": [ { "in": "query", @@ -28430,7 +28430,7 @@ }, "post": { "description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n \"name\": \"my-vector-store\",\n \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```", - "operationId": "vector_store_create_v1_vector_stores_post", + "operationId": "vector_store_create_v1_vector_stores_post_2", "responses": { "200": { "content": { @@ -28455,7 +28455,7 @@ "/v1/vector_stores/{vector_store_id}": { "delete": { "description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete", - "operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete", + "operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete_2", "parameters": [ { "in": "path", @@ -28499,7 +28499,7 @@ }, "get": { "description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve", - "operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get", + "operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get_2", "parameters": [ { "in": "path", @@ -28543,7 +28543,7 @@ }, "post": { "description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify", - "operationId": "vector_store_update_v1_vector_stores__vector_store_id__post", + "operationId": "vector_store_update_v1_vector_stores__vector_store_id__post_2", "parameters": [ { "in": "path", @@ -28588,7 +28588,7 @@ }, "/v1/vector_stores/{vector_store_id}/files": { "get": { - "operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get", + "operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get_2", "parameters": [ { "in": "path", @@ -28631,7 +28631,7 @@ ] }, "post": { - "operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post", + "operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post_2", "parameters": [ { "in": "path", @@ -28676,7 +28676,7 @@ }, "/v1/vector_stores/{vector_store_id}/files/{file_id}": { "delete": { - "operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete", + "operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete_2", "parameters": [ { "in": "path", @@ -28728,7 +28728,7 @@ ] }, "get": { - "operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get", + "operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get_2", "parameters": [ { "in": "path", @@ -28780,7 +28780,7 @@ ] }, "post": { - "operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post", + "operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post_2", "parameters": [ { "in": "path", @@ -28834,7 +28834,7 @@ }, "/v1/vector_stores/{vector_store_id}/files/{file_id}/content": { "get": { - "operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get", + "operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get_2", "parameters": [ { "in": "path", @@ -28889,7 +28889,7 @@ "/v1/vector_stores/{vector_store_id}/search": { "post": { "description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search", - "operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post", + "operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post_2", "parameters": [ { "in": "path", @@ -28935,7 +28935,7 @@ "/vector_stores": { "get": { "description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list", - "operationId": "vector_store_list_vector_stores_get", + "operationId": "vector_store_list_vector_stores_get_2", "parameters": [ { "in": "query", @@ -29036,7 +29036,7 @@ }, "post": { "description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n \"name\": \"my-vector-store\",\n \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```", - "operationId": "vector_store_create_vector_stores_post", + "operationId": "vector_store_create_vector_stores_post_2", "responses": { "200": { "content": { @@ -29061,7 +29061,7 @@ "/vector_stores/{vector_store_id}": { "delete": { "description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete", - "operationId": "vector_store_delete_vector_stores__vector_store_id__delete", + "operationId": "vector_store_delete_vector_stores__vector_store_id__delete_2", "parameters": [ { "in": "path", @@ -29105,7 +29105,7 @@ }, "get": { "description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve", - "operationId": "vector_store_retrieve_vector_stores__vector_store_id__get", + "operationId": "vector_store_retrieve_vector_stores__vector_store_id__get_2", "parameters": [ { "in": "path", @@ -29149,7 +29149,7 @@ }, "post": { "description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify", - "operationId": "vector_store_update_vector_stores__vector_store_id__post", + "operationId": "vector_store_update_vector_stores__vector_store_id__post_2", "parameters": [ { "in": "path", @@ -29194,7 +29194,7 @@ }, "/vector_stores/{vector_store_id}/files": { "get": { - "operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get", + "operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get_2", "parameters": [ { "in": "path", @@ -29237,7 +29237,7 @@ ] }, "post": { - "operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post", + "operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post_2", "parameters": [ { "in": "path", @@ -29282,7 +29282,7 @@ }, "/vector_stores/{vector_store_id}/files/{file_id}": { "delete": { - "operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete", + "operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete_2", "parameters": [ { "in": "path", @@ -29334,7 +29334,7 @@ ] }, "get": { - "operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get", + "operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get_2", "parameters": [ { "in": "path", @@ -29386,7 +29386,7 @@ ] }, "post": { - "operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post", + "operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post_2", "parameters": [ { "in": "path", @@ -29440,7 +29440,7 @@ }, "/vector_stores/{vector_store_id}/files/{file_id}/content": { "get": { - "operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get", + "operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get_2", "parameters": [ { "in": "path", @@ -29495,7 +29495,7 @@ "/vector_stores/{vector_store_id}/search": { "post": { "description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search", - "operationId": "vector_store_search_vector_stores__vector_store_id__search_post", + "operationId": "vector_store_search_vector_stores__vector_store_id__search_post_2", "parameters": [ { "in": "path", diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index fbd1a49aac..de118a66d9 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -11,9 +11,32 @@ import json import re import sys from pathlib import Path -from typing import Dict, Iterable, Optional +from typing import Dict, Optional, Set SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json" +HTTP_METHOD_SUFFIXES = { + "delete", + "get", + "head", + "options", + "patch", + "post", + "put", + "trace", +} + + +def _stabilize_multi_method_route_ids(routes) -> None: + """FastAPI derives route IDs from a set of methods; make snapshots stable.""" + + for route in routes: + methods = sorted(getattr(route, "methods", None) or []) + if len(methods) <= 1 or not getattr(route, "path_format", None): + continue + + operation_id = f"{route.name}{route.path_format}" + operation_id = re.sub(r"\W", "_", operation_id) + route.unique_id = f"{operation_id}_{methods[0].lower()}" def load_snapshot() -> Optional[Dict[str, Dict]]: @@ -26,22 +49,37 @@ def load_snapshot() -> Optional[Dict[str, Dict]]: return None -def _stable_generate_unique_id(route) -> str: - operation_id = f"{route.name}{route.path_format}" - operation_id = re.sub(r"\W", "_", operation_id) - methods = sorted(route.methods or []) - if not methods: - return operation_id - return f"{operation_id}_{methods[0].lower()}" +def _normalize_operation_ids(paths: Dict[str, Dict]) -> None: + """Make FastAPI-generated operation IDs stable for multi-method routes. + 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 -def _set_stable_operation_ids(routes: Iterable) -> None: - for route in routes: - if getattr(route, "operation_id", None) is not None: + methods = {method for method in path_ops if method in HTTP_METHODS} + if not methods: continue - if getattr(route, "methods", None) is None: - continue - route.operation_id = _stable_generate_unique_id(route) + + 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]: @@ -50,7 +88,7 @@ def generate_snapshot() -> Dict[str, Dict]: from fastapi.openapi.utils import get_openapi from litellm.proxy._lazy_features import LAZY_FEATURES - from litellm.proxy.proxy_server import app + from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids for feat in LAZY_FEATURES: if feat.module_path in sys.modules: @@ -62,6 +100,7 @@ def generate_snapshot() -> Dict[str, Dict]: sys.stderr.write(f"warning: skip {feat.name}: {exc}\n") fragments: Dict[str, Dict] = {} + used_operation_ids: Set[str] = set() for feat in LAZY_FEATURES: feat_routes = [ r @@ -70,15 +109,26 @@ def generate_snapshot() -> Dict[str, Dict]: ] if not feat_routes: continue - _set_stable_operation_ids(feat_routes) + _stabilize_multi_method_route_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 op in path_ops.values(): + for method, op in path_ops.items(): if isinstance(op, dict): + operation_id = op.get("operationId") + if isinstance(operation_id, str): + for suffix in HTTP_METHOD_SUFFIXES: + if operation_id.endswith(f"_{suffix}"): + op["operationId"] = ( + operation_id[: -len(suffix)] + method + ) + break op["tags"] = [feat.name] + full = ensure_unique_openapi_operation_ids(full, used_operation_ids) 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..213b6e0987 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, @@ -2574,6 +2579,7 @@ class UserAPIKeyAuth( user_spend: Optional[float] = None user_max_budget: Optional[float] = None request_route: Optional[str] = None + budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True) user: Optional[Any] = None # Expanded user object when expand=user is used created_by_user: Optional[Any] = ( None # Expanded created_by user when expand=user is used diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 65638ed6c1..64709549f8 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, @@ -61,11 +60,17 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.http_parsing_utils import ( + _safe_get_request_headers, + _safe_get_request_query_params, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler 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 @@ -485,7 +490,10 @@ async def common_checks( # noqa: PLR0915 from litellm.proxy.proxy_server import prisma_client, user_api_key_cache _model: Optional[Union[str, List[str]]] = get_model_from_request( - request_body, route + request_data=request_body, + route=route, + request_headers=_safe_get_request_headers(request=request), + request_query_params=_safe_get_request_query_params(request=request), ) # 1. If team is blocked @@ -655,13 +663,7 @@ async def common_checks( # noqa: PLR0915 end_user_object is not None and end_user_object.litellm_budget_table is not None ): - end_user_budget = end_user_object.litellm_budget_table.max_budget - if end_user_budget is not None and end_user_object.spend > end_user_budget: - raise litellm.BudgetExceededError( - current_cost=end_user_object.spend, - max_budget=end_user_budget, - message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}", - ) + await _check_end_user_budget(end_user_obj=end_user_object, route=route) _enforce_user_param_check(general_settings, request, request_body, route) _reject_clientside_metadata_tags_check(general_settings, request_body, route) @@ -852,7 +854,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 +877,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 +896,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 +916,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 +973,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: """ @@ -1006,7 +1013,7 @@ async def _apply_default_budget_to_end_user( return end_user_obj -def _check_end_user_budget( +async def _check_end_user_budget( end_user_obj: LiteLLM_EndUserTable, route: str, ) -> None: @@ -1027,11 +1034,20 @@ def _check_end_user_budget( return end_user_budget = end_user_obj.litellm_budget_table.max_budget - if end_user_budget is not None and end_user_obj.spend > end_user_budget: + if end_user_budget is None: + return + + from litellm.proxy.proxy_server import get_current_spend + + end_user_spend = await get_current_spend( + counter_key=f"spend:end_user:{end_user_obj.user_id}", + fallback_spend=end_user_obj.spend or 0.0, + ) + if end_user_spend > end_user_budget: raise litellm.BudgetExceededError( - current_cost=end_user_obj.spend, + current_cost=end_user_spend, max_budget=end_user_budget, - message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.spend}, Budget={end_user_budget}", + message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}", ) @@ -1039,7 +1055,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 +1086,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, @@ -1083,7 +1101,7 @@ async def get_end_user_object( ) # Check budget limits - _check_end_user_budget(end_user_obj=return_obj, route=route) + await _check_end_user_budget(end_user_obj=return_obj, route=route) return return_obj @@ -1108,13 +1126,15 @@ 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 - _check_end_user_budget(end_user_obj=_response, route=route) + await _check_end_user_budget(end_user_obj=_response, route=route) return _response @@ -1128,7 +1148,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 +1181,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 +1202,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 +1219,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 +1258,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 +1278,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 +1295,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 +1468,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 +1487,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 +1554,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 +1576,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 +1598,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 +1611,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 @@ -1589,17 +1626,21 @@ async def _cache_key_object( ## CACHE REFRESH TIME user_api_key_obj.last_refreshed_at = time.time() + cached_key_obj = _copy_user_api_key_auth_for_cache( + user_api_key_obj=user_api_key_obj + ) await _cache_management_object( key=key, - value=user_api_key_obj, + value=cached_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 +1688,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 +1749,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 +1846,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 +1878,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 +1900,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 +1951,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 +2033,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 +2041,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 +2063,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 +2090,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 +2125,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 +2336,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 +2354,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 _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth) if check_cache_only: raise Exception( @@ -2370,11 +2414,21 @@ async def get_key_object( return _response +def _copy_user_api_key_auth_for_cache( + user_api_key_obj: UserAPIKeyAuth, +) -> UserAPIKeyAuth: + copied_key_obj = user_api_key_obj.model_copy() + copied_key_obj.budget_reservation = None + copied_key_obj.parent_otel_span = None + copied_key_obj.request_route = None + return copied_key_obj + + @log_db_metrics 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 +2444,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 +2460,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 +2477,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 +2497,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 +2528,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 +2541,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 +2572,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 +2591,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 +2613,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 +2671,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 +2690,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 +2709,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 +3433,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 +3501,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 +3808,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 +3823,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 +3847,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 +3857,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 +3951,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], ): @@ -3935,13 +3990,19 @@ async def _tag_max_budget_check( if ( tag_object.litellm_budget_table is not None and tag_object.litellm_budget_table.max_budget is not None - and tag_object.spend is not None - and tag_object.spend > tag_object.litellm_budget_table.max_budget ): + from litellm.proxy.proxy_server import get_current_spend + + tag_spend = await get_current_spend( + counter_key=f"spend:tag:{tag_name}", + fallback_spend=tag_object.spend or 0.0, + ) + if tag_spend <= tag_object.litellm_budget_table.max_budget: + continue raise litellm.BudgetExceededError( - current_cost=tag_object.spend, + current_cost=tag_spend, max_budget=tag_object.litellm_budget_table.max_budget, - message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_object.spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", + message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}", ) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 91c8f2dd7c..51108827f6 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,7 +2,7 @@ import os import re import sys from functools import lru_cache -from typing import Any, List, Optional, Tuple +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union from fastapi import HTTPException, Request, status @@ -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 @@ -942,20 +976,257 @@ def get_end_user_id_from_request_body( return None -def get_model_from_request( - request_data: dict, route: str -) -> Optional[Union[str, List[str]]]: - # First try to get model from request_data - model = request_data.get("model") or request_data.get("target_model_names") +MODEL_ROUTING_HEADER_NAME = "x-litellm-model" +_MODEL_ROUTING_ROUTE_MARKERS = ( + "/files", + "/batches", + "/vector_stores", + "/skills", + "/evals", + "/fine_tuning", + "/videos", +) +_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS = ( + "/files", + "/batches", + "/skills", + "/evals", +) +_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS = ( + "/files", + "/batches", + "/fine_tuning", +) +_MODEL_ROUTING_BODY_TARGET_MODEL_ROUTE_MARKERS = ( + "/files", + "/batches", + "/vector_stores", +) +_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS = ("/evals",) +_MODEL_ROUTING_ID_FIELDS = ( + "file_id", + "input_file_id", + "output_file_id", + "error_file_id", + "batch_id", + "fine_tuning_job_id", + "training_file", + "validation_file", + "vector_store_id", + "video_id", + "character_id", +) - if model is not None: - model_names = model.split(",") - if len(model_names) == 1: - model = model_names[0].strip() + +def _append_model_candidates(candidates: List[str], value: Any) -> None: + if value is None: + return + + values = value if isinstance(value, (list, tuple, set)) else [value] + for item in values: + if item is None: + continue + if isinstance(item, str): + model_names = [model.strip() for model in item.split(",")] else: - model = [m.strip() for m in model_names] + model_names = [str(item).strip()] + candidates.extend(model for model in model_names if model) - # If model not in request_data, try to extract from route + +def _dedupe_model_candidates(candidates: List[str]) -> List[str]: + deduped: List[str] = [] + for model in candidates: + if model not in deduped: + deduped.append(model) + return deduped + + +def _get_case_insensitive_mapping_value( + mapping: Optional[Mapping[str, Any]], key: str +) -> Any: + if not mapping: + return None + if key in mapping: + return mapping[key] + key_lower = key.lower() + for mapping_key, value in mapping.items(): + if str(mapping_key).lower() == key_lower: + return value + return None + + +def _route_matches_any_marker(route: str, markers: Tuple[str, ...]) -> bool: + normalized_route = route.lower() + return any(marker in normalized_route for marker in markers) + + +def _route_uses_model_routing_sources(route: str) -> bool: + return _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_ROUTE_MARKERS) + + +def _extract_models_from_managed_resource_id( + resource_id: Any, resource_id_field: Optional[str] = None +) -> List[str]: + if not isinstance(resource_id, str) or not resource_id: + return [] + + candidates: List[str] = [] + + try: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + decode_model_from_file_id, + get_model_id_from_unified_batch_id, + get_models_from_unified_file_id, + ) + + _append_model_candidates( + candidates=candidates, value=decode_model_from_file_id(resource_id) + ) + unified_file_id = _is_base64_encoded_unified_file_id(resource_id) + if unified_file_id: + _append_model_candidates( + candidates=candidates, + value=get_models_from_unified_file_id(unified_file_id), + ) + _append_model_candidates( + candidates=candidates, + value=get_model_id_from_unified_batch_id(unified_file_id), + ) + except Exception as e: + verbose_proxy_logger.debug( + "Unable to extract model from managed file/batch ID: %s", str(e) + ) + + try: + from litellm.llms.base_llm.managed_resources.utils import parse_unified_id + + parsed_id = parse_unified_id(resource_id) + if parsed_id: + _append_model_candidates( + candidates=candidates, value=parsed_id.get("model_id") + ) + _append_model_candidates( + candidates=candidates, value=parsed_id.get("target_model_names") + ) + except Exception as e: + verbose_proxy_logger.debug( + "Unable to extract model from unified managed resource ID: %s", str(e) + ) + + if resource_id_field in ("video_id", "character_id"): + try: + from litellm.types.videos.utils import ( + decode_character_id_with_provider, + decode_video_id_with_provider, + ) + + if resource_id_field == "video_id": + _append_model_candidates( + candidates=candidates, + value=decode_video_id_with_provider(resource_id).get("model_id"), + ) + else: + _append_model_candidates( + candidates=candidates, + value=decode_character_id_with_provider(resource_id).get( + "model_id" + ), + ) + except Exception as e: + verbose_proxy_logger.debug( + "Unable to extract model from managed video/character ID: %s", str(e) + ) + + return _dedupe_model_candidates(candidates) + + +def _extract_model_candidates_from_request( + request_data: dict, + route: str, + request_headers: Optional[Mapping[str, Any]] = None, + request_query_params: Optional[Mapping[str, Any]] = None, +) -> List[str]: + candidates: List[str] = [] + uses_model_routing_sources = _route_uses_model_routing_sources(route=route) + uses_header_or_query_model_sources = _route_matches_any_marker( + route=route, markers=_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS + ) + uses_query_target_model_sources = _route_matches_any_marker( + route=route, markers=_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS + ) + uses_body_target_model_sources = _route_matches_any_marker( + route=route, markers=_MODEL_ROUTING_BODY_TARGET_MODEL_ROUTE_MARKERS + ) + uses_completion_model_sources = _route_matches_any_marker( + route=route, markers=_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS + ) + + body_model = request_data.get("model") + _append_model_candidates(candidates, body_model) + if uses_body_target_model_sources or not body_model: + _append_model_candidates(candidates, request_data.get("target_model_names")) + if uses_completion_model_sources and isinstance( + request_data.get("completion"), dict + ): + _append_model_candidates(candidates, request_data["completion"].get("model")) + + if uses_model_routing_sources: + if uses_header_or_query_model_sources: + _append_model_candidates( + candidates, + _get_case_insensitive_mapping_value(request_query_params, "model"), + ) + _append_model_candidates( + candidates, + _get_case_insensitive_mapping_value( + request_headers, MODEL_ROUTING_HEADER_NAME + ), + ) + if uses_query_target_model_sources: + _append_model_candidates( + candidates, + _get_case_insensitive_mapping_value( + request_query_params, "target_model_names" + ), + ) + + for field in _MODEL_ROUTING_ID_FIELDS: + _append_model_candidates( + candidates, + _extract_models_from_managed_resource_id( + request_data.get(field), resource_id_field=field + ), + ) + + return _dedupe_model_candidates(candidates) + + +def _format_model_candidates( + candidates: List[str], +) -> Optional[Union[str, List[str]]]: + if not candidates: + return None + if len(candidates) == 1: + return candidates[0] + return candidates + + +def get_model_from_request( + request_data: dict, + route: str, + request_headers: Optional[Mapping[str, Any]] = None, + request_query_params: Optional[Mapping[str, Any]] = None, +) -> Optional[Union[str, List[str]]]: + candidates = _extract_model_candidates_from_request( + request_data=request_data, + route=route, + request_headers=request_headers, + request_query_params=request_query_params, + ) + model = _format_model_candidates(candidates) + + # If no explicit model was found, try to extract from route if model is None: # Parse model from route that follows the pattern /openai/deployments/{model}/* match = re.match(r"/openai/deployments/([^/]+)", route) 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/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f0c2a4514f..2e5140e0e3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -11,7 +11,7 @@ import asyncio import re import secrets from datetime import datetime, timezone -from typing import Any, List, Optional, Tuple, cast +from typing import Any, List, Optional, Tuple, Union, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -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,9 +59,11 @@ 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, + _safe_get_request_query_params, populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body @@ -118,6 +119,29 @@ azure_apim_header = APIKeyHeader( ) +def _get_model_from_request_context( + request_data: dict, + route: str, + request: Optional[Request], +) -> Optional[Union[str, List[str]]]: + return get_model_from_request( + request_data=request_data, + route=route, + request_headers=_safe_get_request_headers(request=request), + request_query_params=_safe_get_request_query_params(request=request), + ) + + +def _get_model_names_for_budget_checks( + model: Optional[Union[str, List[str]]], +) -> List[str]: + if model is None: + return [] + if isinstance(model, str): + return [model] + return model + + def _get_bearer_token_or_received_api_key(api_key: str) -> str: if api_key.startswith("Bearer "): # ensure Bearer token passed in api_key = api_key.replace("Bearer ", "") # extract the token @@ -329,7 +353,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 +369,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 +497,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 +534,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]: @@ -879,7 +908,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) # Check if model has zero cost - if so, skip all budget checks - model = get_model_from_request(request_data, route) + model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + ) skip_budget_checks = False if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero @@ -1107,9 +1140,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 +1210,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( @@ -1247,6 +1282,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 valid_token=valid_token, request_data=request_data, route=route, + request=request, llm_model_list=llm_model_list, llm_router=llm_router, ) @@ -1272,7 +1308,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_obj = None # Check 2a. Check if model has zero cost - if so, skip all budget checks - model = get_model_from_request(request_data, route) + model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + ) skip_budget_checks = False if model is not None and llm_router is not None: from litellm.proxy.auth.auth_checks import _is_model_cost_zero @@ -1291,7 +1331,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 +1340,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 @@ -1390,21 +1436,29 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 # Check 5. Token Model Spend is under Model budget max_budget_per_model = valid_token.model_max_budget - current_model = request_data.get("model", None) + current_model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + ) + current_models = _get_model_names_for_budget_checks( + model=current_model + ) if ( max_budget_per_model is not None and isinstance(max_budget_per_model, dict) and len(max_budget_per_model) > 0 and prisma_client is not None - and current_model is not None + and current_models and valid_token.token is not None ): ## GET THE SPEND FOR THIS MODEL - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=current_model, - ) + for model_name in current_models: + await model_max_budget_limiter.is_key_within_model_budget( + user_api_key_dict=valid_token, + model=model_name, + ) # Check 5b. End-user model max budget end_user_mmb = valid_token.end_user_model_max_budget @@ -1412,14 +1466,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 end_user_mmb is not None and isinstance(end_user_mmb, dict) and len(end_user_mmb) > 0 - and current_model is not None + and current_models and valid_token.end_user_id is not None ): - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_mmb, - model=current_model, - ) + for model_name in current_models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=valid_token.end_user_id, + end_user_model_max_budget=end_user_mmb, + model=model_name, + ) # Check 6: Additional Common Checks across jwt + key auth if valid_token.team_id is not None: @@ -1457,9 +1512,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 @@ -1845,10 +1904,12 @@ async def _run_centralized_common_checks( user_api_key_auth_obj.project_metadata = project_object.metadata user_api_key_auth_obj.project_alias = project_object.project_alias - skip_budget_checks = False - model = get_model_from_request(request_data, route) - if model is not None and llm_router is not None: - skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + skip_budget_checks = _should_skip_budget_checks( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + ) _ = await common_checks( request=request, @@ -1866,6 +1927,21 @@ async def _run_centralized_common_checks( project_object=project_object, ) + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data=request_data, + route=route, + llm_router=llm_router, + team_object=team_object, + user_object=user_object, + end_user_id=end_user_id, + end_user_object=end_user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + skip_budget_checks=skip_budget_checks, + ) + async def _noop_none() -> None: """Sentinel coroutine for asyncio.gather when a fetch is unnecessary @@ -1873,6 +1949,59 @@ async def _noop_none() -> None: return None +async def _reserve_budget_after_common_checks( + user_api_key_auth_obj: UserAPIKeyAuth, + request_data: dict, + route: str, + llm_router: Optional[Any], + team_object: Optional[LiteLLM_TeamTableCachedObj], + user_object: Optional[LiteLLM_UserTable], + prisma_client: Optional[PrismaClient], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + skip_budget_checks: bool, + end_user_id: Optional[str] = None, + end_user_object: Optional[LiteLLM_EndUserTable] = None, +) -> None: + user_api_key_auth_obj.budget_reservation = None + if skip_budget_checks: + return + + from litellm.proxy.spend_tracking.budget_reservation import ( + reserve_budget_for_request, + ) + + user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request( + request_body=request_data, + route=route, + llm_router=llm_router, + valid_token=user_api_key_auth_obj, + team_object=team_object, + user_object=user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_id=end_user_id, + end_user_object=end_user_object, + ) + + +def _should_skip_budget_checks( + request_data: dict, + route: str, + request: Optional[Request], + llm_router: Optional[Any], +) -> bool: + model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + ) + if model is not None and llm_router is not None: + return _is_model_cost_zero(model=model, llm_router=llm_router) + return False + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -1910,6 +2039,7 @@ async def user_api_key_auth( request_data=request_data, custom_litellm_key_header=custom_litellm_key_header, ) + user_api_key_auth_obj.budget_reservation = None ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj) @@ -2117,6 +2247,7 @@ async def _enforce_key_and_fallback_model_access( valid_token: UserAPIKeyAuth, request_data: dict, route: str, + request: Optional[Request], llm_model_list: Optional[list], llm_router: Optional[Any], ) -> None: @@ -2135,7 +2266,11 @@ async def _enforce_key_and_fallback_model_access( ): pass else: - model = get_model_from_request(request_data, route) + model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + ) fallback_models = cast( Optional[List[ALL_FALLBACK_MODEL_VALUES]], request_data.get("fallbacks", None), @@ -2222,11 +2357,17 @@ async def _run_post_custom_auth_checks( valid_token=valid_token, request_data=request_data, route=route, + request=request, llm_model_list=llm_model_list, llm_router=llm_router, ) - current_model = request_data.get("model", None) + current_model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + ) + current_models = _get_model_names_for_budget_checks(model=current_model) # 3. Check key-level model_max_budget max_budget_per_model = valid_token.model_max_budget @@ -2234,13 +2375,14 @@ async def _run_post_custom_auth_checks( max_budget_per_model is not None and isinstance(max_budget_per_model, dict) and len(max_budget_per_model) > 0 - and current_model is not None + and current_models and valid_token.token is not None ): - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=current_model, - ) + for model_name in current_models: + await model_max_budget_limiter.is_key_within_model_budget( + user_api_key_dict=valid_token, + model=model_name, + ) # 4. Check end-user model_max_budget end_user_mmb = valid_token.end_user_model_max_budget @@ -2248,14 +2390,15 @@ async def _run_post_custom_auth_checks( end_user_mmb is not None and isinstance(end_user_mmb, dict) and len(end_user_mmb) > 0 - and current_model is not None + and current_models and valid_token.end_user_id is not None ): - await model_max_budget_limiter.is_end_user_within_model_budget( - end_user_id=valid_token.end_user_id, - end_user_model_max_budget=end_user_mmb, - model=current_model, - ) + for model_name in current_models: + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=valid_token.end_user_id, + end_user_model_max_budget=end_user_mmb, + model=model_name, + ) # team / user / end_user / project context objects are fetched by # the centralized common_checks gate in user_api_key_auth after 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..8fc891d892 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -97,6 +97,55 @@ def _serialize_http_exception_detail( return str(detail), None +def _collect_response_file_search_vector_store_ids(data: Dict[str, Any]) -> set[str]: + vector_store_ids: set[str] = set() + tools = data.get("tools") + if not isinstance(tools, list): + return vector_store_ids + + for tool in tools: + if not isinstance(tool, dict) or tool.get("type") != "file_search": + continue + ids = tool.get("vector_store_ids") or [] + if not isinstance(ids, list): + raise HTTPException( + status_code=400, + detail={ + "error": "file_search.vector_store_ids must be a list of strings" + }, + ) + for vector_store_id in ids: + if not isinstance(vector_store_id, str) or not vector_store_id: + raise HTTPException( + status_code=400, + detail={ + "error": "file_search.vector_store_ids must be a list of strings" + }, + ) + vector_store_ids.add(vector_store_id) + + return vector_store_ids + + +async def _authorize_response_file_search_vector_stores( + data: Dict[str, Any], + user_api_key_dict: UserAPIKeyAuth, +) -> None: + vector_store_ids = _collect_response_file_search_vector_store_ids(data) + if not vector_store_ids: + return + + from litellm.proxy.vector_store_endpoints.utils import ( + assert_user_can_access_vector_store_id, + ) + + for vector_store_id in sorted(vector_store_ids): + await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) + + async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: """Parses an event line and returns an error code if present, else None.""" event_line = ( @@ -744,6 +793,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", @@ -786,6 +840,11 @@ class ProxyBaseLLMRequestProcessing: version=version, proxy_config=proxy_config, ) + if route_type in {"aresponses", "_aresponses_websocket"}: + await _authorize_response_file_search_vector_stores( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) # Calculate request queue time after add_litellm_data_to_request # which sets arrival_time in proxy_server_request @@ -1001,6 +1060,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..19ec669939 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -14,10 +14,12 @@ memory in long-lived deployments. import asyncio from collections import OrderedDict +from datetime import datetime from typing import TYPE_CHECKING, ClassVar, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE +from litellm.litellm_core_utils.duration_parser import duration_in_seconds if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -35,6 +37,10 @@ class SpendCounterReseed: spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend spend:user:{user_id} -> LiteLLM_UserTable.spend spend:org:{org_id} -> LiteLLM_OrganizationTable.spend + + End-user and tag spend counters intentionally do not reseed here. Their + auth paths already load the corresponding objects via get_end_user_object() + and get_tag_objects_batch(); callers pass those values as fallback_spend. """ _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() @@ -69,9 +75,10 @@ class SpendCounterReseed: """ if prisma_client is None: return None - # Per-window counters share prefixes with primary counters but - # don't correspond to a DB row. - if ":window:" in counter_key: + # Per-window key/team counters share prefixes with primary counters + # but don't correspond to a DB row. Do not reject arbitrary entity IDs + # or tag names that merely contain ":window:". + if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None try: if counter_key.startswith("spend:key:"): @@ -97,6 +104,10 @@ class SpendCounterReseed: row = await prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_id} ) + elif counter_key.startswith("spend:end_user:"): + return None + elif counter_key.startswith("spend:tag:"): + return None elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] row = await prisma_client.db.litellm_organizationtable.find_unique( @@ -113,11 +124,27 @@ class SpendCounterReseed: return None return float(getattr(row, "spend", 0.0) or 0.0) + @staticmethod + def _is_key_or_team_window_counter(counter_key: str) -> bool: + for prefix in ("spend:key:", "spend:team:"): + if not counter_key.startswith(prefix): + continue + _, separator, duration = counter_key.rpartition(":window:") + if not separator or not duration: + return False + try: + duration_in_seconds(duration) + except Exception: + return False + return True + return False + @staticmethod async def coalesced( prisma_client: Optional["PrismaClient"], spend_counter_cache: "DualCache", counter_key: str, + require_cache_warm: bool = False, ) -> Optional[float]: """ Reseed a cold spend counter from the DB and warm the cache, @@ -129,7 +156,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,23 +166,169 @@ 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: return None # 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 - ) + if spend_counter_cache.redis_cache is not None: + current_value = ( + await spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=db_spend, + refresh_ttl=True, + ) + ) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, + value=current_value, + ) + else: + await spend_counter_cache.async_increment_cache( + key=counter_key, value=db_spend, refresh_ttl=True + ) except Exception: verbose_proxy_logger.exception( "SpendCounterReseed.coalesced: failed to warm counter %s", counter_key, ) + if require_cache_warm: + raise return db_spend + + @staticmethod + async def window_from_spend_logs( + prisma_client: Optional["PrismaClient"], + entity_type: str, + entity_id: str, + window_start: datetime, + ) -> Optional[float]: + if prisma_client is None: + return None + + if entity_type == "Key": + group_field = "api_key" + where = { + "api_key": entity_id, + "startTime": {"gte": window_start}, + } + elif entity_type == "Team": + group_field = "team_id" + where = { + "team_id": entity_id, + "startTime": {"gte": window_start}, + } + else: + return None + + try: + response = await prisma_client.db.litellm_spendlogs.group_by( + by=[group_field], + where=where, # type: ignore[arg-type] + sum={"spend": True}, + ) + except Exception: + verbose_proxy_logger.exception( + "SpendCounterReseed.window_from_spend_logs: failed for %s=%s", + entity_type, + entity_id, + ) + return None + + if not response: + return 0.0 + first_row = response[0] + sum_row = ( + first_row.get("_sum") + if isinstance(first_row, dict) + else getattr(first_row, "_sum", None) + ) + spend = ( + sum_row.get("spend") + if isinstance(sum_row, dict) + else getattr(sum_row, "spend", None) + ) + return float(spend or 0.0) + + @staticmethod + async def coalesced_window( + prisma_client: Optional["PrismaClient"], + spend_counter_cache: "DualCache", + counter_key: str, + entity_type: str, + entity_id: str, + window_start: datetime, + ) -> Optional[float]: + lock = await SpendCounterReseed._get_lock(counter_key) + async with lock: + 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: + pass + 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) + + window_spend = await SpendCounterReseed.window_from_spend_logs( + prisma_client=prisma_client, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + if window_spend is None: + return None + try: + if spend_counter_cache.redis_cache is not None: + seeded = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=window_spend, + nx=True, + ) + if seeded: + current_value = window_spend + else: + current_cached_value = ( + await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key + ) + ) + if current_cached_value is None: + current_value = ( + await spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=window_spend, + ) + ) + else: + current_value = float(current_cached_value) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, + value=current_value, + ) + else: + await spend_counter_cache.async_increment_cache( + key=counter_key, value=window_spend + ) + except Exception: + verbose_proxy_logger.exception( + "SpendCounterReseed.coalesced_window: failed to warm counter %s", + counter_key, + ) + raise + return window_spend 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..1eda01e5c6 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,90 @@ 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 _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 +857,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 +925,26 @@ async def health_endpoint( detail={"error": f"Model with ID {model_id} not found"}, ) + is_admin = _is_proxy_admin(user_api_key_dict) + + 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. + 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 +955,59 @@ 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 + 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") + } + filtered = _filter_health_check_results_by_model_ids( + health_check_results, allowed_model_ids + ) + if 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) + 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 +1018,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( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c9946f4e26..bd1b8ea79c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -30,16 +30,35 @@ class _ProxyDBLogger(CustomLogger): kwargs, response_obj, start_time, end_time ) - async def async_post_call_failure_hook( - self, - request_data: dict, - original_exception: Exception, - user_api_key_dict: UserAPIKeyAuth, - traceback_str: Optional[str] = None, - ): - request_route = user_api_key_dict.request_route - if _ProxyDBLogger._should_track_errors_in_db() is False: - return + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: Optional[str] = None, + ): + try: + await _release_budget_reservation( + budget_reservation=user_api_key_dict.budget_reservation + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to release budget reservation during failure handling" + ) + try: + await _invalidate_budget_reservation_counters( + budget_reservation=user_api_key_dict.budget_reservation + ) + if user_api_key_dict.budget_reservation is not None: + user_api_key_dict.budget_reservation["finalized"] = True + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate budget reservation counters after failure release failed" + ) + + request_route = user_api_key_dict.request_route + if _ProxyDBLogger._should_track_errors_in_db() is False: + return elif request_route is not None and not ( RouteChecks.is_llm_api_route(route=request_route) or RouteChecks.is_info_route(route=request_route) @@ -155,66 +174,64 @@ class _ProxyDBLogger(CustomLogger): f"kwargs stream: {kwargs.get('stream', None)} + complete streaming response: {kwargs.get('complete_streaming_response', None)}" ) parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs=kwargs) - litellm_params = kwargs.get("litellm_params", {}) or {} - end_user_id = get_end_user_id_for_cost_tracking(litellm_params) - metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) - user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) - team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) - org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) + litellm_params = kwargs.get("litellm_params", {}) or {} + end_user_id = get_end_user_id_for_cost_tracking(litellm_params) + metadata = get_litellm_metadata_from_kwargs(kwargs=kwargs) + budget_reservation = _get_budget_reservation_from_metadata( + metadata=metadata + ) + user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) + team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) + org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) sl_object: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object", None ) - response_cost = ( - sl_object.get("response_cost", None) - if sl_object is not None - else kwargs.get("response_cost", None) - ) - tags: Optional[List[str]] = ( - sl_object.get("request_tags", None) if sl_object is not None else None - ) - - if response_cost is not None: - user_api_key = metadata.get("user_api_key", None) + response_cost = ( + sl_object.get("response_cost", None) + if sl_object is not None + else kwargs.get("response_cost", None) + ) + tags = _get_request_tags_for_cost_tracking( + sl_object=sl_object, + metadata=metadata, + ) + + if response_cost is not None: + user_api_key = metadata.get("user_api_key", None) if kwargs.get("cache_hit", False) is True: response_cost = 0.0 verbose_proxy_logger.debug( f"Cache Hit: response_cost {response_cost}, for user_id {user_id}" ) - verbose_proxy_logger.debug( - f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" - ) - if _should_track_cost_callback( - user_api_key=user_api_key, + verbose_proxy_logger.debug( + f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" + ) + if _should_track_cost_callback( + user_api_key=user_api_key, user_id=user_id, team_id=team_id, - end_user_id=end_user_id, - ): - ## UPDATE DATABASE - await proxy_logging_obj.db_spend_update_writer.update_database( - token=user_api_key, - response_cost=response_cost, - user_id=user_id, - end_user_id=end_user_id, - team_id=team_id, - kwargs=kwargs, - completion_response=completion_response, - start_time=start_time, - end_time=end_time, - org_id=org_id, - ) - - # Atomically update spend counters (in-memory + Redis) - # for cross-pod budget enforcement. - await increment_spend_counters( - token=user_api_key, - team_id=team_id, - user_id=user_id, - response_cost=response_cost, - org_id=org_id, - ) + end_user_id=end_user_id, + ): + ## UPDATE DATABASE + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + team_id=team_id, + org_id=org_id, + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + response_cost=response_cost, + budget_reservation=budget_reservation, + request_tags=tags, + ) # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) @@ -234,10 +251,15 @@ class _ProxyDBLogger(CustomLogger): token=user_api_key, key_alias=key_alias, end_user_id=end_user_id, - response_cost=response_cost, - max_budget=end_user_max_budget, - ) + response_cost=response_cost, + max_budget=end_user_max_budget, + ) + elif budget_reservation is not None: + await _release_budget_reservation( + budget_reservation=budget_reservation + ) else: + await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. if sl_object is None and not kwargs.get("model"): @@ -366,7 +388,7 @@ class _ProxyDBLogger(CustomLogger): return -def _should_track_cost_callback( +def _should_track_cost_callback( user_api_key: Optional[str], user_id: Optional[str], team_id: Optional[str], @@ -387,4 +409,135 @@ def _should_track_cost_callback( or end_user_id is not None ): return True - return False + return False + + +def _get_budget_reservation_from_metadata(metadata: dict) -> Optional[dict]: + metadata_budget_reservation = metadata.get("user_api_key_budget_reservation") + if isinstance(metadata_budget_reservation, dict): + return metadata_budget_reservation + + user_api_key_auth_obj = metadata.get("user_api_key_auth") + if user_api_key_auth_obj is None: + return None + if isinstance(user_api_key_auth_obj, dict): + budget_reservation = user_api_key_auth_obj.get("budget_reservation") + return budget_reservation if isinstance(budget_reservation, dict) else None + return getattr(user_api_key_auth_obj, "budget_reservation", None) + + +def _get_request_tags_for_cost_tracking( + sl_object: Optional[StandardLoggingPayload], + metadata: dict, +) -> Optional[List[str]]: + if sl_object is not None: + request_tags = sl_object.get("request_tags", None) + if isinstance(request_tags, list): + return request_tags + + metadata_tags = metadata.get("tags", None) + if isinstance(metadata_tags, list): + return metadata_tags + + return None + + +async def _update_database_and_spend_counters( + proxy_logging_obj: Any, + increment_spend_counters: Any, + user_api_key: Optional[str], + user_id: Optional[str], + end_user_id: Optional[str], + team_id: Optional[str], + org_id: Optional[str], + kwargs: dict, + completion_response: Optional[Union[litellm.ModelResponse, Any]], + start_time: Any, + end_time: Any, + response_cost: float, + budget_reservation: Optional[dict], + request_tags: Optional[List[str]] = None, +) -> None: + try: + await proxy_logging_obj.db_spend_update_writer.update_database( + token=user_api_key, + response_cost=response_cost, + user_id=user_id, + end_user_id=end_user_id, + team_id=team_id, + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + org_id=org_id, + ) + except Exception: + if budget_reservation is not None: + try: + await _release_budget_reservation(budget_reservation=budget_reservation) + except Exception: + verbose_proxy_logger.exception( + "Failed to release budget reservation after database update failed" + ) + try: + await _invalidate_budget_reservation_counters( + budget_reservation=budget_reservation + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate budget reservation counters after release failed" + ) + raise + + try: + await increment_spend_counters( + token=user_api_key, + team_id=team_id, + user_id=user_id, + response_cost=response_cost, + org_id=org_id, + budget_reservation=budget_reservation, + end_user_id=end_user_id, + tags=request_tags, + ) + except Exception: + if budget_reservation is not None: + try: + await _invalidate_budget_reservation_counters( + budget_reservation=budget_reservation + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate budget reservation counters after spend counter update failed" + ) + finally: + budget_reservation["finalized"] = True + raise + + +async def _release_budget_reservation(budget_reservation: Optional[dict]) -> None: + if budget_reservation is None: + return + + from litellm.proxy.spend_tracking.budget_reservation import ( + release_budget_reservation, + ) + + await release_budget_reservation( + budget_reservation=budget_reservation, + ) + + +async def _invalidate_budget_reservation_counters( + budget_reservation: Optional[dict], +) -> None: + if budget_reservation is None: + return + + from litellm.proxy.spend_tracking.budget_reservation import ( + invalidate_budget_reservation_counters, + ) + + await invalidate_budget_reservation_counters( + budget_reservation=budget_reservation, + ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 3077efe116..853c56856f 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -893,6 +893,10 @@ class LiteLLMProxyRequestSetup: data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr( user_api_key_dict, "end_user_max_budget", None ) + if user_api_key_dict.budget_reservation is not None: + data[_metadata_variable_name][ + "user_api_key_budget_reservation" + ] = user_api_key_dict.budget_reservation # Add the full UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict return data 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 2485aea14f..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, @@ -1059,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. @@ -1834,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, @@ -3298,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) @@ -3341,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: @@ -3415,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]]: @@ -3605,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, @@ -3862,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.""" @@ -4152,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 @@ -5173,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_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6521abffb8..ce103f806e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -47,6 +47,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.proxy.utils import is_known_model from litellm.proxy.vector_store_endpoints.utils import ( + assert_user_can_access_vector_store, + get_litellm_managed_vector_store, is_allowed_to_call_vector_store_endpoint, ) from litellm.secret_managers.main import get_secret_str @@ -533,6 +535,10 @@ async def milvus_proxy_route( ) if vector_store is None: raise Exception(f"Vector store not found for {vector_store_name}") + await assert_user_can_access_vector_store( + vector_store=vector_store, + user_api_key_dict=user_api_key_dict, + ) litellm_params = vector_store.get("litellm_params") or {} auth_credentials = provider_config.get_auth_credentials( litellm_params=litellm_params @@ -1438,6 +1444,10 @@ async def azure_proxy_route( ) if vector_store is None: raise Exception(f"Vector store not found for {vector_store_name}") + await assert_user_can_access_vector_store( + vector_store=vector_store, + user_api_key_dict=user_api_key_dict, + ) litellm_params = vector_store.get("litellm_params") or {} auth_credentials = provider_config.get_auth_credentials( litellm_params=litellm_params @@ -1777,6 +1787,11 @@ async def _base_vertex_proxy_route( request=request, api_key=api_key_to_use, ) + if router_credentials is not None: + await assert_user_can_access_vector_store( + vector_store=router_credentials, + user_api_key_dict=user_api_key_dict, + ) vertex_project: Optional[str] = get_vertex_project_id_from_url(endpoint) vertex_location: Optional[str] = get_vertex_location_from_url(endpoint) @@ -1913,11 +1928,11 @@ async def vertex_discovery_proxy_route( "Extracted vector store ID from endpoint: %s", vector_store_id ) - # Retrieve vector store credentials from the registry - vector_store_credentials = ( - passthrough_endpoint_router.get_vector_store_credentials( - vector_store_id=vector_store_id - ) + # Retrieve LiteLLM-managed vector store credentials if the datastore id + # is registered with LiteLLM. Unknown datastore ids keep the existing + # direct Vertex pass-through behavior. + vector_store_credentials = await get_litellm_managed_vector_store( + vector_store_id=vector_store_id ) if vector_store_credentials: @@ -1925,7 +1940,7 @@ async def vertex_discovery_proxy_route( "Found vector store credentials for ID: %s", vector_store_id ) else: - verbose_proxy_logger.warning( + verbose_proxy_logger.debug( "Vector store ID %s found in endpoint but no credentials found in registry", vector_store_id, ) 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..cc6c26fdf9 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,10 @@ 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. + # That left OSS with no safe configuration: auth=True raised at startup + # unless the operator had a license. The safe option must always be free, + # and unauthenticated forwarding should require explicit opt-in. 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..551c85202b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6,6 +6,7 @@ import inspect import io import os import random +import re import secrets import shutil import subprocess @@ -78,8 +79,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 +98,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 +211,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, @@ -950,6 +956,85 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] +def _generate_stable_operation_id(route: Any) -> str: + operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") + route_methods = sorted(route.methods or []) + if len(route_methods) == 1: + operation_id = f"{operation_id}_{route_methods[0].lower()}" + return operation_id + + +_OPENAPI_HTTP_METHODS = { + "delete", + "get", + "head", + "options", + "patch", + "post", + "put", + "trace", +} + + +def _strip_operation_id_method_suffix(operation_id: str) -> str: + base, separator, suffix = operation_id.rpartition("_") + if separator and suffix in _OPENAPI_HTTP_METHODS: + return base + return operation_id + + +def ensure_unique_openapi_operation_ids( + openapi_schema: Dict[str, Any], + reserved_operation_ids: Optional[Set[str]] = None, +) -> Dict[str, Any]: + operation_entries = [] + operation_id_counts: Dict[str, int] = {} + for path_item in openapi_schema.get("paths", {}).values(): + if not isinstance(path_item, dict): + continue + for method, operation in path_item.items(): + if method not in _OPENAPI_HTTP_METHODS or not isinstance(operation, dict): + continue + operation_id = operation.get("operationId") + if not isinstance(operation_id, str): + continue + operation_entries.append((method, operation, operation_id)) + operation_id_counts[operation_id] = ( + operation_id_counts.get(operation_id, 0) + 1 + ) + + used_operation_ids = set(reserved_operation_ids or set()) + seen_operation_ids: Set[str] = set() + for method, operation, operation_id in operation_entries: + should_rewrite = ( + operation_id_counts[operation_id] > 1 + or operation_id in used_operation_ids + or operation_id in seen_operation_ids + ) + if not should_rewrite: + seen_operation_ids.add(operation_id) + used_operation_ids.add(operation_id) + continue + + base_operation_id = _strip_operation_id_method_suffix(operation_id) + new_operation_id = f"{base_operation_id}_{method}" + suffix = 2 + while ( + new_operation_id in used_operation_ids + or new_operation_id in seen_operation_ids + ): + new_operation_id = f"{base_operation_id}_{method}_{suffix}" + suffix += 1 + operation["operationId"] = new_operation_id + seen_operation_ids.add(new_operation_id) + used_operation_ids.add(new_operation_id) + + if reserved_operation_ids is not None: + reserved_operation_ids.update(used_operation_ids) + + return openapi_schema + + app = FastAPI( docs_url=_get_docs_url(), redoc_url=_get_redoc_url(), @@ -959,6 +1044,7 @@ app = FastAPI( version=version, root_path=server_root_path, lifespan=proxy_startup_event, # type: ignore[reportGeneralTypeIssues] + generate_unique_id_function=_generate_stable_operation_id, ) vertex_live_passthrough_vertex_base = VertexBase() @@ -1038,6 +1124,7 @@ def get_openapi_schema(): from litellm.proxy._lazy_features import inject_lazy_stubs openapi_schema = inject_lazy_stubs(openapi_schema) + openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema) # Fix Swagger UI execute path error when server_root_path is set if server_root_path: @@ -1069,6 +1156,7 @@ def custom_openapi(): from litellm.proxy._lazy_features import inject_lazy_stubs openapi_schema = inject_lazy_stubs(openapi_schema) + openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema) # Fix Swagger UI execute path error when server_root_path is set if server_root_path: @@ -1612,7 +1700,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 +1886,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 +1903,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( @@ -1835,6 +1928,9 @@ async def increment_spend_counters( user_id: Optional[str], response_cost: Optional[float], org_id: Optional[str] = None, + budget_reservation: Optional[dict] = None, + end_user_id: Optional[str] = None, + tags: Optional[List[str]] = None, ): """ Atomically increment spend counters for budget enforcement. @@ -1846,7 +1942,14 @@ async def increment_spend_counters( Awaited (not create_task) in the cost callback, so the counter is updated before the next request's auth check runs. """ + reserved_counter_keys = await _reconcile_budget_reservation_for_counter_update( + budget_reservation=budget_reservation, + response_cost=response_cost, + ) + if response_cost is None or response_cost == 0: + if budget_reservation is not None: + budget_reservation["finalized"] = True return if token is not None: @@ -1861,11 +1964,13 @@ async def increment_spend_counters( if isinstance(token, str) and token.startswith("sk-") else token ) - await _init_and_increment_spend_counter( - counter_key=f"spend:key:{hashed_token}", - source_cache_key=hashed_token, - increment=response_cost, - ) + key_counter_key = f"spend:key:{hashed_token}" + if key_counter_key not in reserved_counter_keys: + await _init_and_increment_spend_counter( + counter_key=key_counter_key, + source_cache_key=hashed_token, + increment=response_cost, + ) # Increment per-window budget counters for multi-budget keys key_obj = await user_api_key_cache.async_get_cache(key=hashed_token) @@ -1882,17 +1987,28 @@ async def increment_spend_counters( if isinstance(window, dict) else window.budget_duration ) - await spend_counter_cache.async_increment_cache( - key=f"spend:key:{hashed_token}:window:{duration}", - value=response_cost, - ) + key_window_counter = f"spend:key:{hashed_token}:window:{duration}" + if key_window_counter not in reserved_counter_keys: + from litellm.proxy.spend_tracking.budget_reservation import ( + get_budget_window_start, + ) + + await _init_and_increment_window_spend_counter( + counter_key=key_window_counter, + entity_type="Key", + entity_id=hashed_token, + window_start=get_budget_window_start(window), + increment=response_cost, + ) if team_id is not None: - await _init_and_increment_spend_counter( - counter_key=f"spend:team:{team_id}", - source_cache_key=f"team_id:{team_id}", - increment=response_cost, - ) + team_counter_key = f"spend:team:{team_id}" + if team_counter_key not in reserved_counter_keys: + await _init_and_increment_spend_counter( + counter_key=team_counter_key, + source_cache_key=f"team_id:{team_id}", + increment=response_cost, + ) # Increment per-window budget counters for multi-budget teams team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{team_id}") @@ -1909,36 +2025,157 @@ async def increment_spend_counters( if isinstance(window, dict) else window.budget_duration ) - await spend_counter_cache.async_increment_cache( - key=f"spend:team:{team_id}:window:{duration}", - value=response_cost, - ) + team_window_counter = f"spend:team:{team_id}:window:{duration}" + if team_window_counter not in reserved_counter_keys: + from litellm.proxy.spend_tracking.budget_reservation import ( + get_budget_window_start, + ) + + await _init_and_increment_window_spend_counter( + counter_key=team_window_counter, + entity_type="Team", + entity_id=team_id, + window_start=get_budget_window_start(window), + increment=response_cost, + ) if user_id is not None and team_id is not None: - await _init_and_increment_spend_counter( - counter_key=f"spend:team_member:{user_id}:{team_id}", - source_cache_key=f"team_membership:{user_id}:{team_id}", - increment=response_cost, - ) + team_member_counter_key = f"spend:team_member:{user_id}:{team_id}" + if team_member_counter_key not in reserved_counter_keys: + await _init_and_increment_spend_counter( + counter_key=team_member_counter_key, + source_cache_key=f"team_membership:{user_id}:{team_id}", + increment=response_cost, + ) if user_id is not None: - await _init_and_increment_spend_counter( - counter_key=f"spend:user:{user_id}", - source_cache_key=user_id, + user_counter_key = f"spend:user:{user_id}" + if user_counter_key not in reserved_counter_keys: + await _init_and_increment_spend_counter( + counter_key=user_counter_key, + source_cache_key=user_id, + increment=response_cost, + ) + + await _increment_end_user_and_tag_spend_counters( + end_user_id=end_user_id, + tags=tags, + response_cost=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + + await _increment_org_spend_counter( + org_id=org_id, + response_cost=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + if budget_reservation is not None: + budget_reservation["finalized"] = True + + +async def _reconcile_budget_reservation_for_counter_update( + budget_reservation: Optional[dict], + response_cost: Optional[float], +) -> Set[str]: + if budget_reservation is None: + return set() + + from litellm.proxy.spend_tracking.budget_reservation import ( + get_reserved_counter_keys, + invalidate_budget_reservation_counters, + reconcile_budget_reservation, + ) + + reserved_counter_keys = get_reserved_counter_keys( + budget_reservation=budget_reservation + ) + try: + await reconcile_budget_reservation( + budget_reservation=budget_reservation, + actual_cost=response_cost or 0.0, + finalize=False, + ) + except Exception: + verbose_proxy_logger.warning( + "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and continuing", + exc_info=True, + ) + try: + await invalidate_budget_reservation_counters( + budget_reservation=budget_reservation + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate reserved counters after reservation reconciliation failed" + ) + return reserved_counter_keys + + +async def _increment_end_user_and_tag_spend_counters( + end_user_id: Optional[str], + tags: Optional[List[str]], + response_cost: float, + reserved_counter_keys: Set[str], +) -> None: + if end_user_id is not None: + await _init_and_increment_unreserved_spend_counter( + counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=f"end_user_id:{end_user_id}", increment=response_cost, + reserved_counter_keys=reserved_counter_keys, ) - if org_id is not None: - await _init_and_increment_spend_counter( - counter_key=f"spend:org:{org_id}", - source_cache_key=f"org_id:{org_id}", + if tags is None: + return + + seen_tags: Set[str] = set() + for tag_name in tags: + if not tag_name or not isinstance(tag_name, str) or tag_name in seen_tags: + continue + seen_tags.add(tag_name) + await _init_and_increment_unreserved_spend_counter( + counter_key=f"spend:tag:{tag_name}", + source_cache_key=f"tag:{tag_name}", increment=response_cost, + reserved_counter_keys=reserved_counter_keys, ) +async def _increment_org_spend_counter( + org_id: Optional[str], + response_cost: float, + reserved_counter_keys: Set[str], +) -> None: + if org_id is None: + return + + await _init_and_increment_unreserved_spend_counter( + counter_key=f"spend:org:{org_id}", + source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"], + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + + +async def _init_and_increment_unreserved_spend_counter( + counter_key: str, + source_cache_key: Union[str, List[str]], + increment: float, + reserved_counter_keys: Set[str], +) -> None: + if counter_key in reserved_counter_keys: + return + + await _init_and_increment_spend_counter( + counter_key=counter_key, + source_cache_key=source_cache_key, + increment=increment, + ) + + async def _init_and_increment_spend_counter( counter_key: str, - source_cache_key: str, + source_cache_key: Union[str, List[str]], increment: float, ): """ @@ -1957,29 +2194,163 @@ async def _init_and_increment_spend_counter( under-counting (would allow overspend). 4. Increment atomically (both in-memory + Redis) """ - current = await spend_counter_cache.async_get_cache(key=counter_key) - if current is None: + await _ensure_spend_counter_initialized( + counter_key=counter_key, + source_cache_key=source_cache_key, + ) + await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + + +async def _init_and_increment_window_spend_counter( + counter_key: str, + entity_type: str, + entity_id: str, + window_start: Optional[datetime], + increment: float, +): + if window_start is None: + verbose_proxy_logger.warning( + "Skipping spend counter increment for invalid budget window %s", + counter_key, + ) + return + + initialized = await _ensure_window_spend_counter_initialized( + counter_key=counter_key, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + if initialized is False: + return + await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) + + +async def _ensure_spend_counter_initialized( + counter_key: str, + source_cache_key: Union[str, List[str]], +): + is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key) + if is_warm is False: # Shares the per-counter lock with get_current_spend. db_spend = await SpendCounterReseed.coalesced( prisma_client=prisma_client, spend_counter_cache=spend_counter_cache, counter_key=counter_key, + require_cache_warm=True, ) if db_spend is None: # DB unavailable - fall back to in-process cache (may be stale). - source = await user_api_key_cache.async_get_cache(key=source_cache_key) - base_spend: float = 0.0 - if source is not None: - if isinstance(source, dict): - base_spend = source.get("spend", 0.0) or 0.0 - else: - base_spend = getattr(source, "spend", 0.0) or 0.0 + base_spend = await _get_source_cache_base_spend( + source_cache_key=source_cache_key + ) if base_spend > 0: - await spend_counter_cache.async_increment_cache( - key=counter_key, value=base_spend + await _increment_spend_counter_cache( + counter_key=counter_key, increment=base_spend ) - await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + +async def _get_source_cache_base_spend( + source_cache_key: Union[str, List[str]], +) -> float: + source_cache_keys = ( + [source_cache_key] if isinstance(source_cache_key, str) else source_cache_key + ) + for cache_key in source_cache_keys: + source = await user_api_key_cache.async_get_cache(key=cache_key) + if source is None: + continue + if isinstance(source, dict): + return float(source.get("spend", 0.0) or 0.0) + return float(getattr(source, "spend", 0.0) or 0.0) + return 0.0 + + +async def _ensure_window_spend_counter_initialized( + counter_key: str, + entity_type: str, + entity_id: str, + window_start: datetime, +) -> bool: + is_warm = await _is_spend_counter_cache_warm(counter_key=counter_key) + if is_warm is True: + return True + + window_spend = await SpendCounterReseed.coalesced_window( + prisma_client=prisma_client, + spend_counter_cache=spend_counter_cache, + counter_key=counter_key, + entity_type=entity_type, + entity_id=entity_id, + window_start=window_start, + ) + if window_spend is None: + verbose_proxy_logger.warning( + "Skipping cold spend counter seed for %s because window spend could not be loaded", + counter_key, + ) + return False + return True + + +async def _is_spend_counter_cache_warm(counter_key: str) -> bool: + if spend_counter_cache.redis_cache is not None: + try: + current_value = await spend_counter_cache.redis_cache.async_get_cache( + key=counter_key, + ) + if current_value is None: + return False + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, + value=current_value, + ) + return True + except Exception as e: + verbose_proxy_logger.debug( + "Unable to read Redis spend counter %s before initialization, falling back to in-memory: %s", + counter_key, + e, + ) + + return spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is not None + + +async def _increment_spend_counter_cache(counter_key: str, increment: float): + if spend_counter_cache.redis_cache is not None: + try: + current_value = await spend_counter_cache.redis_cache.async_increment( + key=counter_key, + value=increment, + refresh_ttl=True, + ) + except Exception: + await _invalidate_spend_counter(counter_key=counter_key) + raise + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, + value=current_value, + ) + return current_value + + return await spend_counter_cache.async_increment_cache( + key=counter_key, + value=increment, + refresh_ttl=True, + ) + + +async def _invalidate_spend_counter(counter_key: str): + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) + except Exception: + verbose_proxy_logger.debug( + "Unable to delete stale spend counter %s after increment failure", + counter_key, + exc_info=True, + ) async def update_cache( # noqa: PLR0915 @@ -2007,14 +2378,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 +2445,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 +2518,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 +2563,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 +2615,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 +3308,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 +3326,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 +3674,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 +3951,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 +6696,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 +6745,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 +7385,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 +10687,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 +10804,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 +12843,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 +12895,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 +12945,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 +13250,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 +13263,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/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 95ca51612f..498d77f753 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -15,6 +15,7 @@ from fastapi.responses import ORJSONResponse import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import ( @@ -22,10 +23,88 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, get_form_data, ) +from litellm.proxy.vector_store_endpoints.utils import ( + assert_user_can_access_vector_store_id, +) router = APIRouter() +def _raise_vector_store_scan_depth_exceeded() -> None: + raise HTTPException( + status_code=400, + detail={ + "error": f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while scanning vector_store_id values" + }, + ) + + +def _append_payload_to_scan_stack( + payload_stack: list[tuple[Any, int]], + value: Any, + next_depth: int, +) -> None: + if isinstance(value, dict): + if next_depth > DEFAULT_MAX_RECURSE_DEPTH: + _raise_vector_store_scan_depth_exceeded() + payload_stack.append((value, next_depth)) + elif isinstance(value, list): + if next_depth > DEFAULT_MAX_RECURSE_DEPTH: + if any(isinstance(item, (dict, list)) for item in value): + _raise_vector_store_scan_depth_exceeded() + return + payload_stack.append((value, next_depth)) + + +def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]: + vector_store_ids: set[str] = set() + payload_stack = [(payload, 0)] + + while payload_stack: + current_payload, depth = payload_stack.pop() + if depth > DEFAULT_MAX_RECURSE_DEPTH: + _raise_vector_store_scan_depth_exceeded() + + if isinstance(current_payload, dict): + for key, value in current_payload.items(): + if key == "vector_store_id": + if not isinstance(value, str) or not value: + raise HTTPException( + status_code=400, + detail={ + "error": "vector_store_id must be a non-empty string" + }, + ) + vector_store_ids.add(value) + continue + if isinstance(value, (dict, list)): + _append_payload_to_scan_stack( + payload_stack=payload_stack, + value=value, + next_depth=depth + 1, + ) + elif isinstance(current_payload, list): + for item in current_payload: + _append_payload_to_scan_stack( + payload_stack=payload_stack, + value=item, + next_depth=depth + 1, + ) + + return vector_store_ids + + +async def _authorize_nested_vector_store_ids( + payload: Any, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + for vector_store_id in sorted(_collect_vector_store_ids_from_payload(payload)): + await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) + + def _build_file_metadata_entry( response: Any, file_data: Optional[Tuple[str, bytes, str]] = None, @@ -385,6 +464,11 @@ async def rag_ingest( }, ) + await _authorize_nested_vector_store_ids( + payload=ingest_options, + user_api_key_dict=user_api_key_dict, + ) + # Add litellm data request_data: Dict[str, Any] = {} request_data = await add_litellm_data_to_request( @@ -537,11 +621,20 @@ async def rag_query( status_code=400, detail={"error": "retrieval_config is required"}, ) + if not isinstance(retrieval_config, dict): + raise HTTPException( + status_code=400, + detail={"error": "retrieval_config must be an object"}, + ) if "vector_store_id" not in retrieval_config: raise HTTPException( status_code=400, detail={"error": "retrieval_config must contain 'vector_store_id'"}, ) + await _authorize_nested_vector_store_ids( + payload=retrieval_config, + user_api_key_dict=user_api_key_dict, + ) # Add litellm data request_data: Dict[str, Any] = {} diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py new file mode 100644 index 0000000000..1d296611bf --- /dev/null +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -0,0 +1,1029 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List, Optional, Sequence, cast + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache +from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_utils import get_model_from_request +from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router import Router + + +@dataclass +class _BudgetCounter: + counter_key: str + max_budget: float + fallback_spend: float + entity_type: str + entity_id: str + source_cache_key: Optional[str] = None + spend_log_entity_id: Optional[str] = None + window_start: Optional[datetime] = None + + +class _CounterReservationUnavailable(Exception): + def __init__( + self, + touched_counter: bool = False, + counter_invalidated: bool = False, + ) -> None: + self.touched_counter = touched_counter + self.counter_invalidated = counter_invalidated + super().__init__("Counter reservation unavailable") + + +def get_reserved_counter_keys(budget_reservation: Optional[dict]) -> set: + if not budget_reservation: + return set() + entries = budget_reservation.get("entries") or [] + return { + entry["counter_key"] + for entry in entries + if isinstance(entry, dict) and entry.get("counter_key") is not None + } + + +async def reserve_budget_for_request( + request_body: dict, + route: str, + llm_router: Optional[Router], + valid_token: Optional[UserAPIKeyAuth], + team_object: Optional[LiteLLM_TeamTable], + user_object: Optional[LiteLLM_UserTable], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, + end_user_id: Optional[str] = None, + end_user_object: Optional[Any] = None, +) -> Optional[dict]: + if valid_token is None or not RouteChecks.is_llm_api_route(route=route): + return None + if route in {"/models", "/v1/models", "/utils/token_counter"}: + return None + if get_model_from_request(request_body, route) is None: + return None + + counters = await _get_budget_counters( + request_body=request_body, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_id=end_user_id, + end_user_object=end_user_object, + ) + if not counters: + return None + + current_spend_by_counter_key: Dict[str, float] = {} + reservation_cost = estimate_request_max_cost( + request_body=request_body, + route=route, + llm_router=llm_router, + ) + if reservation_cost is None: + reservation_cost = await _get_smallest_remaining_budget( + counters=counters, + current_spend_by_counter_key=current_spend_by_counter_key, + ) + if reservation_cost is None or reservation_cost <= 0: + return None + + applied_entries: List[Dict[str, Any]] = [] + try: + for counter in counters: + entry = _counter_to_reservation_entry( + counter=counter, + reserved_cost=reservation_cost, + ) + applied_entries.append(entry) + try: + reserved_value = await _reserve_counter( + counter=counter, + reservation_cost=reservation_cost, + ) + except _CounterReservationUnavailable as exc: + if exc.touched_counter and not exc.counter_invalidated: + await _release_applied_entries_best_effort( + entries=[entry], + default_reserved_cost=reservation_cost, + ) + applied_entries.remove(entry) + continue + + if reserved_value is not None: + current_spend = reserved_value + else: + cached_spend = current_spend_by_counter_key.get(counter.counter_key) + if cached_spend is None: + cached_spend = await _get_current_counter_value(counter=counter) + current_spend = cached_spend + reservation_cost + if current_spend > counter.max_budget: + remaining_before_reservation = counter.max_budget - ( + current_spend - reservation_cost + ) + if remaining_before_reservation > 1e-12: + await _resize_applied_reservation( + entries=applied_entries, + current_reserved_cost=reservation_cost, + new_reserved_cost=remaining_before_reservation, + ) + reservation_cost = remaining_before_reservation + continue + raise litellm.BudgetExceededError( + current_cost=current_spend, + max_budget=counter.max_budget, + message=( + "Budget has been exceeded! " + f"{counter.entity_type}={counter.entity_id} " + f"Current cost: {current_spend}, " + f"Max budget: {counter.max_budget}" + ), + ) + except Exception: + await _release_applied_entries_best_effort( + entries=applied_entries, + default_reserved_cost=reservation_cost, + ) + raise + + if not applied_entries: + return None + + return { + "reserved_cost": reservation_cost, + "entries": applied_entries, + "finalized": False, + } + + +async def reconcile_budget_reservation( + budget_reservation: Optional[dict], + actual_cost: Optional[float], + finalize: bool = True, +) -> None: + if not budget_reservation or budget_reservation.get("finalized") is True: + return + + reserved_cost = float(budget_reservation.get("reserved_cost") or 0.0) + actual = float(actual_cost or 0.0) + await _set_reserved_entries_actual_cost( + entries=budget_reservation.get("entries") or [], + actual_cost=actual, + default_reserved_cost=reserved_cost, + ) + if finalize: + budget_reservation["finalized"] = True + + +async def release_budget_reservation(budget_reservation: Optional[dict]) -> None: + await reconcile_budget_reservation( + budget_reservation=budget_reservation, + actual_cost=0.0, + ) + + +async def invalidate_budget_reservation_counters( + budget_reservation: Optional[dict], +) -> None: + if budget_reservation is None: + return + + from litellm.proxy.proxy_server import _invalidate_spend_counter + + for counter_key in get_reserved_counter_keys(budget_reservation=budget_reservation): + await _invalidate_spend_counter(counter_key=counter_key) + + +async def _get_budget_counters( + request_body: dict, + valid_token: UserAPIKeyAuth, + team_object: Optional[LiteLLM_TeamTable], + user_object: Optional[LiteLLM_UserTable], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, + end_user_id: Optional[str] = None, + end_user_object: Optional[Any] = None, +) -> List[_BudgetCounter]: + counters: List[_BudgetCounter] = [] + + if valid_token.token is not None: + if valid_token.max_budget is not None and valid_token.max_budget > 0: + counters.append( + _BudgetCounter( + counter_key=f"spend:key:{valid_token.token}", + source_cache_key=valid_token.token, + max_budget=float(valid_token.max_budget), + fallback_spend=float(valid_token.spend or 0.0), + entity_type="Key", + entity_id=valid_token.token, + ) + ) + counters.extend( + _get_budget_limit_counters( + entity_prefix=f"spend:key:{valid_token.token}", + entity_type="Key", + entity_id=valid_token.token, + budget_limits=valid_token.budget_limits, + fallback_spend=float(valid_token.spend or 0.0), + ) + ) + + if team_object is not None and team_object.team_id is not None: + team_id = team_object.team_id + if team_object.max_budget is not None and team_object.max_budget > 0: + counters.append( + _BudgetCounter( + counter_key=f"spend:team:{team_id}", + source_cache_key=f"team_id:{team_id}", + max_budget=float(team_object.max_budget), + fallback_spend=float(team_object.spend or 0.0), + entity_type="Team", + entity_id=team_id, + ) + ) + counters.extend( + _get_budget_limit_counters( + entity_prefix=f"spend:team:{team_id}", + entity_type="Team", + entity_id=team_id, + budget_limits=team_object.budget_limits, + fallback_spend=float(team_object.spend or 0.0), + ) + ) + + if ( + (team_object is None or team_object.team_id is None) + and user_object is not None + and user_object.user_id is not None + and user_object.max_budget is not None + and user_object.max_budget > 0 + ): + counters.append( + _BudgetCounter( + counter_key=f"spend:user:{user_object.user_id}", + source_cache_key=user_object.user_id, + max_budget=float(user_object.max_budget), + fallback_spend=float(user_object.spend or 0.0), + entity_type="User", + entity_id=user_object.user_id, + ) + ) + + end_user_counter = await _get_end_user_budget_counter( + valid_token=valid_token, + end_user_id=end_user_id, + end_user_object=end_user_object, + ) + if end_user_counter is not None: + counters.append(end_user_counter) + + counters.extend( + await _get_tag_budget_counters( + request_body=request_body, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ) + + team_member_counter = await _get_team_member_budget_counter( + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + user_api_key_cache=user_api_key_cache, + ) + if team_member_counter is not None: + counters.append(team_member_counter) + + org_counter = await _get_org_budget_counter( + valid_token=valid_token, + team_object=team_object, + user_api_key_cache=user_api_key_cache, + ) + if org_counter is not None: + counters.append(org_counter) + + return counters + + +async def _get_end_user_budget_counter( + valid_token: UserAPIKeyAuth, + end_user_id: Optional[str], + end_user_object: Optional[Any], +) -> Optional[_BudgetCounter]: + end_user_id = end_user_id or valid_token.end_user_id + if end_user_id is None: + return None + + source_cache_key = f"end_user_id:{end_user_id}" + max_budget = _to_float(valid_token.end_user_max_budget) + fallback_spend = 0.0 + if end_user_object is not None: + fallback_spend = _to_float(_get_value(end_user_object, "spend")) or 0.0 + if max_budget is None: + budget_table = _get_value(end_user_object, "litellm_budget_table") + max_budget = _to_float(_get_value(budget_table, "max_budget")) + + if max_budget is None or max_budget <= 0: + return None + + return _BudgetCounter( + counter_key=f"spend:end_user:{end_user_id}", + source_cache_key=source_cache_key, + max_budget=max_budget, + fallback_spend=fallback_spend, + entity_type="EndUser", + entity_id=end_user_id, + ) + + +async def _get_tag_budget_counters( + request_body: dict, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, +) -> List[_BudgetCounter]: + from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body + from litellm.proxy.auth.auth_checks import get_tag_objects_batch + + tag_names = _dedupe_tags(get_tags_from_request_body(request_body=request_body)) + if not tag_names: + return [] + + tag_objects = await get_tag_objects_batch( + tag_names=tag_names, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + counters: List[_BudgetCounter] = [] + for tag_name in tag_names: + tag_object = tag_objects.get(tag_name) + if tag_object is None: + continue + budget_table = _get_value(tag_object, "litellm_budget_table") + max_budget = _to_float(_get_value(budget_table, "max_budget")) + if max_budget is None or max_budget <= 0: + continue + counters.append( + _BudgetCounter( + counter_key=f"spend:tag:{tag_name}", + source_cache_key=f"tag:{tag_name}", + max_budget=max_budget, + fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0, + entity_type="Tag", + entity_id=tag_name, + ) + ) + return counters + + +def _dedupe_tags(tags: List[str]) -> List[str]: + seen = set() + deduped_tags = [] + for tag in tags: + if tag in seen: + continue + seen.add(tag) + deduped_tags.append(tag) + return deduped_tags + + +async def _get_team_member_budget_counter( + valid_token: UserAPIKeyAuth, + team_object: Optional[LiteLLM_TeamTable], + user_object: Optional[LiteLLM_UserTable], + user_api_key_cache: DualCache, +) -> Optional[_BudgetCounter]: + if ( + team_object is None + or team_object.team_id is None + or user_object is None + or valid_token.user_id is None + ): + return None + + membership_cache_key = ( + f"team_membership:{valid_token.user_id}:{team_object.team_id}" + ) + cached_team_membership = await user_api_key_cache.async_get_cache( + key=membership_cache_key + ) + team_membership: Optional[LiteLLM_TeamMembership] = None + if isinstance(cached_team_membership, LiteLLM_TeamMembership): + team_membership = cached_team_membership + elif isinstance(cached_team_membership, dict): + team_membership = LiteLLM_TeamMembership(**cached_team_membership) + + team_member_budget: Optional[float] = None + if team_membership is not None and team_membership.litellm_budget_table is not None: + team_member_budget = team_membership.litellm_budget_table.max_budget + else: + default_budget_id = (team_object.metadata or {}).get("team_member_budget_id") + if isinstance(default_budget_id, str): + default_budget = await user_api_key_cache.async_get_cache( + key=f"team_member_default_budget:{default_budget_id}", + ) + team_member_budget = _to_float(_get_value(default_budget, "max_budget")) + + if team_member_budget is None or team_member_budget <= 0: + return None + + team_member_spend = ( + cast(LiteLLM_TeamMembership, team_membership).spend + if team_membership is not None + else 0.0 + ) + return _BudgetCounter( + counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + source_cache_key=membership_cache_key, + max_budget=float(team_member_budget), + fallback_spend=float(team_member_spend or 0.0), + entity_type="TeamMember", + entity_id=f"{valid_token.user_id}:{team_object.team_id}", + ) + + +async def _get_org_budget_counter( + valid_token: UserAPIKeyAuth, + team_object: Optional[LiteLLM_TeamTable], + user_api_key_cache: DualCache, +) -> Optional[_BudgetCounter]: + org_id: Optional[str] = None + if valid_token.org_id is not None: + org_id = valid_token.org_id + elif team_object is not None and team_object.organization_id is not None: + org_id = team_object.organization_id + if org_id is None: + return None + + org_table = await user_api_key_cache.async_get_cache( + key=f"org_id:{org_id}:with_budget", + ) + if org_table is None: + return None + + org_budget_table = _get_value(org_table, "litellm_budget_table") + if org_budget_table is None: + return None + + org_max_budget = _to_float(_get_value(org_budget_table, "max_budget")) + if org_max_budget is None or org_max_budget <= 0: + return None + + org_spend = _to_float(_get_value(org_table, "spend")) or 0.0 + return _BudgetCounter( + counter_key=f"spend:org:{org_id}", + source_cache_key=f"org_id:{org_id}:with_budget", + max_budget=org_max_budget, + fallback_spend=org_spend, + entity_type="Organization", + entity_id=org_id, + ) + + +def _get_budget_limit_counters( + entity_prefix: str, + entity_type: str, + entity_id: str, + budget_limits: Optional[Sequence[Any]], + fallback_spend: float, +) -> List[_BudgetCounter]: + counters: List[_BudgetCounter] = [] + if not budget_limits: + return counters + + for window in budget_limits: + window_dict = _coerce_window(window) + budget_duration = window_dict.get("budget_duration") + max_budget = window_dict.get("max_budget") + if not budget_duration or max_budget is None or max_budget <= 0: + continue + window_start = get_budget_window_start(window_dict) + if window_start is None: + verbose_proxy_logger.warning( + "Skipping budget window with invalid duration for %s=%s: %s", + entity_type, + entity_id, + budget_duration, + ) + continue + counters.append( + _BudgetCounter( + counter_key=f"{entity_prefix}:window:{budget_duration}", + max_budget=float(max_budget), + fallback_spend=0.0, + entity_type=entity_type, + entity_id=f"{entity_id}:{budget_duration}", + spend_log_entity_id=entity_id, + window_start=window_start, + ) + ) + return counters + + +def _coerce_window(window: Any) -> dict: + if isinstance(window, dict): + return window + if isinstance(window, str): + try: + parsed = json.loads(window) + return parsed if isinstance(parsed, dict) else {} + except Exception: + return {} + if hasattr(window, "model_dump"): + return window.model_dump() + return {} + + +async def _get_smallest_remaining_budget( + counters: List[_BudgetCounter], + current_spend_by_counter_key: Dict[str, float], +) -> Optional[float]: + remaining_budget: Optional[float] = None + for counter in counters: + current_spend = await _get_current_counter_value(counter=counter) + current_spend_by_counter_key[counter.counter_key] = current_spend + remaining = counter.max_budget - current_spend + if remaining <= 0: + raise litellm.BudgetExceededError( + current_cost=current_spend, + max_budget=counter.max_budget, + message=( + "Budget has been exceeded! " + f"{counter.entity_type}={counter.entity_id} " + f"Current cost: {current_spend}, " + f"Max budget: {counter.max_budget}" + ), + ) + remaining_budget = ( + remaining if remaining_budget is None else min(remaining_budget, remaining) + ) + return remaining_budget + + +async def _reserve_counter( + counter: _BudgetCounter, + reservation_cost: float, +) -> Optional[float]: + from litellm.proxy.proxy_server import ( + _ensure_spend_counter_initialized, + _ensure_window_spend_counter_initialized, + _invalidate_spend_counter, + _increment_spend_counter_cache, + ) + + attempted_increment = False + try: + if counter.source_cache_key is not None: + await _ensure_spend_counter_initialized( + counter_key=counter.counter_key, + source_cache_key=counter.source_cache_key, + ) + elif ( + counter.spend_log_entity_id is not None and counter.window_start is not None + ): + initialized = await _ensure_window_spend_counter_initialized( + counter_key=counter.counter_key, + entity_type=counter.entity_type, + entity_id=counter.spend_log_entity_id, + window_start=counter.window_start, + ) + if initialized is False: + verbose_proxy_logger.warning( + "Skipping budget reservation for %s because window spend could not be loaded", + counter.counter_key, + ) + raise _CounterReservationUnavailable + + attempted_increment = True + reserved_value = await _increment_spend_counter_cache( + counter_key=counter.counter_key, + increment=reservation_cost, + ) + return float(reserved_value) if reserved_value is not None else None + except _CounterReservationUnavailable: + raise + except Exception: + verbose_proxy_logger.warning( + "Skipping budget reservation for %s because spend counter reservation failed", + counter.counter_key, + exc_info=True, + ) + counter_invalidated = False + try: + await _invalidate_spend_counter(counter_key=counter.counter_key) + counter_invalidated = True + except Exception: + verbose_proxy_logger.warning( + "Failed to invalidate spend counter after budget reservation failure for %s", + counter.counter_key, + exc_info=True, + ) + raise _CounterReservationUnavailable( + touched_counter=attempted_increment, + counter_invalidated=counter_invalidated, + ) + + +async def _get_current_counter_value(counter: _BudgetCounter) -> float: + from litellm.proxy.proxy_server import get_current_spend + + return await get_current_spend( + counter_key=counter.counter_key, + fallback_spend=counter.fallback_spend, + ) + + +async def _set_reserved_entries_actual_cost( + entries: List[dict], + actual_cost: float, + default_reserved_cost: float, +) -> None: + for entry in entries: + await _set_reserved_entry_actual_cost( + entry=entry, + actual_cost=actual_cost, + default_reserved_cost=default_reserved_cost, + ) + + +async def _set_reserved_entry_actual_cost( + entry: dict, + actual_cost: float, + default_reserved_cost: float, +) -> None: + from litellm.proxy.proxy_server import _increment_spend_counter_cache + + counter_key = entry.get("counter_key") + if counter_key is None: + return + reserved_cost = _get_entry_reserved_cost( + entry=entry, + default_reserved_cost=default_reserved_cost, + ) + target_adjustment = actual_cost - reserved_cost + applied_adjustment = float(entry.get("applied_adjustment") or 0.0) + adjustment = target_adjustment - applied_adjustment + if adjustment == 0: + return + await _ensure_counter_can_apply_adjustment( + counter_key=counter_key, + adjustment=adjustment, + ) + await _increment_spend_counter_cache( + counter_key=counter_key, + increment=adjustment, + ) + entry["applied_adjustment"] = target_adjustment + + +async def _ensure_counter_can_apply_adjustment( + counter_key: str, + adjustment: float, +) -> None: + from litellm.proxy.proxy_server import ( + _invalidate_spend_counter, + spend_counter_cache, + ) + + current_value = await spend_counter_cache.async_get_cache(key=counter_key) + if current_value is None: + await _invalidate_spend_counter(counter_key=counter_key) + raise RuntimeError( + f"Cannot apply budget reservation adjustment to missing counter {counter_key}" + ) + + try: + current_float = float(current_value) + except (TypeError, ValueError): + await _invalidate_spend_counter(counter_key=counter_key) + raise RuntimeError( + f"Cannot apply budget reservation adjustment to non-numeric counter {counter_key}" + ) + + if adjustment < 0 and current_float + adjustment < -1e-12: + await _invalidate_spend_counter(counter_key=counter_key) + raise RuntimeError( + f"Budget reservation adjustment would make counter negative {counter_key}" + ) + + +async def _release_applied_entries_best_effort( + entries: List[dict], + default_reserved_cost: float, +) -> None: + for entry in entries: + try: + await _set_reserved_entry_actual_cost( + entry=entry, + actual_cost=0.0, + default_reserved_cost=default_reserved_cost, + ) + except Exception: + counter_key = entry.get("counter_key") + verbose_proxy_logger.exception( + "Failed to release partial budget reservation during exception cleanup" + ) + if counter_key is None: + continue + try: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter(counter_key=counter_key) + except Exception: + verbose_proxy_logger.exception( + "Failed to invalidate partial budget reservation counter during exception cleanup" + ) + + +async def _resize_applied_reservation( + entries: List[dict], + current_reserved_cost: float, + new_reserved_cost: float, +) -> None: + await _set_reserved_entries_actual_cost( + entries=entries, + actual_cost=new_reserved_cost, + default_reserved_cost=current_reserved_cost, + ) + for entry in entries: + entry["reserved_cost"] = new_reserved_cost + entry["applied_adjustment"] = 0.0 + + +def _counter_to_reservation_entry( + counter: _BudgetCounter, + reserved_cost: float, +) -> Dict[str, Any]: + return { + "counter_key": counter.counter_key, + "entity_type": counter.entity_type, + "entity_id": counter.entity_id, + "reserved_cost": reserved_cost, + "applied_adjustment": 0.0, + } + + +def _get_entry_reserved_cost(entry: dict, default_reserved_cost: float) -> float: + try: + return float(entry.get("reserved_cost", default_reserved_cost) or 0.0) + except (TypeError, ValueError): + return default_reserved_cost + + +def get_budget_window_start(window: Any) -> Optional[datetime]: + window_dict = _coerce_window(window) + budget_duration = window_dict.get("budget_duration") + if budget_duration is None: + return None + try: + duration_seconds = duration_in_seconds(str(budget_duration)) + except Exception: + return None + + reset_at = _coerce_datetime(window_dict.get("reset_at")) + if reset_at is None: + return datetime.now(timezone.utc) - timedelta(seconds=duration_seconds) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) + return reset_at - timedelta(seconds=duration_seconds) + + +def _coerce_datetime(value: Any) -> Optional[datetime]: + if value is None: + return None + if isinstance(value, datetime): + return value + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return None + + +def estimate_request_max_cost( + request_body: dict, + route: str, + llm_router: Optional[Router], +) -> Optional[float]: + model = get_model_from_request(request_body, route) + if model is None: + return None + + models = [model] if isinstance(model, str) else model + estimates = [ + _estimate_request_max_cost_for_model( + request_body=request_body, + route=route, + model=model_name, + llm_router=llm_router, + ) + for model_name in models + ] + estimates = [estimate for estimate in estimates if estimate is not None] + if not estimates: + return None + return max(cast(List[float], estimates)) + + +def _estimate_request_max_cost_for_model( + request_body: dict, + route: str, + model: str, + llm_router: Optional[Router], +) -> Optional[float]: + model_info = _get_model_cost_info(model=model, llm_router=llm_router) + if model_info is None: + return None + + input_cost_per_token = _to_float(model_info.get("input_cost_per_token")) + output_cost_per_token = _to_float(model_info.get("output_cost_per_token")) + input_tokens = _estimate_input_tokens( + request_body=request_body, + route=route, + model=model, + model_info=model_info, + ) + output_tokens = _estimate_output_tokens( + request_body=request_body, + route=route, + model_info=model_info, + ) + if input_tokens is None or output_tokens is None: + return None + + cost = 0.0 + if input_cost_per_token is not None: + cost += input_tokens * input_cost_per_token + elif input_tokens > 0: + return None + + output_multiplier = _get_output_multiplier(request_body=request_body) + if output_cost_per_token is not None: + cost += output_tokens * output_multiplier * output_cost_per_token + elif output_tokens > 0: + return None + + return cost + + +def _get_model_cost_info( + model: str, + llm_router: Optional[Router], +) -> Optional[Dict[str, Any]]: + if llm_router is not None: + try: + model_group_info = llm_router.get_model_group_info(model_group=model) + if model_group_info is not None: + return model_group_info.model_dump() + except Exception: + verbose_proxy_logger.debug( + "Unable to load router model group info for budget reservation", + exc_info=True, + ) + + try: + return dict(litellm.get_model_info(model=model)) + except Exception: + return None + + +def _estimate_input_tokens( + request_body: dict, + route: str, + model: str, + model_info: Dict[str, Any], +) -> Optional[int]: + try: + if "messages" in request_body: + return litellm.token_counter( + model=model, + messages=request_body.get("messages") or [], + tools=request_body.get("tools"), + tool_choice=request_body.get("tool_choice"), + ) + if "prompt" in request_body: + return _count_text_tokens(model=model, text=request_body.get("prompt")) + if "input" in request_body: + return _count_text_tokens(model=model, text=request_body.get("input")) + if "query" in request_body or "documents" in request_body: + query_tokens = _count_text_tokens( + model=model, text=request_body.get("query") + ) + document_tokens = _count_text_tokens( + model=model, + text=request_body.get("documents"), + ) + return query_tokens + document_tokens + except Exception: + verbose_proxy_logger.debug( + "Unable to count input tokens for budget reservation", exc_info=True + ) + + max_input_tokens = _to_int(model_info.get("max_input_tokens")) + if max_input_tokens is not None: + return max_input_tokens + + return None + + +def _estimate_output_tokens( + request_body: dict, + route: str, + model_info: Dict[str, Any], +) -> Optional[int]: + if _is_input_only_route(route=route): + return 0 + + for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"): + max_tokens = _to_int(request_body.get(key)) + if max_tokens is not None: + return max_tokens + + # If the caller did not cap output tokens, avoid reserving a model's + # theoretical maximum context. The caller can still admit one request by + # reserving the smallest remaining budget in reserve_budget_for_request(). + return None + + +def _count_text_tokens(model: str, text: Any) -> int: + if text is None: + return 0 + + token_count = 0 + stack = [text] + while stack: + item = stack.pop() + if item is None: + continue + if isinstance(item, list): + stack.extend(item) + continue + if isinstance(item, dict): + token_count += litellm.token_counter(model=model, text=json.dumps(item)) + continue + token_count += litellm.token_counter(model=model, text=str(item)) + return token_count + + +def _get_output_multiplier(request_body: dict) -> int: + output_multiplier = 1 + for key in ("n", "best_of"): + value = _to_int(request_body.get(key)) + if value is not None: + output_multiplier = max(output_multiplier, value) + return output_multiplier + + +def _is_input_only_route(route: str) -> bool: + return any( + route_part in route + for route_part in ( + "embeddings", + "rerank", + "moderations", + ) + ) + + +def _to_float(value: Any) -> Optional[float]: + if value is None: + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def _to_int(value: Any) -> Optional[int]: + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _get_value(obj: Any, key: str) -> Any: + if isinstance(obj, dict): + return obj.get(key) + return getattr(obj, key, None) 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/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 1fdfad8c96..86e316e7f4 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,8 +1,6 @@ from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response - -import litellm from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -10,7 +8,10 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import jsonify_object -from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.proxy.vector_store_endpoints.utils import ( + assert_user_can_access_vector_store, + get_litellm_managed_vector_store, +) from litellm.types.vector_stores import IndexCreateRequest router = APIRouter() @@ -19,24 +20,6 @@ router = APIRouter() ######################################################## -async def _check_vector_store_access( - vector_store: LiteLLM_ManagedVectorStore, - user_api_key_dict: UserAPIKeyAuth, -) -> bool: - """ - Check if the user has access to the vector store. - - Delegates to :func:`can_user_access_vector_store`, which honors: - - PROXY_ADMIN bypass - - legacy vector stores with no team_id - - key-level and team-level ``object_permission.vector_stores`` allowlists - - team_id match between key and store - """ - return await can_user_access_vector_store( - vector_store=vector_store, user_api_key_dict=user_api_key_dict - ) - - async def _update_request_data_with_litellm_managed_vector_store_registry( data: Dict, vector_store_id: str, @@ -53,35 +36,27 @@ async def _update_request_data_with_litellm_managed_vector_store_registry( Raises: HTTPException: If user doesn't have access to the vector store """ - if litellm.vector_store_registry is not None: - vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = ( - litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( - vector_store_id=vector_store_id + vector_store_to_run: Optional[LiteLLM_ManagedVectorStore] = ( + await get_litellm_managed_vector_store(vector_store_id=vector_store_id) + ) + if vector_store_to_run is not None: + if user_api_key_dict is not None: + await assert_user_can_access_vector_store( + vector_store=vector_store_to_run, + user_api_key_dict=user_api_key_dict, ) - ) - if vector_store_to_run is not None: - if user_api_key_dict is not None: - if not await _check_vector_store_access( - vector_store_to_run, user_api_key_dict - ): - raise HTTPException( - status_code=403, - detail="Access denied: You do not have permission to access this vector store", - ) - if "custom_llm_provider" in vector_store_to_run: - data["custom_llm_provider"] = vector_store_to_run.get( - "custom_llm_provider" - ) + if "custom_llm_provider" in vector_store_to_run: + data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider") - if "litellm_credential_name" in vector_store_to_run: - data["litellm_credential_name"] = vector_store_to_run.get( - "litellm_credential_name" - ) + if "litellm_credential_name" in vector_store_to_run: + data["litellm_credential_name"] = vector_store_to_run.get( + "litellm_credential_name" + ) - if "litellm_params" in vector_store_to_run: - litellm_params = vector_store_to_run.get("litellm_params", {}) or {} - data.update(litellm_params) + if "litellm_params" in vector_store_to_run: + litellm_params = vector_store_to_run.get("litellm_params", {}) or {} + data.update(litellm_params) return data @@ -121,8 +96,7 @@ async def vector_store_search( ) data = await _read_request_body(request=request) - if "vector_store_id" not in data: - data["vector_store_id"] = vector_store_id + data["vector_store_id"] = vector_store_id # Check for legacy vector store registry (non-managed vector stores) data = await _update_request_data_with_litellm_managed_vector_store_registry( diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 061a8aaa24..657b520b27 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -1,7 +1,9 @@ +import json from typing import Any, Dict, Literal, Optional from fastapi import HTTPException, Request +import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, @@ -13,6 +15,21 @@ from litellm.types.vector_stores import LiteLLM_ManagedVectorStore from litellm.utils import ProviderConfigManager +def _normalize_litellm_params( + vector_store: LiteLLM_ManagedVectorStore, +) -> LiteLLM_ManagedVectorStore: + litellm_params = vector_store.get("litellm_params") + if isinstance(litellm_params, str): + normalized = LiteLLM_ManagedVectorStore(**dict(vector_store)) + try: + parsed = json.loads(litellm_params) + normalized["litellm_params"] = parsed if isinstance(parsed, dict) else {} + except (TypeError, ValueError): + normalized["litellm_params"] = {} + return normalized + return vector_store + + def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: return ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN @@ -120,6 +137,104 @@ async def can_user_access_vector_store( return False +async def get_litellm_managed_vector_store( + vector_store_id: str, +) -> Optional[LiteLLM_ManagedVectorStore]: + """ + Resolve a LiteLLM-managed vector store from the registry or shared cache. + + Provider-native vector store IDs will not be present in either location and + return None, preserving direct provider behavior while still protecting + LiteLLM-managed multi-tenant stores. + """ + if not vector_store_id: + return None + + if litellm.vector_store_registry is not None: + try: + vector_store = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( + vector_store_id=vector_store_id + ) + if vector_store is not None: + return _normalize_litellm_params(vector_store) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to resolve vector store id=%s from registry: %s", + vector_store_id, + e, + ) + raise HTTPException( + status_code=500, + detail="Unable to validate vector store access", + ) from e + + try: + from litellm.proxy.auth.auth_checks import ( + get_managed_vector_store_rows_by_uuids, + ) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return None + rows = await get_managed_vector_store_rows_by_uuids( + uuids=[vector_store_id], + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if not rows: + return None + return _normalize_litellm_params( + LiteLLM_ManagedVectorStore(**rows[0].model_dump()) + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to resolve vector store id=%s from shared cache: %s", + vector_store_id, + e, + ) + raise HTTPException( + status_code=500, + detail="Unable to validate vector store access", + ) from e + + +async def assert_user_can_access_vector_store( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth, + detail: str = "Access denied: You do not have permission to access this vector store", +) -> None: + """Raise 403 unless the caller can access the resolved vector store.""" + if not await can_user_access_vector_store(vector_store, user_api_key_dict): + raise HTTPException(status_code=403, detail=detail) + + +async def assert_user_can_access_vector_store_id( + vector_store_id: str, + user_api_key_dict: UserAPIKeyAuth, + detail: str = "Access denied: You do not have permission to access this vector store", +) -> Optional[LiteLLM_ManagedVectorStore]: + """ + Resolve a managed vector store id and enforce ownership if it exists. + + Unknown ids are treated as provider-native ids and are not rejected here. + """ + vector_store = await get_litellm_managed_vector_store( + vector_store_id=vector_store_id + ) + if vector_store is not None: + await assert_user_can_access_vector_store( + vector_store=vector_store, + user_api_key_dict=user_api_key_dict, + detail=detail, + ) + return vector_store + + def _does_endpoint_match(endpoint_path: str, request_path: str) -> bool: if endpoint_path in request_path: return True diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 7cdf865692..346a847c5d 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -17,9 +17,11 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( prepare_data_with_credentials, ) from litellm.proxy.vector_store_endpoints.utils import ( + assert_user_can_access_vector_store_id, is_allowed_to_call_vector_store_files_endpoint, ) from litellm.types.utils import LlmProviders +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore if TYPE_CHECKING: from litellm.router import Router @@ -193,6 +195,8 @@ def _update_request_data_with_litellm_managed_vector_store_registry( data: Dict, vector_store_id: str, llm_router: Optional["Router"] = None, + managed_vector_store: Optional[LiteLLM_ManagedVectorStore] = None, + should_lookup_registry: bool = True, ) -> Dict: """ Update request data with model routing information from managed vector store. @@ -262,23 +266,27 @@ def _update_request_data_with_litellm_managed_vector_store_registry( return data - # Legacy path: Check vector store registry for non-managed vector stores - if litellm.vector_store_registry is not None: + # Legacy path: Check vector store registry for non-managed vector stores. + vector_store_to_run = managed_vector_store + if ( + vector_store_to_run is None + and should_lookup_registry + and litellm.vector_store_registry is not None + ): vector_store_to_run = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( vector_store_id=vector_store_id ) - if vector_store_to_run is not None: - if "custom_llm_provider" in vector_store_to_run: - data["custom_llm_provider"] = vector_store_to_run.get( - "custom_llm_provider" - ) - if "litellm_credential_name" in vector_store_to_run: - data["litellm_credential_name"] = vector_store_to_run.get( - "litellm_credential_name" - ) - if "litellm_params" in vector_store_to_run: - litellm_params = vector_store_to_run.get("litellm_params", {}) or {} - data.update(litellm_params) + + if vector_store_to_run is not None: + if "custom_llm_provider" in vector_store_to_run: + data["custom_llm_provider"] = vector_store_to_run.get("custom_llm_provider") + if "litellm_credential_name" in vector_store_to_run: + data["litellm_credential_name"] = vector_store_to_run.get( + "litellm_credential_name" + ) + if "litellm_params" in vector_store_to_run: + litellm_params = vector_store_to_run.get("litellm_params", {}) or {} + data.update(litellm_params) return data @@ -363,8 +371,11 @@ async def vector_store_file_create( ) data = await _read_request_body(request=request) - if "vector_store_id" not in data: - data["vector_store_id"] = vector_store_id + data["vector_store_id"] = vector_store_id + managed_vector_store = await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) # Handle managed file IDs if present in request body original_managed_file_id = None @@ -375,7 +386,11 @@ async def vector_store_file_create( # Then handle managed vector store IDs data = _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, llm_router=llm_router + data=data, + vector_store_id=vector_store_id, + llm_router=llm_router, + managed_vector_store=managed_vector_store, + should_lookup_registry=False, ) provider_enum = await _resolve_provider(data=data, request=request) @@ -459,9 +474,18 @@ async def vector_store_file_list( query_params = dict(request.query_params) data: Dict[str, Optional[str]] = {"vector_store_id": vector_store_id} data.update(query_params) + data["vector_store_id"] = vector_store_id + managed_vector_store = await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) data = _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, llm_router=llm_router + data=data, + vector_store_id=vector_store_id, + llm_router=llm_router, + managed_vector_store=managed_vector_store, + should_lookup_registry=False, ) provider_enum = await _resolve_provider(data=data, request=request) @@ -541,6 +565,10 @@ async def vector_store_file_retrieve( "vector_store_id": vector_store_id, "file_id": file_id, } + managed_vector_store = await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) # Handle managed file IDs first data, original_managed_file_id = _update_request_data_with_managed_file_id( @@ -549,7 +577,11 @@ async def vector_store_file_retrieve( # Then handle managed vector store IDs data = _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, llm_router=llm_router + data=data, + vector_store_id=vector_store_id, + llm_router=llm_router, + managed_vector_store=managed_vector_store, + should_lookup_registry=False, ) provider_enum = await _resolve_provider(data=data, request=request) @@ -635,6 +667,10 @@ async def vector_store_file_content( "vector_store_id": vector_store_id, "file_id": file_id, } + managed_vector_store = await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) # Handle managed file IDs first data, original_managed_file_id = _update_request_data_with_managed_file_id( @@ -643,7 +679,11 @@ async def vector_store_file_content( # Then handle managed vector store IDs data = _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, llm_router=llm_router + data=data, + vector_store_id=vector_store_id, + llm_router=llm_router, + managed_vector_store=managed_vector_store, + should_lookup_registry=False, ) provider_enum = await _resolve_provider(data=data, request=request) @@ -729,6 +769,10 @@ async def vector_store_file_update( data = await _read_request_body(request=request) data["vector_store_id"] = vector_store_id data["file_id"] = file_id + managed_vector_store = await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) # Handle managed file IDs first data, original_managed_file_id = _update_request_data_with_managed_file_id( @@ -737,7 +781,11 @@ async def vector_store_file_update( # Then handle managed vector store IDs data = _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, llm_router=llm_router + data=data, + vector_store_id=vector_store_id, + llm_router=llm_router, + managed_vector_store=managed_vector_store, + should_lookup_registry=False, ) provider_enum = await _resolve_provider(data=data, request=request) @@ -823,6 +871,10 @@ async def vector_store_file_delete( "vector_store_id": vector_store_id, "file_id": file_id, } + managed_vector_store = await assert_user_can_access_vector_store_id( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + ) # Handle managed file IDs first data, original_managed_file_id = _update_request_data_with_managed_file_id( @@ -831,7 +883,11 @@ async def vector_store_file_delete( # Then handle managed vector store IDs data = _update_request_data_with_litellm_managed_vector_store_registry( - data=data, vector_store_id=vector_store_id, llm_router=llm_router + data=data, + vector_store_id=vector_store_id, + llm_router=llm_router, + managed_vector_store=managed_vector_store, + should_lookup_registry=False, ) provider_enum = await _resolve_provider(data=data, request=request) 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/litellm/vector_stores/main.py b/litellm/vector_stores/main.py index 6d28d67097..13f2f27d3f 100644 --- a/litellm/vector_stores/main.py +++ b/litellm/vector_stores/main.py @@ -377,15 +377,11 @@ def search( _is_async = kwargs.pop("asearch", False) is True # pull credentials from registry if available - vector_store_id_for_credentials = kwargs.get("vector_store_id", vector_store_id) - if ( - litellm.vector_store_registry is not None - and vector_store_id_for_credentials is not None - ): + if litellm.vector_store_registry is not None and vector_store_id is not None: try: registry_credentials = ( litellm.vector_store_registry.get_credentials_for_vector_store( - vector_store_id_for_credentials + vector_store_id ) ) kwargs.update(registry_credentials) 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..0ef0a993dd 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", 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/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index fdf8c24ac9..7b2b6bed6a 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, ) 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..45fa23bcb6 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, @@ -518,3 +519,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/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/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 4579c20321..2fb36bf403 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -394,3 +394,76 @@ class TestHostAllowlist: monkeypatch.setattr(url_utils.socket, "getaddrinfo", fake) validate_url("http://internal.corp/") + + +# ── assert_same_origin ──────────────────────────────────────────────────────── + + +from litellm.litellm_core_utils.url_utils import 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/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/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_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_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/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_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/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a848db27fc..5dedf05215 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -17,7 +17,10 @@ import litellm from litellm.proxy._types import ( CallInfo, Litellm_EntityType, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, LiteLLM_ObjectPermissionTable, + LiteLLM_TagTable, LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, @@ -29,10 +32,12 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _can_object_call_vector_stores, + _check_end_user_budget, _check_team_member_budget, _get_fuzzy_user_object, _get_team_db_check, _log_budget_lookup_failure, + _tag_max_budget_check, _team_max_budget_check, _virtual_key_max_budget_alert_check, _virtual_key_max_budget_check, @@ -922,19 +927,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 +987,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() @@ -1962,6 +1969,67 @@ async def test_team_budget_check_reads_from_spend_counter(): assert exc_info.value.current_cost == 1.5 +@pytest.mark.asyncio +async def test_end_user_budget_check_reads_from_spend_counter(): + """End-user budget check should use get_current_spend when counter exists.""" + end_user_object = LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:end_user:customer-1": + return 1.5 + return fallback_spend + + with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_end_user_budget( + end_user_obj=end_user_object, + route="/chat/completions", + ) + assert exc_info.value.current_cost == 1.5 + assert exc_info.value.max_budget == 1.0 + + +@pytest.mark.asyncio +async def test_tag_budget_check_reads_from_spend_counter(): + """Tag budget check should use get_current_spend when counter exists.""" + from litellm.proxy.utils import ProxyLogging + + tag_object = LiteLLM_TagTable( + tag_name="paid-tag", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:tag:paid-tag": + return 1.5 + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"paid-tag": tag_object}, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body={"metadata": {"tags": ["paid-tag"]}}, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 1.5 + assert exc_info.value.max_budget == 1.0 + + @pytest.mark.asyncio async def test_team_member_budget_check_reads_from_spend_counter(): """Team member budget check should use get_current_spend when counter exists.""" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 91f300b88c..c146b5ded5 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -2,6 +2,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID extraction. """ +import base64 from typing import Optional from unittest.mock import MagicMock, patch @@ -10,11 +11,12 @@ import pytest from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( _get_customer_id_from_standard_headers, + abbreviate_api_key, check_complete_credentials, get_end_user_id_from_request_body, - get_model_from_request, get_key_model_rpm_limit, get_key_model_tpm_limit, + get_model_from_request, get_project_model_rpm_limit, get_project_model_tpm_limit, is_request_body_safe, @@ -258,6 +260,206 @@ def test_get_model_from_request_vertex_passthrough_still_works(): assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro" +def test_get_model_from_request_openai_deployment_route_still_works(): + assert ( + get_model_from_request( + request_data={}, + route="/openai/deployments/my-azure-deployment/chat/completions", + ) + == "my-azure-deployment" + ) + + +def test_get_model_from_request_includes_file_endpoint_header_model(): + assert ( + get_model_from_request( + request_data={}, + route="/v1/files", + request_headers={"X-LiteLLM-Model": "restricted-model"}, + ) + == "restricted-model" + ) + + +def test_get_model_from_request_ignores_routing_header_on_standard_llm_routes(): + assert ( + get_model_from_request( + request_data={"model": "allowed-model"}, + route="/v1/chat/completions", + request_headers={"x-litellm-model": "restricted-model"}, + ) + == "allowed-model" + ) + + +def test_get_model_from_request_authorizes_all_file_routing_model_sources(): + models = get_model_from_request( + request_data={"model": "body-model"}, + route="/v1/files", + request_headers={"x-litellm-model": "header-model"}, + request_query_params={"target_model_names": "query-model-a,query-model-b"}, + ) + assert isinstance(models, list) + assert set(models) == { + "body-model", + "query-model-a", + "query-model-b", + "header-model", + } + + +def test_get_model_from_request_extracts_simple_encoded_file_id_model(): + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + file_id = encode_file_id_with_model( + file_id="file-provider-id", + model="restricted-model", + ) + + assert ( + get_model_from_request( + request_data={"file_id": file_id}, + route="/v1/files/{file_id}", + ) + == "restricted-model" + ) + + +def test_get_model_from_request_extracts_unified_file_id_models(): + raw_unified_file_id = ( + "litellm_proxy:application/octet-stream;unified_id,test-id;" + "target_model_names,model-a,model-b;llm_output_file_id,file-provider-id" + ) + encoded_unified_file_id = ( + base64.urlsafe_b64encode(raw_unified_file_id.encode()).decode().rstrip("=") + ) + + assert get_model_from_request( + request_data={"file_id": encoded_unified_file_id}, + route="/v1/files/{file_id}", + ) == ["model-a", "model-b"] + + +def test_get_model_from_request_extracts_eval_completion_model(): + assert ( + get_model_from_request( + request_data={"completion": {"model": "judge-model"}}, + route="/v1/evals/{eval_id}/runs", + ) + == "judge-model" + ) + + +def test_get_model_from_request_includes_fine_tuning_target_model_query(): + assert ( + get_model_from_request( + request_data={}, + route="/v1/fine_tuning/jobs", + request_query_params={"target_model_names": "fine-tune-model"}, + ) + == "fine-tune-model" + ) + + +def test_get_model_from_request_extracts_video_id_model(): + from litellm.types.videos.utils import encode_video_id_with_provider + + video_id = encode_video_id_with_provider( + video_id="video-provider-id", + provider="openai", + model_id="video-model", + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + ) + == "video-model" + ) + + +def test_get_model_from_request_only_runs_media_decoders_for_matching_fields(): + with ( + patch( + "litellm.types.videos.utils.decode_video_id_with_provider", + return_value={"model_id": "video-model"}, + ) as video_decoder, + patch( + "litellm.types.videos.utils.decode_character_id_with_provider", + return_value={"model_id": "character-model"}, + ) as character_decoder, + ): + assert ( + get_model_from_request( + request_data={"file_id": "file-provider-id"}, + route="/v1/files/{file_id}", + ) + is None + ) + video_decoder.assert_not_called() + character_decoder.assert_not_called() + + assert ( + get_model_from_request( + request_data={"video_id": "video-provider-id"}, + route="/v1/videos/{video_id}", + ) + == "video-model" + ) + video_decoder.assert_called_once_with("video-provider-id") + character_decoder.assert_not_called() + + video_decoder.reset_mock() + character_decoder.reset_mock() + assert ( + get_model_from_request( + request_data={"character_id": "character-provider-id"}, + route="/v1/videos/{character_id}", + ) + == "character-model" + ) + video_decoder.assert_not_called() + character_decoder.assert_called_once_with("character-provider-id") + + +def test_get_model_from_request_handles_managed_id_decoder_failures(): + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils.decode_model_from_file_id", + side_effect=Exception("decode failed"), + ), + patch( + "litellm.llms.base_llm.managed_resources.utils.parse_unified_id", + side_effect=Exception("parse failed"), + ), + patch( + "litellm.types.videos.utils.decode_video_id_with_provider", + side_effect=Exception("video decode failed"), + ), + ): + assert ( + get_model_from_request( + request_data={"file_id": "not-a-managed-resource-id"}, + route="/v1/files/{file_id}", + ) + is None + ) + assert ( + get_model_from_request( + request_data={"video_id": "not-a-managed-resource-id"}, + route="/v1/videos/{video_id}", + ) + is None + ) + + +def test_abbreviate_api_key(): + assert abbreviate_api_key("sk-test-1234") == "sk-...1234" + + def test_get_customer_user_header_returns_none_when_no_customer_role(): from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping @@ -964,3 +1166,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/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 08f4bd0ebf..dd24ac8749 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1,8 +1,7 @@ -import asyncio import json import os import sys -from typing import Tuple +from types import SimpleNamespace from unittest.mock import ANY, AsyncMock, MagicMock, patch sys.path.insert( @@ -15,6 +14,8 @@ import litellm.proxy.proxy_server from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( LiteLLM_JWTAuth, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -23,8 +24,10 @@ from litellm.proxy._types import ( JWTRoutingOverride, ) from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( + _reserve_budget_after_common_checks, _run_centralized_common_checks, _run_post_custom_auth_checks, get_api_key, @@ -32,6 +35,13 @@ from litellm.proxy.auth.user_api_key_auth import ( ) +class _RoutingRequest: + def __init__(self, headers=None, query_params=None): + self.headers = headers or {} + self.query_params = query_params or {} + self.state = SimpleNamespace() + + def test_get_api_key(): bearer_token = "Bearer sk-12345678" api_key = "sk-12345678" @@ -49,6 +59,74 @@ def test_get_api_key(): ) == (api_key, passed_in_key) +@pytest.mark.asyncio +async def test_should_clear_stale_budget_reservation_when_budget_checks_skip(): + user_api_key_auth_obj = UserAPIKeyAuth( + token="test_token", + budget_reservation={ + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_token"}], + }, + ) + + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "free-model"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=True, + ) + + assert user_api_key_auth_obj.budget_reservation is None + + +@pytest.mark.asyncio +async def test_should_not_reuse_cached_key_object_for_request_state(): + key_cache = DualCache() + cached_key = UserAPIKeyAuth( + token="cached-token", + request_route="/old-route", + budget_reservation={ + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:cached-token"}], + }, + ) + + await _cache_key_object( + hashed_token="cached-token", + user_api_key_obj=cached_key, + user_api_key_cache=key_cache, + proxy_logging_obj=None, + ) + + first_request_key = await get_key_object( + hashed_token="cached-token", + prisma_client=MagicMock(), + user_api_key_cache=key_cache, + ) + first_request_key.budget_reservation = { + "reserved_cost": 0.9, + "entries": [{"counter_key": "spend:key:cached-token"}], + } + first_request_key.request_route = "/chat/completions" + + second_request_key = await get_key_object( + hashed_token="cached-token", + prisma_client=MagicMock(), + user_api_key_cache=key_cache, + ) + + assert first_request_key is not cached_key + assert second_request_key is not first_request_key + assert second_request_key.budget_reservation is None + assert second_request_key.request_route is None + + @pytest.mark.asyncio async def test_custom_auth_does_not_enforce_key_model_access_by_default(): valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) @@ -107,6 +185,39 @@ async def test_custom_auth_honors_key_level_model_access_restriction_allowed_wit ) +@pytest.mark.asyncio +async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in(): + valid_token = UserAPIKeyAuth(token="test_token", models=["allowed-model"]) + request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"}) + + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.can_key_call_model", + new_callable=AsyncMock, + ) as mock_can_key, + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"custom_auth_run_common_checks": True}, + ), + ): + await _run_post_custom_auth_checks( + valid_token=valid_token, + request=request, + request_data={}, + route="/v1/files", + parent_otel_span=None, + ) + mock_can_key.assert_awaited_once_with( + model="restricted-model", + llm_model_list=ANY, + valid_token=valid_token, + llm_router=ANY, + ) + + @pytest.mark.asyncio async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in(): valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) @@ -1752,7 +1863,11 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): from starlette.datastructures import URL from starlette.requests import Request - from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy._types import ( + LiteLLM_TeamTableCachedObj, + LitellmUserRoles, + UserAPIKeyAuth, + ) from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder api_key = "sk-test-team-metadata-refresh" @@ -1833,16 +1948,17 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): request_data={}, ) - assert result.team_metadata == {"guardrails": ["test-guardrail-333"]}, ( - f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" - ) + assert result.team_metadata == { + "guardrails": ["test-guardrail-333"] + }, f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" finally: for k, v in _originals.items(): setattr(_proxy_server_mod, k, v) - + + # --------------------------------------------------------------------------- - + # _run_centralized_common_checks — centralized authz gate # --------------------------------------------------------------------------- @@ -1859,7 +1975,7 @@ def _proxy_attrs_for_centralized_checks( """ return { "prisma_client": None, - "user_api_key_cache": MagicMock(), + "user_api_key_cache": DualCache(), "proxy_logging_obj": MagicMock(), "general_settings": ({"custom_auth_run_common_checks": True} if flag else {}), "llm_router": None, @@ -2120,6 +2236,81 @@ async def test_centralized_common_checks_propagates_end_user_budget_error(): setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_centralized_common_checks_reserves_request_end_user_budget(): + """Regression: reservation runs before user_api_key_auth() copies the + request end-user onto the token, so centralized checks must pass the + locally extracted end_user_id/end_user_object into reservation.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u") + request = Request(scope={"type": "http", "headers": []}) + request._url = URL(url="/chat/completions") + request_data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "user": "alice", + } + end_user_object = LiteLLM_EndUserTable( + user_id="alice", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + counter_cache = DualCache() + attrs["spend_counter_cache"] = counter_cache + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_end_user_object", + new_callable=AsyncMock, + return_value=end_user_object, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ), + ): + assert token.end_user_id is None + + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data=request_data, + route="/chat/completions", + ) + + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert token.end_user_id is None + assert token.budget_reservation is not None + assert token.budget_reservation["entries"] == [ + { + "counter_key": "spend:end_user:alice", + "entity_type": "EndUser", + "entity_id": "alice", + "reserved_cost": 0.6, + "applied_adjustment": 0.0, + } + ] + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:alice" + ) == pytest.approx(0.6) + + @pytest.mark.asyncio async def test_centralized_common_checks_short_circuits_when_master_key_unset(): """master_key=None is no-auth dev mode — admin-only routes and 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..d59682c2d5 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,373 @@ 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 + + result = await health_endpoint( + response=Response(), user_api_key_dict=user_api_key_dict + ) + + # 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 + ) + non_admin_result = await health_endpoint( + response=non_admin_response, user_api_key_dict=non_admin_key + ) + 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 + ) + + 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"]) + + +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_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 65e7f744c8..8b5835139b 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,9 +1,7 @@ -import json import os import sys import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -13,8 +11,11 @@ from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger -from litellm.types.utils import StandardLoggingPayload +from litellm.proxy.hooks.proxy_track_cost_callback import ( + _ProxyDBLogger, + _get_budget_reservation_from_metadata, + _update_database_and_spend_counters, +) @pytest.mark.asyncio @@ -62,7 +63,6 @@ async def test_async_post_call_failure_hook(): # Check the arguments passed to update_database call_args = mock_update_database.call_args[1] - print("call_args", json.dumps(call_args, indent=4, default=str)) assert call_args["token"] == "test_api_key" assert call_args["response_cost"] == 0.0 assert call_args["user_id"] == "test_user_id" @@ -128,6 +128,440 @@ async def test_async_post_call_failure_hook_non_llm_route(): mock_update_database.assert_not_called() +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_releases_budget_reservation_before_route_skip(): + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + request_route="/custom/route", + budget_reservation=budget_reservation, + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + ): + await logger.async_post_call_failure_hook( + request_data={}, + original_exception=Exception("Test exception"), + user_api_key_dict=user_api_key_dict, + ) + + assert mock_release_budget_reservation.await_count == 1 + assert ( + mock_release_budget_reservation.await_args.kwargs["budget_reservation"] + is user_api_key_dict.budget_reservation + ) + mock_update_database.assert_not_called() + + +@pytest.mark.asyncio +async def test_should_continue_failure_tracking_when_budget_release_fails(): + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + user_id="test_user_id", + team_id="test_team_id", + request_route="/chat/completions", + budget_reservation=budget_reservation, + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + side_effect=RuntimeError("redis unavailable"), + ) as mock_release_budget_reservation, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback._invalidate_budget_reservation_counters", + new_callable=AsyncMock, + ) as mock_invalidate_budget_reservation_counters, + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception", + ) as mock_log_exception, + ): + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + }, + original_exception=Exception("provider failed"), + user_api_key_dict=user_api_key_dict, + ) + + assert mock_release_budget_reservation.await_count == 1 + assert ( + mock_release_budget_reservation.await_args.kwargs["budget_reservation"] + is user_api_key_dict.budget_reservation + ) + assert mock_invalidate_budget_reservation_counters.await_count == 1 + assert ( + mock_invalidate_budget_reservation_counters.await_args.kwargs[ + "budget_reservation" + ] + is user_api_key_dict.budget_reservation + ) + assert user_api_key_dict.budget_reservation["finalized"] is True + mock_log_exception.assert_called_once() + mock_update_database.assert_called_once() + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_budget_reservation_when_spend_tracking_skips(): + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation) + + kwargs = { + "model": "gpt-4", + "litellm_params": { + "metadata": { + "user_api_key_auth": user_api_key_auth, + }, + }, + "standard_logging_object": { + "response_cost": 0.1, + "request_tags": None, + }, + "stream": False, + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation: + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_release_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_budget_reservation_when_response_cost_missing(): + logger = _ProxyDBLogger() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + user_api_key_auth = UserAPIKeyAuth(budget_reservation=budget_reservation) + + kwargs = { + "model": "gpt-4", + "call_type": "acompletion", + "litellm_params": { + "metadata": { + "user_api_key_auth": user_api_key_auth, + }, + }, + "standard_logging_object": { + "response_cost": None, + "response_cost_failure_debug_info": "missing custom price", + "request_tags": None, + }, + "stream": False, + } + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging, + patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_release_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + + +def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): + budget_reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_api_key"}], + } + + assert ( + _get_budget_reservation_from_metadata( + metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} + ) + is None + ) + assert ( + _get_budget_reservation_from_metadata( + metadata={ + "user_api_key_auth": UserAPIKeyAuth( + budget_reservation=budget_reservation + ) + } + ) + == budget_reservation + ) + assert ( + _get_budget_reservation_from_metadata( + metadata={ + "user_api_key_auth": dict( + UserAPIKeyAuth(budget_reservation=budget_reservation) + ) + } + ) + == budget_reservation + ) + assert ( + _get_budget_reservation_from_metadata( + metadata={"user_api_key_budget_reservation": budget_reservation} + ) + is budget_reservation + ) + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + side_effect=Exception("db unavailable") + ) + increment_spend_counters = AsyncMock() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation: + with pytest.raises(Exception, match="db unavailable"): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + mock_release_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + + increment_spend_counters.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): + proxy_logging_obj = MagicMock() + db_exception = RuntimeError("db unavailable") + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( + side_effect=db_exception + ) + increment_spend_counters = AsyncMock() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + side_effect=RuntimeError("release unavailable"), + ) as mock_release_budget_reservation, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception", + ) as mock_log_exception, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback._invalidate_budget_reservation_counters", + new_callable=AsyncMock, + side_effect=RuntimeError("invalidate unavailable"), + ) as mock_invalidate_budget_reservation_counters, + ): + with pytest.raises(RuntimeError) as exc_info: + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + assert exc_info.value is db_exception + mock_release_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + mock_invalidate_budget_reservation_counters.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + assert mock_log_exception.call_count == 2 + mock_log_exception.assert_any_call( + "Failed to release budget reservation after database update failed" + ) + mock_log_exception.assert_any_call( + "Failed to invalidate budget reservation counters after release failed" + ) + + increment_spend_counters.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_updates_counters_after_db_update(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + increment_spend_counters = AsyncMock() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id="test_end_user_id", + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + request_tags=["tag-a"], + ) + + proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() + increment_spend_counters.assert_awaited_once_with( + token="test_api_key", + team_id="test_team_id", + user_id="test_user_id", + response_cost=0.2, + org_id="test_org_id", + budget_reservation=budget_reservation, + end_user_id="test_end_user_id", + tags=["tag-a"], + ) + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_invalidates_reservation_when_counter_update_fails(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + increment_spend_counters = AsyncMock(side_effect=Exception("counter unavailable")) + budget_reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_api_key"}], + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", + new_callable=AsyncMock, + ) as mock_invalidate_budget_reservation_counters: + with pytest.raises(Exception, match="counter unavailable"): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + mock_invalidate_budget_reservation_counters.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + assert budget_reservation["finalized"] is True + + proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_preserves_counter_exception_when_invalidation_fails(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + counter_exception = RuntimeError("counter unavailable") + increment_spend_counters = AsyncMock(side_effect=counter_exception) + budget_reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_api_key"}], + } + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", + new_callable=AsyncMock, + side_effect=RuntimeError("invalidate unavailable"), + ) as mock_invalidate_budget_reservation_counters, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.exception", + ) as mock_log_exception, + ): + with pytest.raises(RuntimeError) as exc_info: + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + assert exc_info.value is counter_exception + mock_invalidate_budget_reservation_counters.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + mock_log_exception.assert_called_once_with( + "Failed to invalidate budget reservation counters after spend counter update failed" + ) + assert budget_reservation["finalized"] is True + + proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() + + @pytest.mark.asyncio async def test_track_cost_callback_skips_when_no_standard_logging_object(): """ @@ -344,7 +778,7 @@ async def test_enrich_failure_metadata_skips_when_no_api_key(): "user_api_key_team_id": None, "user_api_key_team_alias": None, } - result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) mock_get_key.assert_not_called() 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_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py new file mode 100644 index 0000000000..070b232066 --- /dev/null +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -0,0 +1,1495 @@ +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + LiteLLM_OrganizationTable, + LiteLLM_TagTable, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.spend_tracking.budget_reservation import ( + estimate_request_max_cost, + get_budget_window_start, + invalidate_budget_reservation_counters, + release_budget_reservation, + reserve_budget_for_request, +) +from litellm.proxy.utils import ProxyLogging + + +@pytest.fixture() +def spend_counter_state(): + import litellm.proxy.proxy_server as ps + + original_counter_cache = ps.spend_counter_cache + original_key_cache = ps.user_api_key_cache + original_prisma_client = ps.prisma_client + + counter_cache = DualCache() + key_cache = DualCache() + ps.spend_counter_cache = counter_cache + ps.user_api_key_cache = key_cache + ps.prisma_client = None + + try: + yield counter_cache, key_cache + finally: + ps.spend_counter_cache = original_counter_cache + ps.user_api_key_cache = original_key_cache + ps.prisma_client = original_prisma_client + + +def _request_body() -> dict: + return { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + } + + +def test_should_not_serialize_budget_reservation_on_user_api_key_auth(): + auth = UserAPIKeyAuth( + token="key-budget-runtime-state", + budget_reservation={ + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:key-budget-runtime-state"}], + }, + ) + + assert "budget_reservation" not in auth.model_dump() + assert "budget_reservation" not in auth.model_dump(exclude_none=True) + assert "budget_reservation" not in auth.model_dump_json() + + +@pytest.mark.asyncio +async def test_should_shrink_second_key_reservation_to_remaining_budget( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-race", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert reservation is not None + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-race") + == 0.6 + ) + + second_reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert second_reservation is not None + assert second_reservation["reserved_cost"] == pytest.approx(0.4) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-race" + ) == pytest.approx(1.0) + + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-race" + ) == pytest.approx(1.0) + + await release_budget_reservation(second_reservation) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-race" + ) == pytest.approx(0.6) + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_shrink_second_end_user_reservation_to_remaining_budget( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-end-user", + end_user_id="end-user-budget-race", + ) + end_user_object = LiteLLM_EndUserTable( + user_id="end-user-budget-race", + blocked=False, + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_object=end_user_object, + ) + assert reservation is not None + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:end-user-budget-race" + ) == pytest.approx(0.6) + + second_reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_object=end_user_object, + ) + assert second_reservation is not None + assert second_reservation["reserved_cost"] == pytest.approx(0.4) + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:end-user-budget-race" + ) == pytest.approx(1.0) + + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + end_user_object=end_user_object, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:end-user-budget-race" + ) == pytest.approx(1.0) + + await release_budget_reservation(second_reservation) + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:end-user-budget-race" + ) == pytest.approx(0.6) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + end_user_id="end-user-budget-race", + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:end-user-budget-race" + ) == pytest.approx(0.2) + + +@pytest.mark.asyncio +async def test_should_shrink_second_tag_reservation_to_remaining_budget( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-budget-tag") + request_body = _request_body() + request_body["metadata"] = { + "tags": ["tag-budget-race", "tag-without-budget", "tag-budget-race"] + } + await key_cache.async_set_cache( + key="tag:tag-budget-race", + value=LiteLLM_TagTable( + tag_name="tag-budget-race", + spend=0.0, + budget_id="tag-budget-id", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="tag:tag-without-budget", + value=LiteLLM_TagTable( + tag_name="tag-without-budget", + spend=0.0, + ).model_dump(), + ) + prisma_client = MagicMock() + prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=prisma_client, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert reservation is not None + assert reservation["entries"] == [ + { + "counter_key": "spend:tag:tag-budget-race", + "entity_type": "Tag", + "entity_id": "tag-budget-race", + "reserved_cost": 0.6, + "applied_adjustment": 0.0, + } + ] + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:tag-budget-race" + ) == pytest.approx(0.6) + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:tag:tag-without-budget") + is None + ) + + second_reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=prisma_client, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert second_reservation is not None + assert second_reservation["reserved_cost"] == pytest.approx(0.4) + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:tag-budget-race" + ) == pytest.approx(1.0) + + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=prisma_client, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:tag-budget-race" + ) == pytest.approx(1.0) + + await release_budget_reservation(second_reservation) + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:tag-budget-race" + ) == pytest.approx(0.6) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + tags=["tag-budget-race"], + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:tag-budget-race" + ) == pytest.approx(0.2) + + +@pytest.mark.asyncio +async def test_should_seed_and_update_end_user_and_tag_counters_without_reservation( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + await key_cache.async_set_cache( + key="end_user_id:customer-1", + value=LiteLLM_EndUserTable( + user_id="customer-1", + blocked=False, + spend=4.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=10.0), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="tag:paid-tag", + value=LiteLLM_TagTable( + tag_name="paid-tag", + spend=7.0, + ).model_dump(), + ) + await key_cache.async_set_cache( + key="tag:other-tag", + value=LiteLLM_TagTable( + tag_name="other-tag", + spend=2.0, + ).model_dump(), + ) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.50, + end_user_id="customer-1", + tags=["paid-tag", "paid-tag", "other-tag", ""], + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:end_user:customer-1" + ) == pytest.approx(4.50) + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:paid-tag" + ) == pytest.approx(7.50) + assert counter_cache.in_memory_cache.get_cache( + key="spend:tag:other-tag" + ) == pytest.approx(2.50) + + +@pytest.mark.asyncio +async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_state): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-shared", + spend=0.0, + max_budget=1.0, + user_id="user-budget-shared", + team_id="team-budget-shared", + org_id="org-budget-shared", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-shared", + spend=0.0, + max_budget=1.0, + ) + user_object = LiteLLM_UserTable( + user_id="user-budget-shared", + spend=0.0, + ) + await key_cache.async_set_cache( + key="team_membership:user-budget-shared:team-budget-shared", + value=LiteLLM_TeamMembership( + user_id="user-budget-shared", + team_id="team-budget-shared", + spend=0.1, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ).model_dump(), + ) + await key_cache.async_set_cache( + key="org_id:org-budget-shared:with_budget", + value=LiteLLM_OrganizationTable( + organization_id="org-budget-shared", + organization_alias="shared-org", + budget_id="org-budget-id", + spend=0.1, + models=[], + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ).model_dump(), + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:team_member:user-budget-shared:team-budget-shared" + ) == pytest.approx(0.4) + assert counter_cache.in_memory_cache.get_cache( + key="spend:org:org-budget-shared" + ) == pytest.approx(0.4) + + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state): + counter_cache, key_cache = spend_counter_state + await key_cache.async_set_cache( + key="org_id:org-counter-with-budget:with_budget", + value=LiteLLM_OrganizationTable( + organization_id="org-counter-with-budget", + organization_alias="shared-org", + budget_id="org-budget-id", + spend=2.0, + models=[], + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=10.0), + ).model_dump(), + ) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + org_id="org-counter-with-budget", + response_cost=0.25, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:org:org-counter-with-budget" + ) == pytest.approx(2.25) + + +@pytest.mark.asyncio +async def test_should_seed_org_counter_from_plain_org_cache(spend_counter_state): + counter_cache, key_cache = spend_counter_state + await key_cache.async_set_cache( + key="org_id:org-counter-plain", + value=LiteLLM_OrganizationTable( + organization_id="org-counter-plain", + organization_alias="shared-org", + budget_id="org-budget-id", + spend=2.0, + models=[], + created_by="test", + updated_by="test", + ).model_dump(), + ) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + org_id="org-counter-plain", + response_cost=0.25, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:org:org-counter-plain" + ) == pytest.approx(2.25) + + +@pytest.mark.asyncio +async def test_should_cap_known_estimate_to_remaining_budget( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-known-estimate-cap", + spend=0.9, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-known-estimate-cap", + value=0.9, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.1) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-known-estimate-cap" + ) == pytest.approx(1.0) + + await release_budget_reservation(reservation) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-known-estimate-cap" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_should_reserve_remaining_budget_when_output_cap_missing( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-uncapped", + spend=0.2, + max_budget=1.0, + ) + await key_cache.async_set_cache( + key="key-budget-uncapped", + value=valid_token, + ) + request_body = _request_body() + request_body.pop("max_tokens") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info", + return_value={ + "input_cost_per_token": 0.0, + "output_cost_per_token": 100.0, + "max_output_tokens": 200000, + }, + ): + assert ( + estimate_request_max_cost( + request_body=request_body, + route="/chat/completions", + llm_router=None, + ) + is None + ) + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.8) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-uncapped" + ) == pytest.approx(1.0) + + await release_budget_reservation(reservation) + + +@pytest.mark.asyncio +async def test_should_shrink_uncapped_reservation_when_counter_advances( + spend_counter_state, + monkeypatch, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-uncapped-race", + spend=0.2, + max_budget=1.0, + ) + request_body = _request_body() + request_body.pop("max_tokens") + + from litellm.proxy.spend_tracking import budget_reservation + + async def stale_counter_read(counter): + await counter_cache.async_increment_cache( + key=counter.counter_key, + value=0.3, + ) + return 0.2 + + monkeypatch.setattr( + budget_reservation, + "_get_current_counter_value", + stale_counter_read, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=None, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.7) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-uncapped-race" + ) == pytest.approx(1.0) + + await release_budget_reservation(reservation) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-uncapped-race" + ) == pytest.approx(0.3) + + +@pytest.mark.asyncio +async def test_should_shrink_uncapped_reservation_multiple_times( + spend_counter_state, + monkeypatch, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-double-resize", + spend=0.2, + max_budget=1.0, + team_id="team-budget-double-resize", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-double-resize", + spend=0.2, + max_budget=1.0, + ) + request_body = _request_body() + request_body.pop("max_tokens") + + from litellm.proxy.spend_tracking import budget_reservation + + stale_spend_by_counter_key = { + "spend:key:key-budget-double-resize": 0.3, + "spend:team:team-budget-double-resize": 0.4, + } + + async def stale_counter_read(counter): + await counter_cache.async_increment_cache( + key=counter.counter_key, + value=stale_spend_by_counter_key[counter.counter_key], + ) + return 0.2 + + monkeypatch.setattr( + budget_reservation, + "_get_current_counter_value", + stale_counter_read, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=None, + ): + reservation = await reserve_budget_for_request( + request_body=request_body, + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.6) + assert [entry["reserved_cost"] for entry in reservation["entries"]] == [ + pytest.approx(0.6), + pytest.approx(0.6), + ] + assert [entry["applied_adjustment"] for entry in reservation["entries"]] == [ + pytest.approx(0.0), + pytest.approx(0.0), + ] + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-double-resize" + ) == pytest.approx(0.9) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-double-resize" + ) == pytest.approx(1.0) + + await release_budget_reservation(reservation) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-double-resize" + ) == pytest.approx(0.3) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-double-resize" + ) == pytest.approx(0.4) + + +def test_should_start_window_without_reset_at_at_duration_boundary(): + before = datetime.now(timezone.utc) - timedelta(hours=1) + + window_start = get_budget_window_start({"budget_duration": "1h"}) + + after = datetime.now(timezone.utc) - timedelta(hours=1) + assert window_start is not None + assert before <= window_start <= after + + +@pytest.mark.asyncio +async def test_should_skip_budget_window_with_unparseable_duration( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-malformed-window", + spend=0.9, + max_budget=10.0, + budget_limits=[ + { + "budget_duration": "not-a-duration", + "max_budget": 1.0, + } + ], + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-malformed-window", + value=0.9, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.2, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert [entry["counter_key"] for entry in reservation["entries"]] == [ + "spend:key:key-budget-malformed-window" + ] + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-malformed-window" + ) == pytest.approx(1.1) + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-malformed-window:window:not-a-duration" + ) + is None + ) + + await release_budget_reservation(reservation) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-malformed-window" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_should_skip_window_reservation_when_db_baseline_unavailable( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-window-db-unavailable", + budget_limits=[ + { + "budget_duration": "1h", + "max_budget": 1.0, + } + ], + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is None + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-window-db-unavailable:window:1h" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_should_skip_reservation_when_counter_increment_fails( + spend_counter_state, + monkeypatch, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reserve-unavailable", + spend=0.0, + max_budget=1.0, + ) + + async def fail_increment_cache(*args, **kwargs): + raise RuntimeError("counter unavailable") + + monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.warning" + ) as mock_warning, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is None + assert mock_warning.call_count >= 1 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reserve-unavailable" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_should_skip_reservation_when_counter_initialization_fails( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reserve-init-unavailable", + spend=0.0, + max_budget=1.0, + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ), + patch( + "litellm.proxy.proxy_server._ensure_spend_counter_initialized", + side_effect=RuntimeError("redis unavailable"), + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.warning" + ) as mock_warning, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is None + assert mock_warning.call_count >= 1 + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reserve-init-unavailable" + ) + is None + ) + + +@pytest.mark.asyncio +async def test_should_release_tracked_entry_when_reservation_fails_after_increment( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reserve-after-increment-failure", + spend=0.0, + max_budget=1.0, + ) + + import litellm.proxy.proxy_server as ps + + original_increment_counter = ps._increment_spend_counter_cache + first_increment = True + + async def fail_after_increment(counter_key: str, increment: float): + nonlocal first_increment + if first_increment: + first_increment = False + await counter_cache.async_increment_cache(key=counter_key, value=increment) + raise RuntimeError("lost increment response") + return await original_increment_counter( + counter_key=counter_key, + increment=increment, + ) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ), + patch( + "litellm.proxy.proxy_server._increment_spend_counter_cache", + side_effect=fail_after_increment, + ), + patch( + "litellm.proxy.proxy_server._invalidate_spend_counter", + side_effect=RuntimeError("invalidate unavailable"), + ), + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is None + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reserve-after-increment-failure" + ) == pytest.approx(0.0) + + +@pytest.mark.asyncio +async def test_should_not_re_read_uncapped_budget_after_reservation_fallback( + spend_counter_state, + monkeypatch, +): + _, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-uncapped-read-once", + spend=0.2, + max_budget=1.0, + ) + + from litellm.proxy.spend_tracking import budget_reservation + + current_counter_reads = [] + + async def mock_get_current_counter_value(counter): + current_counter_reads.append(counter.counter_key) + return counter.fallback_spend + + async def mock_reserve_counter(counter, reservation_cost): + return None + + monkeypatch.setattr( + budget_reservation, + "_get_current_counter_value", + mock_get_current_counter_value, + ) + monkeypatch.setattr( + budget_reservation, + "_reserve_counter", + mock_reserve_counter, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=None, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.8) + assert current_counter_reads == ["spend:key:key-budget-uncapped-read-once"] + + +@pytest.mark.asyncio +async def test_should_reconcile_reserved_counter_to_actual_spend( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-reconcile", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token="key-budget-reconcile", + team_id="team-without-budget", + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + ) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-reconcile" + ) == pytest.approx(0.2) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-without-budget" + ) == pytest.approx(0.2) + + +@pytest.mark.asyncio +async def test_should_release_reservation_on_failure(spend_counter_state): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-release", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.4, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + await release_budget_reservation(reservation) + await release_budget_reservation(reservation) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-release" + ) == pytest.approx(0.0) + + +@pytest.mark.asyncio +async def test_should_retry_partial_release_without_double_decrement( + spend_counter_state, + monkeypatch, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-partial-release", + spend=0.0, + max_budget=1.0, + team_id="team-budget-partial-release", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-partial-release", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.4, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + original_increment_cache = counter_cache.async_increment_cache + fail_next_team_release = True + + async def flaky_increment_cache(key, value, *args, **kwargs): + nonlocal fail_next_team_release + if ( + key == "spend:team:team-budget-partial-release" + and value < 0 + and fail_next_team_release + ): + fail_next_team_release = False + raise RuntimeError("simulated counter failure") + return await original_increment_cache(key=key, value=value, *args, **kwargs) + + monkeypatch.setattr(counter_cache, "async_increment_cache", flaky_increment_cache) + + with pytest.raises(RuntimeError): + await release_budget_reservation(reservation) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-partial-release" + ) == pytest.approx(0.0) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-partial-release" + ) == pytest.approx(0.4) + + await release_budget_reservation(reservation) + + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-partial-release" + ) == pytest.approx(0.0) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-partial-release" + ) == pytest.approx(0.0) + + +@pytest.mark.asyncio +async def test_should_preserve_budget_error_and_continue_partial_cleanup( + spend_counter_state, + monkeypatch, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-cleanup-failure", + spend=0.0, + max_budget=1.0, + team_id="team-budget-cleanup-failure", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-cleanup-failure", + spend=0.3, + max_budget=0.3, + ) + await key_cache.async_set_cache( + key="team_id:team-budget-cleanup-failure", + value=team_object, + ) + + original_increment_cache = counter_cache.async_increment_cache + fail_key_cleanup = True + + async def flaky_increment_cache(key, value, *args, **kwargs): + nonlocal fail_key_cleanup + if key == "spend:key:key-budget-cleanup-failure" and value < 0: + if fail_key_cleanup: + fail_key_cleanup = False + raise RuntimeError("simulated cleanup failure") + return await original_increment_cache(key=key, value=value, *args, **kwargs) + + monkeypatch.setattr(counter_cache, "async_increment_cache", flaky_increment_cache) + + with ( + patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.4, + ), + patch( + "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.exception" + ) as mock_log_exception, + ): + with pytest.raises(litellm.BudgetExceededError): + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-cleanup-failure" + ) + is None + ) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-cleanup-failure" + ) == pytest.approx(0.3) + mock_log_exception.assert_called() + + +@pytest.mark.asyncio +async def test_should_not_create_negative_counter_when_release_counter_is_missing( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + reservation = { + "reserved_cost": 0.4, + "entries": [ + { + "counter_key": "spend:key:key-budget-missing-release", + "reserved_cost": 0.4, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with pytest.raises(RuntimeError, match="missing counter"): + await release_budget_reservation(reservation) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-missing-release" + ) + is None + ) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_should_invalidate_counter_when_release_would_underflow( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + await counter_cache.async_increment_cache( + key="spend:key:key-budget-underflow-release", + value=0.1, + ) + reservation = { + "reserved_cost": 0.4, + "entries": [ + { + "counter_key": "spend:key:key-budget-underflow-release", + "reserved_cost": 0.4, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with pytest.raises(RuntimeError, match="negative"): + await release_budget_reservation(reservation) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-underflow-release" + ) + is None + ) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_should_invalidate_non_numeric_counter_during_release( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-nonnumeric-release", + value="stale", + ) + reservation = { + "reserved_cost": 0.4, + "entries": [ + { + "counter_key": "spend:key:key-budget-nonnumeric-release", + "reserved_cost": 0.4, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + with pytest.raises(RuntimeError, match="non-numeric"): + await release_budget_reservation(reservation) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-nonnumeric-release" + ) + is None + ) + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_should_invalidate_reserved_counters_after_persisted_spend_failure( + spend_counter_state, +): + counter_cache, _ = spend_counter_state + await counter_cache.async_increment_cache( + key="spend:key:key-budget-invalidate", + value=0.4, + ) + await counter_cache.async_increment_cache( + key="spend:team:team-budget-invalidate", + value=0.4, + ) + + await invalidate_budget_reservation_counters( + { + "reserved_cost": 0.4, + "entries": [ + {"counter_key": "spend:key:key-budget-invalidate"}, + {"counter_key": "spend:team:team-budget-invalidate"}, + ], + } + ) + + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-invalidate") + is None + ) + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:team:team-budget-invalidate") + is None + ) + + +@pytest.mark.asyncio +async def test_should_reserve_all_budgeted_counters(spend_counter_state): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-all", + spend=0.0, + max_budget=1.0, + team_id="team-budget-all", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-all", + spend=0.0, + max_budget=1.0, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:key:key-budget-all") == 0.3 + ) + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:team:team-budget-all") == 0.3 + ) + + await release_budget_reservation(reservation) 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 33f8ded84c..64cb931888 100644 --- a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py +++ b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py @@ -1,16 +1,118 @@ -from types import SimpleNamespace +import sys +from types import ModuleType, SimpleNamespace -from litellm.proxy._lazy_openapi_snapshot import _stable_generate_unique_id +from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids -def test_stable_generate_unique_id_sorts_route_methods(): - route = SimpleNamespace( - name="langfuse_proxy_route", - path_format="/langfuse/{endpoint}", - methods={"POST", "GET", "DELETE", "PATCH", "PUT"}, +def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch): + from litellm.proxy import _lazy_openapi_snapshot + + route_a = SimpleNamespace(path="/feature-a/items") + route_b = SimpleNamespace(path="/feature-b/items") + fake_app = SimpleNamespace( + title="LiteLLM test", + version="0.0.0", + routes=[route_a, route_b], ) + fake_feature_a_module = ModuleType("fake_feature_a") + fake_feature_b_module = ModuleType("fake_feature_b") + monkeypatch.setitem(sys.modules, "fake_feature_a", fake_feature_a_module) + monkeypatch.setitem(sys.modules, "fake_feature_b", fake_feature_b_module) + + fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features") + fake_lazy_features_module.LAZY_FEATURES = [ + SimpleNamespace( + name="feature-a", + module_path="fake_feature_a", + path_prefixes=("/feature-a",), + register_fn=lambda app, module: None, + ), + SimpleNamespace( + name="feature-b", + module_path="fake_feature_b", + path_prefixes=("/feature-b",), + register_fn=lambda app, module: None, + ), + ] + monkeypatch.setitem( + sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module + ) + + def fake_get_openapi(title, version, routes): + path = routes[0].path + return { + "paths": {path: {"get": {"operationId": "shared_operation_id_get"}}}, + "components": {"schemas": {"Example": {"type": "object"}}}, + } + + def fake_ensure_unique_openapi_operation_ids(schema, reserved_operation_ids): + for path_item in schema["paths"].values(): + operation = path_item["get"] + operation_id = operation["operationId"] + if operation_id in reserved_operation_ids: + operation_id = f"{operation_id}_2" + operation["operationId"] = operation_id + reserved_operation_ids.add(operation_id) + return schema + + fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server") + fake_proxy_server_module.app = fake_app + fake_proxy_server_module.ensure_unique_openapi_operation_ids = ( + fake_ensure_unique_openapi_operation_ids + ) + monkeypatch.setitem( + sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module + ) + monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi) + + fragments = _lazy_openapi_snapshot.generate_snapshot() + assert ( - _stable_generate_unique_id(route) - == "langfuse_proxy_route_langfuse__endpoint__delete" + fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"] + == "shared_operation_id_get" ) + assert ( + fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"] + == "shared_operation_id_get_2" + ) + assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["tags"] == [ + "feature-a" + ] + assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["tags"] == [ + "feature-b" + ] + + +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"}, + } + } + + _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..3f19db36c3 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5,7 +5,7 @@ import os import socket import subprocess import sys -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from unittest import mock from unittest.mock import AsyncMock, MagicMock, mock_open, patch @@ -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 +6426,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/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 6dd0e0e68a..e67a04c749 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -156,8 +156,11 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry(): vector_store_id="test_store_id" ) - # Test with no vector store registry - with patch.object(litellm, "vector_store_registry", None): + # Test with no vector store registry or DB fallback + with ( + patch.object(litellm, "vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): original_data = {"existing_key": "existing_value"} result = await _update_request_data_with_litellm_managed_vector_store_registry( data=original_data, vector_store_id=vector_store_id diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py new file mode 100644 index 0000000000..48262afd36 --- /dev/null +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py @@ -0,0 +1,540 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException, Request, Response + +import litellm +from litellm.proxy._types import LiteLLM_ManagedVectorStoresTable, UserAPIKeyAuth + + +def _mock_request() -> MagicMock: + request = MagicMock(spec=Request) + request.headers = {} + request.method = "POST" + request.query_params = {} + request.url.path = "/v1/vector_stores/vs_path/search" + return request + + +@pytest.mark.asyncio +async def test_vector_store_search_forces_path_id_over_body_id(): + from litellm.proxy.vector_store_endpoints.endpoints import vector_store_search + + captured_data = {} + + async def fake_base_process(self, **kwargs): + captured_data.update(self.data) + return {"ok": True} + + request = _mock_request() + with ( + patch( + "litellm.proxy.proxy_server._read_request_body", + new=AsyncMock( + return_value={ + "vector_store_id": "vs_body_victim", + "query": "test", + } + ), + ), + patch.object(litellm, "vector_store_registry", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch( + "litellm.proxy.vector_store_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=fake_base_process, + ), + ): + response = await vector_store_search( + request=request, + vector_store_id="vs_path_allowed", + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert response == {"ok": True} + assert captured_data["vector_store_id"] == "vs_path_allowed" + + +@pytest.mark.asyncio +async def test_vector_store_file_create_forces_path_id_over_body_id(): + from litellm.proxy.vector_store_files_endpoints.endpoints import ( + vector_store_file_create, + ) + + captured_data = {} + + async def fake_base_process(self, **kwargs): + captured_data.update(self.data) + return {"ok": True} + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = { + "vector_store_id": "vs_path_allowed", + "custom_llm_provider": "openai", + "team_id": "team-a", + } + + request = _mock_request() + with ( + patch( + "litellm.proxy.proxy_server._read_request_body", + new=AsyncMock( + return_value={ + "vector_store_id": "vs_body_victim", + "file_id": "file_123", + } + ), + ), + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.proxy.vector_store_files_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=fake_base_process, + ), + ): + response = await vector_store_file_create( + vector_store_id="vs_path_allowed", + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert response == {"ok": True} + assert captured_data["vector_store_id"] == "vs_path_allowed" + assert captured_data["custom_llm_provider"] == "openai" + mock_registry.get_litellm_managed_vector_store_from_registry.assert_called_once_with( + vector_store_id="vs_path_allowed" + ) + + +@pytest.mark.asyncio +async def test_vector_store_file_create_denies_other_team_path_store(): + from litellm.proxy.vector_store_files_endpoints.endpoints import ( + vector_store_file_create, + ) + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "openai", + "team_id": "team-b", + } + + request = _mock_request() + with ( + patch( + "litellm.proxy.proxy_server._read_request_body", + new=AsyncMock(return_value={"file_id": "file_123"}), + ), + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.proxy.vector_store_files_endpoints.endpoints.ProxyBaseLLMRequestProcessing.base_process_llm_request", + new=AsyncMock(), + ) as mock_base_process, + ): + with pytest.raises(HTTPException) as exc_info: + await vector_store_file_create( + vector_store_id="vs_other_team", + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert exc_info.value.status_code == 403 + mock_base_process.assert_not_called() + + +@pytest.mark.asyncio +async def test_rag_query_denies_nested_other_team_vector_store(): + from litellm.proxy.rag_endpoints.endpoints import rag_query + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "openai", + "team_id": "team-b", + } + + request = _mock_request() + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "vs_other_team"}, + } + ), + ), + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aquery", + new=AsyncMock(), + ) as mock_aquery, + ): + with pytest.raises(HTTPException) as exc_info: + await rag_query( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert exc_info.value.status_code == 403 + mock_aquery.assert_not_called() + + +@pytest.mark.asyncio +async def test_rag_ingest_denies_nested_other_team_vector_store(): + from litellm.proxy.rag_endpoints.endpoints import rag_ingest + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "openai", + "team_id": "team-b", + } + + request = _mock_request() + with ( + patch( + "litellm.proxy.rag_endpoints.endpoints.parse_rag_ingest_request", + new=AsyncMock( + return_value=( + { + "vector_store": { + "custom_llm_provider": "openai", + "vector_store_id": "vs_other_team", + } + }, + None, + "https://example.com/file.txt", + None, + ) + ), + ), + patch.object(litellm, "vector_store_registry", mock_registry), + patch( + "litellm.proxy.rag_endpoints.endpoints.litellm.aingest", + new=AsyncMock(), + ) as mock_aingest, + ): + with pytest.raises(HTTPException) as exc_info: + await rag_ingest( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert exc_info.value.status_code == 403 + mock_aingest.assert_not_called() + + +def test_rag_payload_scan_rejects_excessive_nesting(): + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + from litellm.proxy.rag_endpoints.endpoints import ( + _collect_vector_store_ids_from_payload, + ) + + payload = {} + current = payload + for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1): + current["nested"] = {} + current = current["nested"] + current["vector_store_id"] = "vs_too_deep" + + with pytest.raises(HTTPException) as exc_info: + _collect_vector_store_ids_from_payload(payload) + + assert exc_info.value.status_code == 400 + + +def test_rag_payload_scan_accepts_vector_store_id_at_depth_limit(): + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + from litellm.proxy.rag_endpoints.endpoints import ( + _collect_vector_store_ids_from_payload, + ) + + payload = {} + current = payload + for _ in range(DEFAULT_MAX_RECURSE_DEPTH): + current["nested"] = {} + current = current["nested"] + current["vector_store_id"] = "vs_at_limit" + + assert _collect_vector_store_ids_from_payload(payload) == {"vs_at_limit"} + + +def test_rag_payload_scan_ignores_primitive_list_beyond_depth_limit(): + from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + from litellm.proxy.rag_endpoints.endpoints import ( + _collect_vector_store_ids_from_payload, + ) + + payload = {} + current = payload + for _ in range(DEFAULT_MAX_RECURSE_DEPTH): + current["nested"] = {} + current = current["nested"] + current["labels"] = ["alpha", "beta"] + + assert _collect_vector_store_ids_from_payload(payload) == set() + + +@pytest.mark.asyncio +async def test_responses_file_search_denies_other_team_vector_store(): + from litellm.proxy.common_request_processing import ( + _authorize_response_file_search_vector_stores, + ) + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "openai", + "team_id": "team-b", + } + + with patch.object(litellm, "vector_store_registry", mock_registry): + with pytest.raises(HTTPException) as exc_info: + await _authorize_response_file_search_vector_stores( + data={ + "tools": [ + { + "type": "file_search", + "vector_store_ids": ["vs_other_team"], + } + ] + }, + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vertex_discovery_denies_other_team_vector_store_credentials(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _base_vertex_proxy_route, + ) + + request = _mock_request() + request.method = "GET" + vector_store_credentials = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "vertex_ai", + "team_id": "team-b", + } + + with patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", + new=AsyncMock(return_value=UserAPIKeyAuth(team_id="team-a")), + ): + with pytest.raises(HTTPException) as exc_info: + await _base_vertex_proxy_route( + endpoint="projects/p/locations/us-central1/dataStores/vs_other_team", + request=request, + fastapi_response=Response(), + get_vertex_pass_through_handler=MagicMock(), + router_credentials=vector_store_credentials, + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_managed_vector_store_uses_shared_cache_helper_for_db_fallback(): + from litellm.proxy.vector_store_endpoints.utils import ( + get_litellm_managed_vector_store, + ) + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = None + cache_helper = AsyncMock( + return_value=[ + LiteLLM_ManagedVectorStoresTable( + vector_store_id="vs_cached", + custom_llm_provider="openai", + vector_store_name=None, + vector_store_description=None, + vector_store_metadata=None, + created_at=None, + updated_at=None, + litellm_credential_name=None, + litellm_params={"api_base": "https://example.com"}, + team_id="team-a", + user_id=None, + ) + ] + ) + + with ( + patch.object(litellm, "vector_store_registry", mock_registry), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.auth.auth_checks.get_managed_vector_store_rows_by_uuids", + new=cache_helper, + ), + ): + vector_store = await get_litellm_managed_vector_store( + vector_store_id="vs_cached" + ) + + assert vector_store is not None + assert vector_store["vector_store_id"] == "vs_cached" + assert vector_store["team_id"] == "team-a" + cache_helper.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_managed_vector_store_fails_closed_on_lookup_error(): + from litellm.proxy.vector_store_endpoints.utils import ( + get_litellm_managed_vector_store, + ) + + mock_registry = MagicMock() + mock_registry.get_litellm_managed_vector_store_from_registry.side_effect = ( + RuntimeError("registry unavailable") + ) + + with patch.object(litellm, "vector_store_registry", mock_registry): + with pytest.raises(HTTPException) as exc_info: + await get_litellm_managed_vector_store(vector_store_id="vs_registry_only") + + assert exc_info.value.status_code == 500 + + +@pytest.mark.asyncio +async def test_vertex_discovery_allows_unregistered_provider_native_datastore_id(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + vertex_discovery_proxy_route, + ) + + request = _mock_request() + request.method = "GET" + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_litellm_managed_vector_store", + new=AsyncMock(return_value=None), + ) as mock_lookup, + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._base_vertex_proxy_route", + new=AsyncMock(return_value={"ok": True}), + ) as mock_base_route, + ): + response = await vertex_discovery_proxy_route( + endpoint="projects/p/locations/us-central1/dataStores/vs_unknown", + request=request, + fastapi_response=Response(), + ) + + assert response == {"ok": True} + mock_lookup.assert_awaited_once_with(vector_store_id="vs_unknown") + assert mock_base_route.call_args.kwargs["router_credentials"] is None + + +@pytest.mark.asyncio +async def test_milvus_passthrough_denies_other_team_vector_store_index(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + milvus_proxy_route, + ) + + request = _mock_request() + request.url.path = "/milvus/v2/vectordb/entities/search" + + index_object = MagicMock() + index_object.litellm_params.vector_store_name = "tenant-b-store" + index_object.litellm_params.vector_store_index = "tenant_b_collection" + + mock_index_registry = MagicMock() + mock_index_registry.is_vector_store_index.return_value = True + mock_index_registry.get_vector_store_index_by_name.return_value = index_object + + mock_vector_registry = MagicMock() + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "milvus", + "team_id": "team-b", + "litellm_params": {"api_base": "https://milvus.example.com"}, + } + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_request_body", + new=AsyncMock(return_value={"collectionName": "managed_index"}), + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint", + return_value=True, + ), + patch.object(litellm, "vector_store_index_registry", mock_index_registry), + patch.object(litellm, "vector_store_registry", mock_vector_registry), + ): + with pytest.raises(HTTPException) as exc_info: + await milvus_proxy_route( + endpoint="v2/vectordb/entities/search", + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_azure_passthrough_denies_other_team_vector_store_index(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + azure_proxy_route, + ) + + request = _mock_request() + request.url.path = "/azure/indexes/managed_index/docs/search" + + index_object = MagicMock() + index_object.litellm_params.vector_store_name = "tenant-b-store" + + mock_index_registry = MagicMock() + mock_index_registry.is_vector_store_index.side_effect = ( + lambda vector_store_index_name: vector_store_index_name == "managed_index" + ) + mock_index_registry.get_vector_store_index_by_name.return_value = index_object + + mock_vector_registry = MagicMock() + mock_vector_registry.get_litellm_managed_vector_store_from_registry_by_name.return_value = { + "vector_store_id": "vs_other_team", + "custom_llm_provider": "azure_ai", + "team_id": "team-b", + "litellm_params": {"api_base": "https://azure.example.com"}, + } + + with ( + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_passthrough_request_using_router_model", + return_value=False, + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.ProviderConfigManager.get_provider_vector_stores_config", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_allowed_to_call_vector_store_endpoint", + return_value=True, + ), + patch.object(litellm, "vector_store_index_registry", mock_index_registry), + patch.object(litellm, "vector_store_registry", mock_vector_registry), + ): + with pytest.raises(HTTPException) as exc_info: + await azure_proxy_route( + endpoint="indexes/managed_index/docs/search", + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(team_id="team-a"), + ) + + assert exc_info.value.status_code == 403 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/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