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/Makefile b/Makefile index b6b674ff3b..5dbd308a3e 100644 --- a/Makefile +++ b/Makefile @@ -185,3 +185,6 @@ test-llm-translation-single: install-test-deps $(UV_RUN) pytest tests/llm_translation/$(FILE) \ --junitxml=test-results/junit.xml \ -v --tb=short --maxfail=100 --timeout=300 + +test-llm-translation-flush-vcr-cache: + $(UV_RUN) python tests/_flush_vcr_cache.py diff --git a/README.md b/README.md index d72fb746ed..72fd43925c 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,7 @@ Managing LLM calls across providers gets complicated fast — different SDKs, au Stripe image Google ADK - Greptile + Greptile OpenHands

Netflix

OpenAI Agents SDK diff --git a/docs/my-website/docs/providers/crusoe.md b/docs/my-website/docs/providers/crusoe.md new file mode 100644 index 0000000000..aa737cbdcd --- /dev/null +++ b/docs/my-website/docs/providers/crusoe.md @@ -0,0 +1,196 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Crusoe + +## Overview + +| Property | Details | +|-------|-------| +| Description | Crusoe Cloud provides GPU-accelerated inference for open-source large language models, optimized for performance and cost efficiency. | +| Provider Route on LiteLLM | `crusoe/` | +| Link to Provider Doc | [Crusoe Managed Inference Documentation ↗](https://docs.crusoecloud.com/managed-inference/overview/index.html) | +| Base URL | `https://managed-inference-api-proxy.crusoecloud.com/v1` | +| Supported Operations | [`/chat/completions`](#sample-usage) | + +
+
+ +**We support ALL Crusoe models, just set `crusoe/` as a prefix when sending completion requests** + +## Available Models + +| Model | Description | Context Window | +|-------|-------------|----------------| +| `crusoe/deepseek-ai/DeepSeek-R1-0528` | DeepSeek R1 reasoning model (May 2025) | 163,840 tokens | +| `crusoe/deepseek-ai/DeepSeek-V3-0324` | DeepSeek V3 chat model (March 2025) | 163,840 tokens | +| `crusoe/google/gemma-3-12b-it` | Google Gemma 3 12B instruction-tuned | 131,072 tokens | +| `crusoe/meta-llama/Llama-3.3-70B-Instruct` | Llama 3.3 70B instruction-tuned | 131,072 tokens | +| `crusoe/moonshotai/Kimi-K2-Thinking` | Kimi K2 extended thinking model | 262,144 tokens | +| `crusoe/openai/gpt-oss-120b` | OpenAI 120B open-source model | 131,072 tokens | +| `crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507` | Qwen3 235B MoE instruction-tuned | 262,144 tokens | + +## Required Variables + +```python showLineNumbers title="Environment Variables" +os.environ["CRUSOE_API_KEY"] = "" # your Crusoe API key +``` + +## Usage - LiteLLM Python SDK + +### Non-streaming + +```python showLineNumbers title="Crusoe Non-streaming Completion" +import os +import litellm +from litellm import completion + +os.environ["CRUSOE_API_KEY"] = "" # your Crusoe API key + +messages = [{"content": "Hello, how are you?", "role": "user"}] + +# Crusoe call +response = completion( + model="crusoe/meta-llama/Llama-3.3-70B-Instruct", + messages=messages +) + +print(response) +``` + +### Streaming + +```python showLineNumbers title="Crusoe Streaming Completion" +import os +import litellm +from litellm import completion + +os.environ["CRUSOE_API_KEY"] = "" # your Crusoe API key + +messages = [{"content": "Write a short story about AI", "role": "user"}] + +# Crusoe call with streaming +response = completion( + model="crusoe/meta-llama/Llama-3.3-70B-Instruct", + messages=messages, + stream=True +) + +for chunk in response: + print(chunk) +``` + +### Function Calling + +```python showLineNumbers title="Crusoe Function Calling" +import os +import litellm +from litellm import completion + +os.environ["CRUSOE_API_KEY"] = "" # your Crusoe API key + +tools = [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA" + } + }, + "required": ["location"] + } + } +}] + +messages = [{"role": "user", "content": "What's the weather in Boston?"}] + +response = completion( + model="crusoe/meta-llama/Llama-3.3-70B-Instruct", + messages=messages, + tools=tools, + tool_choice="auto" +) + +print(response) +``` + +## Usage - LiteLLM Proxy Server + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: llama-3.3-70b + litellm_params: + model: crusoe/meta-llama/Llama-3.3-70B-Instruct + api_key: os.environ/CRUSOE_API_KEY + - model_name: deepseek-r1 + litellm_params: + model: crusoe/deepseek-ai/DeepSeek-R1-0528 + api_key: os.environ/CRUSOE_API_KEY + - model_name: deepseek-v3 + litellm_params: + model: crusoe/deepseek-ai/DeepSeek-V3-0324 + api_key: os.environ/CRUSOE_API_KEY + - model_name: qwen3-235b + litellm_params: + model: crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507 + api_key: os.environ/CRUSOE_API_KEY + - model_name: kimi-k2 + litellm_params: + model: crusoe/moonshotai/Kimi-K2-Thinking + api_key: os.environ/CRUSOE_API_KEY +``` + +## Custom API Base + +**Option 1: Environment variable** + +```python showLineNumbers title="Custom API Base via env var" +import os +from litellm import completion + +os.environ["CRUSOE_API_BASE"] = "https://custom.crusoecloud.com/v1" +os.environ["CRUSOE_API_KEY"] = "" # your API key + +response = completion( + model="crusoe/meta-llama/Llama-3.3-70B-Instruct", + messages=[{"content": "Hello!", "role": "user"}], +) +``` + +**Option 2: Pass directly** + +```python showLineNumbers title="Custom API Base via parameter" +from litellm import completion + +response = completion( + model="crusoe/meta-llama/Llama-3.3-70B-Instruct", + messages=[{"content": "Hello!", "role": "user"}], + api_base="https://custom.crusoecloud.com/v1", + api_key="your-api-key", +) +``` + +## Supported OpenAI Parameters + +- `temperature` +- `max_tokens` +- `max_completion_tokens` +- `top_p` +- `frequency_penalty` +- `presence_penalty` +- `stop` +- `n` +- `stream` +- `tools` +- `tool_choice` +- `response_format` +- `seed` +- `user` +- `logit_bias` +- `logprobs` +- `top_logprobs` diff --git a/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py b/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py index a368232038..e8f104c262 100644 --- a/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py +++ b/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py @@ -10,28 +10,21 @@ has already authenticated the user) and you need to extract user information fro custom headers or other request attributes. """ -from typing import TYPE_CHECKING, Dict, Optional, Union, cast +from typing import cast from fastapi import Request from fastapi.responses import RedirectResponse -if TYPE_CHECKING: - from fastapi_sso.sso.base import OpenID -else: - from typing import Any as OpenID - -from litellm.proxy.management_endpoints.types import CustomOpenID - class EnterpriseCustomSSOHandler: """ Enterprise Custom SSO Handler for LiteLLM Proxy - + This class provides methods for handling custom SSO authentication flows where users can implement their own authentication logic by processing request headers and returning user information in OpenID format. """ - + @staticmethod async def handle_custom_ui_sso_sign_in( request: Request, @@ -40,16 +33,16 @@ class EnterpriseCustomSSOHandler: Allow a user to execute their custom code to parse incoming request headers and return a OpenID object Use this when you have an OAuth proxy in front of LiteLLM (where the OAuth proxy has already authenticated the user) - + Args: request: The FastAPI request object containing headers and other request data - + Returns: RedirectResponse: Redirect response that sends the user to the LiteLLM UI with authentication token - + Raises: ValueError: If custom_ui_sso_sign_in_handler is not configured - + Example: This method is typically called when a user has already been authenticated by an external OAuth proxy and the proxy has added custom headers containing user information. @@ -60,27 +53,44 @@ class EnterpriseCustomSSOHandler: from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler from litellm.proxy.proxy_server import ( CommonProxyErrors, + general_settings, premium_user, user_custom_ui_sso_sign_in_handler, ) + from litellm.proxy.auth.trusted_proxy_utils import ( + require_trusted_proxy_request, + ) + if premium_user is not True: raise ValueError(CommonProxyErrors.not_premium_user.value) - + if user_custom_ui_sso_sign_in_handler is None: - raise ValueError("custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings.") - - custom_sso_login_handler = cast(CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler) - openid_response: OpenID = await custom_sso_login_handler.handle_custom_ui_sso_sign_in( + raise ValueError( + "custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings." + ) + + require_trusted_proxy_request( request=request, + general_settings=general_settings, + feature_name="Custom UI SSO", ) - + + custom_sso_login_handler = cast( + CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler + ) + openid_response: OpenID = ( + await custom_sso_login_handler.handle_custom_ui_sso_sign_in( + request=request, + ) + ) + # Import here to avoid circular imports from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler - + return await SSOAuthenticationHandler.get_redirect_response_from_openid( result=openid_response, request=request, received_response=None, generic_client_id=None, ui_access_mode=None, - ) \ No newline at end of file + ) 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/_logging.py b/litellm/_logging.py index d072cc549d..5ddafd6c6a 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,12 +1,12 @@ import ast import logging import os -import re import sys from datetime import datetime from logging import Formatter -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional +from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -21,74 +21,11 @@ _ENABLE_SECRET_REDACTION = ( os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true" ) -_REDACTED = "REDACTED" - - -def _build_secret_patterns() -> re.Pattern: - patterns: List[str] = [ - # ── PEM private key / certificate blocks ── - r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----", - # ── GCP OAuth2 access tokens (ya29.*) ── - r"\bya29\.[A-Za-z0-9_.~+/-]+", - # ── Credential %s formatting (space separator, no key= prefix) ── - r"(?:client_secret|azure_password|azure_username)\s+[^\s,'\"})\]{}>]+", - # AWS access key IDs - r"(?:AKIA|ASIA)[0-9A-Z]{16}", - # AWS secrets / session tokens / access key IDs (key=value) - r"(?:aws_secret_access_key|aws_session_token|aws_access_key_id)" - r"\s*[:=]\s*[A-Za-z0-9/+=]{20,}", - # Bearer tokens (OAuth, JWT, etc.) - r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*", - # Basic auth headers - r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", - # OpenAI / Anthropic sk- prefixed keys - r"sk-[A-Za-z0-9\-_]{20,}", - # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr) - r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}", - # x-api-key / api-key header values (handles 'key': 'value' dict repr) - r"(?:x-api-key|api-key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+", - # Anthropic internal header keys - r"x-ak-[A-Za-z0-9\-_]{20,}", - # Google API keys - r"AIza[0-9A-Za-z\-_]{35}", - # Password / secret params (handles key=value and 'key': 'value') - # Word boundary prevents O(n^2) backtracking on long word-char runs. - r"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)" - r"['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+", - # Database connection string credentials (scheme://user:pass@host) - r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)", - # Databricks personal access tokens - r"dapi[0-9a-f]{32}", - # ── Key-name-based redaction ── - # Catches secrets inside dicts/config dumps by matching on the KEY name - # regardless of what the value looks like. - # e.g. 'master_key': 'any-value-here', "database_url": "postgres://..." - # private_key with PEM-aware value capture - r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""", - r"(?:master_key|database_url|db_url|connection_string|" - r"signing_key|encryption_key|" - r"auth_token|access_token|refresh_token|" - r"slack_webhook_url|webhook_url|" - r"database_connection_string|" - r"huggingface_token|jwt_secret)" - r"""['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+""", - # ── Raw JWTs (without Bearer prefix) ── - r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", - # ── Azure SAS tokens in URLs ── - r"[?&]sig=[A-Za-z0-9%+/=]+", - # ── Full JSON service-account blobs (single-line and multi-line) ── - r'\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', - ] - return re.compile("|".join(patterns), re.IGNORECASE) - - -_SECRET_RE = _build_secret_patterns() - def _redact_string(value: str) -> str: if not _ENABLE_SECRET_REDACTION: return value - return _SECRET_RE.sub(_REDACTED, value) + return redact_string(value) def redact_secrets(value: str) -> str: diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4b965d4e63..aaf083e75d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -387,6 +387,27 @@ def _get_batch_job_total_usage_from_file_content( ) +def _get_models_from_batch_input_file_content( + file_content_dictionary: List[dict], +) -> List[str]: + """Extract the distinct ``body.model`` values from a batch *input* file. + + Used by the proxy's batch pre-call hook to enforce that the caller is + authorized for every model named inside the JSONL — not just the one + on the outer request — so the proxy's per-key model allowlist isn't + bypassed by smuggling expensive models into the batch file. + """ + models: List[str] = [] + seen: set = set() + for _item in file_content_dictionary: + body = _item.get("body") or {} + model = body.get("model") + if model and model not in seen: + seen.add(model) + models.append(model) + return models + + def _get_batch_job_input_file_usage( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", @@ -403,11 +424,25 @@ def _get_batch_job_input_file_usage( for _item in file_content_dictionary: body = _item.get("body", {}) model = body.get("model", model_name or "") - messages = body.get("messages", []) + # Chat completion payloads. + messages = body.get("messages") if messages: - item_tokens = token_counter(model=model, messages=messages) - prompt_tokens += item_tokens + prompt_tokens += token_counter(model=model, messages=messages) + continue + + # Text completion payloads (`prompt`). + prompt = body.get("prompt") + if prompt: + prompt_tokens += _count_prompt_or_input_tokens(model=model, value=prompt) + continue + + # Embedding payloads (`input`). + input_data = body.get("input") + if input_data: + prompt_tokens += _count_prompt_or_input_tokens( + model=model, value=input_data + ) return Usage( total_tokens=prompt_tokens + completion_tokens, @@ -416,6 +451,43 @@ def _get_batch_job_input_file_usage( ) +def _count_prompt_or_input_tokens(model: str, value: Any) -> int: + """Token-count a ``prompt`` / ``input`` field that the OpenAI batch + schema allows in four shapes: + + - ``str``: a single text prompt. + - ``list[str]``: multiple text prompts. + - ``list[int]``: a pre-tokenized prompt (each int counts as 1 token). + - ``list[list[int]]``: multiple pre-tokenized prompts. + + Pre-fix only the string shapes were counted, so a caller could send + a large ``list[list[int]]`` payload and slip past TPM rate limits + with a recorded cost of zero tokens. + """ + if isinstance(value, str): + return token_counter(model=model, text=value) + if isinstance(value, list): + total = 0 + for chunk in value: + if isinstance(chunk, str): + total += token_counter(model=model, text=chunk) + elif isinstance(chunk, int): + # Single pre-tokenized prompt at the top level: each + # int counts as one token. + total += 1 + elif isinstance(chunk, list): + # Nested pre-tokenized prompt: every int contributes a + # token. Mixed string/int items still count. + total += sum(1 if isinstance(t, int) else 0 for t in chunk) + total += sum( + token_counter(model=model, text=t) + for t in chunk + if isinstance(t, str) + ) + return total + return 0 + + def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage: """ Get the tokens of a batch job from the response body diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8a68d74be5..9b4dd80265 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -513,7 +513,10 @@ def cost_per_token( # noqa: PLR0915 return fireworks_ai_cost_per_token(model=model, usage=usage_block) elif custom_llm_provider == "azure": return azure_openai_cost_per_token( - model=model, usage=usage_block, response_time_ms=response_time_ms + model=model, + usage=usage_block, + response_time_ms=response_time_ms, + service_tier=service_tier, ) elif custom_llm_provider == "gemini": return gemini_cost_per_token( @@ -539,6 +542,7 @@ def cost_per_token( # noqa: PLR0915 usage=usage_block, response_time_ms=response_time_ms, request_model=request_model, + service_tier=service_tier, ) else: model_info = _cached_get_model_info_helper( 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/custom_sso_handler.py b/litellm/integrations/custom_sso_handler.py index 7f60decabc..202e488e0e 100644 --- a/litellm/integrations/custom_sso_handler.py +++ b/litellm/integrations/custom_sso_handler.py @@ -18,6 +18,17 @@ class CustomSSOLoginHandler(CustomLogger): self, request: Request, ) -> OpenID: + from litellm.proxy.auth.trusted_proxy_utils import ( + require_trusted_proxy_request, + ) + from litellm.proxy.proxy_server import general_settings + + require_trusted_proxy_request( + request=request, + general_settings=general_settings, + feature_name="Custom UI SSO", + ) + request_headers_dict = dict(request.headers) return OpenID( id=request_headers_dict.get("x-litellm-user-id"), diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index e691c490c8..0efc7d6687 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -90,6 +90,29 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: return cache_read_input_tokens +def resolve_langfuse_credentials( + langfuse_public_key=None, + langfuse_secret=None, + langfuse_secret_key=None, + langfuse_host=None, + allow_env_credentials: bool = True, +): + if allow_env_credentials is False and langfuse_host is not None: + secret_key = langfuse_secret or langfuse_secret_key + public_key = langfuse_public_key + else: + secret_key = ( + langfuse_secret or langfuse_secret_key or os.getenv("LANGFUSE_SECRET_KEY") + ) + public_key = langfuse_public_key or os.getenv("LANGFUSE_PUBLIC_KEY") + + resolved_host = langfuse_host or os.getenv( + "LANGFUSE_HOST", "https://cloud.langfuse.com" + ) + + return public_key, secret_key, resolved_host + + class LangFuseLogger: # Class variables or attributes def __init__( @@ -98,6 +121,7 @@ class LangFuseLogger: langfuse_secret=None, langfuse_host=None, flush_interval=1, + allow_env_credentials: bool = True, ): try: import langfuse @@ -106,11 +130,13 @@ class LangFuseLogger: raise Exception( f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\n{traceback.format_exc()}\033[0m" ) - # Instance variables - self.secret_key = langfuse_secret or os.getenv("LANGFUSE_SECRET_KEY") - self.public_key = langfuse_public_key or os.getenv("LANGFUSE_PUBLIC_KEY") - self.langfuse_host = langfuse_host or os.getenv( - "LANGFUSE_HOST", "https://cloud.langfuse.com" + self.public_key, self.secret_key, self.langfuse_host = ( + resolve_langfuse_credentials( + langfuse_public_key=langfuse_public_key, + langfuse_secret=langfuse_secret, + langfuse_host=langfuse_host, + allow_env_credentials=allow_env_credentials, + ) ) if not ( self.langfuse_host.startswith("http://") @@ -160,9 +186,10 @@ class LangFuseLogger: project_id = None if os.getenv("UPSTREAM_LANGFUSE_SECRET_KEY") is not None: + upstream_langfuse_debug_env = os.getenv("UPSTREAM_LANGFUSE_DEBUG") upstream_langfuse_debug = ( - str_to_bool(self.upstream_langfuse_debug) - if self.upstream_langfuse_debug is not None + str_to_bool(upstream_langfuse_debug_env) + if upstream_langfuse_debug_env is not None else None ) self.upstream_langfuse_secret_key = os.getenv( @@ -173,7 +200,7 @@ class LangFuseLogger: ) self.upstream_langfuse_host = os.getenv("UPSTREAM_LANGFUSE_HOST") self.upstream_langfuse_release = os.getenv("UPSTREAM_LANGFUSE_RELEASE") - self.upstream_langfuse_debug = os.getenv("UPSTREAM_LANGFUSE_DEBUG") + self.upstream_langfuse_debug = upstream_langfuse_debug_env self.upstream_langfuse = Langfuse( public_key=self.upstream_langfuse_public_key, secret_key=self.upstream_langfuse_secret_key, diff --git a/litellm/integrations/langfuse/langfuse_handler.py b/litellm/integrations/langfuse/langfuse_handler.py index fbadf1a2fc..4a80972642 100644 --- a/litellm/integrations/langfuse/langfuse_handler.py +++ b/litellm/integrations/langfuse/langfuse_handler.py @@ -115,8 +115,10 @@ class LangFuseHandler: langfuse_logger = LangFuseLogger( langfuse_public_key=credentials.get("langfuse_public_key"), - langfuse_secret=credentials.get("langfuse_secret"), + langfuse_secret=credentials.get("langfuse_secret") + or credentials.get("langfuse_secret_key"), langfuse_host=credentials.get("langfuse_host"), + allow_env_credentials=credentials.get("langfuse_host") is None, ) in_memory_dynamic_logger_cache.set_cache( credentials=credentials, diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 5f4ced3a5c..b7a565512c 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -20,7 +20,7 @@ from ...litellm_core_utils.specialty_caches.dynamic_logging_cache import ( DynamicLoggingCache, ) from ..prompt_management_base import PromptManagementBase -from .langfuse import LangFuseLogger +from .langfuse import LangFuseLogger, resolve_langfuse_credentials from .langfuse_handler import LangFuseHandler if TYPE_CHECKING: @@ -46,6 +46,7 @@ def langfuse_client_init( langfuse_secret_key=None, langfuse_host=None, flush_interval=1, + allow_env_credentials: bool = True, ) -> LangfuseClass: """ Initialize Langfuse client with caching to prevent multiple initializations. @@ -70,14 +71,12 @@ def langfuse_client_init( f"\033[91mLangfuse not installed, try running 'pip install langfuse' to fix this error: {e}\n\033[0m" ) - # Instance variables - - secret_key = ( - langfuse_secret or langfuse_secret_key or os.getenv("LANGFUSE_SECRET_KEY") - ) - public_key = langfuse_public_key or os.getenv("LANGFUSE_PUBLIC_KEY") - langfuse_host = langfuse_host or os.getenv( - "LANGFUSE_HOST", "https://cloud.langfuse.com" + public_key, secret_key, langfuse_host = resolve_langfuse_credentials( + langfuse_public_key=langfuse_public_key, + langfuse_secret=langfuse_secret, + langfuse_secret_key=langfuse_secret_key, + langfuse_host=langfuse_host, + allow_env_credentials=allow_env_credentials, ) if not ( @@ -222,6 +221,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge langfuse_secret=dynamic_callback_params.get("langfuse_secret"), langfuse_secret_key=dynamic_callback_params.get("langfuse_secret_key"), langfuse_host=dynamic_callback_params.get("langfuse_host"), + allow_env_credentials=dynamic_callback_params.get("langfuse_host") is None, ) langfuse_prompt_client = self._get_prompt_from_id( langfuse_prompt_id=prompt_id, @@ -246,6 +246,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge langfuse_secret=dynamic_callback_params.get("langfuse_secret"), langfuse_secret_key=dynamic_callback_params.get("langfuse_secret_key"), langfuse_host=dynamic_callback_params.get("langfuse_host"), + allow_env_credentials=dynamic_callback_params.get("langfuse_host") is None, ) langfuse_prompt_client = self._get_prompt_from_id( langfuse_prompt_id=prompt_id, diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 3d4fd39ebe..3a20612237 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -112,17 +112,28 @@ class LangsmithLogger(CustomBatchLogger): langsmith_project: Optional[str] = None, langsmith_base_url: Optional[str] = None, langsmith_tenant_id: Optional[str] = None, + allow_env_credentials: bool = True, ) -> LangsmithCredentialsObject: - _credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY") - _credentials_project = ( - langsmith_project or os.getenv("LANGSMITH_PROJECT") or "litellm-completion" - ) - _credentials_base_url = ( - langsmith_base_url - or os.getenv("LANGSMITH_BASE_URL") - or "https://api.smith.langchain.com" - ) - _credentials_tenant_id = langsmith_tenant_id or os.getenv("LANGSMITH_TENANT_ID") + if allow_env_credentials is False and langsmith_base_url is not None: + _credentials_api_key = langsmith_api_key + _credentials_project = langsmith_project or "litellm-completion" + _credentials_base_url = langsmith_base_url + _credentials_tenant_id = langsmith_tenant_id + else: + _credentials_api_key = langsmith_api_key or os.getenv("LANGSMITH_API_KEY") + _credentials_project = ( + langsmith_project + or os.getenv("LANGSMITH_PROJECT") + or "litellm-completion" + ) + _credentials_base_url = ( + langsmith_base_url + or os.getenv("LANGSMITH_BASE_URL") + or "https://api.smith.langchain.com" + ) + _credentials_tenant_id = langsmith_tenant_id or os.getenv( + "LANGSMITH_TENANT_ID" + ) return LangsmithCredentialsObject( LANGSMITH_API_KEY=_credentials_api_key, @@ -540,6 +551,10 @@ class LangsmithLogger(CustomBatchLogger): langsmith_tenant_id=standard_callback_dynamic_params.get( "langsmith_tenant_id", None ), + allow_env_credentials=standard_callback_dynamic_params.get( + "langsmith_base_url", None + ) + is None, ) else: credentials = self.default_credentials diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index b6d91d0b76..77833e5de0 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -69,6 +69,8 @@ class OpenTelemetryConfig: deployment_environment: Optional[str] = None model_id: Optional[str] = None ignore_context_propagation: Optional[bool] = None + # When True, create a private TracerProvider instead of reusing or setting the global one. + skip_set_global: bool = False def __post_init__(self) -> None: # If endpoint is specified but exporter is still the default "console", @@ -259,16 +261,21 @@ class OpenTelemetry(CustomLogger): try: existing_provider = get_existing_provider_fn() - # If a real SDK provider exists (set by another SDK like Langfuse), use it - # This uses a positive check for SDK providers instead of a negative check for proxy providers if isinstance(existing_provider, sdk_provider_class): - verbose_logger.debug( - "OpenTelemetry: Using existing %s: %s", - provider_name, - type(existing_provider).__name__, - ) - provider = existing_provider - # Don't call set_provider to preserve existing context + if skip_set_global: + verbose_logger.debug( + "OpenTelemetry: existing %s found but skip_set_global=True; creating private %s for isolation", + provider_name, + provider_name, + ) + provider = create_new_provider_fn() + else: + verbose_logger.debug( + "OpenTelemetry: Using existing %s: %s", + provider_name, + type(existing_provider).__name__, + ) + provider = existing_provider else: # Default proxy provider or unknown type, create our own verbose_logger.debug("OpenTelemetry: Creating new %s", provider_name) @@ -293,6 +300,12 @@ class OpenTelemetry(CustomLogger): return provider + def _skip_set_global(self) -> bool: + # langfuse_otel relies on the Langfuse SDK's providers; don't overwrite them. + return self.config.skip_set_global or ( + hasattr(self, "callback_name") and self.callback_name == "langfuse_otel" + ) + def _init_tracing(self, tracer_provider): from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider @@ -303,11 +316,6 @@ class OpenTelemetry(CustomLogger): provider.add_span_processor(self._get_span_processor()) return provider - # CRITICAL FIX: For Langfuse OTEL, skip setting global provider to prevent interference - skip_global = ( - hasattr(self, "callback_name") and self.callback_name == "langfuse_otel" - ) - tracer_provider = self._get_or_create_provider( provider=tracer_provider, provider_name="TracerProvider", @@ -315,16 +323,18 @@ class OpenTelemetry(CustomLogger): sdk_provider_class=TracerProvider, create_new_provider_fn=create_tracer_provider, set_provider_fn=trace.set_tracer_provider, - skip_set_global=skip_global, + skip_set_global=self._skip_set_global(), ) # Grab our tracer from the TracerProvider (not from global context) # This ensures we use the provided TracerProvider (e.g., for testing) self.tracer = tracer_provider.get_tracer(LITELLM_TRACER_NAME) + self._tracer_provider = tracer_provider self.span_kind = SpanKind def _init_metrics(self, meter_provider): if not self.config.enable_metrics: + self._meter_provider = None self._operation_duration_histogram = None self._token_usage_histogram = None self._cost_histogram = None @@ -350,7 +360,9 @@ class OpenTelemetry(CustomLogger): sdk_provider_class=MeterProvider, create_new_provider_fn=create_meter_provider, set_provider_fn=metrics.set_meter_provider, + skip_set_global=self._skip_set_global(), ) + self._meter_provider = meter_provider meter = meter_provider.get_meter(__name__) @@ -388,6 +400,7 @@ class OpenTelemetry(CustomLogger): def _init_logs(self, logger_provider): # nothing to do if events disabled if not self.config.enable_events: + self._logger_provider = None return from opentelemetry._logs import get_logger_provider, set_logger_provider @@ -404,13 +417,14 @@ class OpenTelemetry(CustomLogger): ) return provider - self._get_or_create_provider( + self._logger_provider = self._get_or_create_provider( provider=logger_provider, provider_name="LoggerProvider", get_existing_provider_fn=get_logger_provider, sdk_provider_class=OTLoggerProvider, create_new_provider_fn=create_logger_provider, set_provider_fn=set_logger_provider, + skip_set_global=self._skip_set_global(), ) def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -1073,7 +1087,7 @@ class OpenTelemetry(CustomLogger): # See: https://github.com/open-telemetry/opentelemetry-python/pull/4676 # TODO: Refactor to use the proper OTEL Logs API instead of directly creating SDK LogRecords - from opentelemetry._logs import SeverityNumber, get_logger + from opentelemetry._logs import SeverityNumber try: from opentelemetry.sdk._logs import ( # type: ignore[attr-defined] # OTEL < 1.39.0 @@ -1084,7 +1098,10 @@ class OpenTelemetry(CustomLogger): LogRecord as SdkLogRecord, # type: ignore[attr-defined] # OTEL >= 1.39.0 ) - otel_logger = get_logger(LITELLM_LOGGER_NAME) + # Resolve through the handler's own LoggerProvider (which may be a + # private one when skip_set_global=True) rather than the module-level + # get_logger() which always goes through the global provider. + otel_logger = self._logger_provider.get_logger(LITELLM_LOGGER_NAME) parent_ctx = span.get_span_context() provider = (kwargs.get("litellm_params") or {}).get( diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index b25da57723..0901d7b680 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/litellm/integrations/prometheus_helpers/prometheus_api.py @@ -2,6 +2,7 @@ Helper functions to query prometheus API """ +import json import time from datetime import datetime, timedelta from typing import Optional @@ -81,6 +82,24 @@ def is_prometheus_connected() -> bool: return False +def _quote_promql_string_literal(value: str) -> str: + """Render ``value`` as a PromQL double-quoted string literal. + + PromQL string literals follow Go's escape rules + (https://prometheus.io/docs/prometheus/latest/querying/basics/): a + backslash begins an escape sequence and a bare ``"`` ends the literal. + Without escaping, callers that accept arbitrary user-supplied values + (like the ``api_key`` filter on ``/global/spend/logs``) can inject extra + label matchers or selectors and read cross-tenant metrics. + + JSON's quoting rules are a strict subset of Go's, so ``json.dumps`` of + a Python string produces a literal Prometheus accepts: ``\\``, ``\\"``, + and the standard ``\\n`` / ``\\t`` / ``\\uNNNN`` control-character + escapes. The returned value already includes the surrounding quotes. + """ + return json.dumps(value, ensure_ascii=False) + + async def get_daily_spend_from_prometheus(api_key: Optional[str]): """ Expected Response Format: @@ -109,8 +128,11 @@ async def get_daily_spend_from_prometheus(api_key: Optional[str]): if api_key is None: query = "sum(delta(litellm_spend_metric_total[1d]))" else: + quoted_api_key = _quote_promql_string_literal(api_key) query = ( - f'sum(delta(litellm_spend_metric_total{{hashed_api_key="{api_key}"}}[1d]))' + "sum(delta(litellm_spend_metric_total{" + f"hashed_api_key={quoted_api_key}" + "}[1d]))" ) params = { 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/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 5a7d4e33b6..2c1d92920a 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -6,7 +6,8 @@ from typing import Any, Optional import httpx import litellm -from litellm._logging import _redact_string, verbose_logger +from litellm._logging import _ENABLE_SECRET_REDACTION, _redact_string, verbose_logger +from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.types.utils import LlmProviders from ..exceptions import ( @@ -261,10 +262,18 @@ def exception_type( # type: ignore # noqa: PLR0915 original_exception=original_exception ) try: - error_str = str(original_exception) + error_str = ( + redact_string(str(original_exception)) + if _ENABLE_SECRET_REDACTION + else str(original_exception) + ) if model: if hasattr(original_exception, "message"): - error_str = str(original_exception.message) + error_str = ( + redact_string(str(original_exception.message)) + if _ENABLE_SECRET_REDACTION + else str(original_exception.message) + ) if isinstance(original_exception, BaseException): exception_type = type(original_exception).__name__ else: @@ -2431,7 +2440,8 @@ def exception_type( # type: ignore # noqa: PLR0915 else: raise APIConnectionError( message="{}\n{}".format( - str(original_exception), _redact_string(traceback.format_exc()) + str(original_exception), + _redact_string(traceback.format_exc()), ), llm_provider=custom_llm_provider, model=model, @@ -2461,7 +2471,8 @@ def exception_type( # type: ignore # noqa: PLR0915 raise e # it's already mapped raised_exc = APIConnectionError( message="{}\n{}".format( - original_exception, _redact_string(traceback.format_exc()) + original_exception, + _redact_string(traceback.format_exc()), ), llm_provider="", model="", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 829c1c9ca0..a815442c2f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3242,10 +3242,15 @@ class Logging(LiteLLMLoggingBaseClass): ), langfuse_secret=self.standard_callback_dynamic_params.get( "langfuse_secret" - ), + ) + or self.standard_callback_dynamic_params.get("langfuse_secret_key"), langfuse_host=self.standard_callback_dynamic_params.get( "langfuse_host" ), + allow_env_credentials=self.standard_callback_dynamic_params.get( + "langfuse_host" + ) + is None, ) return langFuseLogger @@ -4720,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/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index f5f28822ca..7be7085297 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -77,8 +77,8 @@ def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: if litellm_params is None: return {} - proxy_request_headers = ( - litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} - ) + proxy_request_headers = (litellm_params.get("proxy_server_request") or {}).get( + "headers" + ) or {} return proxy_request_headers diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py new file mode 100644 index 0000000000..5c4e3e3dac --- /dev/null +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -0,0 +1,81 @@ +""" +Credential/secret redaction utilities. + +This module owns the compiled regex and the public `redact_string` helper so +that any part of the codebase (logging, exception mapping, etc.) can scrub +secrets from strings without depending on the logging-configuration module. +""" + +import re +from typing import List + +_REDACTED = "REDACTED" + + +def _build_secret_patterns() -> "re.Pattern[str]": + patterns: List[str] = [ + # PEM private key / certificate blocks + r"-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----", + # GCP OAuth2 access tokens (ya29.*) + r"\bya29\.[A-Za-z0-9_.~+/-]+", + # Credential %s formatting (space separator, no key= prefix) + r"(?:client_secret|azure_password|azure_username)\s+[^\s,'\"})\]{}>]+", + # AWS access key IDs + r"(?:AKIA|ASIA)[0-9A-Z]{16}", + # AWS secrets / session tokens / access key IDs (key=value) + r"(?:aws_secret_access_key|aws_session_token|aws_access_key_id)" + r"\s*[:=]\s*[A-Za-z0-9/+=]{20,}", + # Bearer tokens (OAuth, JWT, etc.) + r"Bearer\s+[A-Za-z0-9\-._~+/]{10,}=*", + # Basic auth headers + r"Basic\s+[A-Za-z0-9+/]{10,}={0,2}", + # OpenAI / Anthropic sk- prefixed keys + r"sk-[A-Za-z0-9\-_]{20,}", + # Generic api_key / api-key / apikey (handles 'key': 'value' dict repr) + r"(?:api[_-]?key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]{8,}", + # x-api-key / api-key header values (handles 'key': 'value' dict repr) + r"(?:x-api-key|api-key)['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+", + # Anthropic internal header keys + r"x-ak-[A-Za-z0-9\-_]{20,}", + # Google API keys (bare key value) + r"AIza[0-9A-Za-z\-_]{35}", + # URL query-param key=VALUE (e.g. ?key=AIza... or &key=...) — catches the + # full "key=" fragment so the value is redacted regardless of format. + r"(?<=[?&])key=[^\s&'\"]{8,}", + # Password / secret params (handles key=value and 'key': 'value') + # Word boundary prevents O(n^2) backtracking on long word-char runs. + r"(?:^|(?<=\W))\w*(?:password|passwd|client_secret|secret_key|_secret)" + r"['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+", + # Database connection string credentials (scheme://user:pass@host) + r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)", + # Databricks personal access tokens + r"dapi[0-9a-f]{32}", + # ── Key-name-based redaction ── + # Catches secrets inside dicts/config dumps by matching on the KEY name + # regardless of what the value looks like. + # e.g. 'master_key': 'any-value-here', "database_url": "postgres://..." + # private_key with PEM-aware value capture + r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""", + r"(?:master_key|database_url|db_url|connection_string|" + r"signing_key|encryption_key|" + r"auth_token|access_token|refresh_token|" + r"slack_webhook_url|webhook_url|" + r"database_connection_string|" + r"huggingface_token|jwt_secret)" + r"""['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+""", + # Raw JWTs (without Bearer prefix) + r"\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]*", + # Azure SAS tokens in URLs + r"[?&]sig=[A-Za-z0-9%+/=]+", + # Full JSON service-account blobs (single-line and multi-line) + r'\{[^{}]*"type"\s*:\s*"service_account"[^{}]*(?:\{[^{}]*\}[^{}]*)*\}', + ] + return re.compile("|".join(patterns), re.IGNORECASE) + + +_SECRET_RE = _build_secret_patterns() + + +def redact_string(value: str) -> str: + """Scrub known secret/credential patterns from *value* and return the result.""" + return _SECRET_RE.sub(_REDACTED, value) 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..224927e5ac 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -22,7 +22,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): import socket from ipaddress import ip_address, ip_network from typing import Any, List, Set, Tuple -from urllib.parse import urlparse, urlunparse +from urllib.parse import quote, urlparse, urlunparse import httpx @@ -46,6 +46,46 @@ class SSRFError(ValueError): pass +def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: + """Percent-encode one user-controlled URL path segment. + + ``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986 + unreserved characters such as ``.`` unescaped, so reject standalone dot + segments before they can be appended to an upstream URL and normalized by + the HTTP client. + """ + if value is None: + raise ValueError(f"{field_name} is required") + + value_str = str(value) + if value_str == "": + raise ValueError(f"{field_name} is required") + if value_str in {".", ".."}: + raise ValueError(f"{field_name} cannot be a dot path segment") + + return quote(value_str, safe="") + + +def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: + """Percent-encode a user-controlled URL path made of multiple segments. + + Empty segments are rejected, so leading, trailing, or consecutive slashes + fail closed instead of being normalized by the HTTP client. + """ + if value is None: + raise ValueError(f"{field_name} is required") + + value_str = str(value) + if value_str == "": + raise ValueError(f"{field_name} is required") + + encoded_segments = [] + for segment in value_str.split("/"): + encoded_segments.append(encode_url_path_segment(segment, field_name=field_name)) + + return "/".join(encoded_segments) + + def _is_blocked_ip(addr: str) -> bool: """Return True for any IP not safe to reach from a user-supplied URL. @@ -199,6 +239,47 @@ def validate_url(url: str) -> Tuple[str, str]: return rewritten, host_header +def assert_same_origin(candidate_url: str, expected_url: str) -> None: + """Verify ``candidate_url`` shares scheme, host, and port with ``expected_url``. + + Use when an upstream API returns a URL meant for follow-up requests + (e.g. an async-job polling URL that will be hit with the operator's + API key in the headers). The upstream is trusted because the operator + configured ``api_base``, but the URL it hands back must actually point + back at the same origin or we'd be blindly forwarding credentials + wherever the upstream told us to. + + Hostnames are compared case-insensitively. Default ports are made + explicit (HTTP→80, HTTPS→443) so ``https://api.example.com:443/...`` + and ``https://api.example.com/...`` are treated as the same origin. + + Error messages identify *which* component mismatched but never echo + the operator's ``expected`` host or the candidate's hostname back to + the caller — in the SSRF threat model the caller is the attacker, + and reflecting host info would be a secondary leak of operator + infrastructure details. + """ + candidate = urlparse(candidate_url) + expected = urlparse(expected_url) + + if candidate.scheme not in _ALLOWED_SCHEMES: + raise SSRFError("URL scheme is not allowed") + + if candidate.scheme != expected.scheme: + raise SSRFError("Origin mismatch on scheme") + + candidate_host = _normalize_host(candidate.hostname or "") + expected_host = _normalize_host(expected.hostname or "") + if not candidate_host or candidate_host != expected_host: + raise SSRFError("Origin mismatch on host") + + default_port = 443 if candidate.scheme == "https" else 80 + candidate_port = candidate.port if candidate.port is not None else default_port + expected_port = expected.port if expected.port is not None else default_port + if candidate_port != expected_port: + raise SSRFError("Origin mismatch on port") + + _MAX_REDIRECTS = 10 diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 3f03c744ef..fd67a7fbaf 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cas import httpx from httpx import Headers, Response +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest @@ -122,7 +123,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig): Complete URL for Anthropic batch retrieval: {api_base}/v1/messages/batches/{batch_id} """ api_base = api_base or self.anthropic_model_info.get_api_base(api_base) - return f"{api_base.rstrip('/')}/v1/messages/batches/{batch_id}" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + return f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}" def transform_retrieve_batch_request( self, diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index c56799f30c..56296df94a 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( FileContentRequest, @@ -89,7 +90,10 @@ class AnthropicFilesHandler: raise ValueError("Missing Anthropic API Key") # Construct the Anthropic batch results URL - results_url = f"{api_base.rstrip('/')}/v1/messages/batches/{batch_id}/results" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + results_url = ( + f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}/results" + ) # Prepare headers headers = { diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index aeaab4e57b..ea9bf00f50 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -19,6 +19,7 @@ from typing import Any, Dict, List, Optional, Union, cast import httpx from openai.types.file_deleted import FileDeleted +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( @@ -185,7 +186,8 @@ class AnthropicFilesConfig(BaseFilesConfig): AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE ) - return f"{api_base.rstrip('/')}/v1/files/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} def transform_retrieve_file_response( self, @@ -206,7 +208,8 @@ class AnthropicFilesConfig(BaseFilesConfig): AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE ) - return f"{api_base.rstrip('/')}/v1/files/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} def transform_delete_file_response( self, @@ -268,7 +271,8 @@ class AnthropicFilesConfig(BaseFilesConfig): AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE ) - return f"{api_base.rstrip('/')}/v1/files/{file_id}/content", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}/content", {} def transform_file_content_response( self, diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index a992d84d45..4ea768b02a 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Tuple import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.skills.transformation import ( BaseSkillsAPIConfig, LiteLLMLoggingObj, @@ -81,7 +82,8 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): api_base = AnthropicModelInfo.get_api_base() if skill_id: - return f"{api_base}/v1/skills/{skill_id}" + encoded_skill_id = encode_url_path_segment(skill_id, field_name="skill_id") + return f"{api_base}/v1/skills/{encoded_skill_id}" return f"{api_base}/v1/{endpoint}" def transform_create_skill_request( diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 61cfd54b56..c0e070b6c1 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, @@ -792,6 +793,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): client=client, litellm_params=litellm_params, api_base=api_base, + api_version=api_version, ) azure_client = self.get_azure_openai_client( api_version=api_version, @@ -898,6 +900,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 +921,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 +1027,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 +1044,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/cost_calculation.py b/litellm/llms/azure/cost_calculation.py index 5b411095ea..2a20c55a6c 100644 --- a/litellm/llms/azure/cost_calculation.py +++ b/litellm/llms/azure/cost_calculation.py @@ -12,7 +12,10 @@ from litellm.utils import get_model_info def cost_per_token( - model: str, usage: Usage, response_time_ms: Optional[float] = 0.0 + model: str, + usage: Usage, + response_time_ms: Optional[float] = 0.0, + service_tier: Optional[str] = None, ) -> Tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -47,4 +50,5 @@ def cost_per_token( model=model, usage=usage, custom_llm_provider="azure", + service_tier=service_tier, ) diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 76a6d485bc..ca9293325f 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -5,6 +5,7 @@ import httpx from openai.types.responses import ResponseReasoningItem from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import * @@ -201,7 +202,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Insert the response_id at the end of the path component # Remove trailing slash if present to avoid double slashes path = parsed_url.path.rstrip("/") - new_path = f"{path}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + new_path = f"{path}/{encoded_response_id}" # Reconstruct the URL with all original components but with the modified path constructed_url = urlunparse( @@ -322,7 +326,10 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig): # Insert the response_id and /cancel at the end of the path component # Remove trailing slash if present to avoid double slashes path = parsed_url.path.rstrip("/") - new_path = f"{path}/{response_id}/cancel" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + new_path = f"{path}/{encoded_response_id}/cancel" # Reconstruct the URL with all original components but with the modified path cancel_url = urlunparse( diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index c3cd06ab4d..9bae8abce8 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -36,6 +36,7 @@ from typing import ( import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, @@ -75,20 +76,29 @@ class AzureAIAgentsHandler: def _build_messages_url( self, api_base: str, thread_id: str, api_version: str ) -> str: - return f"{api_base}/threads/{thread_id}/messages?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + return ( + f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + ) def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: - return f"{api_base}/threads/{thread_id}/runs?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + return f"{api_base}/threads/{encoded_thread_id}/runs?api-version={api_version}" def _build_run_status_url( self, api_base: str, thread_id: str, run_id: str, api_version: str ) -> str: - return f"{api_base}/threads/{thread_id}/runs/{run_id}?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + return f"{api_base}/threads/{encoded_thread_id}/runs/{encoded_run_id}?api-version={api_version}" def _build_list_messages_url( self, api_base: str, thread_id: str, api_version: str ) -> str: - return f"{api_base}/threads/{thread_id}/messages?api-version={api_version}" + encoded_thread_id = encode_url_path_segment(thread_id, field_name="thread_id") + return ( + f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + ) def _build_create_thread_and_run_url(self, api_base: str, api_version: str) -> str: """URL for the create-thread-and-run endpoint (supports streaming).""" diff --git a/litellm/llms/azure_ai/cost_calculator.py b/litellm/llms/azure_ai/cost_calculator.py index 067181b946..755d44fdef 100644 --- a/litellm/llms/azure_ai/cost_calculator.py +++ b/litellm/llms/azure_ai/cost_calculator.py @@ -65,6 +65,7 @@ def cost_per_token( usage: Usage, response_time_ms: Optional[float] = 0.0, request_model: Optional[str] = None, + service_tier: Optional[str] = None, ) -> Tuple[float, float]: """ Calculate the cost per token for Azure AI models. @@ -102,6 +103,7 @@ def cost_per_token( model=model, usage=usage, custom_llm_provider="azure_ai", + service_tier=service_tier, ) except Exception as e: # For Model Router, the model name (e.g., "azure-model-router") may not be in the cost map diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index 76c247aea8..d4144a7571 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -17,11 +17,13 @@ from urllib.parse import quote import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.constants import ( AZURE_DOCUMENT_INTELLIGENCE_API_VERSION, AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, AZURE_OPERATION_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.ocr.transformation import ( BaseOCRConfig, DocumentType, @@ -217,11 +219,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): if "/" in model: # Extract the last part after the last slash model_id = model.split("/")[-1] + encoded_model_id = encode_url_path_segment(model_id, field_name="model_id") # Azure Document Intelligence analyze endpoint # Note: API version 2024-11-30+ uses /documentintelligence/ (not /formrecognizer/) url = ( - f"{api_base}/documentintelligence/documentModels/{model_id}:analyze" + f"{api_base}/documentintelligence/documentModels/{encoded_model_id}:analyze" f"?api-version={AZURE_DOCUMENT_INTELLIGENCE_API_VERSION}" ) @@ -599,6 +602,16 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "Azure Document Intelligence returned 202 but no Operation-Location header found" ) + # Reject cross-origin polling URLs — the auth headers + # below would otherwise leak to whatever URL the upstream + # (or an attacker-controlled upstream) returns. VERIA-51. + try: + assert_same_origin(operation_url, str(raw_response.request.url)) + except SSRFError as ssrf_err: + raise ValueError( + f"Azure Document Intelligence: rejected polling URL ({ssrf_err})" + ) + # Get headers for polling (need auth) poll_headers = { "Ocp-Apim-Subscription-Key": raw_response.request.headers.get( @@ -711,6 +724,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): "Azure Document Intelligence returned 202 but no Operation-Location header found" ) + # Reject cross-origin polling URLs (see sync path). VERIA-51. + try: + assert_same_origin(operation_url, str(raw_response.request.url)) + except SSRFError as ssrf_err: + raise ValueError( + f"Azure Document Intelligence: rejected polling URL ({ssrf_err})" + ) + # Get headers for polling (need auth) poll_headers = { "Ocp-Apim-Subscription-Key": raw_response.request.headers.get( diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index 2c7135f4d8..4c667b0ce3 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -12,6 +12,7 @@ import httpx from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -97,8 +98,15 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): agent_id, agent_alias_id = self._get_agent_id_and_alias_id(model) session_id = self._get_session_id(optional_params) + encoded_agent_id = encode_url_path_segment(agent_id, field_name="agent_id") + encoded_agent_alias_id = encode_url_path_segment( + agent_alias_id, field_name="agent_alias_id" + ) + encoded_session_id = encode_url_path_segment( + session_id, field_name="session_id" + ) - endpoint_url = f"{endpoint_url}/agents/{agent_id}/agentAliases/{agent_alias_id}/sessions/{session_id}/text" + endpoint_url = f"{endpoint_url}/agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text" return endpoint_url diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index a37af13162..c967fd334b 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -201,13 +201,14 @@ class BedrockCountTokensConfig(BaseAWSLLM): # Remove bedrock/ prefix if present if model_id.startswith("bedrock/"): model_id = model_id[8:] # Remove "bedrock/" prefix + encoded_model_id = self.encode_model_id(model_id=model_id) base_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=aws_region_name, ) - endpoint = f"{base_url}/model/{model_id}/count-tokens" + endpoint = f"{base_url}/model/{encoded_model_id}/count-tokens" return endpoint diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index f028503c6a..ec20d76102 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -5,6 +5,7 @@ from urllib.parse import urlparse import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( @@ -209,7 +210,10 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): if isinstance(query, list): query = " ".join(query) - url = f"{api_base}/{vector_store_id}/retrieve" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}/retrieve" request_body: Dict[str, Any] = { "retrievalQuery": BedrockKBRetrievalQuery(text=query), diff --git a/litellm/llms/black_forest_labs/image_edit/handler.py b/litellm/llms/black_forest_labs/image_edit/handler.py index dea2683a04..f5784e0836 100644 --- a/litellm/llms/black_forest_labs/image_edit/handler.py +++ b/litellm/llms/black_forest_labs/image_edit/handler.py @@ -15,6 +15,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -331,6 +332,17 @@ class BlackForestLabsImageEdit: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} @@ -416,6 +428,17 @@ class BlackForestLabsImageEdit: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} diff --git a/litellm/llms/black_forest_labs/image_generation/handler.py b/litellm/llms/black_forest_labs/image_generation/handler.py index 5a1d885e52..8af4a236fd 100644 --- a/litellm/llms/black_forest_labs/image_generation/handler.py +++ b/litellm/llms/black_forest_labs/image_generation/handler.py @@ -15,6 +15,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -317,6 +318,17 @@ class BlackForestLabsImageGeneration: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} @@ -402,6 +414,17 @@ class BlackForestLabsImageGeneration: message="No polling_url in BFL response", ) + # Reject cross-origin polling URLs — the ``x-key`` auth header + # would otherwise leak to whatever URL the upstream returns. + # VERIA-51. + try: + assert_same_origin(polling_url, str(initial_response.request.url)) + except SSRFError as ssrf_err: + raise BlackForestLabsError( + status_code=502, + message=f"Rejected polling URL: {ssrf_err}", + ) + # Get just the auth header for polling polling_headers = {"x-key": headers.get("x-key", "")} diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index a72f732a30..5b08670f9f 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import httpx +from litellm.litellm_core_utils.url_utils import encode_url_path_segments from litellm.litellm_core_utils.exception_mapping_utils import exception_type from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -149,7 +150,8 @@ class BytezChatConfig(BaseConfig): litellm_params: dict, stream: Optional[bool] = None, ) -> str: - return f"{API_BASE}/{model}" + encoded_model = encode_url_path_segments(model, field_name="model") + return f"{API_BASE}/{encoded_model}" def transform_request( self, diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 9e59782bf7..b9e219f5cb 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -5,6 +5,7 @@ from typing import AsyncIterator, Iterator, List, Optional, Union import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segments from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import ( BaseConfig, @@ -89,7 +90,8 @@ class CloudflareChatConfig(BaseConfig): api_base = ( f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/" ) - return api_base + model + encoded_model = encode_url_path_segments(model, field_name="model") + return api_base + encoded_model def get_supported_openai_params(self, model: str) -> List[str]: return [ diff --git a/litellm/llms/custom_httpx/container_handler.py b/litellm/llms/custom_httpx/container_handler.py index afdd7bc6a8..599cd705eb 100644 --- a/litellm/llms/custom_httpx/container_handler.py +++ b/litellm/llms/custom_httpx/container_handler.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Coroutine, Dict, Optional, Type, Union import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -72,7 +73,8 @@ def _build_url( # Substitute path parameters for param, value in path_params.items(): - path_template = path_template.replace(f"{{{param}}}", value) + encoded_value = encode_url_path_segment(value, field_name=param) + path_template = path_template.replace(f"{{{param}}}", encoded_value) # Parse the api_base to extract existing query params parsed_base = httpx.URL(api_base) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index dc625918b9..0c4816fcda 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -26,6 +26,7 @@ from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -8948,7 +8949,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", @@ -9015,7 +9019,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", @@ -9214,7 +9221,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" request_body: Dict[str, Any] = dict(vector_store_update_optional_params) @@ -9297,7 +9307,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" request_body: Dict[str, Any] = dict(vector_store_update_optional_params) @@ -9363,7 +9376,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", @@ -9428,7 +9444,10 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url = f"{api_base}/{vector_store_id}" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}" logging_obj.pre_call( input="", diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 4dac2b8ba9..6a59911701 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -11,13 +11,14 @@ import httpx from httpx import Headers import litellm -from litellm.types.utils import all_litellm_params +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, ) from litellm.secret_managers.main import get_secret_str +from litellm.types.utils import all_litellm_params from ..common_utils import ElevenLabsException @@ -321,7 +322,8 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): "ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`." ) - url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{voice_id}" + encoded_voice_id = encode_url_path_segment(voice_id, field_name="voice_id") + url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}" query_params = litellm_params.get(self.ELEVENLABS_QUERY_PARAMS_KEY, {}) if query_params: diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 401d7bb9f4..63a383ebd3 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -12,6 +12,7 @@ import httpx from openai.types.file_deleted import FileDeleted from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, @@ -258,10 +259,14 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): normalized_file_id = file_id normalized_file_id = normalized_file_id.strip("/") - if not normalized_file_id.startswith("files/"): - normalized_file_id = f"files/{normalized_file_id}" + if normalized_file_id.startswith("files/"): + normalized_file_id = normalized_file_id.removeprefix("files/") - return normalized_file_id + encoded_file_id = encode_url_path_segment( + normalized_file_id, field_name="file_id" + ) + + return f"files/{encoded_file_id}" def transform_retrieve_file_response( self, @@ -337,13 +342,8 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not api_key: raise ValueError("api_key is required") - # Extract file name from URI if full URI is provided - # file_id could be "files/abc123" or "https://generativelanguage.googleapis.com/v1beta/files/abc123" - if file_id.startswith("http"): - # Extract the file path from full URI - file_name = file_id.split("/v1beta/")[-1] - else: - file_name = file_id if file_id.startswith("files/") else f"files/{file_id}" + # Normalize and encode the file name before interpolating it into the URL. + file_name = self._normalize_gemini_file_id(file_id) # Construct the delete URL url = f"{api_base}/v1beta/{file_name}" diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index c34da83cb8..593cbf7c2c 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -15,6 +15,7 @@ import httpx from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.interactions import ( @@ -205,8 +206,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") + encoded_interaction_id = encode_url_path_segment( + interaction_id, field_name="interaction_id" + ) return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", {}, ) @@ -238,8 +242,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") + encoded_interaction_id = encode_url_path_segment( + interaction_id, field_name="interaction_id" + ) return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", {}, ) @@ -268,8 +275,11 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): resolved_api_base = GeminiModelInfo.get_api_base(api_base) if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") + encoded_interaction_id = encode_url_path_segment( + interaction_id, field_name="interaction_id" + ) return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel", + f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}:cancel", {}, ) diff --git a/litellm/llms/hosted_vllm/embedding/README.md b/litellm/llms/hosted_vllm/embedding/README.md index f82b3c77a6..2c58e16fc2 100644 --- a/litellm/llms/hosted_vllm/embedding/README.md +++ b/litellm/llms/hosted_vllm/embedding/README.md @@ -2,4 +2,15 @@ No transformation is required for hosted_vllm embedding. VLLM is a superset of OpenAI's `embedding` endpoint. -To pass provider-specific parameters, see [this](https://docs.litellm.ai/docs/completion/provider_specific_params) \ No newline at end of file +## `encoding_format` + +For OpenAI-compatible embedding calls (including `openai/...` with a custom `api_base` pointing at vLLM), LiteLLM resolves `encoding_format` when it is not set on the request: + +1. Explicit value on the embedding call (`encoding_format=...`). +2. Model config (`litellm_params.encoding_format` on the proxy `model_list` entry). +3. Environment variable `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT` (e.g. in `.env` or container env). +4. Default **`float`**. + +That avoids forwarding `encoding_format=None` to the provider/SDK where some servers behave poorly. + +To pass provider-specific parameters, see [provider-specific params](https://docs.litellm.ai/docs/completion/provider_specific_params). \ No newline at end of file diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index 3381a5327e..3416616139 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -18,6 +18,7 @@ from openai.types.file_deleted import FileDeleted import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( @@ -306,7 +307,8 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_retrieve_file_response( self, @@ -336,7 +338,8 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_delete_file_response( self, @@ -422,7 +425,8 @@ class ManusFilesConfig(BaseFilesConfig): optional_params=optional_params, litellm_params=litellm_params, ) - return f"{api_base}/{file_id}/content", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}/content", {} def transform_file_content_response( self, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index 510c41304a..b3a0073a5c 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -6,6 +6,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -270,7 +271,10 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): Reference: https://open.manus.im/docs/openai-compatibility """ - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 955b9f760d..7f874ffd3b 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -6,6 +6,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, @@ -198,7 +199,10 @@ class OpenAIContainerConfig(BaseContainerConfig): ) -> Tuple[str, Dict]: """Transform the OpenAI container retrieve request.""" # For container retrieve, we just need to construct the URL - url = join_container_api_base_path(api_base, f"/{container_id}") + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No additional data needed for GET request data: Dict[str, Any] = {} @@ -230,7 +234,10 @@ class OpenAIContainerConfig(BaseContainerConfig): - DELETE /v1/containers/{container_id} """ # Construct the URL for container delete - url = join_container_api_base_path(api_base, f"/{container_id}") + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}") # No data needed for DELETE request data: Dict[str, Any] = {} @@ -267,7 +274,10 @@ class OpenAIContainerConfig(BaseContainerConfig): - GET /v1/containers/{container_id}/files """ # Construct the URL for container files - url = join_container_api_base_path(api_base, f"/{container_id}/files") + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + url = join_container_api_base_path(api_base, f"/{encoded_container_id}/files") # Prepare query parameters params: Dict[str, Any] = {} @@ -311,8 +321,12 @@ class OpenAIContainerConfig(BaseContainerConfig): - GET /v1/containers/{container_id}/files/{file_id}/content """ # Construct the URL for container file content + encoded_container_id = encode_url_path_segment( + container_id, field_name="container_id" + ) + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") url = join_container_api_base_path( - api_base, f"/{container_id}/files/{file_id}/content" + api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content" ) # No query parameters needed diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index c24dbf8637..66537e56a6 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Tuple import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.evals.transformation import ( BaseEvalsAPIConfig, LiteLLMLoggingObj, @@ -76,7 +77,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base = "https://api.openai.com" if eval_id: - return f"{api_base}/v1/evals/{eval_id}" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + return f"{api_base}/v1/evals/{encoded_eval_id}" return f"{api_base}/v1/{endpoint}" def transform_create_eval_request( @@ -276,7 +278,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if litellm_params and litellm_params.api_base: api_base = litellm_params.api_base - url = f"{api_base}/v1/evals/{eval_id}/runs" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs" # Build request body request_body = {k: v for k, v in create_request.items() if v is not None} @@ -310,7 +313,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if litellm_params and litellm_params.api_base: api_base = litellm_params.api_base - url = f"{api_base}/v1/evals/{eval_id}/runs" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs" # Build query parameters query_params: Dict[str, Any] = {} @@ -350,7 +354,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict]: """Transform get run request for OpenAI""" - url = f"{api_base}/v1/evals/{eval_id}/runs/{run_id}" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" verbose_logger.debug("Get run request - URL: %s", url) @@ -376,7 +382,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict, Dict]: """Transform cancel run request for OpenAI""" - url = f"{api_base}/v1/evals/{eval_id}/runs/{run_id}/cancel" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel" # Empty body for cancel request request_body: Dict[str, Any] = {} @@ -405,7 +413,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> Tuple[str, Dict, Dict]: """Transform delete run request for OpenAI""" - url = f"{api_base}/v1/evals/{eval_id}/runs/{run_id}" + encoded_eval_id = encode_url_path_segment(eval_id, field_name="eval_id") + encoded_run_id = encode_url_path_segment(run_id, field_name="run_id") + url = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" # Empty body for delete request request_body: Dict[str, Any] = {} diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 87c502032c..b7d5340d8d 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -7,6 +7,7 @@ from pydantic import BaseModel, ValidationError import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -421,7 +422,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - DELETE /v1/responses/{response_id} """ - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -457,7 +461,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - GET /v1/responses/{response_id} """ - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -498,7 +505,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}/input_items" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/input_items" params: Dict[str, Any] = {} if after is not None: params["after"] = after @@ -540,7 +550,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): OpenAI API expects the following request - POST /v1/responses/{response_id}/cancel """ - url = f"{api_base}/{response_id}/cancel" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/cancel" data: Dict = {} return url, data diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index cd5f10251b..52202f57fd 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -3,6 +3,7 @@ from typing import Any, Dict, Optional, Tuple, cast import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) @@ -98,7 +99,10 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): or "https://api.openai.com/v1" ) base_url = base_url.rstrip("/") - return f"{base_url}/vector_stores/{vector_store_id}/files" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + return f"{base_url}/vector_stores/{encoded_vector_store_id}/files" def transform_create_vector_store_file_request( self, @@ -163,7 +167,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, api_base: str, ) -> Tuple[str, Dict[str, Any]]: - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_retrieve_vector_store_file_response( self, @@ -186,7 +191,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, api_base: str, ) -> Tuple[str, Dict[str, Any]]: - return f"{api_base}/{file_id}/content", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}/content", {} def transform_retrieve_vector_store_file_content_response( self, @@ -218,7 +224,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): payload["attributes"] = filtered_attributes else: payload.pop("attributes", None) - return f"{api_base}/{file_id}", payload + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", payload def transform_update_vector_store_file_response( self, @@ -241,7 +248,8 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): file_id: str, api_base: str, ) -> Tuple[str, Dict[str, Any]]: - return f"{api_base}/{file_id}", {} + encoded_file_id = encode_url_path_segment(file_id, field_name="file_id") + return f"{api_base}/{encoded_file_id}", {} def transform_delete_vector_store_file_response( self, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 2c11d13748..bd095a0a1b 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast import httpx import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -108,7 +109,10 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: - url = f"{api_base}/{vector_store_id}/search" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( query=query, filters=vector_store_search_optional_params.get("filters", None), diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 61baa56949..2d165a7d7d 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -1,11 +1,13 @@ import mimetypes from io import BufferedReader, BytesIO from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from urllib.parse import quote import httpx from httpx._types import RequestFiles import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils from litellm.secret_managers.main import get_secret_str @@ -220,11 +222,18 @@ class OpenAIVideoConfig(BaseVideoConfig): - GET /v1/videos/{video_id}/content?variant=thumbnail """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for video content download - url = f"{api_base.rstrip('/')}/{original_video_id}/content" + url = f"{api_base.rstrip('/')}/{encoded_video_id}/content" if variant is not None: - url = f"{url}?variant={variant}" + # Encode the user-controlled ``variant`` so a value like + # ``thumbnail&extra=1`` cannot inject additional query params + # into the upstream request — same hardening rationale as the + # path-segment encoding above. + url = f"{url}?variant={quote(variant, safe='')}" # No additional data needed for GET content request data: Dict[str, Any] = {} @@ -247,9 +256,12 @@ class OpenAIVideoConfig(BaseVideoConfig): - POST /v1/videos/{video_id}/remix """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for video remix - url = f"{api_base.rstrip('/')}/{original_video_id}/remix" + url = f"{api_base.rstrip('/')}/{encoded_video_id}/remix" # Prepare the request data data = {"prompt": prompt} @@ -391,9 +403,12 @@ class OpenAIVideoConfig(BaseVideoConfig): - DELETE /v1/videos/{video_id} """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for video delete - url = f"{api_base.rstrip('/')}/{original_video_id}" + url = f"{api_base.rstrip('/')}/{encoded_video_id}" # No data needed for DELETE request data: Dict[str, Any] = {} @@ -427,9 +442,12 @@ class OpenAIVideoConfig(BaseVideoConfig): """ # Extract the original video_id (remove provider encoding if present) original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # For video retrieve, we just need to construct the URL - url = f"{api_base.rstrip('/')}/{original_video_id}" + url = f"{api_base.rstrip('/')}/{encoded_video_id}" # No additional data needed for GET request data: Dict[str, Any] = {} @@ -494,7 +512,11 @@ class OpenAIVideoConfig(BaseVideoConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base.rstrip('/')}/characters/{character_id}" + original_character_id = extract_original_character_id(character_id) + encoded_character_id = encode_url_path_segment( + original_character_id, field_name="character_id" + ) + url = f"{api_base.rstrip('/')}/characters/{encoded_character_id}" return url, {} def transform_video_get_character_response( diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 5dd1247001..b5e5aa4ea2 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -106,5 +106,13 @@ "base_url": "https://aihubmix.com/v1", "api_key_env": "AIHUBMIX_API_KEY", "api_base_env": "AIHUBMIX_API_BASE" + }, + "crusoe": { + "base_url": "https://managed-inference-api-proxy.crusoecloud.com/v1", + "api_key_env": "CRUSOE_API_KEY", + "api_base_env": "CRUSOE_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + } } } diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 7b22edd867..fc4cfc7b08 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -1,5 +1,6 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -82,7 +83,10 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): litellm_params: dict, extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: - url = f"{api_base}/{vector_store_id}/search" + encoded_vector_store_id = encode_url_path_segment( + vector_store_id, field_name="vector_store_id" + ) + url = f"{api_base}/{encoded_vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py index d49a5fd370..990fc2b2e6 100644 --- a/litellm/llms/ragflow/chat/transformation.py +++ b/litellm/llms/ragflow/chat/transformation.py @@ -13,6 +13,7 @@ Model name format: from typing import List, Optional, Tuple import litellm +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.openai.openai import OpenAIConfig from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -126,10 +127,11 @@ class RAGFlowConfig(OpenAIConfig): api_base = api_base[:-3] # Remove /v1 # Construct the RAGFlow-specific path + encoded_entity_id = encode_url_path_segment(entity_id, field_name="entity_id") if endpoint_type == "chat": - path = f"/api/v1/chats_openai/{entity_id}/chat/completions" + path = f"/api/v1/chats_openai/{encoded_entity_id}/chat/completions" else: # agent - path = f"/api/v1/agents_openai/{entity_id}/chat/completions" + path = f"/api/v1/agents_openai/{encoded_entity_id}/chat/completions" # Ensure path starts with / if not path.startswith("/"): diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 8377dea952..4f84816a2b 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -6,6 +6,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -334,9 +335,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): We'll retrieve the task and extract the video URL. """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Get task status to retrieve video URL - url = f"{api_base}/tasks/{original_video_id}" + url = f"{api_base}/tasks/{encoded_video_id}" params: Dict[str, Any] = {} @@ -495,9 +499,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): RunwayML uses task cancellation. """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the URL for task cancellation - url = f"{api_base}/tasks/{original_video_id}/cancel" + url = f"{api_base}/tasks/{encoded_video_id}/cancel" data: Dict[str, Any] = {} @@ -533,9 +540,12 @@ class RunwayMLVideoConfig(BaseVideoConfig): RunwayML uses GET /v1/tasks/{task_id} to retrieve task status. """ original_video_id = extract_original_video_id(video_id) + encoded_video_id = encode_url_path_segment( + original_video_id, field_name="video_id" + ) # Construct the full URL for task status retrieval - url = f"{api_base}/tasks/{original_video_id}" + url = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) data: Dict[str, Any] = {} diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7436bfef58..c627599da8 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -4,7 +4,11 @@ from typing import Any, Coroutine, Dict, Optional, Union import httpx import litellm -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import ( + async_safe_get, + encode_url_path_segment, + safe_get, +) from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -170,7 +174,8 @@ class VertexAIBatchPrediction(VertexLLM): ) # Append batch_id to the URL - default_api_base = f"{default_api_base}/{batch_id}" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + default_api_base = f"{default_api_base}/{encoded_batch_id}" if len(default_api_base.split(":")) > 1: endpoint = default_api_base.split(":")[-1] @@ -413,7 +418,8 @@ class VertexAIBatchPrediction(VertexLLM): vertex_project=vertex_project or project_id, ) - retrieve_api_base_default = f"{default_api_base}/{batch_id}" + encoded_batch_id = encode_url_path_segment(batch_id, field_name="batch_id") + retrieve_api_base_default = f"{default_api_base}/{encoded_batch_id}" cancel_api_base_default = f"{retrieve_api_base_default}:cancel" _, api_base = self._check_custom_proxy( diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 87bd484382..9afa5dec46 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -212,6 +212,22 @@ def _process_gemini_media( return _apply_gemini_metadata( part, model, media_resolution_enum, video_metadata ) + elif image_url.startswith( + "https://generativelanguage.googleapis.com/v1beta/files/" + ): + # Gemini Files API URIs — the file is already uploaded to Google's + # servers; pass the URI through as file_data without fetching it. + # These URLs return 403 when accessed directly, so we must not try + # to resolve their MIME type via HTTP. + if format: + file_data = FileDataType(mime_type=format, file_uri=image_url) + else: + # Gemini Files API references can be passed through as URI-only. + file_data = cast(FileDataType, {"file_uri": image_url}) + part = {"file_data": file_data} + return _apply_gemini_metadata( + part, model, media_resolution_enum, video_metadata + ) elif ( "https://" in image_url and (image_type := format or _get_image_mime_type_from_url(image_url)) @@ -743,16 +759,22 @@ def _transform_request_body( # noqa: PLR0915 ] data = RequestBody(contents=content) - if system_instructions is not None: - data["system_instruction"] = system_instructions - if tools is not None: - data["tools"] = tools - if tool_choice is not None: - data["toolConfig"] = tool_choice - if include_server_side_tool_invocations: - if "toolConfig" not in data: - data["toolConfig"] = {} - data["toolConfig"]["includeServerSideToolInvocations"] = True + # Vertex rejects system_instruction/tools/toolConfig alongside cachedContent. + # Treat dropping these fields as a request mutation guarded by modify_params. + can_send_cache_incompatible_fields = ( + cached_content is None or litellm.modify_params is False + ) + if can_send_cache_incompatible_fields: + if system_instructions is not None: + data["system_instruction"] = system_instructions + if tools is not None: + data["tools"] = tools + if tool_choice is not None: + data["toolConfig"] = tool_choice + if include_server_side_tool_invocations: + if "toolConfig" not in data: + data["toolConfig"] = {} + data["toolConfig"]["includeServerSideToolInvocations"] = True if safety_settings is not None: data["safetySettings"] = safety_settings if generation_config is not None and len(generation_config) > 0: 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/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 6cb7a86bea..61fb848b40 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import httpx from litellm import get_model_info +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.router import GenericLiteLLMParams @@ -91,12 +92,18 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): raise ValueError("vector_store_id is required") if api_base: return api_base.rstrip("/") + encoded_collection_id = encode_url_path_segment( + collection_id, field_name="vertex_collection_id" + ) + encoded_datastore_id = encode_url_path_segment( + datastore_id, field_name="vector_store_id" + ) # Vertex AI Search API endpoint for search return ( f"https://discoveryengine.googleapis.com/v1/" f"projects/{vertex_project}/locations/{vertex_location}/" - f"collections/{collection_id}/dataStores/{datastore_id}/servingConfigs/default_config" + f"collections/{encoded_collection_id}/dataStores/{encoded_datastore_id}/servingConfigs/default_config" ) def transform_search_vector_store_request( diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index f6dda4dd25..99e0a958ef 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -17,6 +17,7 @@ from pydantic import fields as pyd_fields import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( _safe_convert_created_field, ) @@ -300,7 +301,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -333,7 +337,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}" data: Dict = {} return url, data @@ -372,7 +379,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): limit: int = 20, order: Literal["asc", "desc"] = "desc", ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}/input_items" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/input_items" params: Dict[str, Any] = {} if after is not None: params["after"] = after @@ -408,7 +418,10 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> Tuple[str, Dict]: - url = f"{api_base}/{response_id}/cancel" + encoded_response_id = encode_url_path_segment( + response_id, field_name="response_id" + ) + url = f"{api_base}/{encoded_response_id}/cancel" data: Dict = {} return url, data diff --git a/litellm/main.py b/litellm/main.py index 0079bd750c..0553cf9d42 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4923,8 +4923,17 @@ def embedding( # noqa: PLR0915 if encoding_format is not None: optional_params["encoding_format"] = encoding_format else: - # Omiting causes openai sdk to add default value of "float" - optional_params["encoding_format"] = None + env_fmt = get_secret_str("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT") + if env_fmt is not None and env_fmt.strip().lower() == "none": + optional_params.pop("encoding_format", None) + else: + _default_fmt = ( + optional_params.get("encoding_format") or env_fmt or "float" + ) + if _default_fmt.strip().lower() == "none": + optional_params.pop("encoding_format", None) + else: + optional_params["encoding_format"] = _default_fmt api_version = None diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a1e3e42a9c..6078d7e690 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -22061,6 +22061,98 @@ "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, + "crusoe/deepseek-ai/DeepSeek-R1-0528": { + "input_cost_per_token": 3e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 7e-06, + "supports_function_calling": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": false + }, + "crusoe/deepseek-ai/DeepSeek-V3-0324": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "crusoe/google/gemma-3-12b-it": { + "input_cost_per_token": 1e-07, + "litellm_provider": "crusoe", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "crusoe/meta-llama/Llama-3.3-70B-Instruct": { + "input_cost_per_token": 2e-07, + "litellm_provider": "crusoe", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "crusoe/moonshotai/Kimi-K2-Thinking": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": false + }, + "crusoe/openai/gpt-oss-120b": { + "input_cost_per_token": 8e-07, + "litellm_provider": "crusoe", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507": { + "input_cost_per_token": 3e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "lambda_ai/deepseek-llama3.3-70b": { "input_cost_per_token": 2e-07, "litellm_provider": "lambda_ai", diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index cebd224a1a..1794cd1438 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -33,10 +33,12 @@ def get_request_base_url(request: Request) -> str: """ Get the base URL for the request, considering X-Forwarded-* headers. - When behind a proxy (like nginx), the proxy may set: - - X-Forwarded-Proto: The original protocol (http/https) - - X-Forwarded-Host: The original host (may include port) - - X-Forwarded-Port: The original port (if not in Host header) + X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured + when the request comes from a configured trusted proxy + (``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``). + Otherwise the request's literal ``base_url`` is returned, so an + untrusted caller cannot poison OAuth-discovery / redirect_uri values + by injecting headers. Args: request: FastAPI Request object @@ -47,34 +49,28 @@ def get_request_base_url(request: Request) -> str: base_url = str(request.base_url).rstrip("/") parsed = urlparse(base_url) - # Get forwarded headers + if not IPAddressUtils.is_request_from_trusted_proxy(request): + return base_url + x_forwarded_proto = request.headers.get("X-Forwarded-Proto") x_forwarded_host = request.headers.get("X-Forwarded-Host") x_forwarded_port = request.headers.get("X-Forwarded-Port") - # Start with the original scheme scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme - # Handle host and port if x_forwarded_host: # X-Forwarded-Host may already include port (e.g., "example.com:8080") if ":" in x_forwarded_host and not x_forwarded_host.startswith("["): - # Host includes port netloc = x_forwarded_host elif x_forwarded_port: - # Port is separate netloc = f"{x_forwarded_host}:{x_forwarded_port}" else: - # Just host, no explicit port netloc = x_forwarded_host else: - # No X-Forwarded-Host, use original netloc netloc = parsed.netloc if x_forwarded_port and ":" not in netloc: - # Add forwarded port if not already in netloc netloc = f"{netloc}:{x_forwarded_port}" - # Reconstruct the URL return urlunparse((scheme, netloc, parsed.path, "", "", "")) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index abb4b5cfa6..54d9bbe6e2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2138,6 +2138,47 @@ if MCP_AVAILABLE: ######################################################### local_tool = global_mcp_tool_registry.get_tool(name) if local_tool: + # OpenAPI-backed tools used to bypass `pre_call_tool_check` — + # only the managed path ran allowed/banned-tool checks, key/team + # tool permissions, and parameter validation. Run the same checks + # before dispatching to the local registry. Refuse the call if + # we cannot resolve a server: tools registered via + # openapi_to_mcp_generator are always tied to a server, so a + # missing mcp_server here means the tool->server mapping has + # not finished initializing or the registry entry is orphaned. + # Skipping the check would re-open the same authorization gap. + if mcp_server is None: + raise HTTPException( + status_code=503, + detail=( + f"MCP server for tool '{name}' is not available; " + "refusing to dispatch without authorization checks. " + "Retry once the server is registered." + ), + ) + + # `pre_call_tool_check` calls into `proxy_logging_obj` for the + # pre-call guardrail hooks, so source it from the canonical + # `proxy_server` module the same way `_handle_managed_mcp_tool` + # does. `kwargs.get("proxy_logging_obj")` is None on the MCP + # entry path and would crash with AttributeError after the + # security checks pass. + from litellm.proxy.proxy_server import proxy_logging_obj + + hook_result = await global_mcp_server_manager.pre_call_tool_check( + name=original_tool_name, + arguments=arguments or {}, + server_name=server_name or mcp_server.name, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, + server=mcp_server, + raw_headers=raw_headers, + ) + # `pre_call_tool_check` may return guardrail-modified + # arguments; honor them on the local path too. + if isinstance(hook_result, dict) and "arguments" in hook_result: + arguments = hook_result["arguments"] + verbose_logger.debug(f"Executing local registry tool: {name}") # For BYOK servers the credential must be injected via a ContextVar # because the tool function has headers baked into its closure. 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 8331f748c6..eb35dd6cb3 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3572,7 +3572,7 @@ "/anthropic/{endpoint}": { "delete": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3616,7 +3616,7 @@ }, "get": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3660,7 +3660,7 @@ }, "patch": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3704,7 +3704,7 @@ }, "post": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3748,7 +3748,7 @@ }, "put": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -13260,7 +13260,7 @@ "/langfuse/{endpoint}": { "delete": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13299,7 +13299,7 @@ }, "get": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13338,7 +13338,7 @@ }, "patch": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13377,7 +13377,7 @@ }, "post": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13416,7 +13416,7 @@ }, "put": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -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.", @@ -26883,7 +26883,7 @@ "/toolset/{toolset_name}/mcp": { "delete": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -26922,7 +26922,7 @@ }, "get": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -26961,7 +26961,7 @@ }, "head": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27000,7 +27000,7 @@ }, "options": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27039,7 +27039,7 @@ }, "patch": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27078,7 +27078,7 @@ }, "post": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27117,7 +27117,7 @@ }, "put": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -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 315f6a9742..c63ff8d073 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -8,11 +8,35 @@ any drift as a neutral check. """ import json +import re import sys from pathlib import Path -from typing import Dict, 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]]: @@ -25,13 +49,46 @@ def load_snapshot() -> Optional[Dict[str, Dict]]: return None +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 + + methods = {method for method in path_ops if method in HTTP_METHOD_SUFFIXES} + if not methods: + continue + + for method, operation in path_ops.items(): + if method not in HTTP_METHOD_SUFFIXES 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]: import importlib 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: @@ -43,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 @@ -51,14 +109,26 @@ def generate_snapshot() -> Dict[str, Dict]: ] if not feat_routes: continue + _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 10ffb7403f..fceaa331ea 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -667,6 +667,8 @@ class LiteLLMRoutes(enum.Enum): "/models/{model_id}", "/guardrails/list", "/v2/guardrails/list", + "/project/list", + "/project/info", ] + spend_tracking_routes + key_management_routes @@ -691,6 +693,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", @@ -718,21 +723,73 @@ class LiteLLMRoutes(enum.Enum): "/organization/member_delete", ] - # Routes accessible by Admin Viewer (read-only admin access) - admin_viewer_routes = [ - "/user/list", - "/user/available_users", - "/user/available_roles", - "/user/daily/activity", - "/team/daily/activity", - "/tag/daily/activity", - "/tag/list", - "/audit", - "/audit/{id}", - "/global/activity", - "/global/activity/model", - "/global/activity/cache_hits", - ] + info_routes + # Routes accessible by Admin Viewer (read-only admin access). + # + # Admin Viewer follows a read-parity-with-Proxy-Admin rule: anything Proxy + # Admin can read/list/get, Admin Viewer can too (no writes, no cost-incurring + # actions). + # + # NOTE: This list is no longer the primary mechanism for granting access — + # `_check_proxy_admin_viewer_access()` in route_checks.py default-allows + # any safe HTTP method (GET/HEAD/OPTIONS) on non-inference routes. This + # list now matters only for non-GET routes that are semantically reads + # (e.g. POST /spend/calculate). Adding a new GET endpoint does not require + # updating this list — the default-allow behavior covers it automatically. + admin_viewer_routes = ( + [ + "/user/list", + "/user/available_users", + "/user/available_roles", + "/user/daily/activity", + "/team/daily/activity", + "/tag/daily/activity", + "/tag/list", + "/audit", + "/audit/{id}", + "/global/activity", + "/global/activity/model", + "/global/activity/cache_hits", + # Customer / end-user listing (handlers already gate on + # PROXY_ADMIN_VIEW_ONLY — the route gate must match). + "/customer/list", + "/customer/info", + # UI Logs page detail drawer (single + session). The list endpoint + # `/spend/logs/ui` is covered via spend_tracking_routes below. + "/spend/logs/ui/{logId}", + "/spend/logs/session/ui", + # Settings / observability read endpoints exposed in admin-only + # sidebar groups (Logging & Alerts, Admin Settings, Budgets, + # Invitations). + "/callbacks/list", + "/callbacks/configs", + "/get/config/callbacks", + "/alerting/settings", + "/config/list", + "/config/field/info", + "/budget/list", + "/budget/settings", + # Invitation viewing (admin viewer cannot create/delete; can read). + "/invitation/info", + # Guardrails / Policies pages (read-only views). + "/guardrails/list", + "/v2/guardrails/list", + "/guardrails/submissions", + "/guardrails/submissions/{guardrail_id}", + "/guardrails/usage/overview", + "/policies/attachments/list", + # MCP semantic filter settings (read). + "/get/mcp_semantic_filter_settings", + # Model cost map maintenance views (read-only status / source). + "/schedule/model_cost_map_reload/status", + "/model/cost_map/source", + ] + # Spend tracking reads (/spend/logs, /spend/logs/ui, /spend/keys, + # /spend/users, /spend/tags, /spend/calculate, /cost/estimate). Admin + # Viewer can already read /global/spend/* via global_spend_tracking_routes; + # the per-tenant /spend/* views were the missing peer. + + spend_tracking_routes + + info_routes + ) # All routes accesible by an Org Admin org_admin_allowed_routes = ( @@ -2380,6 +2437,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For headers are only trusted from these IPs.", ) + trusted_proxy_ranges: Optional[List[str]] = Field( + None, + description="CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler.", + ) store_model_in_db: Optional[bool] = Field( None, description="If True, models and config are stored in and loaded from the database. Default is False.", @@ -2573,6 +2634,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 113a8f538c..754488367e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -60,6 +60,10 @@ 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, @@ -486,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 @@ -495,23 +502,28 @@ async def common_checks( # noqa: PLR0915 f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin." ) - # 2. If team can call model + # 2. If team can call model (or key's access_group_ids grant it) if _model and team_object: with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"): - if not await can_team_access_model( - model=_model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=( - valid_token.team_model_aliases if valid_token else None - ), - ): - raise ProxyException( - message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=status.HTTP_401_UNAUTHORIZED, + try: + await can_team_access_model( + model=_model, + team_object=team_object, + llm_router=llm_router, + team_model_aliases=( + valid_token.team_model_aliases if valid_token else None + ), ) + except ProxyException as team_denial: + if team_denial.type != ProxyErrorTypes.team_model_access_denied: + raise + if not await _key_access_group_grants_model( + model=_model, + valid_token=valid_token, + team_object=team_object, + llm_router=llm_router, + ): + raise # 2.2. If team member has per-member model scope, enforce it if _model and team_object and valid_token and valid_token.user_id: @@ -656,13 +668,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) @@ -1012,7 +1018,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: @@ -1033,11 +1039,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}", ) @@ -1091,7 +1106,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 @@ -1124,7 +1139,7 @@ async def get_end_user_object( ) # 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 @@ -1616,9 +1631,12 @@ 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, @@ -2348,7 +2366,7 @@ async def get_key_object( model_type=UserAPIKeyAuth, ) if user_api_key_auth is not None: - return user_api_key_auth + return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth) if check_cache_only: raise Exception( @@ -2401,6 +2419,16 @@ 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, @@ -2952,6 +2980,77 @@ async def can_team_access_model( raise +async def _key_access_group_grants_model( + model: Union[str, List[str]], + valid_token: Optional[UserAPIKeyAuth], + team_object: Optional[LiteLLM_TeamTable], + llm_router: Optional[Router], +) -> bool: + """ + Returns True if the key's `access_group_ids` expand to models that grant + access to `model`. Used to let a key's access group override a team's + model restriction in `common_checks`. + + A key's access group only counts if the access group itself authorizes the + caller as an owner — that is, the group's `assigned_team_ids` includes the + key's `team_id`, or the group's `assigned_key_ids` includes the key's + token. This preserves the team-as-owner boundary (a team member cannot + escalate by naming a group assigned to a different team) while still + letting a group reach the key without first being added to the team's + `access_group_ids` list. + """ + if valid_token is None: + return False + key_access_group_ids = list(valid_token.access_group_ids or []) + if not key_access_group_ids: + return False + + from litellm.proxy.proxy_server import prisma_client as _prisma_client + from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache + + if _prisma_client is None or _user_api_key_cache is None: + return False + + key_team_id = valid_token.team_id or ( + team_object.team_id if team_object is not None else None + ) + key_token = valid_token.token + + authorized_models: List[str] = [] + for ag_id in key_access_group_ids: + try: + ag = await get_access_object( + access_group_id=ag_id, + prisma_client=_prisma_client, + user_api_key_cache=_user_api_key_cache, + proxy_logging_obj=_proxy_logging_obj, + ) + except Exception: + continue + team_authorized = bool( + key_team_id and key_team_id in (ag.assigned_team_ids or []) + ) + key_authorized = bool(key_token and key_token in (ag.assigned_key_ids or [])) + if team_authorized or key_authorized: + authorized_models.extend(ag.access_model_names or []) + + if not authorized_models: + return False + try: + _can_object_call_model( + model=model, + llm_router=llm_router, + models=list(set(authorized_models)), + team_model_aliases=valid_token.team_model_aliases, + team_id=valid_token.team_id, + object_type="key", + ) + return True + except ProxyException: + return False + + def can_project_access_model( model: Union[str, List[str]], project_object: LiteLLM_ProjectTableCachedObj, @@ -3967,13 +4066,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/ip_address_utils.py b/litellm/proxy/auth/ip_address_utils.py index 34fab4849e..39d3282942 100644 --- a/litellm/proxy/auth/ip_address_utils.py +++ b/litellm/proxy/auth/ip_address_utils.py @@ -13,6 +13,10 @@ from fastapi import Request from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.auth_utils import _get_request_ip_address +# One-shot warning so operators upgrading from the prior "always trust X-Forwarded-*" +# behaviour see an actionable message in their logs the first time it triggers. +_warned_xff_without_trusted_ranges = False + class IPAddressUtils: """Static utilities for IP-based MCP access control.""" @@ -106,6 +110,61 @@ class IPAddressUtils: return any(addr in network for network in networks) + @staticmethod + def is_request_from_trusted_proxy( + request: Request, + general_settings: Optional[Dict[str, Any]] = None, + ) -> bool: + """ + Return True if X-Forwarded-* headers on this request should be trusted. + + Trusts the headers iff both: + 1. ``use_x_forwarded_for`` is enabled in proxy settings, AND + 2. ``mcp_trusted_proxy_ranges`` is configured AND the direct + connection IP (``request.client.host``) falls inside one of + those CIDRs. + + When ``use_x_forwarded_for`` is enabled but ``mcp_trusted_proxy_ranges`` + is missing, the headers are NOT trusted: there is no way to + distinguish a trusted reverse proxy from a direct attacker, so callers + that build URLs (OAuth issuer / redirect_uri / etc.) must fall back + to the request's literal base URL instead of risking a poisoned host. + """ + if general_settings is None: + try: + from litellm.proxy.proxy_server import ( + general_settings as proxy_general_settings, + ) + + general_settings = proxy_general_settings + except ImportError: + general_settings = {} + + if general_settings is None: + general_settings = {} + + if not general_settings.get("use_x_forwarded_for", False): + return False + + trusted_ranges = general_settings.get("mcp_trusted_proxy_ranges") + if not trusted_ranges: + global _warned_xff_without_trusted_ranges + if not _warned_xff_without_trusted_ranges: + verbose_proxy_logger.warning( + "use_x_forwarded_for is enabled but mcp_trusted_proxy_ranges " + "is not configured. X-Forwarded-* headers will NOT be " + "trusted, so MCP OAuth discovery URLs will use the proxy's " + "literal base URL. Set mcp_trusted_proxy_ranges in " + "general_settings to your reverse-proxy CIDR(s) to allow " + "X-Forwarded-* through." + ) + _warned_xff_without_trusted_ranges = True + return False + + direct_ip = request.client.host if request.client else None + trusted_networks = IPAddressUtils.parse_trusted_proxy_networks(trusted_ranges) + return IPAddressUtils.is_trusted_proxy(direct_ip, trusted_networks) + @staticmethod def get_mcp_client_ip( request: Request, diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index 0dc696bc45..9fc4c4fb53 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -1,19 +1,69 @@ -from typing import Any, Dict +from typing import Any, Dict, FrozenSet from fastapi import Request from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.trusted_proxy_utils import require_trusted_proxy_request + +# OAuth2-proxy header trust is for **identity assertion** from a trusted +# upstream auth proxy (oauth2-proxy, Authelia, etc.). The allowlist below +# is the only safe surface — anything else (``user_role``, ``api_key``, +# ``permissions``, ``max_budget``, ``user_max_budget``, +# ``team_tpm_limit``, ``end_user_max_budget``, ``allowed_model_region``, +# and dozens of similar policy fields scattered across the +# ``LiteLLM_VerificationTokenView`` hierarchy) is a privilege grant that +# would let a caller forge their own enforcement parameters by sending +# the matching header. +# +# A denylist of "privileged fields" is unmaintainable in this codebase: +# the auth model has ~50 budget/spend/limit/permission fields and gains +# more with each release. An allowlist scoped to identity assertion is +# default-secure — new fields are blocked automatically. +# +# Operators who need a trusted upstream to assert anything beyond +# identity should switch to JWT authentication, which validates a +# signature on the assertion rather than blindly trusting headers. +ALLOWED_OAUTH2_PROXY_FIELDS: FrozenSet[str] = frozenset( + { + "user_id", + "user_email", + "team_id", + "team_alias", + "org_id", + "models", + } +) async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: """ - Handle request from oauth2 proxy. + Resolve a ``UserAPIKeyAuth`` from request headers per the admin-set + ``oauth2_config_mappings``. + + The auth model assumes the proxy is deployed behind a trusted OAuth2 + reverse proxy that injects authenticated identity headers (e.g. + oauth2-proxy, Authelia). + + **Identity-only allowlist.** ``oauth2_config_mappings`` maps header + names to ``UserAPIKeyAuth`` fields. Without an allowlist, an admin + who maps the wrong header to ``user_role`` lets any caller send + ``X-User-Role: proxy_admin`` and gain full admin privileges + (Pydantic coerces the string into the enum). Only fields in + ``ALLOWED_OAUTH2_PROXY_FIELDS`` (identity assertion only — see the + constant's comment) may be mapped; any other mapping is rejected at + request time so the misconfiguration surfaces loudly rather than as + a silent privesc. """ from litellm.proxy.proxy_server import general_settings verbose_proxy_logger.debug("Handling oauth2 proxy request") - # Define the OAuth2 config mappings + require_trusted_proxy_request( + request=request, + general_settings=general_settings, + feature_name="OAuth2 proxy auth", + ) + oauth2_config_mappings: Dict[str, str] = ( general_settings.get("oauth2_config_mappings") or {} ) @@ -21,21 +71,32 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: if not oauth2_config_mappings: raise ValueError("Oauth2 config mappings not found in general_settings") - # Initialize a dictionary to store the mapped values - auth_data: Dict[str, Any] = {} - # Extract values from headers based on the mappings + disallowed = sorted( + set(oauth2_config_mappings.keys()) - ALLOWED_OAUTH2_PROXY_FIELDS + ) + if disallowed: + raise ValueError( + "Oauth2 proxy auth refuses to map non-identity UserAPIKeyAuth " + f"fields from request headers: {disallowed}. Only identity " + f"fields are accepted ({sorted(ALLOWED_OAUTH2_PROXY_FIELDS)}); " + "anything else (privileges, budgets, rate limits, metadata) " + "would let a caller forge enforcement parameters by spoofing " + "the matching header. If you need a trusted upstream to " + "assert anything beyond identity, use JWT auth " + "(signature-validated) instead of header-trust." + ) + + auth_data: Dict[str, Any] = {} for key, header in oauth2_config_mappings.items(): value = request.headers.get(header) - if value: - # Convert max_budget to float if present - if key == "max_budget": - auth_data[key] = float(value) - # Convert models to list if present - elif key == "models": - auth_data[key] = [model.strip() for model in value.split(",")] - else: - auth_data[key] = value + if not value: + continue + if key == "models": + auth_data[key] = [model.strip() for model in value.split(",")] + else: + auth_data[key] = value + verbose_proxy_logger.debug( "Auth data before creating UserAPIKeyAuth object: keys=%s", list(auth_data.keys()), @@ -45,5 +106,4 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: "UserAPIKeyAuth object created with keys: %s", list(user_api_key_auth.__fields_set__), ) - # Create and return UserAPIKeyAuth object return user_api_key_auth diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 6417307f69..dba29f8413 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -202,6 +202,7 @@ class RouteChecks: route=route, _user_role=_user_role, request_data=request_data, + request=request, ) elif ( _user_role == LitellmUserRoles.INTERNAL_USER.value @@ -596,14 +597,66 @@ class RouteChecks: return True return False + # HTTP methods that are intrinsically read-only and therefore safe to + # default-allow for PROXY_ADMIN_VIEW_ONLY. Anything else (POST/PUT/PATCH/ + # DELETE) is treated as a write attempt and goes through the explicit + # write-allowlist below. + _SAFE_HTTP_METHODS = frozenset({"GET", "HEAD", "OPTIONS"}) + + # Explicit write routes that PROXY_ADMIN_VIEW_ONLY must NEVER call. The + # role-principle is "no writes, ever" — the management_routes list is the + # authoritative source for which non-llm routes are writes; we just need + # to filter out the read endpoints (info / list) that share the prefix. + # A cleaner approach is to denylist by HTTP verb (POST/PUT/PATCH/DELETE); + # this block stays as a backstop in case a write is implemented as GET. + _ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset( + [ + "/user/new", + "/user/delete", + "/user/bulk_update", + "/team/new", + "/team/update", + "/team/delete", + "/model/new", + "/model/update", + "/model/delete", + "/key/generate", + "/key/delete", + "/key/update", + "/key/regenerate", + "/key/service-account/generate", + "/key/block", + "/key/unblock", + ] + ) + @staticmethod def _check_proxy_admin_viewer_access( route: str, _user_role: str, request_data: dict, + request: Optional[Request] = None, ) -> None: """ - Check access for PROXY_ADMIN_VIEW_ONLY role + Check access for PROXY_ADMIN_VIEW_ONLY role. + + Admin Viewer follows a read-parity-with-Proxy-Admin rule: anything Proxy + Admin can read/list/get, Admin Viewer can read/list/get. The only + exclusions are cost-incurring inference routes (Playground, /chat/ + completions, etc.) and any state-mutating request. + + Implementation: + 1. LLM/inference routes → 403 (cost-incurring). + 2. Safe HTTP method (GET/HEAD/OPTIONS) → allow by default. This is + the read-parity guarantee — every new GET endpoint added anywhere + in the codebase is automatically readable by Admin Viewer + without needing to remember to add it to an allowlist. + 3. Unsafe HTTP method (POST/PUT/PATCH/DELETE): + - Allow `/user/update` only when restricted to user_email/password. + - Block all explicit writes in `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. + - Otherwise allow only if the route is in admin_viewer_routes / + global_spend_tracking_routes (legacy explicit-allow set). + - Else 403. """ if RouteChecks.is_llm_api_route(route=route): raise HTTPException( @@ -611,65 +664,60 @@ class RouteChecks: detail=f"user not allowed to access this OpenAI routes, role= {_user_role}", ) - # Check if this is a write operation on management routes - if RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.management_routes.value - ): - # For management routes, only allow read operations or specific allowed updates - if route == "/user/update": - # Check the Request params are valid for PROXY_ADMIN_VIEW_ONLY - if request_data is not None and isinstance(request_data, dict): - _params_updated = request_data.keys() - for param in _params_updated: - if param not in ["user_email", "password"]: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated", - ) - elif ( - route - in [ - "/user/new", - "/user/delete", - "/user/bulk_update", - "/team/new", - "/team/update", - "/team/delete", - "/model/new", - "/model/update", - "/model/delete", - "/key/generate", - "/key/delete", - "/key/update", - "/key/regenerate", - "/key/service-account/generate", - "/key/block", - "/key/unblock", - ] - or route.startswith("/key/") - and route.endswith("/regenerate") - ): - # Block write operations for PROXY_ADMIN_VIEW_ONLY - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", - ) - # Allow read operations on management routes (like /user/info, /team/info, /model/info) + method = request.method.upper() if request is not None else "GET" + is_safe_method = method in RouteChecks._SAFE_HTTP_METHODS + + # ── Safe HTTP method: default-allow ────────────────────────────── + if is_safe_method: return - elif RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value - ): - # Allow access to admin viewer routes (read-only admin endpoints) + + # ── Unsafe HTTP method: explicit checks ────────────────────────── + # Allow `/user/update` for self-service email / password change. + if route == "/user/update": + if request_data is not None and isinstance(request_data, dict): + for param in request_data.keys(): + if param not in ["user_email", "password"]: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"user not allowed to access this route, role= {_user_role}. " + f"Trying to access: {route} and updating invalid param: {param}. " + "only user_email and password can be updated" + ), + ) return - elif RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value + + # Hard-block known write routes regardless of HTTP method (defensive + # — these are POSTs in practice, but pinning them here protects + # against future GET-shaped writes). + if route in RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES or ( + route.startswith("/key/") and route.endswith("/regenerate") ): - # Allow access to global spend tracking routes (read-only spend endpoints) - # proxy_admin_viewer role description: "view all keys, view all spend" - return - else: - # For other routes, block access raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", ) + + # Legacy explicit-allow sets (kept for routes that are POST but + # semantically read-only, e.g. /spend/calculate). Both admin_viewer_routes + # and global_spend_tracking_routes are reads/listings. + if RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.admin_viewer_routes.value + ): + return + if RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.global_spend_tracking_routes.value + ): + return + + # NOTE: We intentionally do NOT fall back to allowing all + # `management_routes`. That set is a mix of reads (info/list — handled + # via the safe-method branch above) and writes (`/team/block`, + # `/team/permissions_update`, `/jwt/key/mapping/{new,update,delete}`, + # `/key/bulk_update`, `/key/{id}/reset_spend`). A blanket allow would + # let Admin Viewer POST these write endpoints — violating the + # "no writes, ever" rule. Default-deny instead. + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}", + ) diff --git a/litellm/proxy/auth/trusted_proxy_utils.py b/litellm/proxy/auth/trusted_proxy_utils.py new file mode 100644 index 0000000000..df7b3080f2 --- /dev/null +++ b/litellm/proxy/auth/trusted_proxy_utils.py @@ -0,0 +1,118 @@ +import ipaddress +from typing import Any, Dict, List, Optional, Union + +from fastapi import Request + +from litellm._logging import verbose_proxy_logger + +TRUSTED_PROXY_RANGES_KEY = "trusted_proxy_ranges" +TrustedProxyNetwork = Union[ipaddress.IPv4Network, ipaddress.IPv6Network] + + +def _get_proxy_general_settings() -> Dict[str, Any]: + try: + from litellm.proxy.proxy_server import general_settings + + return general_settings or {} + except ImportError: + return {} + + +def _normalize_cidr_ranges(configured_ranges: Any, *, setting_name: str) -> List[str]: + if not configured_ranges: + return [] + if isinstance(configured_ranges, str): + return [ + raw_range.strip() + for raw_range in configured_ranges.split(",") + if raw_range.strip() + ] + if isinstance(configured_ranges, (list, tuple, set)): + return [ + str(raw_range).strip() + for raw_range in configured_ranges + if str(raw_range).strip() + ] + verbose_proxy_logger.warning( + "Invalid %s value: expected a list of CIDR ranges, got %s", + setting_name, + type(configured_ranges).__name__, + ) + return [] + + +def parse_trusted_proxy_ranges( + configured_ranges: Any, + *, + setting_name: str = TRUSTED_PROXY_RANGES_KEY, +) -> List[TrustedProxyNetwork]: + networks: List[TrustedProxyNetwork] = [] + for cidr in _normalize_cidr_ranges(configured_ranges, setting_name=setting_name): + try: + networks.append(ipaddress.ip_network(cidr, strict=False)) + except ValueError: + verbose_proxy_logger.warning( + "Invalid CIDR in %s: %s, skipping", setting_name, cidr + ) + return networks + + +def _get_direct_client_ip(request: Request) -> Optional[str]: + client = getattr(request, "client", None) + client_host = getattr(client, "host", None) + if isinstance(client_host, str): + return client_host + return None + + +def _is_ip_in_networks( + client_ip: Optional[str], networks: List[TrustedProxyNetwork] +) -> bool: + if not client_ip or not networks: + return False + try: + addr = ipaddress.ip_address(client_ip.strip()) + except ValueError: + return False + return any(addr in network for network in networks) + + +def require_trusted_proxy_request( + *, + request: Request, + general_settings: Optional[Dict[str, Any]] = None, + feature_name: str, + setting_name: str = TRUSTED_PROXY_RANGES_KEY, +) -> None: + """ + Fail closed unless the direct TCP peer is one of the configured + trusted reverse proxies. + + Header-based auth paths must validate the direct peer, not + X-Forwarded-For, because the direct peer is the actor supplying the + identity headers. + """ + if general_settings is None: + general_settings = _get_proxy_general_settings() + + trusted_networks = parse_trusted_proxy_ranges( + general_settings.get(setting_name), setting_name=setting_name + ) + if not trusted_networks: + raise ValueError( + f"{feature_name} requires general_settings.{setting_name} before " + "trusting identity headers from an upstream proxy." + ) + + direct_client_ip = _get_direct_client_ip(request) + if not _is_ip_in_networks(direct_client_ip, trusted_networks): + verbose_proxy_logger.warning( + "%s rejected identity headers from untrusted direct client IP %r", + feature_name, + direct_client_ip, + ) + raise ValueError( + f"{feature_name} only accepts identity headers from configured " + f"trusted proxy ranges. Direct client IP {direct_client_ip!r} " + "is not trusted." + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4fb82e2254..ef5bdc4fb7 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 @@ -63,6 +63,7 @@ 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 @@ -135,6 +136,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 @@ -903,7 +927,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 @@ -1273,6 +1301,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, ) @@ -1298,7 +1327,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 @@ -1422,21 +1455,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 @@ -1444,14 +1485,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: @@ -1881,10 +1923,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, @@ -1902,6 +1946,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 @@ -1909,6 +1968,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, @@ -1946,6 +2058,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) @@ -2153,6 +2266,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: @@ -2171,7 +2285,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), @@ -2258,11 +2376,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 @@ -2270,13 +2394,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 @@ -2284,14 +2409,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/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index a9ea7a84e1..447837c35e 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -53,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 @@ -572,9 +576,11 @@ def login(ctx: click.Context): 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 f3138f10da..baa0853700 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 = ( @@ -791,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 @@ -1604,6 +1658,12 @@ class ProxyBaseLLMRequestProcessing: # here would duplicate the guardrail API call # (e.g. double OpenAI Moderation charges). continue + if "async_post_call_streaming_iterator_hook" in type(cb).__dict__: + # Skip — the guardrail already scanned the assembled + # response via its own streaming iterator hook in the + # streaming pipeline. re running this function async_post_call_success_hook + # here would duplicate the scan and can spuriously block the guardrail that already passed / failed. + continue else: guardrail_result = await cb.async_post_call_success_hook( user_api_key_dict=captured_user_api_key_dict, diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index a979471dc8..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, @@ -152,12 +179,156 @@ class SpendCounterReseed: 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, refresh_ttl=True - ) + 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/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index ac487fb06d..5351391e5e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( build_sandbox_globals, compile_sandboxed, @@ -842,7 +843,10 @@ async def list_guardrail_submissions( if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") - is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + # Admin Viewer follows the read-parity rule: see all submissions like a + # Proxy Admin would (no writes — registration / approval still gated + # elsewhere by their own per-action checks). + is_admin = _user_has_admin_view(user_api_key_dict) visible_team_ids: Optional[List[str]] = None if not is_admin: visible_team_ids = await _get_user_team_ids(user_api_key_dict) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qohash/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/qohash/__init__.py new file mode 100644 index 0000000000..465f52db3d --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/qohash/__init__.py @@ -0,0 +1,35 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .qohash import QostodianNexus + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _instance = QostodianNexus( + api_base=litellm_params.api_base, + guardrail_name=guardrail.get("guardrail_name", ""), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + additional_provider_specific_params=litellm_params.additional_provider_specific_params, + extra_headers=getattr(litellm_params, "extra_headers", None), + ) + + litellm.logging_callback_manager.add_litellm_callback(_instance) + + return _instance + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.QOSTODIAN_NEXUS.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.QOSTODIAN_NEXUS.value: QostodianNexus, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py b/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py new file mode 100644 index 0000000000..a1bab6dbac --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/qohash/qohash.py @@ -0,0 +1,81 @@ +""" +Qostodian Nexus (by Qohash) — LiteLLM guardrail integration. +""" + +import os +from typing import TYPE_CHECKING, Literal, Optional, Type + +from litellm.integrations.custom_guardrail import log_guardrail_information +from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( + GenericGuardrailAPI, +) +from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +GUARDRAIL_NAME = "qostodian_nexus" + + +class QostodianNexus(GenericGuardrailAPI): + def __init__( + self, + api_base: Optional[str] = None, + **kwargs, + ): + api_base = api_base or os.environ.get( + "QOSTODIAN_NEXUS_API_BASE", "http://nexus:8800" + ) + + kwargs["guardrail_name"] = kwargs.get("guardrail_name", GUARDRAIL_NAME) + + # Merge built-in Qostodian Nexus identifier headers with any caller-supplied extra_headers + nexus_headers = [ + "x-qostodian-nexus-identifiers-trace", + "x-qostodian-nexus-identifiers-source", + "x-qostodian-nexus-identifiers-container", + "x-qostodian-nexus-identifiers-identity", + ] + + existing = kwargs.get("extra_headers") or [] + kwargs["extra_headers"] = nexus_headers + [ + h for h in existing if h not in nexus_headers + ] + + super().__init__( + api_base=api_base, + **kwargs, + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Apply Qostodian Nexus to the given inputs. + + NOTE: This override is intentionally a pass-through. It must be present + directly in this class's __dict__ so that LiteLLM's unified guardrail + routing check (`"apply_guardrail" in type(callback).__dict__` in + litellm/proxy/utils.py) routes calls correctly. Do not remove. + """ + return await super().apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type=input_type, + logging_obj=logging_obj, + ) + + @classmethod + def get_config_model(cls) -> Optional[Type[QostodianNexusConfigModel]]: + """ + Returns the config model for Qostodian Nexus. + """ + return QostodianNexusConfigModel 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 0f7067b1c2..c9859ed25b 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( CallInfo, EnterpriseLicenseData, Litellm_EntityType, + LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth, @@ -28,6 +29,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.health_check import ( + ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, _update_litellm_params_for_health_check, perform_health_check, @@ -723,6 +725,129 @@ async def _save_background_health_checks_to_db( # Continue execution - don't let database save failure break health checks +_PROXY_ADMIN_ROLES = frozenset( + { + LitellmUserRoles.PROXY_ADMIN.value, + # View-only admins are operators (oncall, support); they need the + # routing fields (api_base, api_version) to diagnose health and tell + # which provider region a check is hitting. They cannot mutate config + # so granting them the read-only view is safe. + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + } +) + + +def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Return True if the caller has a proxy-admin role (full or view-only). + + user_role on UserAPIKeyAuth can be either a LitellmUserRoles enum or its + string value depending on how the auth path constructed the object, so we + compare against the raw value rather than the enum identity. + """ + role = user_api_key_dict.user_role + if role is None: + return False + role_value = role.value if hasattr(role, "value") else role + return role_value in _PROXY_ADMIN_ROLES + + +def _strip_admin_only_fields_from_health_result(result: dict) -> dict: + """ + Return a copy of the /health response with provider routing fields + (``api_base``, ``api_version``) removed from each healthy/unhealthy + endpoint entry. Used to hide those fields from non-admin callers while + still showing them which deployments they own and whether each one is + healthy. Proxy admins receive the unmodified result. + """ + out = dict(result) + drop = set(ADMIN_ONLY_HEALTH_DISPLAY_PARAMS) + for key in ("healthy_endpoints", "unhealthy_endpoints"): + eps = out.get(key) + if isinstance(eps, list): + out[key] = [ + ( + {k: v for k, v in ep.items() if k not in drop} + if isinstance(ep, dict) + else ep + ) + for ep in eps + ] + return out + + +def _resolve_targeted_model_ids( + model_list: list, model: Optional[str], model_id: Optional[str] +) -> Optional[set]: + """ + Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of + deployment IDs the response should be scoped to. + + Mirrors the live-path semantics in ``perform_health_check()``: ``model`` + matches either the deployment's ``model_name`` alias or its + ``litellm_params.model`` provider string. ``model_id`` matches + ``model_info.id``. + + Both query params are validated against the supplied ``model_list``. + Callers pass an already-scoped list (filtered to the caller's allowed + models for non-admins, full list for admins), so a ``model_id`` that + isn't present resolves to an empty set rather than a single-element + set — preventing a non-admin from reading another deployment's cached + health entry by guessing its ID. + + Returns ``None`` when no targeting is requested — callers should treat + that as "no filter." + """ + if not model and not model_id: + return None + target_ids: set = set() + for m in model_list: + deployment_id = (m.get("model_info") or {}).get("id") + if not deployment_id: + continue + if model_id and deployment_id == model_id: + target_ids.add(deployment_id) + continue + if model: + litellm_model = (m.get("litellm_params") or {}).get("model") + if m.get("model_name") == model or litellm_model == model: + target_ids.add(deployment_id) + return target_ids + + +def _filter_health_check_results_by_model_ids( + results: dict, allowed_model_ids: set +) -> dict: + """ + Restrict a cached background health-check result dict to endpoints whose + model_id is in ``allowed_model_ids``. + + Endpoints without a model_id (e.g. CLI-model entries that predate the + model_id wiring) are dropped conservatively — we cannot prove they belong + to the caller, so they are excluded rather than leaked. + + Each retained endpoint is shallow-copied before being returned, so any + downstream transform (e.g. _strip_admin_only_fields_from_health_result) + cannot accidentally mutate the shared ``health_check_results`` cache. + """ + healthy = [ + dict(ep) + for ep in (results.get("healthy_endpoints") or []) + if ep.get("model_id") in allowed_model_ids + ] + unhealthy = [ + dict(ep) + for ep in (results.get("unhealthy_endpoints") or []) + if ep.get("model_id") in allowed_model_ids + ] + return { + "healthy_endpoints": healthy, + "unhealthy_endpoints": unhealthy, + "healthy_count": len(healthy), + "unhealthy_count": len(unhealthy), + } + + async def _perform_health_check_and_save( model_list, target_model, @@ -771,6 +896,7 @@ async def _perform_health_check_and_save( @router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def health_endpoint( + response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), model: Optional[str] = fastapi.Query( None, description="Specify the model name (optional)" @@ -838,11 +964,33 @@ async def health_endpoint( detail={"error": f"Model with ID {model_id} not found"}, ) + is_admin = _is_proxy_admin(user_api_key_dict) + model_specific_request = bool(model or model_id) + + def _post_process(result: dict) -> dict: + # api_base / api_version reveal which provider/region/internal host the + # deployment talks to; only proxy admins receive them. Non-admin keys + # still see model/model_id and the healthy/unhealthy status. We also + # set a header so non-admin clients that previously parsed those + # fields can detect the change programmatically. + # When a caller asked about a specific model/model_id and zero + # endpoints came back healthy, surface that as a 503 so monitoring + # systems can rely on the HTTP status instead of having to parse the + # body. The body shape is unchanged. + if model_specific_request and result.get("healthy_count", 0) == 0: + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE + if is_admin: + return result + response.headers["Litellm-Health-Field-Notice"] = ( + "api_base and api_version are admin-only on this endpoint" + ) + return _strip_admin_only_fields_from_health_result(result) + try: if llm_model_list is None: # if no router set, check if user set a model using litellm --model ollama/llama2 if user_model is not None: - return await _perform_health_check_and_save( + cli_result = await _perform_health_check_and_save( model_list=[], target_model=None, cli_model=user_model, @@ -853,20 +1001,81 @@ async def health_endpoint( model_id=None, # CLI model doesn't have model_id max_concurrency=health_check_concurrency, ) + return _post_process(cli_result) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "Model list not initialized"}, ) _llm_model_list = copy.deepcopy(llm_model_list) ### FILTER MODELS FOR ONLY THOSE USER HAS ACCESS TO ### + # Live path: scope by model_name (every deployment has one). + # Cache path: scope by model_id (the cache is keyed on model_id). + # Consequence: a deployment whose model_name the caller can access + # but which lacks model_info.id will appear in the live /health + # response but NOT in the background-cache /health response. This is + # surfaced via the "warnings" field below so operators can fix the + # missing model_info.id rather than guess at the discrepancy. if len(user_api_key_dict.models) > 0: - pass - else: - pass # + allowed_models = set(user_api_key_dict.models) + _llm_model_list = [ + m for m in _llm_model_list if m.get("model_name") in allowed_models + ] if use_background_health_checks: - return health_check_results + # The cached background result covers every model. When the + # caller targets a specific model/model_id we have to narrow the + # cache to that deployment before _post_process evaluates + # healthy_count, otherwise an unhealthy "foo" combined with any + # other healthy model would still report healthy_count > 0 and + # the targeted-503 path would never fire. + targeted_ids = _resolve_targeted_model_ids(_llm_model_list, model, model_id) + if len(user_api_key_dict.models) > 0: + allowed_model_ids = { + (m.get("model_info") or {}).get("id") + for m in _llm_model_list + if (m.get("model_info") or {}).get("id") + } + # _llm_model_list is already scoped to the caller's allowed + # model_names above, so targeted_ids is implicitly the + # intersection of "targeted" and "allowed." + filter_ids = ( + targeted_ids if targeted_ids is not None else allowed_model_ids + ) + filtered = _filter_health_check_results_by_model_ids( + health_check_results, filter_ids + ) + if targeted_ids is None and not allowed_model_ids: + # Caller has accessible model_names but none of the + # matching deployments expose a model_info.id, so the + # cache filter (which keys on model_id) drops every + # entry. Surface this both as a warning log and a + # structured "warnings" field on the response so the + # caller can distinguish "no deployments found" from + # "deployments excluded due to missing model_info.id". + verbose_proxy_logger.warning( + "health_endpoint: scoped key %s has accessible models %s " + "but none of the matching deployments carry a model_info.id; " + "background health-check cache will return an empty result.", + user_api_key_dict.user_id, + list(user_api_key_dict.models), + ) + filtered["warnings"] = [ + "Some accessible deployments are missing model_info.id " + "and were excluded from this response. Ask a proxy admin " + "to populate model_info.id for these models." + ] + return _post_process(filtered) + if targeted_ids is not None: + # Admin caller targeting a specific model: filter the cache + # so the response (and the targeted-503 check) reflects only + # that deployment, not the global aggregate. + return _post_process( + _filter_health_check_results_by_model_ids( + health_check_results, targeted_ids + ) + ) + return _post_process(health_check_results) else: - return await _perform_health_check_and_save( + router_result = await _perform_health_check_and_save( model_list=_llm_model_list, target_model=target_model, cli_model=None, @@ -877,6 +1086,7 @@ async def health_endpoint( model_id=model_id, max_concurrency=health_check_concurrency, ) + return _post_process(router_result) except Exception as e: verbose_proxy_logger.error( "litellm.proxy.proxy_server.py::health_endpoint(): Exception occured - {}".format( @@ -1237,7 +1447,9 @@ def callback_name(callback): return str(callback) -async def _get_health_readiness_details() -> Dict[str, Any]: +async def _get_health_readiness_details( + response: Optional[Response] = None, +) -> Dict[str, Any]: """ Detailed health payload for authenticated diagnostics. """ @@ -1271,8 +1483,8 @@ async def _get_health_readiness_details() -> Dict[str, Any]: try: index_info = await litellm.cache.cache._index_info() except Exception as e: - index_info = "index does not exist - error: " + str(e) - cache_type = {"type": cache_type, "index_info": index_info} + index_info = "index does not exist - error: " + str(e) # type: ignore[assignment] + cache_type = {"type": cache_type, "index_info": index_info} # type: ignore[assignment] # check log level log_level_name = logging.getLevelName(verbose_logger.getEffectiveLevel()) @@ -1281,6 +1493,12 @@ async def _get_health_readiness_details() -> Dict[str, Any]: # check DB if prisma_client is not None: # if db passed in, check if it's connected db_health_status = await _db_health_readiness_check() + # A configured DB that is not reachable means the worker cannot + # serve requests that depend on persisted state (keys, budgets, + # spend logs). Return 503 so orchestrators take this pod out of + # rotation; "Not connected" (no DB configured at all) stays 200. + if response is not None and db_health_status["status"] != "connected": + response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", "db": db_health_status["status"], @@ -1316,14 +1534,14 @@ def _allow_public_health_readiness_details() -> bool: "/health/readiness", tags=["health"], ) -async def health_readiness(): +async def health_readiness(response: Response): """ Public readiness probe. Keep this low-detail for unauthenticated load balancers by default. Admins can opt into the legacy detailed public payload with general_settings.allow_public_health_readiness_details. """ if _allow_public_health_readiness_details(): - return await _get_health_readiness_details() + return await _get_health_readiness_details(response=response) return {"status": "healthy"} @@ -1332,11 +1550,11 @@ async def health_readiness(): tags=["health"], dependencies=[Depends(user_api_key_auth)], ) -async def health_readiness_details(): +async def health_readiness_details(response: Response): """ Authenticated readiness diagnostics with DB/cache/callback metadata. """ - return await _get_health_readiness_details() + return await _get_health_readiness_details(response=response) @router.get( diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 06b7d85789..f740d5dd40 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import ( _get_batch_job_input_file_usage, _get_file_content_as_dictionary, + _get_models_from_batch_input_file_content, ) from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -164,13 +165,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage: BatchFileUsage, ) -> None: """ - Check rate limits and increment counters by the batch amounts. + Atomically check + increment rate-limit counters by the batch amounts. - Raises HTTPException if any limit would be exceeded. + Raises HTTPException if any descriptor would exceed its limit; in that + case no counter is modified. Backed by `atomic_check_and_increment_by_n` + which uses a Redis Lua script when available (multi-process atomic) and + falls back to a per-process asyncio.Lock + in-memory operation. """ - from litellm.types.caching import RedisPipelineIncrementOperation - - # Create descriptors and check if batch would exceed limits descriptors = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, @@ -179,73 +180,31 @@ class _PROXY_BatchRateLimiter(CustomLogger): model_has_failures=False, ) - # Check current usage without incrementing - rate_limit_response = await self.parallel_request_limiter.should_rate_limit( - descriptors=descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, - ) + increment: Dict[Literal["requests", "tokens"], int] = { + "requests": batch_usage.request_count, + "tokens": batch_usage.total_tokens, + } + increments: List[Dict[Literal["requests", "tokens"], int]] = [ + increment for _ in descriptors + ] - # Verify batch won't exceed any limits - for status in rate_limit_response["statuses"]: - rate_limit_type = status["rate_limit_type"] - limit_remaining = status["limit_remaining"] - - required_capacity = ( - batch_usage.request_count - if rate_limit_type == "requests" - else batch_usage.total_tokens if rate_limit_type == "tokens" else 0 - ) - - if required_capacity > limit_remaining: - self._raise_rate_limit_error( - status, descriptors, batch_usage, rate_limit_type - ) - - # Build pipeline operations for batch increments - # Reuse the same keys that descriptors check - pipeline_operations: List[RedisPipelineIncrementOperation] = [] - - for descriptor in descriptors: - key = descriptor["key"] - value = descriptor["value"] - rate_limit = descriptor.get("rate_limit") - - if rate_limit is None: - continue - - # Add RPM increment if limit is set - if rate_limit.get("requests_per_unit") is not None: - rpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="requests" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=rpm_key, - increment_value=batch_usage.request_count, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Add TPM increment if limit is set - if rate_limit.get("tokens_per_unit") is not None: - tpm_key = self.parallel_request_limiter.create_rate_limit_keys( - key=key, value=value, rate_limit_type="tokens" - ) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=tpm_key, - increment_value=batch_usage.total_tokens, - ttl=self.parallel_request_limiter.window_size, - ) - ) - - # Execute increments - if pipeline_operations: - await self.parallel_request_limiter.async_increment_tokens_with_ttl_preservation( - pipeline_operations=pipeline_operations, + rate_limit_response = ( + await self.parallel_request_limiter.atomic_check_and_increment_by_n( + descriptors=descriptors, + increments=increments, parent_otel_span=user_api_key_dict.parent_otel_span, ) + ) + + if rate_limit_response["overall_code"] == "OVER_LIMIT": + for status in rate_limit_response["statuses"]: + if status["code"] == "OVER_LIMIT": + self._raise_rate_limit_error( + status, + descriptors, + batch_usage, + status["rate_limit_type"], + ) async def count_input_file_usage( self, @@ -288,6 +247,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content_as_dict = _get_file_content_as_dictionary(file_content.content) + # Validate every model named in the batch JSONL against the + # caller's per-key model allowlist. Without this, a caller + # could smuggle restricted/expensive models inside the file + # and the upstream provider would execute the batch under + # the proxy's shared API key. + if user_api_key_dict is not None: + await self._enforce_batch_file_model_access( + user_api_key_dict=user_api_key_dict, + file_content_as_dict=file_content_as_dict, + ) + input_file_usage = _get_batch_job_input_file_usage( file_content_dictionary=file_content_as_dict, custom_llm_provider=custom_llm_provider, @@ -298,12 +268,69 @@ class _PROXY_BatchRateLimiter(CustomLogger): request_count=request_count, ) + except HTTPException as e: + # Distinguish intentional 403s from `_enforce_batch_file_model_access` + # from genuine I/O failures so security-relevant rejections show up + # in the access log instead of getting buried in error noise. + if e.status_code == 403: + verbose_proxy_logger.warning( + f"Batch rejected: caller not authorized for a model named in {file_id}: {e.detail}" + ) + else: + verbose_proxy_logger.error( + f"Batch input file rejected for {file_id}: status={e.status_code} detail={e.detail}" + ) + raise except Exception as e: verbose_proxy_logger.error( f"Error counting input file usage for {file_id}: {str(e)}" ) raise + async def _enforce_batch_file_model_access( + self, + user_api_key_dict: UserAPIKeyAuth, + file_content_as_dict: List[dict], + ) -> None: + """Reject the batch if the caller is not authorized for every + ``body.model`` named inside the JSONL. + + Reuses ``can_key_call_model`` so the same allowlist semantics + (wildcards, access groups, ``all-proxy-models``, team aliases) + the proxy enforces on `/chat/completions` apply here. + """ + from litellm.proxy.auth.auth_checks import can_key_call_model + from litellm.proxy.proxy_server import llm_router + + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return + + llm_model_list = llm_router.model_list if llm_router is not None else None + for model in models: + try: + await can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + except HTTPException: + raise + except Exception as e: + # `can_key_call_model` raises ProxyException on denial; + # re-shape to a 403 so the batch endpoint returns a + # consistent rejection without leaking internal types. + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Batch input file references a model the caller is " + f"not authorized to use: model={model}, reason={str(e)}" + ) + }, + ) + async def _fetch_managed_file_content( self, file_id: str, diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 72483d29cd..f7c0592992 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting import os from datetime import datetime -from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, Union from fastapi import HTTPException @@ -460,92 +460,128 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if priority_descriptors: descriptors_to_check.extend(priority_descriptors) - # PHASE 1: Read-only check of ALL limits (no increments) - check_response = await self.v3_limiter.should_rate_limit( - descriptors=descriptors_to_check, + # Atomic check-and-increment for the ENFORCED descriptor set: + # - model_saturation_check is always enforced + # - priority_model is enforced only when saturation crosses threshold + # + # Backed by a Redis Lua script (multi-process atomic) with an + # asyncio.Lock + in-memory fallback for single-process deployments. + # All-or-nothing: if any enforced descriptor would exceed its limit, + # no counter is modified and the response carries "OVER_LIMIT". + enforced_descriptors: List[RateLimitDescriptor] = [model_wide_descriptor] + if priority_descriptors and should_enforce_priority: + enforced_descriptors.extend(priority_descriptors) + + per_request_increment: Dict[Literal["requests", "tokens"], int] = { + "requests": 1, + "tokens": 0, + } + atomic_response = await self.v3_limiter.atomic_check_and_increment_by_n( + descriptors=enforced_descriptors, + increments=[per_request_increment for _ in enforced_descriptors], parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=True, # CRITICAL: Don't increment counters yet ) verbose_proxy_logger.debug( - f"Read-only check: {json.dumps(check_response, indent=2)}" + f"Atomic check+increment response: {json.dumps(atomic_response, indent=2)}" ) - # PHASE 2: Decide which limits to enforce - if check_response["overall_code"] == "OVER_LIMIT": - for status in check_response["statuses"]: - if status["code"] == "OVER_LIMIT": - descriptor_key = status["descriptor_key"] + if atomic_response["overall_code"] == "OVER_LIMIT": + for status in atomic_response["statuses"]: + if status["code"] != "OVER_LIMIT": + continue + descriptor_key = status["descriptor_key"] + if descriptor_key == "model_saturation_check": + raise HTTPException( + status_code=429, + detail={ + "error": f"Model capacity reached for {model}. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + }, + ) + if descriptor_key == "priority_model": + verbose_proxy_logger.debug( + f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " + f"priority: {priority}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": f"Priority-based rate limit exceeded. " + f"Priority: {priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + ) - # Model-wide limit exceeded (ALWAYS enforce) - if descriptor_key == "model_saturation_check": - raise HTTPException( - status_code=429, - detail={ - "error": f"Model capacity reached for {model}. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - }, - ) + # Fail-closed guard: overall_code says OVER_LIMIT but no status + # matched a descriptor key we know how to translate into a 429. + # Refuse the request rather than silently fall through and let an + # over-limit request proceed to the model. Without this, a future + # caller wiring an unfamiliar descriptor into enforced_descriptors + # would silently bypass the rate limit. + offending = next( + (s for s in atomic_response["statuses"] if s["code"] == "OVER_LIMIT"), + None, + ) + verbose_proxy_logger.error( + f"Dynamic rate limiter: OVER_LIMIT response with unknown " + f"descriptor_key(s) — refusing request. response={atomic_response}" + ) + raise HTTPException( + status_code=429, + detail={ + "error": "Rate limit exceeded", + "descriptor_key": ( + offending["descriptor_key"] if offending else "unknown" + ), + "rate_limit_type": ( + str(offending["rate_limit_type"]) if offending else "unknown" + ), + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "x-litellm-priority": priority or "default", + }, + ) - # Priority limit exceeded (ONLY enforce when saturated) - elif descriptor_key == "priority_model" and should_enforce_priority: - verbose_proxy_logger.debug( - f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " - f"priority: {priority}" - ) - raise HTTPException( - status_code=429, - detail={ - "error": f"Priority-based rate limit exceeded. " - f"Priority: {priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}, " - f"Model saturation: {saturation:.1%}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": priority or "default", - "x-litellm-saturation": f"{saturation:.2%}", - }, - ) - - # PHASE 3: Increment counters separately to avoid early-exit issues - # Model counter must ALWAYS increment, but priority counter might be over limit - # If we increment them together, v3_limiter's in-memory check will exit early - # and skip incrementing the model counter - - # Step 3a: Increment model-wide counter (always) - model_increment_response = await self.v3_limiter.should_rate_limit( - descriptors=[model_wide_descriptor], - parent_otel_span=user_api_key_dict.parent_otel_span, - read_only=False, - ) - - # Step 3b: Increment priority counter (may be over limit, but we still track it) - if priority_descriptors: - priority_increment_response = await self.v3_limiter.should_rate_limit( + # If priority is NOT enforced (saturation below threshold) but + # priority_descriptors exist, increment them for tracking only — no + # check, no rollback. This matches the prior tracking semantics. + # + # Using the non-atomic should_rate_limit (instead of + # atomic_check_and_increment_by_n) is intentional here: we don't want + # to enforce the limit, we only want to bump the counter so the + # priority allocation has accurate usage when it later becomes + # enforced. The increment-then-check semantics of should_rate_limit + # are fine because we ignore the OVER_LIMIT response. + if priority_descriptors and not should_enforce_priority: + priority_tracking_response = await self.v3_limiter.should_rate_limit( descriptors=priority_descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=False, ) - - # Combine responses for post-call hook - combined_response = { - "overall_code": model_increment_response["overall_code"], - "statuses": model_increment_response["statuses"] - + priority_increment_response["statuses"], + data["litellm_proxy_rate_limit_response"] = { + "overall_code": atomic_response["overall_code"], + "statuses": atomic_response["statuses"] + + priority_tracking_response["statuses"], } - data["litellm_proxy_rate_limit_response"] = combined_response else: - data["litellm_proxy_rate_limit_response"] = model_increment_response + data["litellm_proxy_rate_limit_response"] = atomic_response async def async_pre_call_hook( self, diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 7789fa6a34..9a7e511794 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -32,10 +32,25 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): if user_api_key_dict.team_id is not None: return + # The reservation path admits at the strict-`<` boundary and + # atomically pre-fills the same counter we'd read here. Re-checking + # with `>=` would reject a request the reservation already admitted + # when the reservation fills the counter to exactly max_budget. + # Imported lazily to avoid a circular import via proxy.utils. + from litellm.proxy.spend_tracking.budget_reservation import ( + get_reserved_counter_keys, + ) + + user_counter_key = f"spend:user:{user_id}" + if user_counter_key in get_reserved_counter_keys( + user_api_key_dict.budget_reservation + ): + return + from litellm.proxy.proxy_server import get_current_spend curr_spend = await get_current_spend( - counter_key=f"spend:user:{user_id}", + counter_key=user_counter_key, fallback_spend=user_api_key_dict.user_spend or 0.0, ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index f29bbd2d9d..4497e64c17 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4,6 +4,7 @@ This is a rate limiter implementation based on a similar one by Envoy proxy. This is currently in development and not yet ready for production. """ +import asyncio import binascii import os from datetime import datetime @@ -80,6 +81,90 @@ end return results """ +CHECK_AND_INCREMENT_BY_N_SCRIPT = """ +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) +-- +-- Return on success: { 0, new_counter_1, new_counter_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local descriptor_count = #KEYS / 2 + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + if current_counter + increment > limit then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_expired = descriptor_state[i][1] + + if window_expired then + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end +end + +return results +""" + TOKEN_INCREMENT_SCRIPT = """ local results = {} @@ -162,15 +247,37 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): TOKEN_INCREMENT_SCRIPT ) ) + self.check_and_increment_by_n_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + CHECK_AND_INCREMENT_BY_N_SCRIPT + ) + ) else: self.batch_rate_limiter_script = None self.token_increment_script = None + self.check_and_increment_by_n_script = None self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) # Batch rate limiter (lazy loaded) self._batch_rate_limiter: Optional[Any] = None + # Serializes multi-phase check+increment sequences (batch + dynamic + # limiters) within this process to close the TOCTOU window between + # read-only check and counter increment. Multi-replica deployments + # additionally rely on Redis Lua atomicity for cross-process safety. + # + # Coarse granularity: this single lock serializes ALL atomic check+ + # increment operations across batch and dynamic limiters on this + # instance. A slow batch input-file fetch (which happens upstream of + # the lock) does not block here, but Redis Lua latency does. If + # contention shows up under load (visible as p99 latency spikes + # correlated with batch traffic), shard to a per-descriptor-key lock + # via a `weakref.WeakValueDictionary[str, asyncio.Lock]`. Punted as a + # follow-up because Lua dominates wall-time and the lock is held for + # one round-trip. + self._check_and_increment_lock = asyncio.Lock() + def _get_batch_rate_limiter(self) -> Optional[Any]: """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: @@ -588,6 +695,281 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return rate_limit_response + async def atomic_check_and_increment_by_n( + self, + descriptors: List[RateLimitDescriptor], + increments: List[Dict[Literal["requests", "tokens"], int]], + parent_otel_span: Optional[Span] = None, + ) -> RateLimitResponse: + """ + Atomic check-and-increment-by-N across one or more descriptors. + + All-or-nothing: if any descriptor would exceed its limit, no counter is + modified and the response carries `overall_code = "OVER_LIMIT"` with + the offending descriptor's status. Closes the TOCTOU window between + read and increment in both single-process and multi-process (Redis) + deployments. + + Args: + descriptors: rate-limit descriptors to check + increments: per-descriptor increment amounts, indexed parallel to + `descriptors`. Each entry is `{"requests": int, "tokens": int}` + — values default to 0 when a descriptor has no matching limit. + + Returns: + RateLimitResponse with one status per (descriptor, rate_limit_type) + counter, mirroring `should_rate_limit`'s shape. + """ + if len(descriptors) != len(increments): + raise ValueError( + "atomic_check_and_increment_by_n: descriptors and increments " + "must have the same length" + ) + + keys: List[str] = [] + per_counter_meta: List[Dict[str, Any]] = [] + script_args: List[Any] = [] + + for descriptor, increment_amounts in zip(descriptors, increments): + descriptor_key = descriptor["key"] + descriptor_value = descriptor["value"] + rate_limit: RateLimitDescriptorRateLimitObject = ( + descriptor.get("rate_limit") or RateLimitDescriptorRateLimitObject() + ) + window_size = rate_limit.get("window_size") or self.window_size + window_key = f"{{{descriptor_key}:{descriptor_value}}}:window" + + for rate_limit_type in ("requests", "tokens"): + rlt: Literal["requests", "tokens"] = cast( + Literal["requests", "tokens"], rate_limit_type + ) + if rlt == "requests": + limit_value = rate_limit.get("requests_per_unit") + inc_amount = int(increment_amounts.get("requests", 0) or 0) + else: + limit_value = rate_limit.get("tokens_per_unit") + inc_amount = int(increment_amounts.get("tokens", 0) or 0) + if limit_value is None or inc_amount <= 0: + continue + counter_key = self.create_rate_limit_keys( + descriptor_key, descriptor_value, rlt + ) + # Counter-key TTL and window_size are conceptually distinct + # ("how long the counter Redis key lives" vs "how long the + # sliding window is"). They happen to be equal today because + # we have no descriptor type that needs them apart, but they + # are kept as separate variables here so a future custom-TTL + # descriptor doesn't reintroduce a silent expiry bug. Both + # the Lua script and the in-memory fallback read these from + # their respective ARGV / meta slots. + ttl_seconds = int(window_size) + window_size_seconds = int(window_size) + keys.extend([window_key, counter_key]) + # Per-counter 4-tuple matches the Lua ARGV layout exactly: + # [limit, increment, ttl_seconds, window_size_seconds]. + script_args.extend( + [ + int(limit_value), + inc_amount, + ttl_seconds, + window_size_seconds, + ] + ) + per_counter_meta.append( + { + "descriptor_key": descriptor_key, + "current_limit": int(limit_value), + "rate_limit_type": rlt, + "window_key": window_key, + "counter_key": counter_key, + "increment": inc_amount, + "ttl": ttl_seconds, + "window_size": window_size_seconds, + } + ) + + if not keys: + return RateLimitResponse(overall_code="OK", statuses=[]) + + # Multi-process atomicity via Redis Lua. Single-process atomicity + # falls back to the asyncio.Lock + in-memory sliding window below. + # Note: in-memory state diverges from Redis state — if Lua fails + # mid-write, retrying via in-memory may double-count. See fallback + # warning below. + if self.check_and_increment_by_n_script is not None: + try: + raw = await self.check_and_increment_by_n_script( + keys=keys, + args=script_args, + ) + return self._build_atomic_response(raw, per_counter_meta) + except Exception as e: + # Escalated from warning to error: Lua failures (script timeout, + # Redis OOM, network partition) leave counter state ambiguous. + # The fallback path below uses LOCAL DualCache, which is a + # different store from Redis — counters here will diverge from + # Redis until that key's window expires (TTL bounds divergence). + # Operators should alert on this log line; sustained occurrences + # indicate Redis health degradation that may erode rate-limit + # accuracy. + verbose_proxy_logger.error( + f"atomic_check_and_increment_by_n: Redis Lua execution " + f"failed ({type(e).__name__}: {e}). Falling back to " + f"in-memory enforcement — counters will diverge from Redis " + f"state until window expires (window_size={self.window_size}s)." + ) + + async with self._check_and_increment_lock: + return await self._atomic_check_and_increment_in_memory( + per_counter_meta=per_counter_meta, + parent_otel_span=parent_otel_span, + ) + + def _build_atomic_response( + self, + raw: List[Any], + per_counter_meta: List[Dict[str, Any]], + ) -> RateLimitResponse: + """Convert Lua script return value to RateLimitResponse. + + Indexing invariant: `per_counter_meta` and `KEYS` are parallel-indexed + at the COUNTER level, not the descriptor level. A descriptor with both + RPM and TPM limits emits two `(window_key, counter_key)` pairs and + two meta entries — one per counter. The Lua script's loop variable + `i` therefore enumerates counters, and the over-limit return tuple + `{1, i, ...}` carries a counter index that maps directly to + `per_counter_meta[i - 1]`. Keep these arrays parallel at the counter + level when modifying this code. + """ + if not raw: + return RateLimitResponse(overall_code="OK", statuses=[]) + + status_code = int(raw[0]) + if status_code == 1: + # Over limit: { 1, counter_index (1-based), current_counter, limit } + descriptor_index = int(raw[1]) - 1 + current_counter = int(raw[2]) + limit = int(raw[3]) + meta = per_counter_meta[descriptor_index] + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ + RateLimitStatus( + code="OVER_LIMIT", + current_limit=limit, + limit_remaining=max(0, limit - current_counter), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ], + ) + + statuses: List[RateLimitStatus] = [] + for meta, new_counter in zip(per_counter_meta, raw[1:]): + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - int(new_counter)), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + return RateLimitResponse(overall_code="OK", statuses=statuses) + + async def _atomic_check_and_increment_in_memory( + self, + per_counter_meta: List[Dict[str, Any]], + parent_otel_span: Optional[Span] = None, + ) -> RateLimitResponse: + """In-memory all-or-nothing check-and-increment. Caller holds lock. + + Reads/writes the LOCAL DualCache (`local_only=True`) — note this is + a different store from Redis. When this fallback fires after a Lua + failure, in-memory counters will diverge from Redis until each key's + window expires (TTL bounds divergence). + """ + # Use a single 'now' for the duration of this critical section so all + # descriptors evaluate window expiry consistently. + now_int = int(self._get_current_time().timestamp()) + + # Pass 1: read state, validate. + descriptor_state: List[Dict[str, Any]] = [] + for meta in per_counter_meta: + window_size = meta["window_size"] + window_start = await self.internal_usage_cache.async_get_cache( + key=meta["window_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + window_expired = ( + window_start is None or (now_int - int(window_start)) >= window_size + ) + current_counter = ( + 0 + if window_expired + else int( + await self.internal_usage_cache.async_get_cache( + key=meta["counter_key"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + or 0 + ) + ) + if current_counter + meta["increment"] > meta["current_limit"]: + return RateLimitResponse( + overall_code="OVER_LIMIT", + statuses=[ + RateLimitStatus( + code="OVER_LIMIT", + current_limit=meta["current_limit"], + limit_remaining=max( + 0, meta["current_limit"] - current_counter + ), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ], + ) + descriptor_state.append( + {"window_expired": window_expired, "current": current_counter} + ) + + # Pass 2: apply increments. + statuses: List[RateLimitStatus] = [] + for meta, state in zip(per_counter_meta, descriptor_state): + new_counter = ( + meta["increment"] + if state["window_expired"] + else state["current"] + meta["increment"] + ) + if state["window_expired"]: + await self.internal_usage_cache.async_set_cache( + key=meta["window_key"], + value=str(now_int), + ttl=meta["window_size"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + await self.internal_usage_cache.async_set_cache( + key=meta["counter_key"], + value=new_counter, + ttl=meta["ttl"], + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - new_counter), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + return RateLimitResponse(overall_code="OK", statuses=statuses) + def create_organization_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: Optional[str] = None ) -> List[RateLimitDescriptor]: diff --git a/litellm/proxy/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 ceaef20a8d..62a770f46a 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -38,6 +38,17 @@ def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: ) +def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None: + """Admin Viewer parity: PROXY_ADMIN or PROXY_ADMIN_VIEW_ONLY may read.""" + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view + + if not _user_has_admin_view(user_api_key_dict): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + def _record_to_response(record) -> AccessGroupResponse: return AccessGroupResponse( access_group_id=record.access_group_id, @@ -370,7 +381,7 @@ async def create_access_group( async def list_access_groups( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> List[AccessGroupResponse]: - _require_proxy_admin(user_api_key_dict) + _require_admin_view(user_api_key_dict) prisma_client = get_prisma_client_or_throw( CommonProxyErrors.db_not_connected_error.value ) @@ -389,7 +400,7 @@ async def get_access_group( access_group_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> AccessGroupResponse: - _require_proxy_admin(user_api_key_dict) + _require_admin_view(user_api_key_dict) prisma_client = get_prisma_client_or_throw( CommonProxyErrors.db_not_connected_error.value ) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 90c0d02d1e..81b133e6c8 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -17,6 +17,7 @@ from fastapi import APIRouter, Depends, HTTPException from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.utils import jsonify_object router = APIRouter() @@ -238,7 +239,7 @@ async def budget_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, detail={ @@ -305,7 +306,7 @@ async def list_budget( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, detail={ diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index b0ea6b41ac..8ad44b5300 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,5 +1,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from fastapi import HTTPException, status + from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( @@ -29,6 +31,34 @@ def _user_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool: ) +def require_caller_user_id_for_non_admin( + user_api_key_dict: UserAPIKeyAuth, +) -> str: + """Return the caller's user_id, or raise 403 if missing. + + Non-admin analytics endpoints scope queries by the caller's own user_id. + Service-account keys are deliberately created with user_id=None + (key_management_endpoints.py forces ``data.user_id = None`` at key + creation). Without this guard, that None value flows through to the + daily-activity builder, which treats ``entity_id is None`` as "no filter" + and returns every tenant's data. + + Callers must check is_admin first; this helper is only valid on the + non-admin scoping branch. + """ + if user_api_key_dict.user_id is None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + "Service-account keys cannot query user analytics. " + "Use a user-bound key, or call as a proxy admin." + ) + }, + ) + return user_api_key_dict.user_id + + def _is_user_team_admin( user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable ) -> bool: diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index d78c5526e6..b736ba1081 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -267,9 +267,11 @@ async def get_hashicorp_vault_config( Get current Hashicorp Vault configuration. Returns decrypted values from DB, or falls back to current env vars. """ + from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.proxy_server import prisma_client, proxy_config - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail="Only admin users can view config overrides", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 921d24da04..6f73c6a632 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -35,6 +35,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _user_has_admin_view, + require_caller_user_id_for_non_admin, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -618,6 +619,40 @@ def _normalize_user_info_user_id( return user_id +def _enforce_user_info_access( + user_id: Optional[str], user_api_key_dict: UserAPIKeyAuth +) -> None: + """Re-validate that the caller may read the resolved ``user_id`` after + URL-decoding has been finalized. + + The route-level check in ``RouteChecks.non_proxy_admin_allowed_routes_check`` + runs against ``request.query_params``, which decodes a literal ``+`` to a + space. ``_normalize_user_info_user_id`` then re-parses the raw query with + ``unquote`` so the endpoint can return rows for user_ids that contain ``+`` + (e.g. plus-addressed emails). That asymmetry let an attacker who registered + a username with a literal space pass the route check and then read another + user's row by sending the encoded ``+`` form. Re-checking ownership here + closes the gap without changing the supported user_id grammar. + """ + if user_id is None: + return + # Only true proxy admin bypasses ownership. PROXY_ADMIN_VIEW_ONLY is + # subject to the same `user_id == valid_token.user_id` rule that + # `RouteChecks.non_proxy_admin_allowed_routes_check` applies upstream + # for the `/user/info` route. + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if user_id == user_api_key_dict.user_id: + return + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=( + f"key not allowed to access this user's info. user_id={user_id}, " + f"key's user_id={user_api_key_dict.user_id}" + ), + ) + + async def _get_user_info_teams( prisma_client: Any, user_id: Optional[str], @@ -732,6 +767,7 @@ async def user_info( # noqa: PLR0915 try: user_id = _normalize_user_info_user_id(request=request, user_id=user_id) + _enforce_user_info_access(user_id=user_id, user_api_key_dict=user_api_key_dict) if prisma_client is None: raise Exception( @@ -2587,9 +2623,10 @@ async def get_user_daily_activity( if is_admin: entity_id = user_id # None means global view, otherwise filter by user else: + caller_user_id = require_caller_user_id_for_non_admin(user_api_key_dict) if user_id is None: - user_id = user_api_key_dict.user_id - if user_id != user_api_key_dict.user_id: + user_id = caller_user_id + if user_id != caller_user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={ @@ -2684,9 +2721,10 @@ async def get_user_daily_activity_aggregated( if is_admin: entity_id = user_id # None means global view, otherwise filter by user else: + caller_user_id = require_caller_user_id_for_non_admin(user_api_key_dict) if user_id is None: - user_id = user_api_key_dict.user_id - if user_id != user_api_key_dict.user_id: + user_id = caller_user_id + if user_id != caller_user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={ diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index e474cb7d15..1ee5bfb022 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -10,6 +10,7 @@ from litellm.proxy._types import ( hash_token, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view router = APIRouter() @@ -194,7 +195,8 @@ async def list_jwt_key_mappings( ): from litellm.proxy.proxy_server import prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail="Only proxy admins can list JWT key mappings" ) @@ -233,7 +235,8 @@ async def info_jwt_key_mapping( ): from litellm.proxy.proxy_server import prisma_client - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail="Only proxy admins can get JWT key mapping info" ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a01f5e6321..1145495e7d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_org_admin_for_team, _is_user_team_admin, _set_object_metadata_field, + _team_member_has_permission, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, @@ -4436,6 +4437,26 @@ def _get_admin_team_ids_from_objects( ] +def _get_team_ids_with_key_list_permission_from_objects( + user_api_key_dict: UserAPIKeyAuth, + team_objects: List[LiteLLM_TeamTable], +) -> List[str]: + """Filter team objects to non-admin teams where the caller has /key/list + permission via team_member_permissions. These teams should grant the + caller full key visibility (same as a team admin), so other members' + keys and service account keys (user_id=NULL) are returned.""" + return [ + team.team_id + for team in team_objects + if not _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team) + and _team_member_has_permission( + user_api_key_dict=user_api_key_dict, + team_obj=team, + permission=KeyManagementRoutes.KEY_LIST.value, + ) + ] + + def _get_member_team_ids_from_objects( user_api_key_dict: UserAPIKeyAuth, team_objects: List[LiteLLM_TeamTable], @@ -4589,6 +4610,17 @@ async def list_keys( user_api_key_dict=user_api_key_dict, team_objects=team_objects, ) + # Non-admin members with /key/list permission get full team-key + # visibility for that team — matching the UI contract that + # granting this permission lets them see all keys within the team. + list_permission_team_ids = ( + _get_team_ids_with_key_list_permission_from_objects( + user_api_key_dict=user_api_key_dict, + team_objects=team_objects, + ) + ) + if list_permission_team_ids: + admin_team_ids = list({*admin_team_ids, *list_permission_team_ids}) else: admin_team_ids = None diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9c510a568e..729493e1df 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -2120,7 +2120,8 @@ if MCP_AVAILABLE: Used by the UI to show a discovery grid when adding new MCP servers. """ - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail={ @@ -2177,7 +2178,8 @@ if MCP_AVAILABLE: async def get_openapi_registry( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail={ diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 61247ce5de..466ce47a6f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4683,9 +4683,11 @@ async def team_member_permissions( complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + # Admin Viewer follows the read-parity rule: see team permissions like + # a Proxy Admin would. Team / org admins keep their existing scope. if ( hasattr(user_api_key_dict, "user_role") - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value + and not _user_has_admin_view(user_api_key_dict) and not _is_user_team_admin( user_api_key_dict=user_api_key_dict, team_obj=complete_team_data ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 9dfc67370f..74ee7c7220 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -678,6 +678,7 @@ async def google_login( google_client_id=google_client_id, generic_client_id=generic_client_id, state=cli_state, + request=request, ) if return_to is not None and sso_redirect is not None: if SSOAuthenticationHandler._validate_return_to(return_to): @@ -1159,6 +1160,30 @@ async def get_generic_sso_response( authorization_code = request.query_params.get("code") if code_verifier: + # State-to-session-cookie binding. The non-PKCE branch below + # delegates to fastapi-sso's ``verify_and_process``, which + # performs its own session-cookie check. The PKCE branch + # bypasses that helper, so we validate the URL ``state`` + # against the ``litellm_oauth_state`` cookie set on the + # redirect response — without this an attacker can pre-mint + # a state + cached PKCE verifier and hijack a victim's auth + # code (Login-CSRF / token theft). + url_state = request.query_params.get("state") + cookie_state = request.cookies.get("litellm_oauth_state") + if ( + not url_state + or not cookie_state + or not secrets.compare_digest(url_state, cookie_state) + ): + raise ProxyException( + message=( + "Invalid OAuth state parameter — does not match " + "the browser-bound state cookie." + ), + type=ProxyErrorTypes.auth_error, + param="state", + code=status.HTTP_400_BAD_REQUEST, + ) if not authorization_code: raise ProxyException( message="Missing authorization code in callback", @@ -2147,6 +2172,7 @@ class SSOAuthenticationHandler: microsoft_client_id: Optional[str] = None, generic_client_id: Optional[str] = None, state: Optional[str] = None, + request: Optional[Request] = None, ) -> Optional[RedirectResponse]: """ Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url` @@ -2156,6 +2182,8 @@ class SSOAuthenticationHandler: google_client_id (Optional[str], optional): The Google Client ID. Defaults to None. microsoft_client_id (Optional[str], optional): The Microsoft Client ID. Defaults to None. generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None. + request: Optional FastAPI request, used to drive the ``Secure`` + attribute on the ``litellm_oauth_state`` CSRF cookie. Returns: RedirectResponse: The redirect response from the SSO provider. @@ -2266,6 +2294,7 @@ class SSOAuthenticationHandler: generic_sso=generic_sso, state=state, generic_authorization_endpoint=generic_authorization_endpoint, + request=request, ) raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" @@ -2276,6 +2305,7 @@ class SSOAuthenticationHandler: generic_sso: Any, state: Optional[str] = None, generic_authorization_endpoint: Optional[str] = None, + request: Optional[Request] = None, ) -> Optional[RedirectResponse]: """ Get the redirect response for Generic SSO @@ -2285,10 +2315,13 @@ class SSOAuthenticationHandler: from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache with generic_sso: - # TODO: state should be a random string and added to the user session with cookie - # or a cryptographicly signed state that we can verify stateless - # For simplification we are using a static state, this is not perfect but some - # SSO providers do not allow stateless verification + # State is bound to the caller's browser via a ``litellm_oauth_state`` + # HttpOnly cookie set on the redirect response below; the SSO + # callback validates the URL ``state`` against that cookie before + # completing the PKCE token exchange. Without this binding, an + # attacker who pre-mints a state + a cached PKCE verifier can hand + # the link to a victim and capture the resulting access token + # (Login CSRF / token theft). ( redirect_params, code_verifier, @@ -2355,6 +2388,31 @@ class SSOAuthenticationHandler: # Update the redirect response redirect_response.headers["location"] = new_url + + # Bind state to the user's browser session. The /callback + # handler validates the URL ``state`` against this cookie via + # ``secrets.compare_digest`` before exchanging the PKCE + # code_verifier. Only set the cookie when PKCE is in use + # (i.e. inside this ``code_verifier`` branch) so two + # concurrent SSO sessions — one PKCE, one plain — cannot + # overwrite each other's state cookie. + state_value = redirect_params.get("state") + if state_value and redirect_response is not None: + # Production-safe default: require HTTPS for the + # CSRF-protection cookie unless we can prove the + # incoming request is HTTP (local dev). Without + # ``Secure`` the cookie is sent over plain HTTP, + # letting a network observer read and replay the + # state value and bypass this protection. + secure_flag = request is None or request.url.scheme == "https" + redirect_response.set_cookie( + key="litellm_oauth_state", + value=state_value, + max_age=600, + httponly=True, + samesite="lax", + secure=secure_flag, + ) return redirect_response @staticmethod @@ -3972,6 +4030,7 @@ async def debug_sso_login(request: Request): microsoft_client_id=microsoft_client_id, google_client_id=google_client_id, generic_client_id=generic_client_id, + request=request, ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 4de29e0409..a50ce1d3c4 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -440,6 +440,15 @@ def _resolve_fetch_kwargs( kwargs: Dict[str, Any] = {"start_date": start_date, "end_date": end_date} if fn_name == "get_usage_data": if not is_admin: + if user_id is None: + # Defense-in-depth: the endpoint guard in usage_endpoints/endpoints.py + # should have already rejected this. If we ever reach here it means + # a future caller invoked the helper without scoping — fail loudly + # rather than issuing an unfiltered global query. + raise ValueError( + "Non-admin caller has user_id=None; refusing to issue an " + "unscoped query. Endpoint-level guard missing." + ) kwargs["user_id"] = user_id elif fn_args.get("user_id"): kwargs["user_id"] = fn_args["user_id"] diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py index 0dbe518afb..d0df80fed0 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -44,13 +44,17 @@ async def usage_ai_chat( """ from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_view, + require_caller_user_id_for_non_admin, ) from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( stream_usage_ai_chat, ) is_admin = _user_has_admin_view(user_api_key_dict) - user_id = user_api_key_dict.user_id + if is_admin: + user_id = user_api_key_dict.user_id + else: + user_id = require_caller_user_id_for_non_admin(user_api_key_dict) messages = [{"role": m.role, "content": m.content} for m in data.messages] return StreamingResponse( 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 714b5f3c7b..cc6c26fdf9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2324,14 +2324,10 @@ async def _register_pass_through_endpoint( dependencies = None if auth is not None and str(auth).lower() == "true": - # Authentication on a pass-through endpoint used to be enterprise- - # only — which left the OSS tier with no safe configuration: the - # default was ``auth=False`` (unauthenticated forwarder) and the - # safe ``auth=True`` raised at startup unless the operator had a - # license. The default is now ``True`` (safe-by-default), and - # turning it on no longer requires a license: an unauthenticated - # forwarder is a deployment choice the operator should be allowed - # to make explicitly, but the safe option must always be free. + # 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/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 6d1096d5ee..8d5d811691 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -220,6 +220,7 @@ class AttachmentRegistry: attachment: PolicyAttachment object to add """ self._attachments.append(attachment) + self._initialized = True verbose_proxy_logger.debug(f"Added attachment for policy: {attachment.policy}") def remove_attachments_for_policy(self, policy_name: str) -> int: diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index d3df16afde..75017c4660 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -226,6 +226,7 @@ class PolicyRegistry: policy: Policy object to add """ self._policies[policy_name] = policy + self._initialized = True verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}") def remove_policy(self, policy_name: str) -> bool: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 29d5bf8f6f..7893673ec5 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 @@ -334,6 +335,7 @@ from litellm.proxy.management_endpoints.callback_management_endpoints import ( ) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, + _user_has_admin_view, admin_can_invite_user, ) from litellm.proxy.management_endpoints.cost_tracking_settings import ( @@ -955,6 +957,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(), @@ -964,6 +1045,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() @@ -1043,6 +1125,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: @@ -1074,6 +1157,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: @@ -1845,6 +1929,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. @@ -1856,7 +1943,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: @@ -1871,11 +1965,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) @@ -1892,17 +1988,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}") @@ -1919,36 +2026,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, ): """ @@ -1967,31 +2195,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, refresh_ttl=True + await _increment_spend_counter_cache( + counter_key=counter_key, increment=base_spend ) - await spend_counter_cache.async_increment_cache( - key=counter_key, value=increment, refresh_ttl=True + +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 @@ -11603,7 +11963,7 @@ async def alerting_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, detail={ @@ -12715,7 +13075,7 @@ async def invitation_info( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, detail={ @@ -12891,9 +13251,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: @@ -12901,109 +13264,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 @@ -13147,7 +13497,7 @@ async def get_config_general_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, detail={"error": CommonProxyErrors.not_allowed_access.value}, @@ -13211,7 +13561,7 @@ async def get_config_list( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=400, detail={ @@ -13923,8 +14273,8 @@ async def get_model_cost_map_reload_status( Get the status of the scheduled model cost map reload job. """ - # Check if user is admin - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Read-only status check — admin viewers can read. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}", @@ -14026,7 +14376,8 @@ async def get_model_cost_map_source( - fallback_reason: human-readable reason why remote failed (null on success) - model_count: number of models in the currently loaded cost map """ - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Read-only source info — admin viewers can read. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}", @@ -14283,8 +14634,8 @@ async def get_anthropic_beta_headers_reload_status( Get the status of the scheduled Anthropic beta headers reload job. """ - # Check if user is admin - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Read-only status — admin viewers can read. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail=f"Access denied. Admin role required. Current role: {user_api_key_dict.user_role}", @@ -14390,7 +14741,8 @@ async def get_adaptive_router_state( adaptive-router deployment. Each snapshot's `router_name` field identifies which deployment it came from. """ - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Read-only state — admin viewers can read. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail={"error": CommonProxyErrors.not_allowed_access.value}, 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/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index c7bff7ec64..1f551d5ffe 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -6,6 +6,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -127,10 +128,10 @@ async def get_cloudzero_settings( Only the first 4 and last 4 characters of the API key are shown. Returns null/empty values when settings are not configured (consistent with other settings endpoints). - Only admin users can view CloudZero settings. + Only admin users (Proxy Admin or Admin Viewer) can view CloudZero settings. """ - # Validation - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Validation — Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail={"error": CommonProxyErrors.not_allowed_access.value}, diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 7d8fbf7461..60e54d005b 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -7,6 +7,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -140,9 +141,10 @@ async def get_vantage_settings( View current Vantage settings. Returns the current Vantage configuration with the API key masked for security. - Only admin users can view Vantage settings. + Only admin users (Proxy Admin or Admin Viewer) can view Vantage settings. """ - if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + # Admin Viewer follows the read-parity rule. + if not _user_has_admin_view(user_api_key_dict): raise HTTPException( status_code=403, detail={"error": CommonProxyErrors.not_allowed_access.value}, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8c5fce8409..69cd7b983e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3188,6 +3188,8 @@ class PrismaClient: t.organization_id as org_id, p.project_alias AS project_alias, tm.spend AS team_member_spend, + b_tm.tpm_limit AS team_member_tpm_limit, + b_tm.rpm_limit AS team_member_rpm_limit, m.aliases AS team_model_aliases, -- Added comma to separate b.* columns b.max_budget AS litellm_budget_table_max_budget, @@ -3203,6 +3205,7 @@ class PrismaClient: FROM "LiteLLM_VerificationToken" AS v LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id + LEFT JOIN "LiteLLM_BudgetTable" AS b_tm ON tm.budget_id = b_tm.budget_id LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id 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/proxy/vertex_ai_endpoints/langfuse_endpoints.py b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py index b6454bf077..8ce1bedcf9 100644 --- a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py @@ -12,11 +12,13 @@ import base64 import os from base64 import b64encode from typing import Optional +from urllib.parse import unquote import httpx -from fastapi import APIRouter, Request, Response +from fastapi import APIRouter, HTTPException, Request, Response, status import litellm +from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers @@ -27,6 +29,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router = APIRouter() default_vertex_config = None +_DEFAULT_LANGFUSE_HOST = "https://cloud.langfuse.com" def create_request_copy(request: Request): @@ -39,6 +42,116 @@ def create_request_copy(request: Request): } +def _decode_to_convergence(value: str) -> str: + previous = value + while True: + decoded = unquote(previous) + if decoded == previous: + return decoded + previous = decoded + + +def _normalize_langfuse_base_url(base_target_url: str) -> str: + if not ( + base_target_url.startswith("http://") or base_target_url.startswith("https://") + ): + # Existing behavior allows host-only Langfuse settings. + base_target_url = "http://" + base_target_url + + try: + base_url = httpx.URL(base_target_url) + except Exception as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": f"Invalid Langfuse host: {str(e)}"}, + ) + + if base_url.scheme not in ("http", "https") or not base_url.host: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid Langfuse host"}, + ) + + if base_url.userinfo: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Langfuse host must not include credentials"}, + ) + + return str(base_url) + + +def _validate_langfuse_proxy_path(endpoint: str) -> str: + decoded_endpoint = _decode_to_convergence(endpoint) + if any(ord(char) < 32 for char in decoded_endpoint): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid Langfuse endpoint path"}, + ) + if "\\" in decoded_endpoint or decoded_endpoint.startswith("//"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid Langfuse endpoint path"}, + ) + + endpoint_path = "/" + decoded_endpoint.lstrip("/") + if any(segment in (".", "..") for segment in endpoint_path.split("/")): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid Langfuse endpoint path"}, + ) + return endpoint_path + + +def _get_langfuse_proxy_credentials( + *, + dynamic_host_supplied: bool, + dynamic_langfuse_public_key: Optional[str], + dynamic_langfuse_secret_key: Optional[str], +): + if dynamic_host_supplied: + if not dynamic_langfuse_public_key or not dynamic_langfuse_secret_key: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "Dynamic Langfuse hosts must include dynamic Langfuse credentials" + }, + ) + return dynamic_langfuse_public_key, dynamic_langfuse_secret_key + + return ( + dynamic_langfuse_public_key + or litellm.utils.get_secret(secret_name="LANGFUSE_PUBLIC_KEY"), + dynamic_langfuse_secret_key + or litellm.utils.get_secret(secret_name="LANGFUSE_SECRET_KEY"), + ) + + +def _build_langfuse_proxy_target( + *, + endpoint: str, + base_target_url: str, + dynamic_host_supplied: bool, +): + endpoint_path = _validate_langfuse_proxy_path(endpoint) + base_url = httpx.URL(_normalize_langfuse_base_url(base_target_url)) + updated_url = base_url.copy_with(path=endpoint_path) + custom_headers = {} + + if dynamic_host_supplied and getattr(litellm, "user_url_validation", True): + try: + target_url, host_header = validate_url(str(updated_url)) + except SSRFError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": f"Invalid Langfuse host: {str(e)}"}, + ) + custom_headers["Host"] = host_header + return target_url, custom_headers + + return str(updated_url), custom_headers + + @router.api_route( "/langfuse/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], @@ -91,44 +204,33 @@ async def langfuse_proxy_route( elif k == "langfuse_host": dynamic_langfuse_host = v + dynamic_host_supplied = dynamic_langfuse_host is not None base_target_url: str = ( dynamic_langfuse_host - or os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com") - or "https://cloud.langfuse.com" + or os.getenv("LANGFUSE_HOST", _DEFAULT_LANGFUSE_HOST) + or _DEFAULT_LANGFUSE_HOST ) - if not ( - base_target_url.startswith("http://") or base_target_url.startswith("https://") - ): - # add http:// if unset, assume communicating over private network - e.g. render - base_target_url = "http://" + base_target_url - - encoded_endpoint = httpx.URL(endpoint).path - - # Ensure endpoint starts with '/' for proper URL construction - if not encoded_endpoint.startswith("/"): - encoded_endpoint = "/" + encoded_endpoint - - # Construct the full target URL using httpx - base_url = httpx.URL(base_target_url) - updated_url = base_url.copy_with(path=encoded_endpoint) - - # Add or update query parameters - langfuse_public_key = dynamic_langfuse_public_key or litellm.utils.get_secret( - secret_name="LANGFUSE_PUBLIC_KEY" + langfuse_public_key, langfuse_secret_key = _get_langfuse_proxy_credentials( + dynamic_host_supplied=dynamic_host_supplied, + dynamic_langfuse_public_key=dynamic_langfuse_public_key, + dynamic_langfuse_secret_key=dynamic_langfuse_secret_key, ) - langfuse_secret_key = dynamic_langfuse_secret_key or litellm.utils.get_secret( - secret_name="LANGFUSE_SECRET_KEY" + target_url, target_headers = _build_langfuse_proxy_target( + endpoint=endpoint, + base_target_url=base_target_url, + dynamic_host_supplied=dynamic_host_supplied, ) langfuse_combined_key = "Basic " + b64encode( f"{langfuse_public_key}:{langfuse_secret_key}".encode("utf-8") ).decode("ascii") + target_headers["Authorization"] = langfuse_combined_key ## CREATE PASS-THROUGH endpoint_func = create_pass_through_route( endpoint=endpoint, - target=str(updated_url), - custom_headers={"Authorization": langfuse_combined_key}, + target=target_url, + custom_headers=target_headers, query_params=dict(request.query_params), # type: ignore ) # dynamically construct pass-through endpoint based on incoming path received_value = await endpoint_func( diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index a98f9d666a..04347aebe3 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -38,6 +38,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( HiddenlayerGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, +) """ Pydantic object defining how to set guardrails on litellm proxy @@ -96,6 +99,7 @@ class SupportedGuardrailIntegrations(Enum): AKTO = "akto" MCP_JWT_SIGNER = "mcp_jwt_signer" LLM_AS_A_JUDGE = "llm_as_a_judge" + QOSTODIAN_NEXUS = "qostodian_nexus" class Role(Enum): @@ -773,6 +777,7 @@ class LitellmParams( QualifireGuardrailConfigModel, BlockCodeExecutionGuardrailConfigModel, HiddenlayerGuardrailConfigModel, + QostodianNexusConfigModel, ): guardrail: str = Field(description="The type of guardrail integration to use") mode: Union[str, List[str], Mode] = Field( 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/types/proxy/guardrails/guardrail_hooks/qohash.py b/litellm/types/proxy/guardrails/guardrail_hooks/qohash.py new file mode 100644 index 0000000000..5abfa69148 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/qohash.py @@ -0,0 +1,16 @@ +from typing import Optional + +from pydantic import Field + +from .base import GuardrailConfigModel + + +class QostodianNexusConfigModel(GuardrailConfigModel): + api_base: Optional[str] = Field( + default=None, + description="The API base URL for Qostodian Nexus. If not provided, the `QOSTODIAN_NEXUS_API_BASE` environment variable is checked. Defaults to http://nexus:8800.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Qostodian Nexus" 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 d21a4bf11d..c38c14a1f7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -22109,6 +22109,98 @@ "tool_use_system_prompt_tokens": 346, "supports_native_structured_output": true }, + "crusoe/deepseek-ai/DeepSeek-R1-0528": { + "input_cost_per_token": 3e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 7e-06, + "supports_function_calling": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": false + }, + "crusoe/deepseek-ai/DeepSeek-V3-0324": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "crusoe/google/gemma-3-12b-it": { + "input_cost_per_token": 1e-07, + "litellm_provider": "crusoe", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "crusoe/meta-llama/Llama-3.3-70B-Instruct": { + "input_cost_per_token": 2e-07, + "litellm_provider": "crusoe", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "crusoe/moonshotai/Kimi-K2-Thinking": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_tool_choice": false + }, + "crusoe/openai/gpt-oss-120b": { + "input_cost_per_token": 8e-07, + "litellm_provider": "crusoe", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507": { + "input_cost_per_token": 3e-06, + "litellm_provider": "crusoe", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, "lambda_ai/deepseek-llama3.3-70b": { "input_cost_per_token": 2e-07, "litellm_provider": "lambda_ai", @@ -34894,6 +34986,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/provider_endpoints_support.json b/provider_endpoints_support.json index ed49c14621..3fc7cd4318 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -635,6 +635,24 @@ "interactions": true } }, + "crusoe": { + "display_name": "Crusoe (`crusoe`)", + "url": "https://docs.litellm.ai/docs/providers/crusoe", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": true, + "interactions": true + } + }, "custom": { "display_name": "Custom (`custom`)", "url": "https://docs.litellm.ai/docs/providers/custom_llm_server", diff --git a/pyproject.toml b/pyproject.toml index 0ef0a993dd..65b9bd2c98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -149,6 +149,8 @@ dev = [ "parameterized==0.9.0", "openapi-core==0.22.0; python_version < '3.14'", "pytest-timeout==2.4.0", + "vcrpy==8.1.1", + "pytest-recording==0.13.4", ] proxy-dev = [ "prisma==0.11.0", diff --git a/tests/_flush_vcr_cache.py b/tests/_flush_vcr_cache.py new file mode 100644 index 0000000000..d236c88fa3 --- /dev/null +++ b/tests/_flush_vcr_cache.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import os +import sys + +import redis + +from tests._vcr_redis_persister import CASSETTE_REDIS_URL_ENV, _redis_url_from_env + +PREFIX = "litellm:vcr:cassette:" +SCAN_BATCH = 500 + + +def _client() -> redis.Redis: + url = _redis_url_from_env() + if not url: + sys.exit(f"Set {CASSETTE_REDIS_URL_ENV} to flush the VCR cache") + return redis.Redis.from_url( + url, + socket_timeout=5, + socket_connect_timeout=5, + decode_responses=False, + ) + + +def main() -> None: + client = _client() + deleted = 0 + pipeline = client.pipeline(transaction=False) + pending = 0 + for key in client.scan_iter(match=f"{PREFIX}*", count=SCAN_BATCH): + pipeline.delete(key) + pending += 1 + if pending >= SCAN_BATCH: + deleted += sum(pipeline.execute()) + pipeline = client.pipeline(transaction=False) + pending = 0 + if pending: + deleted += sum(pipeline.execute()) + print(f"Deleted {deleted} VCR cassette key(s) under {PREFIX!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/_vcr_redis_persister.py b/tests/_vcr_redis_persister.py new file mode 100644 index 0000000000..a6ed448f1c --- /dev/null +++ b/tests/_vcr_redis_persister.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import logging +import os +from typing import Any, Optional + +from vcr.persisters.filesystem import CassetteNotFoundError +from vcr.serialize import deserialize, serialize + +CASSETTE_TTL_SECONDS = 24 * 60 * 60 +REDIS_KEY_PREFIX = "litellm:vcr:cassette:" +CASSETTE_REDIS_URL_ENV = "CASSETTE_REDIS_URL" +VCR_VERBOSE_ENV = "LITELLM_VCR_VERBOSE" +MAX_EPISODES_PER_CASSETTE = 50 + +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +_log = logging.getLogger(__name__) +_passed_by_cassette_key: dict[str, bool] = {} + + +def mark_test_outcome_for_cassette(cassette_path: str, passed: bool) -> None: + _passed_by_cassette_key[redis_key_for(cassette_path)] = passed + + +def redis_key_for(cassette_path: str) -> str: + abs_path = os.path.abspath(str(cassette_path)) + try: + rel = os.path.relpath(abs_path, start=_REPO_ROOT) + except ValueError: + rel = os.path.basename(abs_path) + if rel.endswith(".yaml"): + rel = rel[: -len(".yaml")] + rel = rel.replace("/cassettes/", "/").lstrip("./") + return f"{REDIS_KEY_PREFIX}{rel}" + + +def _redis_url_from_env() -> Optional[str]: + return os.environ.get(CASSETTE_REDIS_URL_ENV) or None + + +def _build_default_client(): + import redis + from redis.backoff import ExponentialBackoff + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + from redis.retry import Retry + + url = _redis_url_from_env() + if not url: + raise RuntimeError( + f"Set {CASSETTE_REDIS_URL_ENV} to enable the VCR persister. " + "Cassette Redis is intentionally separate from the application " + "Redis (REDIS_URL/REDIS_HOST) to avoid being flushed by tests." + ) + return redis.Redis.from_url( + url, + socket_timeout=5, + socket_connect_timeout=5, + decode_responses=False, + retry=Retry(ExponentialBackoff(cap=2, base=0.1), retries=2), + retry_on_error=[RedisConnectionError, RedisTimeoutError], + ) + + +def make_redis_persister( + client: Optional[Any] = None, + ttl_seconds: int = CASSETTE_TTL_SECONDS, +): + redis_client = client if client is not None else _build_default_client() + + try: + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import TimeoutError as RedisTimeoutError + + _transient_errors: tuple = (RedisConnectionError, RedisTimeoutError) + except ImportError: # pragma: no cover - redis is a hard test dep + _transient_errors = () + + class _RedisPersister: + @staticmethod + def load_cassette(cassette_path, serializer): + try: + data = redis_client.get(redis_key_for(cassette_path)) + except _transient_errors as exc: + _log.warning( + "VCR redis load failed for %s; treating as cache miss: %s", + cassette_path, + exc, + ) + raise CassetteNotFoundError() from exc + if data is None: + raise CassetteNotFoundError() + if isinstance(data, bytes): + data = data.decode("utf-8") + return deserialize(data, serializer) + + @staticmethod + def save_cassette(cassette_path, cassette_dict, serializer): + key = redis_key_for(cassette_path) + passed = _passed_by_cassette_key.pop(key, True) + episode_count = len(cassette_dict.get("requests", []) or []) + if episode_count > MAX_EPISODES_PER_CASSETTE: + _log.warning( + "VCR redis save refused for %s; cassette has %d episodes " + "(> MAX_EPISODES_PER_CASSETTE=%d). The test likely produces " + "non-deterministic request bodies (e.g. uuid) and is " + "appending instead of replaying. Opt it out with the " + "no-vcr list in conftest, or stabilize its request body.", + cassette_path, + episode_count, + MAX_EPISODES_PER_CASSETTE, + ) + return + if not passed: + _log.info( + "VCR redis save skipped for %s; test did not pass — " + "leaving any prior cassette intact", + cassette_path, + ) + return + data = serialize(cassette_dict, serializer) + payload = data.encode("utf-8") if isinstance(data, str) else data + try: + redis_client.set(key, payload, ex=ttl_seconds) + except _transient_errors as exc: + _log.warning( + "VCR redis save failed for %s; cassette not persisted: %s", + cassette_path, + exc, + ) + + return _RedisPersister + + +def filter_non_2xx_response(response): + if not isinstance(response, dict): + return response + status = response.get("status") + code = status.get("code") if isinstance(status, dict) else status + if not isinstance(code, int): + return response + return response if 200 <= code < 300 else None + + +_PATCHED_AIOHTTP_RECORD = False + + +def patch_vcrpy_aiohttp_record_path() -> None: + """Re-feed the response body into aiohttp's StreamReader after vcrpy's + record_response drains it, so downstream consumers (e.g. + LiteLLMAiohttpTransport.AiohttpResponseStream) can still read it.""" + global _PATCHED_AIOHTTP_RECORD + if _PATCHED_AIOHTTP_RECORD: + return + import vcr.stubs.aiohttp_stubs as _aiohttp_stubs + + _orig_record_response = _aiohttp_stubs.record_response + + async def _record_response_preserving_body(cassette, vcr_request, response): + await _orig_record_response(cassette, vcr_request, response) + body = getattr(response, "_body", None) or b"" + if body: + response.content.unread_data(body) + + _aiohttp_stubs.record_response = _record_response_preserving_body + _PATCHED_AIOHTTP_RECORD = True + + +def vcr_verbose_enabled() -> bool: + return os.environ.get(VCR_VERBOSE_ENV) == "1" + + +def format_vcr_verdict(cassette: Any) -> str: + if cassette is None: + return "[VCR NOOP]" + played = getattr(cassette, "play_count", 0) or 0 + dirty = getattr(cassette, "dirty", False) + total = len(cassette) if hasattr(cassette, "__len__") else 0 + if played == 0 and not dirty: + return "[VCR NOOP] (no http traffic)" + if played > 0 and not dirty: + return f"[VCR HIT] {played} replayed, 0 new ({total} cassette entries)" + if played == 0 and dirty: + return f"[VCR MISS] 0 replayed, recorded new ({total} cassette entries)" + return ( + f"[VCR PARTIAL] {played} replayed + new recordings ({total} cassette entries)" + ) diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index 4ecc0d0bc9..344e38da83 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -169,3 +169,4 @@ langchain-mcp-adapters: >=0.2.1 # MIT License langgraph: >=1.0.10 # MIT License langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE pytest-rerunfailures: >=15.1 # MPL 2.0 license +pytest-recording: >=0.13.4 # MIT license diff --git a/tests/litellm/llms/azure/__init__.py b/tests/litellm/llms/azure/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/litellm/llms/azure/test_azure_embedding.py b/tests/litellm/llms/azure/test_azure_embedding.py new file mode 100644 index 0000000000..22ee503ef0 --- /dev/null +++ b/tests/litellm/llms/azure/test_azure_embedding.py @@ -0,0 +1,94 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) +) + +from litellm.llms.azure.azure import AzureChatCompletion +from litellm.types.utils import EmbeddingResponse, Usage + + +def _make_embedding_response() -> EmbeddingResponse: + return EmbeddingResponse( + model="text-embedding-3-large", + usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), + data=[{"embedding": [0.1, 0.2, 0.3], "index": 0, "object": "embedding"}], + ) + + +def _make_logging_obj() -> MagicMock: + return MagicMock() + + +class TestAzureV1AsyncEmbedding: + def test_aembedding_receives_api_version(self): + """Regression: api_version must be forwarded to aembedding() when aembedding=True. + Without the fix, it was silently dropped, causing AsyncAzureOpenAI to be used + instead of AsyncOpenAI for Azure AI Foundry (v1) endpoints. Fixes #24848.""" + handler = AzureChatCompletion() + + with patch.object(handler, "aembedding") as mock_aembedding: + handler.embedding( + model="text-embedding-3-large", + input=["hello world"], + api_base="https://my-endpoint.openai.azure.com", + api_version="v1", + timeout=60.0, + logging_obj=_make_logging_obj(), + model_response=_make_embedding_response(), + optional_params={}, + api_key="fake-key", + aembedding=True, + litellm_params={}, + ) + + mock_aembedding.assert_called_once() + _, kwargs = mock_aembedding.call_args + assert kwargs.get("api_version") == "v1" + + def test_get_azure_openai_client_returns_async_openai_for_v1(self): + from openai import AsyncAzureOpenAI, AsyncOpenAI + + handler = AzureChatCompletion() + client = handler.get_azure_openai_client( + api_key="fake-key", + api_base="https://my-endpoint.openai.azure.com", + api_version="v1", + _is_async=True, + litellm_params={}, + ) + + assert isinstance(client, AsyncOpenAI) + assert not isinstance(client, AsyncAzureOpenAI) + + def test_get_azure_openai_client_uses_v1_base_url(self): + handler = AzureChatCompletion() + client = handler.get_azure_openai_client( + api_key="fake-key", + api_base="https://my-endpoint.openai.azure.com", + api_version="v1", + _is_async=True, + litellm_params={}, + ) + + assert client is not None + assert "/openai/v1/" in str(client.base_url) + + @pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) + def test_all_v1_variants_use_openai_client(self, api_version: str): + from openai import AsyncOpenAI + + handler = AzureChatCompletion() + client = handler.get_azure_openai_client( + api_key="fake-key", + api_base="https://my-endpoint.openai.azure.com", + api_version=api_version, + _is_async=True, + litellm_params={}, + ) + + assert isinstance(client, AsyncOpenAI) diff --git a/tests/litellm_utils_tests/test_bedrock_token_counter.py b/tests/litellm_utils_tests/test_bedrock_token_counter.py index 62a595baaa..9fb2463e8b 100644 --- a/tests/litellm_utils_tests/test_bedrock_token_counter.py +++ b/tests/litellm_utils_tests/test_bedrock_token_counter.py @@ -130,7 +130,7 @@ class TestBedrockCountTokensEndpoint: ) assert ( url - == "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-lite-v1:0/count-tokens" + == "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.nova-lite-v1%3A0/count-tokens" ) def test_api_base_overrides_default(self): @@ -141,7 +141,7 @@ class TestBedrockCountTokensEndpoint: aws_region_name="us-east-1", api_base=custom_base, ) - assert url == f"{custom_base}/model/amazon.nova-lite-v1:0/count-tokens" + assert url == f"{custom_base}/model/amazon.nova-lite-v1%3A0/count-tokens" def test_aws_bedrock_runtime_endpoint_overrides_default(self): handler = self._make_handler() @@ -153,7 +153,7 @@ class TestBedrockCountTokensEndpoint: aws_region_name="eu-west-1", aws_bedrock_runtime_endpoint=custom_endpoint, ) - assert url == f"{custom_endpoint}/model/amazon.nova-lite-v1:0/count-tokens" + assert url == f"{custom_endpoint}/model/amazon.nova-lite-v1%3A0/count-tokens" def test_api_base_takes_priority_over_aws_bedrock_runtime_endpoint(self): handler = self._make_handler() @@ -165,7 +165,7 @@ class TestBedrockCountTokensEndpoint: api_base=api_base, aws_bedrock_runtime_endpoint=runtime_endpoint, ) - assert url == f"{api_base}/model/amazon.nova-lite-v1:0/count-tokens" + assert url == f"{api_base}/model/amazon.nova-lite-v1%3A0/count-tokens" def test_env_var_overrides_default(self, monkeypatch): monkeypatch.setenv( diff --git a/tests/llm_responses_api_testing/conftest.py b/tests/llm_responses_api_testing/conftest.py index 0b03348190..80f36e159a 100644 --- a/tests/llm_responses_api_testing/conftest.py +++ b/tests/llm_responses_api_testing/conftest.py @@ -1,5 +1,6 @@ # conftest.py +import asyncio import importlib import os import sys @@ -9,9 +10,161 @@ import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import litellm -import asyncio +import litellm # noqa: E402 + +from tests._vcr_redis_persister import ( # noqa: E402 + filter_non_2xx_response, + format_vcr_verdict, + make_redis_persister, + mark_test_outcome_for_cassette, + patch_vcrpy_aiohttp_record_path, + vcr_verbose_enabled, +) + + +_controller_pluginmanager = None +_controller_terminal_reporter = None + + +_FILTERED_REQUEST_HEADERS = ( + "authorization", + "x-api-key", + "anthropic-api-key", + "anthropic-version", + "openai-api-key", + "azure-api-key", + "api-key", + "cookie", + "x-amz-security-token", + "x-amz-date", + "x-amz-content-sha256", + "amz-sdk-invocation-id", + "amz-sdk-request", + "x-goog-api-key", + "x-goog-user-project", +) + +_FILTERED_RESPONSE_HEADERS = ( + "set-cookie", + "x-request-id", + "request-id", + "cf-ray", + "anthropic-organization-id", + "openai-organization", + "x-amzn-requestid", + "x-amzn-trace-id", + "date", +) + + +def _scrub_response(response): + if not isinstance(response, dict): + return response + headers = response.get("headers") or {} + if isinstance(headers, dict): + for header in list(headers): + if header.lower() in _FILTERED_RESPONSE_HEADERS: + headers.pop(header, None) + return response + + +def _before_record_response(response): + return filter_non_2xx_response(_scrub_response(response)) + + +@pytest.fixture(scope="module") +def vcr_config(): + return { + "filter_headers": list(_FILTERED_REQUEST_HEADERS), + "decode_compressed_response": True, + "record_mode": "new_episodes", + "allow_playback_repeats": True, + "match_on": ( + "method", + "scheme", + "host", + "port", + "path", + "query", + "body", + ), + "before_record_response": _before_record_response, + } + + +def _vcr_disabled() -> bool: + if os.environ.get("LITELLM_VCR_DISABLE") == "1": + return True + return not os.environ.get("CASSETTE_REDIS_URL") + + +def pytest_recording_configure(config, vcr): + if _vcr_disabled(): + return + vcr.register_persister(make_redis_persister()) + patch_vcrpy_aiohttp_record_path() + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + outcome = yield + rep = outcome.get_result() + setattr(item, f"rep_{rep.when}", rep) + + +@pytest.fixture(autouse=True) +def _vcr_outcome_gate(request, vcr): + yield + cassette = vcr + rep_call = getattr(request.node, "rep_call", None) + test_passed = bool(rep_call and rep_call.passed) + cassette_path = getattr(cassette, "_path", None) if cassette is not None else None + if cassette_path: + mark_test_outcome_for_cassette(cassette_path, test_passed) + + if not vcr_verbose_enabled(): + return + verdict = format_vcr_verdict(cassette) + request.node.user_properties.append(("vcr_verdict", verdict)) + + +def pytest_configure(config): + global _controller_pluginmanager + if os.environ.get("PYTEST_XDIST_WORKER"): + return + _controller_pluginmanager = config.pluginmanager + + +def _resolve_terminal_reporter(): + global _controller_terminal_reporter + if _controller_terminal_reporter is not None: + return _controller_terminal_reporter + if _controller_pluginmanager is None: + return None + _controller_terminal_reporter = _controller_pluginmanager.getplugin( + "terminalreporter" + ) + return _controller_terminal_reporter + + +def pytest_runtest_logreport(report): + if report.when != "teardown": + return + if os.environ.get("PYTEST_XDIST_WORKER"): + return + if not vcr_verbose_enabled(): + return + reporter = _resolve_terminal_reporter() + if reporter is None: + return + verdict = next( + (v for k, v in (report.user_properties or []) if k == "vcr_verdict"), + None, + ) + if not verdict: + return + reporter.write_line(f"{verdict} :: {report.nodeid}") @pytest.fixture(scope="session") @@ -61,15 +214,18 @@ def setup_and_teardown(): def pytest_collection_modifyitems(config, items): - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + if not _vcr_disabled(): + for item in items: + if item.get_closest_marker("vcr") is not None: + continue + item.add_marker(pytest.mark.vcr) + custom_logger_tests = [ item for item in items if "custom_logger" in item.parent.name ] other_tests = [item for item in items if "custom_logger" not in item.parent.name] - # Sort tests based on their names custom_logger_tests.sort(key=lambda x: x.name) other_tests.sort(key=lambda x: x.name) - # Reorder the items list items[:] = custom_logger_tests + other_tests diff --git a/tests/llm_translation/Readme.md b/tests/llm_translation/Readme.md index db84e7c33c..958adbd975 100644 --- a/tests/llm_translation/Readme.md +++ b/tests/llm_translation/Readme.md @@ -1,3 +1,41 @@ -Unit tests for individual LLM providers. +Unit tests for individual LLM providers. -Name of the test file is the name of the LLM provider - e.g. `test_openai.py` is for OpenAI. \ No newline at end of file +Name of the test file is the name of the LLM provider - e.g. `test_openai.py` is for OpenAI. + +## Redis-backed VCR cache + +Every test in this directory is auto-decorated with `@pytest.mark.vcr` (via +`conftest.py`). The first time a test runs we hit the live provider and +record the HTTP exchange into Redis under +`litellm:vcr:cassette:`. Every subsequent run within 24h replays +from Redis without touching the network. The 24h TTL means each new day's +first run records again, so upstream API drift surfaces within a day. + +The persister, header scrubbing, and 2xx-only filtering are defined in +`tests/_vcr_redis_persister.py`. Files that already use `respx` (which +patches the same httpx transport vcrpy does) are excluded from the +auto-marker — see `_RESPX_CONFLICTING_FILES` in `conftest.py`. + +### Required environment + +`REDIS_HOST`, `REDIS_PORT`, `REDIS_PASSWORD` — same vars CircleCI uses for +its other Redis-backed jobs. Provider credentials +(`ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `AWS_*`, etc.) are needed only on +cache-miss (the daily re-record), not on replay. + +### Flushing the cache + +When you want the next run to re-record immediately instead of waiting +for the 24h TTL: + +```bash +make test-llm-translation-flush-vcr-cache +``` + +### Disabling VCR + +Skip the cache entirely (every call goes live, no recording): + +```bash +LITELLM_VCR_DISABLE=1 uv run pytest tests/llm_translation/test_.py +``` diff --git a/tests/llm_translation/base_llm_unit_tests.py b/tests/llm_translation/base_llm_unit_tests.py index f3b1895323..aedc4f810c 100644 --- a/tests/llm_translation/base_llm_unit_tests.py +++ b/tests/llm_translation/base_llm_unit_tests.py @@ -401,7 +401,7 @@ class BaseLLMChatTest(ABC): { "type": "file", "file": { - "file_id": "https://upload.wikimedia.org/wikipedia/commons/2/20/Re_example.pdf" + "file_id": "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/tests/llm_translation/fixtures/dummy.pdf" }, }, ] diff --git a/tests/llm_translation/conftest.py b/tests/llm_translation/conftest.py index d315dc63bc..09da0520be 100644 --- a/tests/llm_translation/conftest.py +++ b/tests/llm_translation/conftest.py @@ -5,6 +5,7 @@ # - Function-scoped fixture resets litellm globals to true defaults # - Module-scoped reload only in single-process mode +import asyncio import importlib import os import sys @@ -14,9 +15,195 @@ import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import litellm -import asyncio +import litellm # noqa: E402 + +from tests._vcr_redis_persister import ( # noqa: E402 + filter_non_2xx_response, + format_vcr_verdict, + make_redis_persister, + mark_test_outcome_for_cassette, + patch_vcrpy_aiohttp_record_path, + vcr_verbose_enabled, +) + + +_controller_pluginmanager = None +_controller_terminal_reporter = None + + +# vcrpy and respx both patch the httpx transport — applying both makes one +# silently win, so respx-using files opt out of the auto-marker. +_RESPX_CONFLICTING_FILES = frozenset( + { + "test_azure_o_series.py", + "test_gpt4o_audio.py", + "test_nvidia_nim.py", + "test_openai.py", + "test_openai_o1.py", + "test_prompt_caching.py", + "test_text_completion_unit_tests.py", + "test_xai.py", + } +) +_VCR_AUTO_MARKER_SKIP_FILES = _RESPX_CONFLICTING_FILES | frozenset( + {"test_vcr_redis_persister.py"} +) + +# Tests that observe live cross-call provider state (e.g. prompt-cache +# warm-up between two consecutive calls); replay can't reproduce that state. +_VCR_INCOMPATIBLE_NODEID_SUFFIXES = frozenset( + { + "::test_prompt_caching", + "TestBedrockInvokeNovaJson::test_json_response_pydantic_obj", + "::test_bedrock_converse__streaming_passthrough", + } +) + + +def _is_vcr_incompatible(nodeid: str) -> bool: + return any(nodeid.endswith(suffix) for suffix in _VCR_INCOMPATIBLE_NODEID_SUFFIXES) + + +_FILTERED_REQUEST_HEADERS = ( + "authorization", + "x-api-key", + "anthropic-api-key", + "anthropic-version", + "openai-api-key", + "azure-api-key", + "api-key", + "cookie", + "x-amz-security-token", + "x-amz-date", + "x-amz-content-sha256", + "amz-sdk-invocation-id", + "amz-sdk-request", + "x-goog-api-key", + "x-goog-user-project", +) + +_FILTERED_RESPONSE_HEADERS = ( + "set-cookie", + "x-request-id", + "request-id", + "cf-ray", + "anthropic-organization-id", + "openai-organization", + "x-amzn-requestid", + "x-amzn-trace-id", + "date", +) + + +def _scrub_response(response): + if not isinstance(response, dict): + return response + headers = response.get("headers") or {} + if isinstance(headers, dict): + for header in list(headers): + if header.lower() in _FILTERED_RESPONSE_HEADERS: + headers.pop(header, None) + return response + + +def _before_record_response(response): + return filter_non_2xx_response(_scrub_response(response)) + + +@pytest.fixture(scope="module") +def vcr_config(): + return { + "filter_headers": list(_FILTERED_REQUEST_HEADERS), + "decode_compressed_response": True, + "record_mode": "new_episodes", + "allow_playback_repeats": True, + "match_on": ( + "method", + "scheme", + "host", + "port", + "path", + "query", + "body", + ), + "before_record_response": _before_record_response, + } + + +def _vcr_disabled() -> bool: + if os.environ.get("LITELLM_VCR_DISABLE") == "1": + return True + return not os.environ.get("CASSETTE_REDIS_URL") + + +def pytest_recording_configure(config, vcr): + if _vcr_disabled(): + return + vcr.register_persister(make_redis_persister()) + patch_vcrpy_aiohttp_record_path() + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + outcome = yield + rep = outcome.get_result() + setattr(item, f"rep_{rep.when}", rep) + + +@pytest.fixture(autouse=True) +def _vcr_outcome_gate(request, vcr): + yield + cassette = vcr + rep_call = getattr(request.node, "rep_call", None) + test_passed = bool(rep_call and rep_call.passed) + cassette_path = getattr(cassette, "_path", None) if cassette is not None else None + if cassette_path: + mark_test_outcome_for_cassette(cassette_path, test_passed) + + if not vcr_verbose_enabled(): + return + verdict = format_vcr_verdict(cassette) + request.node.user_properties.append(("vcr_verdict", verdict)) + + +def pytest_configure(config): + global _controller_pluginmanager + if os.environ.get("PYTEST_XDIST_WORKER"): + return + _controller_pluginmanager = config.pluginmanager + + +def _resolve_terminal_reporter(): + global _controller_terminal_reporter + if _controller_terminal_reporter is not None: + return _controller_terminal_reporter + if _controller_pluginmanager is None: + return None + _controller_terminal_reporter = _controller_pluginmanager.getplugin( + "terminalreporter" + ) + return _controller_terminal_reporter + + +def pytest_runtest_logreport(report): + if report.when != "teardown": + return + if os.environ.get("PYTEST_XDIST_WORKER"): + return + if not vcr_verbose_enabled(): + return + reporter = _resolve_terminal_reporter() + if reporter is None: + return + verdict = next( + (v for k, v in (report.user_properties or []) if k == "vcr_verdict"), + None, + ) + if not verdict: + return + reporter.write_line(f"{verdict} :: {report.nodeid}") + # --------------------------------------------------------------------------- # Capture TRUE defaults at conftest import time (before test modules pollute). @@ -48,7 +235,6 @@ def event_loop(): @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(event_loop): # Add event_loop as a dependency - curr_dir = os.getcwd() sys.path.insert(0, os.path.abspath("../..")) import litellm @@ -97,15 +283,23 @@ def setup_and_teardown(event_loop): # Add event_loop as a dependency def pytest_collection_modifyitems(config, items): - # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + if not _vcr_disabled(): + for item in items: + filename = os.path.basename(str(item.fspath)) + if filename in _VCR_AUTO_MARKER_SKIP_FILES: + continue + if _is_vcr_incompatible(item.nodeid): + continue + if item.get_closest_marker("vcr") is not None: + continue + item.add_marker(pytest.mark.vcr) + custom_logger_tests = [ item for item in items if "custom_logger" in item.parent.name ] other_tests = [item for item in items if "custom_logger" not in item.parent.name] - # Sort tests based on their names custom_logger_tests.sort(key=lambda x: x.name) other_tests.sort(key=lambda x: x.name) - # Reorder the items list items[:] = custom_logger_tests + other_tests diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 7b2b6bed6a..371b27c5b2 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -1885,3 +1885,42 @@ def test_metadata_filter_applies_to_azure_anthropic(): headers={}, ) assert data.get("metadata") == {"user_id": "u2"} + + +def test_anthropic_basic_completion_replay(): + response = litellm.completion( + model="anthropic/claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello!"}], + ) + + assert response is not None + content = response.choices[0].message.content + assert isinstance(content, str) and content.strip(), content + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.choices[0].finish_reason in {"stop", "length"} + + +def test_anthropic_streaming_completion_replay(): + stream = litellm.completion( + model="anthropic/claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello!"}], + stream=True, + ) + + collected_text = "" + finish_reason = None + chunk_count = 0 + for chunk in stream: + chunk_count += 1 + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if delta and delta.content: + collected_text += delta.content + if chunk.choices[0].finish_reason: + finish_reason = chunk.choices[0].finish_reason + + assert chunk_count > 1, "expected multiple SSE chunks from streaming response" + assert collected_text.strip(), collected_text + assert finish_reason in {"stop", "length"} diff --git a/tests/llm_translation/test_crusoe.py b/tests/llm_translation/test_crusoe.py new file mode 100644 index 0000000000..56aa4e4cd4 --- /dev/null +++ b/tests/llm_translation/test_crusoe.py @@ -0,0 +1,108 @@ +""" +Tests for Crusoe provider integration +""" +import os +from unittest import mock + +import litellm + +CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" + + +def test_crusoe_json_registry(): + """Test CrusoeChatConfig is loaded from JSON provider registry""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert JSONProviderRegistry.exists("crusoe") + config = JSONProviderRegistry.get("crusoe") + assert config is not None + assert config.base_url == CRUSOE_API_BASE + assert config.api_key_env == "CRUSOE_API_KEY" + assert config.api_base_env == "CRUSOE_API_BASE" + + +def test_crusoe_get_openai_compatible_provider_info(): + """Test Crusoe provider info retrieval""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("crusoe"))() + + # Test with default values (no env vars set) + with mock.patch.dict(os.environ, {}, clear=True): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == CRUSOE_API_BASE + assert api_key is None + + # Test with environment variables + with mock.patch.dict( + os.environ, + { + "CRUSOE_API_KEY": "test-key", + "CRUSOE_API_BASE": "https://custom.crusoecloud.com/v1", + }, + ): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == "https://custom.crusoecloud.com/v1" + assert api_key == "test-key" + + # Test with explicit parameters (should override env vars) + with mock.patch.dict( + os.environ, + { + "CRUSOE_API_KEY": "env-key", + "CRUSOE_API_BASE": "https://env.crusoecloud.com/v1", + }, + ): + api_base, api_key = config._get_openai_compatible_provider_info( + "https://param.crusoecloud.com/v1", "param-key" + ) + assert api_base == "https://param.crusoecloud.com/v1" + assert api_key == "param-key" + + +def test_get_llm_provider_crusoe(): + """Test that get_llm_provider correctly identifies Crusoe""" + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + # Test with crusoe/model-name format + model, provider, api_key, api_base = get_llm_provider( + "crusoe/meta-llama/Llama-3.3-70B-Instruct" + ) + assert model == "meta-llama/Llama-3.3-70B-Instruct" + assert provider == "crusoe" + + +def test_crusoe_models_configuration(): + """Test that Crusoe models are configured correctly""" + from litellm import get_model_info + + original_model_cost = litellm.model_cost + original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + try: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + crusoe_models = [ + "crusoe/meta-llama/Llama-3.3-70B-Instruct", + "crusoe/deepseek-ai/DeepSeek-R1-0528", + "crusoe/deepseek-ai/DeepSeek-V3-0324", + "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507", + "crusoe/moonshotai/Kimi-K2-Thinking", + "crusoe/openai/gpt-oss-120b", + "crusoe/google/gemma-3-12b-it", + ] + + for model in crusoe_models: + model_info = get_model_info(model) + assert model_info is not None, f"Model info not found for {model}" + assert model_info.get("litellm_provider") == "crusoe", ( + f"{model} should have crusoe as provider" + ) + assert model_info.get("mode") == "chat", f"{model} should be in chat mode" + finally: + litellm.model_cost = original_model_cost + if original_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env diff --git a/tests/llm_translation/test_vcr_redis_persister.py b/tests/llm_translation/test_vcr_redis_persister.py new file mode 100644 index 0000000000..6e62e4491c --- /dev/null +++ b/tests/llm_translation/test_vcr_redis_persister.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +import os +import sys + +import fakeredis +import pytest +from redis.exceptions import ConnectionError as RedisConnectionError +from vcr.persisters.filesystem import CassetteNotFoundError +from vcr.request import Request +from vcr.serializers import yamlserializer + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))) + +from tests._vcr_redis_persister import ( # noqa: E402 + CASSETTE_TTL_SECONDS, + MAX_EPISODES_PER_CASSETTE, + filter_non_2xx_response, + make_redis_persister, + mark_test_outcome_for_cassette, + redis_key_for, +) + + +def _sample_cassette_dict(): + request = Request( + method="POST", + uri="https://api.anthropic.com/v1/messages", + body=b'{"model":"claude","messages":[{"role":"user","content":"hi"}]}', + headers={"content-type": "application/json"}, + ) + response = { + "status": {"code": 200, "message": "OK"}, + "headers": {"content-type": ["application/json"]}, + "body": {"string": b'{"id":"msg_1","type":"message"}'}, + } + return {"requests": [request], "responses": [response]} + + +def _persister_with_fake_redis(): + fake = fakeredis.FakeStrictRedis() + return fake, make_redis_persister(client=fake) + + +def test_save_then_load_roundtrips_cassette_content(): + _, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_y" + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + requests, responses = persister.load_cassette(cassette_id, yamlserializer) + + assert len(requests) == 1 + assert len(responses) == 1 + assert requests[0].method == "POST" + assert requests[0].uri == "https://api.anthropic.com/v1/messages" + assert responses[0]["status"]["code"] == 200 + assert responses[0]["body"]["string"] == b'{"id":"msg_1","type":"message"}' + + +def test_saved_key_has_24h_ttl(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_ttl" + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + + ttl = fake.ttl(redis_key_for(cassette_id)) + assert CASSETTE_TTL_SECONDS - 5 <= ttl <= CASSETTE_TTL_SECONDS + + +def test_load_missing_key_raises_cassette_not_found(): + _, persister = _persister_with_fake_redis() + with pytest.raises(CassetteNotFoundError): + persister.load_cassette("never/recorded", yamlserializer) + + +def test_redis_key_normalizes_path_passed_by_pytest_recording(): + raw = "tests/llm_translation/cassettes/test_anthropic/test_streaming.yaml" + assert ( + redis_key_for(raw) + == "litellm:vcr:cassette:tests/llm_translation/test_anthropic/test_streaming" + ) + + +def test_redis_key_is_stable_across_working_directories(tmp_path, monkeypatch): + repo_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + abs_cassette = os.path.join( + repo_root, + "tests/llm_translation/cassettes/test_anthropic/test_streaming.yaml", + ) + + monkeypatch.chdir(repo_root) + key_from_root = redis_key_for(abs_cassette) + + monkeypatch.chdir(os.path.join(repo_root, "tests", "llm_translation")) + key_from_subdir = redis_key_for(abs_cassette) + + monkeypatch.chdir(tmp_path) + key_from_tmp = redis_key_for(abs_cassette) + + assert key_from_root == key_from_subdir == key_from_tmp + assert ( + key_from_root + == "litellm:vcr:cassette:tests/llm_translation/test_anthropic/test_streaming" + ) + + +class _FlakyRedis: + def __init__(self, inner, fail_on: str): + self._inner = inner + self._fail_on = fail_on + + def get(self, *args, **kwargs): + if self._fail_on == "get": + raise RedisConnectionError("simulated outage") + return self._inner.get(*args, **kwargs) + + def set(self, *args, **kwargs): + if self._fail_on == "set": + raise RedisConnectionError("simulated outage") + return self._inner.set(*args, **kwargs) + + +def test_save_swallows_connection_errors_so_teardown_does_not_fail(): + flaky = _FlakyRedis(fakeredis.FakeStrictRedis(), fail_on="set") + persister = make_redis_persister(client=flaky) + + persister.save_cassette( + "tests/llm_translation/test_x/test_save_outage", + _sample_cassette_dict(), + yamlserializer, + ) + + +def test_save_skipped_when_test_marked_failed_and_prior_cassette_preserved(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_flaky" + key = redis_key_for(cassette_id) + + good = _sample_cassette_dict() + persister.save_cassette(cassette_id, good, yamlserializer) + good_payload = fake.get(key) + assert good_payload is not None + + mark_test_outcome_for_cassette(cassette_id, passed=False) + bad_response = { + "status": {"code": 200, "message": "OK"}, + "headers": {}, + "body": {"string": b'{"id":"BAD","type":"message"}'}, + } + bad = {"requests": good["requests"], "responses": [bad_response]} + persister.save_cassette(cassette_id, bad, yamlserializer) + + assert fake.get(key) == good_payload + + +def test_save_proceeds_when_test_marked_passed(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_passed" + key = redis_key_for(cassette_id) + + mark_test_outcome_for_cassette(cassette_id, passed=True) + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + + assert fake.get(key) is not None + + +def test_save_refused_when_cassette_exceeds_max_episodes(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_runaway" + key = redis_key_for(cassette_id) + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + seed_payload = fake.get(key) + + request = Request( + method="POST", + uri="https://api.anthropic.com/v1/messages", + body=b"x", + headers={"content-type": "application/json"}, + ) + response = { + "status": {"code": 200, "message": "OK"}, + "headers": {}, + "body": {"string": b"{}"}, + } + bloated = { + "requests": [request] * (MAX_EPISODES_PER_CASSETTE + 1), + "responses": [response] * (MAX_EPISODES_PER_CASSETTE + 1), + } + persister.save_cassette(cassette_id, bloated, yamlserializer) + + assert fake.get(key) == seed_payload + + +def test_save_proceeds_at_max_episodes_threshold(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_at_threshold" + key = redis_key_for(cassette_id) + + request = Request( + method="POST", + uri="https://api.anthropic.com/v1/messages", + body=b"x", + headers={"content-type": "application/json"}, + ) + response = { + "status": {"code": 200, "message": "OK"}, + "headers": {}, + "body": {"string": b"{}"}, + } + at_threshold = { + "requests": [request] * MAX_EPISODES_PER_CASSETTE, + "responses": [response] * MAX_EPISODES_PER_CASSETTE, + } + persister.save_cassette(cassette_id, at_threshold, yamlserializer) + + assert fake.get(key) is not None + + +def test_save_proceeds_when_outcome_unknown(): + fake, persister = _persister_with_fake_redis() + cassette_id = "tests/llm_translation/test_x/test_no_marker" + key = redis_key_for(cassette_id) + + persister.save_cassette(cassette_id, _sample_cassette_dict(), yamlserializer) + + assert fake.get(key) is not None + + +def test_load_treats_connection_errors_as_cassette_miss(): + flaky = _FlakyRedis(fakeredis.FakeStrictRedis(), fail_on="get") + persister = make_redis_persister(client=flaky) + + with pytest.raises(CassetteNotFoundError): + persister.load_cassette( + "tests/llm_translation/test_x/test_load_outage", yamlserializer + ) + + +@pytest.mark.parametrize( + ("status_code", "expect_dropped"), + [ + (200, False), + (201, False), + (204, False), + (299, False), + (300, True), + (400, True), + (401, True), + (404, True), + (429, True), + (500, True), + (502, True), + (503, True), + ], +) +def test_only_2xx_responses_are_cached(status_code, expect_dropped): + response = { + "status": {"code": status_code, "message": "X"}, + "headers": {}, + "body": {"string": ""}, + } + result = filter_non_2xx_response(response) + assert (result is None) == expect_dropped + if not expect_dropped: + assert result is response diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 3f1a397ebc..82510b6f4f 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -1257,22 +1257,13 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input): assert sent_data["input"] == expected_payload_input -def test_encoding_format_none_not_omitted_from_openai_sdk(): +def test_encoding_format_defaults_to_float_for_openai_sdk(monkeypatch): """ - Test that encoding_format=None is explicitly sent to OpenAI SDK. + When encoding_format is not provided, LiteLLM sends `float` for OpenAI-path embeddings. - This test verifies that when encoding_format is not provided by the user, - liteLLM explicitly sets it to None rather than omitting it. This prevents - the OpenAI SDK from adding its default value of 'base64'. - - Without this fix: - - OpenAI SDK adds encoding_format='base64' as default when parameter is missing - - This causes issues with providers that don't support encoding_format (like Gemini) - - With this fix: - - encoding_format=None is explicitly passed - - OpenAI SDK respects the explicit None and doesn't add defaults + Optional global override: `LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT`. """ + monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) with patch( "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" ) as mock_get_client: @@ -1310,17 +1301,12 @@ def test_encoding_format_none_not_omitted_from_openai_sdk(): call_kwargs = call_args[1] # Get kwargs - # The key assertion: encoding_format should be in the request with value None - # This prevents OpenAI SDK from adding its default 'base64' value - assert "encoding_format" in call_kwargs, ( - "encoding_format should be explicitly passed to OpenAI SDK " - "(even if None) to prevent SDK from adding default value" - ) + assert "encoding_format" in call_kwargs assert ( - call_kwargs["encoding_format"] is None - ), "encoding_format should be None when not provided by user" + call_kwargs["encoding_format"] == "float" + ), "encoding_format should default to float when not provided by user" - print("✅ PASS: encoding_format=None is correctly passed to OpenAI SDK") + print("✅ PASS: encoding_format='float' is correctly passed to OpenAI SDK") def test_encoding_format_explicit_value_preserved(): 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/logging_callback_tests/test_langfuse_dynamic_credentials.py b/tests/logging_callback_tests/test_langfuse_dynamic_credentials.py new file mode 100644 index 0000000000..1b19862338 --- /dev/null +++ b/tests/logging_callback_tests/test_langfuse_dynamic_credentials.py @@ -0,0 +1,129 @@ +import sys +from types import ModuleType, SimpleNamespace + +import litellm +from litellm.integrations.langfuse.langfuse import resolve_langfuse_credentials +from litellm.integrations.langfuse.langfuse_handler import LangFuseHandler + + +def test_resolve_langfuse_credentials_does_not_use_env_for_dynamic_host(monkeypatch): + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "global-public") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "global-secret") + + public_key, secret_key, host = resolve_langfuse_credentials( + langfuse_host="https://attacker.example", + allow_env_credentials=False, + ) + + assert public_key is None + assert secret_key is None + assert host == "https://attacker.example" + + +def test_resolve_langfuse_credentials_accepts_secret_key_alias_for_dynamic_host( + monkeypatch, +): + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "global-secret") + + public_key, secret_key, host = resolve_langfuse_credentials( + langfuse_public_key="dynamic-public", + langfuse_secret_key="dynamic-secret", + langfuse_host="https://team-langfuse.example", + allow_env_credentials=False, + ) + + assert public_key == "dynamic-public" + assert secret_key == "dynamic-secret" + assert host == "https://team-langfuse.example" + + +def test_resolve_langfuse_credentials_keeps_env_for_global_config(monkeypatch): + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "global-public") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "global-secret") + + public_key, secret_key, host = resolve_langfuse_credentials( + langfuse_host="https://admin-configured.example", + allow_env_credentials=True, + ) + + assert public_key == "global-public" + assert secret_key == "global-secret" + assert host == "https://admin-configured.example" + + +def test_upstream_langfuse_debug_env_is_passed(monkeypatch): + from litellm.integrations.langfuse.langfuse import LangFuseLogger + + class FakeLangfuse: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + FakeLangfuse.instances.append(self) + + fake_langfuse_module = ModuleType("langfuse") + fake_langfuse_module.Langfuse = FakeLangfuse + fake_langfuse_module.version = SimpleNamespace(__version__="2.6.0") + + monkeypatch.setitem(sys.modules, "langfuse", fake_langfuse_module) + monkeypatch.setattr(litellm, "initialized_langfuse_clients", 0) + monkeypatch.setenv("LANGFUSE_MOCK", "true") + monkeypatch.setenv("UPSTREAM_LANGFUSE_SECRET_KEY", "upstream-secret") + monkeypatch.setenv("UPSTREAM_LANGFUSE_PUBLIC_KEY", "upstream-public") + monkeypatch.setenv("UPSTREAM_LANGFUSE_HOST", "https://upstream.example") + monkeypatch.setenv("UPSTREAM_LANGFUSE_RELEASE", "release") + monkeypatch.setenv("UPSTREAM_LANGFUSE_DEBUG", "true") + + logger = LangFuseLogger( + langfuse_public_key="public", + langfuse_secret="secret", + langfuse_host="https://langfuse.example", + ) + + assert logger.upstream_langfuse_debug == "true" + assert FakeLangfuse.instances[-1].kwargs["debug"] is True + + +def test_langfuse_handler_accepts_secret_key_alias(monkeypatch): + captured = {} + + class FakeLangFuseLogger: + def __init__( + self, + *, + langfuse_public_key=None, + langfuse_secret=None, + langfuse_host=None, + allow_env_credentials=True, + ): + captured["langfuse_public_key"] = langfuse_public_key + captured["langfuse_secret"] = langfuse_secret + captured["langfuse_host"] = langfuse_host + captured["allow_env_credentials"] = allow_env_credentials + + class FakeDynamicLoggingCache: + def set_cache(self, *, credentials, service_name, logging_obj): + captured["cached_credentials"] = credentials + captured["cached_service_name"] = service_name + captured["cached_logging_obj"] = logging_obj + + monkeypatch.setattr( + "litellm.integrations.langfuse.langfuse_handler.LangFuseLogger", + FakeLangFuseLogger, + ) + + logger = LangFuseHandler._create_langfuse_logger_from_credentials( + credentials={ + "langfuse_public_key": "dynamic-public", + "langfuse_secret_key": "dynamic-secret", + "langfuse_host": "https://langfuse.example", + }, + in_memory_dynamic_logger_cache=FakeDynamicLoggingCache(), + ) + + assert captured["langfuse_public_key"] == "dynamic-public" + assert captured["langfuse_secret"] == "dynamic-secret" + assert captured["langfuse_host"] == "https://langfuse.example" + assert captured["allow_env_credentials"] is False + assert captured["cached_service_name"] == "langfuse" + assert captured["cached_logging_obj"] is logger diff --git a/tests/logging_callback_tests/test_langsmith_dynamic_credentials.py b/tests/logging_callback_tests/test_langsmith_dynamic_credentials.py new file mode 100644 index 0000000000..f1912c5846 --- /dev/null +++ b/tests/logging_callback_tests/test_langsmith_dynamic_credentials.py @@ -0,0 +1,50 @@ +import pytest + +from litellm.integrations.langsmith import LangsmithLogger + + +@pytest.mark.asyncio +async def test_get_credentials_from_env_does_not_use_env_for_dynamic_base_url( + monkeypatch, +): + monkeypatch.setenv("LANGSMITH_API_KEY", "global-key") + monkeypatch.setenv("LANGSMITH_PROJECT", "global-project") + monkeypatch.setenv("LANGSMITH_TENANT_ID", "global-tenant") + logger = LangsmithLogger( + langsmith_api_key="default-key", + langsmith_project="default-project", + langsmith_base_url="https://default.example", + ) + + credentials = logger.get_credentials_from_env( + langsmith_base_url="https://attacker.example", + allow_env_credentials=False, + ) + + assert credentials["LANGSMITH_API_KEY"] is None + assert credentials["LANGSMITH_PROJECT"] == "litellm-completion" + assert credentials["LANGSMITH_BASE_URL"] == "https://attacker.example" + assert credentials["LANGSMITH_TENANT_ID"] is None + + +@pytest.mark.asyncio +async def test_dynamic_langsmith_base_url_does_not_inherit_default_api_key( + monkeypatch, +): + monkeypatch.setenv("LANGSMITH_API_KEY", "global-key") + logger = LangsmithLogger( + langsmith_api_key="default-key", + langsmith_project="default-project", + langsmith_base_url="https://default.example", + ) + + credentials = logger._get_credentials_to_use_for_request( + kwargs={ + "standard_callback_dynamic_params": { + "langsmith_base_url": "https://attacker.example" + } + } + ) + + assert credentials["LANGSMITH_API_KEY"] is None + assert credentials["LANGSMITH_BASE_URL"] == "https://attacker.example" 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 5636a55c95..d9f4a6e56b 100644 --- a/tests/proxy_unit_tests/test_auth_checks.py +++ b/tests/proxy_unit_tests/test_auth_checks.py @@ -1153,3 +1153,320 @@ async def test_can_key_call_model_via_access_group_ids(): valid_token=user_api_key_object, llm_router=router, ) + + +# --------------------------------------------------------------------------- +# _key_access_group_grants_model (key access group overriding team restriction) +# --------------------------------------------------------------------------- + + +def _patch_proxy_server_globals(): + """Patch proxy_server's prisma_client and user_api_key_cache to non-None mocks + so the helper's None-guard doesn't short-circuit. The actual values don't + matter because get_access_object is patched separately to return fixtures.""" + from unittest.mock import MagicMock, patch + + return [ + 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()), + ] + + +def _fake_access_group( + access_group_id: str, + access_model_names=None, + assigned_team_ids=None, + assigned_key_ids=None, +): + from litellm.proxy._types import LiteLLM_AccessGroupTable + + return LiteLLM_AccessGroupTable( + access_group_id=access_group_id, + access_group_name=access_group_id, + access_model_names=access_model_names or [], + assigned_team_ids=assigned_team_ids or [], + assigned_key_ids=assigned_key_ids or [], + ) + + +@pytest.mark.asyncio +async def test_key_access_group_grants_model_when_team_authorized(): + """Group's assigned_team_ids includes the key's team and grants the model → True. + + This is the happy path equivalent of Andres's report: admin creates an + access group with assigned_team_ids=[team-a], grants claude-haiku-4-5, + attaches it to a key on team-a. Override fires. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import _key_access_group_grants_model + + valid_token = UserAPIKeyAuth( + token="test-token", + models=[], + access_group_ids=["premium-group"], + team_id="team-a", + ) + team_object = LiteLLM_TeamTable( + team_id="team-a", + models=["mock-success"], + access_group_ids=[], # deliberately not synced — the access group itself authorizes + ) + + fake_ag = _fake_access_group( + access_group_id="premium-group", + access_model_names=["claude-haiku-4-5"], + assigned_team_ids=["team-a"], + ) + + patches = _patch_proxy_server_globals() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + for p in patches: + p.start() + try: + assert ( + await _key_access_group_grants_model( + model="claude-haiku-4-5", + valid_token=valid_token, + team_object=team_object, + llm_router=None, + ) + is True + ) + finally: + for p in patches: + p.stop() + + +@pytest.mark.asyncio +async def test_key_access_group_grants_model_when_key_directly_authorized(): + """Group's assigned_key_ids includes the key's token and grants the model → True. + + Per-key authorization path: an admin scopes a group directly to a key + (assigned_key_ids) without listing the team. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import _key_access_group_grants_model + + valid_token = UserAPIKeyAuth( + token="test-token-hashed", + models=[], + access_group_ids=["per-key-group"], + team_id="team-a", + ) + team_object = LiteLLM_TeamTable( + team_id="team-a", + models=["mock-success"], + access_group_ids=[], + ) + + fake_ag = _fake_access_group( + access_group_id="per-key-group", + access_model_names=["claude-haiku-4-5"], + assigned_team_ids=[], + assigned_key_ids=["test-token-hashed"], + ) + + patches = _patch_proxy_server_globals() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + for p in patches: + p.start() + try: + assert ( + await _key_access_group_grants_model( + model="claude-haiku-4-5", + valid_token=valid_token, + team_object=team_object, + llm_router=None, + ) + is True + ) + finally: + for p in patches: + p.stop() + + +@pytest.mark.asyncio +async def test_key_access_group_grants_model_when_key_has_no_groups(): + """Key with no access_group_ids → False (early return, no DB read).""" + from litellm.proxy.auth.auth_checks import _key_access_group_grants_model + + valid_token = UserAPIKeyAuth( + token="test-token", + models=[], + access_group_ids=[], + team_id="team-a", + ) + team_object = LiteLLM_TeamTable( + team_id="team-a", + models=["mock-success"], + access_group_ids=["any-group"], + ) + assert ( + await _key_access_group_grants_model( + model="claude-haiku-4-5", + valid_token=valid_token, + team_object=team_object, + llm_router=None, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_key_access_group_grants_model_when_group_does_not_cover_model(): + """Group authorizes the team but does not grant the requested model → False.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import _key_access_group_grants_model + + valid_token = UserAPIKeyAuth( + token="test-token", + models=[], + access_group_ids=["basic-group"], + team_id="team-a", + ) + team_object = LiteLLM_TeamTable( + team_id="team-a", + models=["mock-success"], + access_group_ids=[], + ) + + fake_ag = _fake_access_group( + access_group_id="basic-group", + access_model_names=["gpt-4o-mini"], + assigned_team_ids=["team-a"], + ) + + patches = _patch_proxy_server_globals() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + for p in patches: + p.start() + try: + assert ( + await _key_access_group_grants_model( + model="claude-haiku-4-5", + valid_token=valid_token, + team_object=team_object, + llm_router=None, + ) + is False + ) + finally: + for p in patches: + p.stop() + + +@pytest.mark.asyncio +async def test_key_access_group_grants_model_when_group_authorizes_neither(): + """ + Bypass regression test: a team member sets a foreign access group on their + key. The group grants the requested model but its assigned_team_ids / + assigned_key_ids do not include this caller's team or token. Override is + denied — the team's 401 propagates. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import _key_access_group_grants_model + + valid_token = UserAPIKeyAuth( + token="team-a-token", + models=[], + access_group_ids=["team-b-premium"], + team_id="team-a", + ) + team_object = LiteLLM_TeamTable( + team_id="team-a", + models=["mock-success"], + access_group_ids=[], + ) + + fake_ag = _fake_access_group( + access_group_id="team-b-premium", + access_model_names=["claude-opus-4-5"], + assigned_team_ids=["team-b"], + assigned_key_ids=["team-b-token"], + ) + + patches = _patch_proxy_server_globals() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + return_value=fake_ag, + ), + ] + for p in patches: + p.start() + try: + assert ( + await _key_access_group_grants_model( + model="claude-opus-4-5", + valid_token=valid_token, + team_object=team_object, + llm_router=None, + ) + is False + ) + finally: + for p in patches: + p.stop() + + +@pytest.mark.asyncio +async def test_key_access_group_grants_model_when_get_access_object_raises(): + """Group lookup failure (404, network, etc.) is treated as no authorization.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import _key_access_group_grants_model + + valid_token = UserAPIKeyAuth( + token="test-token", + models=[], + access_group_ids=["missing-group"], + team_id="team-a", + ) + team_object = LiteLLM_TeamTable( + team_id="team-a", + models=["mock-success"], + access_group_ids=[], + ) + + patches = _patch_proxy_server_globals() + [ + patch( + "litellm.proxy.auth.auth_checks.get_access_object", + new_callable=AsyncMock, + side_effect=Exception("not found"), + ), + ] + for p in patches: + p.start() + try: + assert ( + await _key_access_group_grants_model( + model="claude-haiku-4-5", + valid_token=valid_token, + team_object=team_object, + llm_router=None, + ) + is False + ) + finally: + for p in patches: + p.stop() 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/test_litellm/containers/test_azure_container_transformation.py b/tests/test_litellm/containers/test_azure_container_transformation.py index 45fa23bcb6..32caaaffef 100644 --- a/tests/test_litellm/containers/test_azure_container_transformation.py +++ b/tests/test_litellm/containers/test_azure_container_transformation.py @@ -324,6 +324,29 @@ class TestAzureContainerConfig: assert url_fc == expected_fc assert url_fc.index("/content") < url_fc.index("?") + def test_transform_requests_encode_path_ids_before_query_string(self): + from litellm.types.router import GenericLiteLLMParams + + api_base = ( + "https://my-resource.openai.azure.com/openai/v1/containers" + "?api-version=v1" + ) + + url, _ = self.config.transform_container_file_content_request( + container_id="../../other", + file_id="file?download=1#frag", + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + expected_url = ( + "https://my-resource.openai.azure.com/openai/v1/containers/" + "..%2F..%2Fother/files/file%3Fdownload%3D1%23frag/content" + "?api-version=v1" + ) + assert url == expected_url + def test_provider_config_manager_returns_azure_config(self): from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager diff --git a/tests/test_litellm/containers/test_container_handler_url.py b/tests/test_litellm/containers/test_container_handler_url.py new file mode 100644 index 0000000000..1927580068 --- /dev/null +++ b/tests/test_litellm/containers/test_container_handler_url.py @@ -0,0 +1,28 @@ +import pytest + +from litellm.llms.custom_httpx.container_handler import _build_url + + +def test_build_url_encodes_path_params_and_preserves_query(): + url = _build_url( + api_base="https://example.com/v1/containers?api-version=v1", + path_template="/containers/{container_id}/files/{file_id}/content", + path_params={ + "container_id": "../../containers/other", + "file_id": "file?download=1#frag", + }, + ) + + assert ( + url + == "https://example.com/v1/containers/..%2F..%2Fcontainers%2Fother/files/file%3Fdownload%3D1%23frag/content?api-version=v1" + ) + + +def test_build_url_rejects_dot_segment_path_param(): + with pytest.raises(ValueError, match="container_id cannot be a dot path segment"): + _build_url( + api_base="https://example.com/v1/containers", + path_template="/containers/{container_id}", + path_params={"container_id": ".."}, + ) diff --git a/tests/test_litellm/containers/test_container_transformation.py b/tests/test_litellm/containers/test_container_transformation.py index 47b5b8dc56..555fe7773f 100644 --- a/tests/test_litellm/containers/test_container_transformation.py +++ b/tests/test_litellm/containers/test_container_transformation.py @@ -230,6 +230,23 @@ class TestOpenAIContainerTransformation: assert url == f"{api_base}/{container_id}" assert params == {} # No query params for retrieve + def test_transform_container_retrieve_request_encodes_path_traversal(self): + """Test container IDs are treated as a single upstream path segment.""" + api_base = "https://api.openai.com/v1/containers" + + url, params = self.config.transform_container_retrieve_request( + container_id="../../vector_stores?x=1#frag", + api_base=api_base, + litellm_params={}, + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/containers/..%2F..%2Fvector_stores%3Fx%3D1%23frag" + ) + assert params == {} + def test_transform_container_retrieve_response(self): """Test container retrieve response transformation.""" # Mock HTTP response diff --git a/tests/test_litellm/integrations/arize/test_arize_phoenix.py b/tests/test_litellm/integrations/arize/test_arize_phoenix.py index 01f85af262..4a2eab29e8 100644 --- a/tests/test_litellm/integrations/arize/test_arize_phoenix.py +++ b/tests/test_litellm/integrations/arize/test_arize_phoenix.py @@ -280,3 +280,55 @@ class TestDynamicProjectNameOnSpan: if __name__ == "__main__": unittest.main() + + +# --- Security: SSRF via prompt_version_id path traversal --- + + +def test_arize_phoenix_client_sanitize_id_rejects_traversal(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + # dotdot without slashes + with pytest.raises(ValueError, match="path traversal"): + _sanitize_id("..something") + # full traversal (slash caught first) + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("../../projects") + + +def test_arize_phoenix_client_sanitize_id_rejects_slash(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("valid/extra") + + +def test_arize_phoenix_client_sanitize_id_rejects_fragment(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("abc#suffix") + + +def test_arize_phoenix_client_sanitize_id_rejects_query(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + with pytest.raises(ValueError, match="disallowed characters"): + _sanitize_id("abc?x=1") + + +def test_arize_phoenix_client_sanitize_id_allows_uuid(): + from litellm.integrations.arize.arize_phoenix_client import _sanitize_id + + uid = "550e8400-e29b-41d4-a716-446655440000" + assert _sanitize_id(uid) == uid + + +def test_arize_phoenix_client_get_prompt_version_rejects_traversal(): + from litellm.integrations.arize.arize_phoenix_client import ArizePhoenixClient + + client = ArizePhoenixClient( + api_key="test-key", api_base="https://app.phoenix.arize.com" + ) + with pytest.raises(ValueError, match="disallowed characters"): + client.get_prompt_version("../../projects") diff --git a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py index a7b2d362ed..46cd1d6e76 100644 --- a/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py +++ b/tests/test_litellm/integrations/bitbucket/test_bitbucket_integration.py @@ -11,6 +11,7 @@ sys.path.insert( import litellm from litellm.integrations.bitbucket import BitBucketPromptManager +from litellm.integrations.bitbucket.bitbucket_client import _sanitize_file_path @patch("litellm.integrations.bitbucket.bitbucket_prompt_manager.BitBucketClient") @@ -370,3 +371,45 @@ def test_bitbucket_prompt_manager_list_templates(mock_client_class): templates = manager.prompt_manager.list_templates() assert isinstance(templates, list) assert "test_prompt" in templates + + +# --- Security: path traversal / SSRF --- + + +def test_sanitize_file_path_rejects_traversal(): + with pytest.raises(ValueError, match="path traversal"): + _sanitize_file_path("../../etc/passwd") + + +def test_sanitize_file_path_rejects_fragment(): + with pytest.raises(ValueError, match="URL special characters"): + _sanitize_file_path("secret#.prompt") + + +def test_sanitize_file_path_rejects_query(): + with pytest.raises(ValueError, match="URL special characters"): + _sanitize_file_path("secret?.prompt") + + +def test_sanitize_file_path_encodes_special_chars(): + result = _sanitize_file_path("prompts/my prompt.prompt") + assert result == "prompts/my%20prompt.prompt" + + +def test_sanitize_file_path_allows_normal_paths(): + assert _sanitize_file_path("prompts/my-prompt") == "prompts/my-prompt" + assert _sanitize_file_path("simple") == "simple" + + +def test_bitbucket_client_rejects_traversal_in_get_file_content(): + from litellm.integrations.bitbucket.bitbucket_client import BitBucketClient + + client = BitBucketClient( + { + "workspace": "ws", + "repository": "repo", + "access_token": "tok", + } + ) + with pytest.raises(ValueError, match="path traversal"): + client.get_file_content("../../admin/credentials") diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index b31bbca889..962806c5b5 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -259,6 +259,189 @@ class TestOpenTelemetryProviderInitialization(unittest.TestCase): ), "Existing LoggerProvider should be respected and not overridden" +class TestOpenTelemetryDualHandlerIsolation(unittest.TestCase): + """Two OpenTelemetry handlers coexisting via skip_set_global=True + must each get their own provider for every signal (tracer/meter/logger).""" + + @staticmethod + def _wire_span_processor(exporter): + """Context manager: while active, the next OpenTelemetry instance + wires its TracerProvider to `exporter`.""" + return patch.object( + OpenTelemetry, + "_get_span_processor", + lambda self, dynamic_headers=None: SimpleSpanProcessor(exporter), + ) + + def test_skip_set_global_creates_isolated_tracer_provider(self): + from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider + + fake_existing = SDKTracerProvider() + own_exporter = InMemorySpanExporter() + cfg = OpenTelemetryConfig( + exporter="console", service_name="iso-test", skip_set_global=True + ) + with ( + patch.object(trace, "get_tracer_provider", return_value=fake_existing), + patch.object(trace, "set_tracer_provider") as mock_set, + self._wire_span_processor(own_exporter), + ): + handler = OpenTelemetry(config=cfg) + + self.assertIsNot(handler._tracer_provider, fake_existing) + mock_set.assert_not_called() + + handler.tracer.start_span("isolation_check").end() + handler._tracer_provider.force_flush(2000) + self.assertEqual( + [s.name for s in own_exporter.get_finished_spans()], + ["isolation_check"], + ) + + def test_skip_set_global_via_callback_name_back_compat(self): + from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider + + fake_existing = SDKTracerProvider() + cfg = OpenTelemetryConfig(exporter="console", service_name="lf-back-compat") + with ( + patch.object(trace, "get_tracer_provider", return_value=fake_existing), + patch.object(trace, "set_tracer_provider"), + self._wire_span_processor(InMemorySpanExporter()), + ): + handler = OpenTelemetry(config=cfg, callback_name="langfuse_otel") + + self.assertIsNot(handler._tracer_provider, fake_existing) + + def test_default_behavior_reuses_existing_sdk_tracer_provider(self): + from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider + + fake_existing = SDKTracerProvider() + with patch.object(trace, "get_tracer_provider", return_value=fake_existing): + handler = OpenTelemetry(config=OpenTelemetryConfig(service_name="shared")) + self.assertIs(handler._tracer_provider, fake_existing) + + def test_skip_set_global_creates_isolated_meter_provider(self): + from opentelemetry import metrics + from opentelemetry.sdk.metrics import MeterProvider as SDKMeterProvider + + fake_existing = SDKMeterProvider() + cfg = OpenTelemetryConfig( + exporter="console", + service_name="meter-iso-test", + enable_metrics=True, + skip_set_global=True, + ) + with ( + patch.object(metrics, "get_meter_provider", return_value=fake_existing), + patch.object(metrics, "set_meter_provider") as mock_set, + self._wire_span_processor(InMemorySpanExporter()), + ): + handler = OpenTelemetry(config=cfg) + + self.assertIsNot(handler._meter_provider, fake_existing) + mock_set.assert_not_called() + + def test_skip_set_global_creates_isolated_logger_provider(self): + from opentelemetry import _logs + from opentelemetry.sdk._logs import LoggerProvider as SDKLoggerProvider + + fake_existing = SDKLoggerProvider() + cfg = OpenTelemetryConfig( + exporter="console", + service_name="logger-iso-test", + enable_events=True, + skip_set_global=True, + ) + with ( + patch.object(_logs, "get_logger_provider", return_value=fake_existing), + patch.object(_logs, "set_logger_provider") as mock_set, + self._wire_span_processor(InMemorySpanExporter()), + ): + handler = OpenTelemetry(config=cfg) + + self.assertIsNot(handler._logger_provider, fake_existing) + mock_set.assert_not_called() + + def test_emitted_logs_route_to_isolated_logger_provider(self): + # End-to-end: emitted logs land in the handler's private LoggerProvider, + # not the global one. Guards against get_logger() bypassing self._logger_provider. + from opentelemetry import _logs + from opentelemetry.sdk._logs import LoggerProvider as SDKLoggerProvider + + global_exporter = InMemoryLogExporter() + fake_existing = SDKLoggerProvider() + fake_existing.add_log_record_processor( + SimpleLogRecordProcessor(global_exporter) + ) + + private_exporter = InMemoryLogExporter() + cfg = OpenTelemetryConfig( + exporter="console", + service_name="logger-emit-test", + enable_events=True, + skip_set_global=True, + ) + with ( + patch.object(_logs, "get_logger_provider", return_value=fake_existing), + patch.object(_logs, "set_logger_provider"), + patch.object( + OpenTelemetry, "_get_log_exporter", return_value=private_exporter + ), + self._wire_span_processor(InMemorySpanExporter()), + ): + handler = OpenTelemetry(config=cfg) + + span = handler.tracer.start_span("emit-test") + handler._emit_semantic_logs( + kwargs={"messages": [{"role": "user", "content": "hi"}]}, + response_obj={"choices": []}, + span=span, + ) + span.end() + handler._logger_provider.force_flush(2000) + + self.assertGreater(len(private_exporter.get_finished_logs()), 0) + self.assertEqual(len(global_exporter.get_finished_logs()), 0) + + def test_two_handlers_each_receive_their_own_spans(self): + # Handler A gets explicit injection (production-ish: claims the global). + exporter_a = InMemorySpanExporter() + provider_a = TracerProvider() + provider_a.add_span_processor(SimpleSpanProcessor(exporter_a)) + handler_a = OpenTelemetry( + config=OpenTelemetryConfig(service_name="handler-a"), + tracer_provider=provider_a, + ) + + # Handler B comes along with the global appearing to be A's provider. + exporter_b = InMemorySpanExporter() + cfg_b = OpenTelemetryConfig( + exporter="console", service_name="handler-b", skip_set_global=True + ) + with ( + patch.object(trace, "get_tracer_provider", return_value=provider_a), + patch.object(trace, "set_tracer_provider"), + self._wire_span_processor(exporter_b), + ): + handler_b = OpenTelemetry(config=cfg_b) + + self.assertIsNot(handler_a._tracer_provider, handler_b._tracer_provider) + + handler_a.tracer.start_span("from_handler_a").end() + handler_b.tracer.start_span("from_handler_b").end() + provider_a.force_flush(2000) + handler_b._tracer_provider.force_flush(2000) + + self.assertEqual( + sorted(s.name for s in exporter_a.get_finished_spans()), + ["from_handler_a"], + ) + self.assertEqual( + sorted(s.name for s in exporter_b.get_finished_spans()), + ["from_handler_b"], + ) + + class TestOpenTelemetry(unittest.TestCase): POLL_INTERVAL = 0.05 POLL_TIMEOUT = 2.0 diff --git a/tests/test_litellm/integrations/test_prometheus_api_promql_escape.py b/tests/test_litellm/integrations/test_prometheus_api_promql_escape.py new file mode 100644 index 0000000000..262ca6b692 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_api_promql_escape.py @@ -0,0 +1,150 @@ +""" +Tests for VERIA-53: PromQL string-literal quoting in +``get_daily_spend_from_prometheus``. + +PromQL string literals follow Go's escape rules +(https://prometheus.io/docs/prometheus/latest/querying/basics/). JSON's +quoting is a strict subset of Go's, so ``json.dumps`` produces a literal +Prometheus parses identically. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +def test_quote_safe_input_round_trips(): + from litellm.integrations.prometheus_helpers.prometheus_api import ( + _quote_promql_string_literal, + ) + + assert _quote_promql_string_literal("sk-abc123") == '"sk-abc123"' + assert _quote_promql_string_literal("hash:deadbeef") == '"hash:deadbeef"' + + +def test_quote_escapes_double_quote(): + from litellm.integrations.prometheus_helpers.prometheus_api import ( + _quote_promql_string_literal, + ) + + # A bare double quote would otherwise terminate the label matcher and + # let the attacker append `, foo="..."} or sum(...)`. + assert _quote_promql_string_literal('hello"injected') == '"hello\\"injected"' + + +def test_quote_escapes_backslash(): + from litellm.integrations.prometheus_helpers.prometheus_api import ( + _quote_promql_string_literal, + ) + + assert _quote_promql_string_literal('a\\"b') == '"a\\\\\\"b"' + + +def test_quote_escapes_newlines_and_control_chars(): + """Beyond the security minimum, the canonical Go/JSON escape also + handles control characters that would otherwise produce an invalid + PromQL string literal.""" + from litellm.integrations.prometheus_helpers.prometheus_api import ( + _quote_promql_string_literal, + ) + + assert _quote_promql_string_literal("a\nb") == '"a\\nb"' + assert _quote_promql_string_literal("a\tb") == '"a\\tb"' + assert _quote_promql_string_literal("a\rb") == '"a\\rb"' + + +@pytest.mark.asyncio +async def test_get_daily_spend_does_not_pass_raw_quote_into_query(): + from litellm.integrations.prometheus_helpers import prometheus_api + + captured = {} + + class _FakeResponse: + def json(self): + return {"data": {"result": []}} + + async def _capture(url, params): + captured["url"] = url + captured["params"] = params + return _FakeResponse() + + fake_client = MagicMock() + fake_client.get = AsyncMock(side_effect=_capture) + + with patch.object(prometheus_api, "PROMETHEUS_URL", "http://prom.example"): + with patch.object(prometheus_api, "async_http_handler", fake_client): + await prometheus_api.get_daily_spend_from_prometheus( + api_key='sk-victim"} or sum(other_metric{a="b' + ) + + rendered_query = captured["params"]["query"] + # The legitimate matcher framing must still be intact: one outer + # `delta()` window, one inner `hashed_api_key="..."` matcher. + assert rendered_query.startswith( + 'sum(delta(litellm_spend_metric_total{hashed_api_key="' + ) + assert rendered_query.endswith('"}[1d]))') + + # Every injected `"` from the attacker payload appears as `\"` so the + # PromQL parser treats them as literal characters inside the matcher + # value, never as the terminator that would let the rest parse as + # PromQL syntax. + inner = rendered_query[ + len('sum(delta(litellm_spend_metric_total{hashed_api_key="') : -len('"}[1d]))') + ] + assert '"' not in inner.replace('\\"', "") + + +@pytest.mark.asyncio +async def test_get_daily_spend_with_no_api_key_uses_unfiltered_query(): + from litellm.integrations.prometheus_helpers import prometheus_api + + captured = {} + + class _FakeResponse: + def json(self): + return {"data": {"result": []}} + + async def _capture(url, params): + captured["params"] = params + return _FakeResponse() + + fake_client = MagicMock() + fake_client.get = AsyncMock(side_effect=_capture) + + with patch.object(prometheus_api, "PROMETHEUS_URL", "http://prom.example"): + with patch.object(prometheus_api, "async_http_handler", fake_client): + await prometheus_api.get_daily_spend_from_prometheus(api_key=None) + + assert captured["params"]["query"] == "sum(delta(litellm_spend_metric_total[1d]))" + + +@pytest.mark.asyncio +async def test_get_daily_spend_legitimate_hashed_key_unchanged(): + """A normal hex hashed_api_key flows through `json.dumps` as itself + plus the surrounding quotes — no spurious escaping that would break + real lookups.""" + from litellm.integrations.prometheus_helpers import prometheus_api + + captured = {} + + class _FakeResponse: + def json(self): + return {"data": {"result": []}} + + async def _capture(url, params): + captured["params"] = params + return _FakeResponse() + + fake_client = MagicMock() + fake_client.get = AsyncMock(side_effect=_capture) + + legit_key = "a" * 64 # 64-char hex-ish hashed key + with patch.object(prometheus_api, "PROMETHEUS_URL", "http://prom.example"): + with patch.object(prometheus_api, "async_http_handler", fake_client): + await prometheus_api.get_daily_spend_from_prometheus(api_key=legit_key) + + assert ( + captured["params"]["query"] + == f'sum(delta(litellm_spend_metric_total{{hashed_api_key="{legit_key}"}}[1d]))' + ) diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py index 37dc491c26..758ff3ea38 100644 --- a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py +++ b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py @@ -147,6 +147,21 @@ class TestInteractionOperationUrls: assert "secret-key" not in url assert expected_suffix in url + def test_interaction_id_is_encoded_as_one_path_segment(self, config): + with patch(_PATCH_GET_API_KEY, return_value="secret-key"): + url, params = config.transform_cancel_interaction_request( + interaction_id="../../interactions/other?x=1#frag", + api_base="https://generativelanguage.googleapis.com", + litellm_params=GenericLiteLLMParams(api_key="secret-key"), + headers={}, + ) + + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/interactions/..%2F..%2Finteractions%2Fother%3Fx%3D1%23frag:cancel" + ) + assert params == {} + def test_get_interaction_raises_without_key(self, config): with patch(_PATCH_GET_API_KEY, return_value=None): with pytest.raises(ValueError, match="Google API key is required"): diff --git a/tests/test_litellm/litellm_core_utils/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..e363418e77 100644 --- a/tests/test_litellm/litellm_core_utils/test_url_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_url_utils.py @@ -4,7 +4,13 @@ import pytest import litellm from litellm.litellm_core_utils import url_utils -from litellm.litellm_core_utils.url_utils import SSRFError, _is_blocked_ip, validate_url +from litellm.litellm_core_utils.url_utils import ( + SSRFError, + _is_blocked_ip, + encode_url_path_segment, + encode_url_path_segments, + validate_url, +) @pytest.fixture @@ -80,6 +86,28 @@ class TestIsBlockedIp: assert _is_blocked_ip("::ffff:168.63.129.16") is True +class TestEncodeUrlPathSegment: + def test_encodes_path_delimiters_and_query_markers(self): + encoded = encode_url_path_segment("../../v1/files?limit=1#frag") + + assert encoded == "..%2F..%2Fv1%2Ffiles%3Flimit%3D1%23frag" + + def test_encodes_path_segments_without_collapsing_valid_model_paths(self): + encoded = encode_url_path_segments("@cf/meta/model?debug=1") + + assert encoded == "%40cf/meta/model%3Fdebug%3D1" + + @pytest.mark.parametrize("value", ["", ".", "..", None]) + def test_rejects_empty_and_dot_segments(self, value): + with pytest.raises(ValueError): + encode_url_path_segment(value, field_name="resource_id") + + @pytest.mark.parametrize("value", ["../model", "model/../other", "/model"]) + def test_rejects_dot_segments_in_multi_segment_paths(self, value): + with pytest.raises(ValueError): + encode_url_path_segments(value, field_name="model") + + class TestValidateUrl: def test_blocks_loopback(self): with pytest.raises(SSRFError): @@ -394,3 +422,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/files/test_anthropic_files_transformation.py b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py index e9509be9e1..9fc4981510 100644 --- a/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py +++ b/tests/test_litellm/llms/anthropic/files/test_anthropic_files_transformation.py @@ -180,6 +180,19 @@ class TestAnthropicFilesConfig: assert url == "https://custom.api.com/v1/files/file-abc123" assert params == {} + def test_transform_retrieve_file_request_encodes_path_traversal(self): + url, params = self.config.transform_retrieve_file_request( + file_id="../../v1/messages/batches?limit=1#frag", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == f"{ANTHROPIC_FILES_API_BASE}/v1/files/..%2F..%2Fv1%2Fmessages%2Fbatches%3Flimit%3D1%23frag" + ) + assert params == {} + def test_transform_retrieve_file_response(self): mock_response = Mock(spec=httpx.Response) mock_response.json.return_value = { @@ -296,6 +309,14 @@ class TestAnthropicFilesConfig: assert url == f"{ANTHROPIC_FILES_API_BASE}/v1/files/file-abc123/content" assert params == {} + def test_transform_file_content_request_rejects_dot_segment(self): + with pytest.raises(ValueError, match="file_id cannot be a dot path segment"): + self.config.transform_file_content_request( + file_content_request={"file_id": ".."}, + optional_params={}, + litellm_params={}, + ) + def test_transform_file_content_response(self): mock_response = Mock(spec=httpx.Response) result = self.config.transform_file_content_response( diff --git a/tests/test_litellm/llms/azure/response/test_azure_transformation.py b/tests/test_litellm/llms/azure/response/test_azure_transformation.py index 9519b7c8a5..a4bd14d69f 100644 --- a/tests/test_litellm/llms/azure/response/test_azure_transformation.py +++ b/tests/test_litellm/llms/azure/response/test_azure_transformation.py @@ -96,6 +96,28 @@ def test_get_complete_url(): assert result == expected +@pytest.mark.serial +def test_response_id_path_requests_encode_response_id(): + config = AzureOpenAIResponsesAPIConfig() + api_base = ( + "https://litellm8397336933.openai.azure.com/openai/responses" + "?api-version=2024-05-01-preview" + ) + + url, params = config.transform_cancel_response_api_request( + response_id="../../responses/other?x=1#frag", + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://litellm8397336933.openai.azure.com/openai/responses/..%2F..%2Fresponses%2Fother%3Fx%3D1%23frag/cancel?api-version=2024-05-01-preview" + ) + assert params == {} + + @pytest.mark.serial def test_azure_o_series_responses_api_supported_params(): """Test that Azure OpenAI O-series responses API excludes temperature from supported parameters.""" diff --git a/tests/test_litellm/llms/azure/test_azure_cost_calculation.py b/tests/test_litellm/llms/azure/test_azure_cost_calculation.py new file mode 100644 index 0000000000..53c91032b3 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_azure_cost_calculation.py @@ -0,0 +1,75 @@ +""" +Test Azure OpenAI cost calculator — service_tier pricing. +""" + +import pytest + +import litellm +from litellm.llms.azure.cost_calculation import cost_per_token +from litellm.types.utils import Usage + + +# Register a test model with tier-specific pricing +TEST_MODEL = "test-azure-gpt-4.1" +TEST_MODEL_COST = { + TEST_MODEL: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "azure", + "max_tokens": 8192, + } +} + + +class TestAzureServiceTierCostCalculation: + """Test that service_tier is passed through Azure cost calculation.""" + + @pytest.fixture(autouse=True) + def register_test_model(self): + litellm.register_model(model_cost=TEST_MODEL_COST) + + def test_service_tier_priority_higher_cost(self): + """Priority tier should cost more than standard.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model=TEST_MODEL, usage=usage + ) + priority_prompt, priority_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier="priority" + ) + + assert priority_prompt > standard_prompt + assert priority_completion > standard_completion + + def test_service_tier_flex_lower_cost(self): + """Flex tier should cost less than standard.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model=TEST_MODEL, usage=usage + ) + flex_prompt, flex_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier="flex" + ) + + assert flex_prompt < standard_prompt + assert flex_completion < standard_completion + + def test_service_tier_none_returns_standard(self): + """service_tier=None should return standard pricing.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + none_prompt, none_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier=None + ) + standard_prompt, standard_completion = cost_per_token( + model=TEST_MODEL, usage=usage, service_tier="standard" + ) + + assert abs(none_prompt - standard_prompt) < 1e-10 + assert abs(none_completion - standard_completion) < 1e-10 diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py new file mode 100644 index 0000000000..f65573b7ae --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_agents_handler.py @@ -0,0 +1,57 @@ +import pytest + +from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler + + +def test_should_encode_thread_id_in_azure_ai_agent_urls(): + handler = AzureAIAgentsHandler() + + assert ( + handler._build_messages_url( + "https://example.services.ai.azure.com/api/projects/proj", + "../../threads/other?x=1#frag", + "2024-05-01-preview", + ) + == "https://example.services.ai.azure.com/api/projects/proj/threads/..%2F..%2Fthreads%2Fother%3Fx%3D1%23frag/messages?api-version=2024-05-01-preview" + ) + assert ( + handler._build_runs_url( + "https://example.services.ai.azure.com/api/projects/proj", + "thread/abc", + "2024-05-01-preview", + ) + == "https://example.services.ai.azure.com/api/projects/proj/threads/thread%2Fabc/runs?api-version=2024-05-01-preview" + ) + + +def test_should_encode_thread_and_run_ids_in_azure_ai_agent_status_url(): + handler = AzureAIAgentsHandler() + + assert ( + handler._build_run_status_url( + "https://example.services.ai.azure.com/api/projects/proj", + "thread/abc", + "../runs/other#frag", + "2024-05-01-preview", + ) + == "https://example.services.ai.azure.com/api/projects/proj/threads/thread%2Fabc/runs/..%2Fruns%2Fother%23frag?api-version=2024-05-01-preview" + ) + + +def test_should_reject_dot_segments_in_azure_ai_agent_urls(): + handler = AzureAIAgentsHandler() + + with pytest.raises(ValueError, match="thread_id cannot be a dot path segment"): + handler._build_messages_url( + "https://example.services.ai.azure.com/api/projects/proj", + "..", + "2024-05-01-preview", + ) + + with pytest.raises(ValueError, match="run_id cannot be a dot path segment"): + handler._build_run_status_url( + "https://example.services.ai.azure.com/api/projects/proj", + "thread_123", + "..", + "2024-05-01-preview", + ) diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py index 37add41b83..20260c744f 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_cost_calculator.py @@ -451,3 +451,51 @@ class TestAzureModelRouterCostBreakdown: assert logging_obj.cost_breakdown["additional_costs"][ "Azure Model Router Flat Cost" ] == pytest.approx(expected_flat_cost, rel=1e-9) + + +class TestAzureAIServiceTierCostCalculation: + """Test that service_tier is passed through Azure AI cost calculation.""" + + @pytest.fixture(autouse=True) + def register_test_model(self): + import litellm + litellm.register_model(model_cost={ + "test-azure-ai-model": { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "input_cost_per_token_priority": 0.01, + "output_cost_per_token_priority": 0.02, + "input_cost_per_token_flex": 0.0005, + "output_cost_per_token_flex": 0.001, + "litellm_provider": "azure_ai", + "max_tokens": 8192, + } + }) + + def test_service_tier_priority_higher_cost(self): + """Priority tier should cost more than standard for azure_ai.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model="test-azure-ai-model", usage=usage + ) + priority_prompt, priority_completion = cost_per_token( + model="test-azure-ai-model", usage=usage, service_tier="priority" + ) + + assert priority_prompt > standard_prompt + assert priority_completion > standard_completion + + def test_service_tier_flex_lower_cost(self): + """Flex tier should cost less than standard for azure_ai.""" + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + standard_prompt, standard_completion = cost_per_token( + model="test-azure-ai-model", usage=usage + ) + flex_prompt, flex_completion = cost_per_token( + model="test-azure-ai-model", usage=usage, service_tier="flex" + ) + + assert flex_prompt < standard_prompt + assert flex_completion < standard_completion diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py new file mode 100644 index 0000000000..e638be68ec --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -0,0 +1,33 @@ +import pytest + +from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, +) + + +def test_should_encode_azure_document_intelligence_model_id(): + config = AzureDocumentIntelligenceOCRConfig() + + url = config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="prebuilt-layout?x=1#frag", + optional_params={}, + litellm_params={}, + ) + + assert ( + url + == "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout%3Fx%3D1%23frag:analyze?api-version=2024-11-30" + ) + + +def test_should_reject_dot_segment_azure_document_intelligence_model_id(): + config = AzureDocumentIntelligenceOCRConfig() + + with pytest.raises(ValueError, match="model_id cannot be a dot path segment"): + config.get_complete_url( + api_base="https://example.cognitiveservices.azure.com", + model="azure_ai/doc-intelligence/..", + optional_params={}, + litellm_params={}, + ) diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index eac022ec23..0de3f833a3 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -167,3 +167,24 @@ def test_tool_name_sanitization(): ] # Should be sanitized: only [a-zA-Z0-9_] assert tool_name == "my_tool_" + + +def test_count_tokens_endpoint_encodes_model_id(monkeypatch): + """Test model IDs are treated as a single Bedrock path segment.""" + config = BedrockCountTokensConfig() + + monkeypatch.setattr( + config, + "get_runtime_endpoint", + lambda **kwargs: ("https://bedrock-runtime.us-east-1.amazonaws.com", None), + ) + + endpoint = config.get_bedrock_count_tokens_endpoint( + model="bedrock/../../model/other?x=1#frag", + aws_region_name="us-east-1", + ) + + assert ( + endpoint + == "https://bedrock-runtime.us-east-1.amazonaws.com/model/..%2F..%2Fmodel%2Fother%3Fx%3D1%23frag/count-tokens" + ) diff --git a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py b/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py index ffd3526698..9e526e4778 100644 --- a/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py +++ b/tests/test_litellm/llms/bedrock/invoke_agent/test_bedrock_agent_transformation.py @@ -1,12 +1,9 @@ import base64 -import json import os import sys -from litellm._uuid import uuid -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../../..") @@ -15,12 +12,7 @@ sys.path.insert( from litellm.llms.bedrock.chat.invoke_agent.transformation import ( AmazonInvokeAgentConfig, ) -from litellm.types.llms.bedrock_invoke_agents import ( - InvokeAgentEvent, - InvokeAgentEventHeaders, - InvokeAgentUsage, -) -from litellm.types.utils import Message, ModelResponse, Usage +from litellm.types.utils import ModelResponse class TestAmazonInvokeAgentConfig: @@ -270,3 +262,28 @@ class TestAmazonInvokeAgentConfig: "https://bedrock-runtime.us-east-1.amazonaws.com/agents/L1RT58GYRW/agentAliases/MFPSBCXYTW/sessions" in result ) + + @patch( + "litellm.llms.bedrock.chat.invoke_agent.transformation.convert_content_list_to_str" + ) + @patch.object(AmazonInvokeAgentConfig, "get_runtime_endpoint") + @patch.object(AmazonInvokeAgentConfig, "_get_aws_region_name") + def test_get_complete_url_encodes_session_id( + self, mock_region, mock_endpoint, mock_convert, config + ): + """Test get_complete_url encodes session ID path segment.""" + mock_endpoint.return_value = ( + "https://bedrock-runtime.us-east-1.amazonaws.com", + None, + ) + mock_region.return_value = "us-east-1" + + result = config.get_complete_url( + api_base=None, + api_key=None, + model="agent/L1RT58GYRW/MFPSBCXYTW", + optional_params={"sessionID": "../../sessions/other?x=1#frag"}, + litellm_params={}, + ) + + assert "sessions/..%2F..%2Fsessions%2Fother%3Fx%3D1%23frag/text" in result diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index d60d0487d0..7b04efa17d 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -28,6 +28,28 @@ def test_transform_search_request(): assert body["retrievalQuery"].get("text") == "hello" +def test_transform_search_request_encodes_vector_store_id(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + + url, body = config.transform_search_vector_store_request( + vector_store_id="../../knowledgebases/other?x=1#frag", + query="hello", + vector_store_search_optional_params={}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + extra_body=None, + ) + + assert ( + url + == "https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases/..%2F..%2Fknowledgebases%2Fother%3Fx%3D1%23frag/retrieve" + ) + assert body["retrievalQuery"].get("text") == "hello" + + def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): config = BedrockVectorStoreConfig() mock_log = MagicMock() diff --git a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py index dd388fedc6..2f8cc5484b 100644 --- a/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py +++ b/tests/test_litellm/llms/bytez/chat/test_bytez_chat_transformation.py @@ -87,6 +87,29 @@ class TestBytezChatConfig: assert response.choices[0].message.content == output_content # type: ignore + def test_get_complete_url_encodes_model_path_segment(self): + config = BytezChatConfig() + + assert ( + config.get_complete_url( + api_base=API_BASE, + api_key=TEST_API_KEY, + model="google/gemma?x=1#frag", + optional_params={}, + litellm_params={}, + ) + == f"{API_BASE}/google/gemma%3Fx%3D1%23frag" + ) + + with pytest.raises(ValueError, match="dot path segment"): + config.get_complete_url( + api_base=API_BASE, + api_key=TEST_API_KEY, + model="../../models/other", + optional_params={}, + litellm_params={}, + ) + def test_bytez_messages_adaptation(self): cases = [ dict( diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py new file mode 100644 index 0000000000..cecb6024de --- /dev/null +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_transformation.py @@ -0,0 +1,27 @@ +import pytest + +from litellm.llms.cloudflare.chat.transformation import CloudflareChatConfig + + +def test_get_complete_url_encodes_model_path_segment(): + config = CloudflareChatConfig() + + assert ( + config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="@cf/meta/llama?x=1#frag", + optional_params={}, + litellm_params={}, + ) + == "https://api.cloudflare.com/client/v4/accounts/acct/ai/run/%40cf/meta/llama%3Fx%3D1%23frag" + ) + + with pytest.raises(ValueError, match="dot path segment"): + config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="../../accounts/other", + optional_params={}, + litellm_params={}, + ) diff --git a/tests/test_litellm/llms/crusoe/__init__.py b/tests/test_litellm/llms/crusoe/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/crusoe/test_crusoe.py b/tests/test_litellm/llms/crusoe/test_crusoe.py new file mode 100644 index 0000000000..0a05126919 --- /dev/null +++ b/tests/test_litellm/llms/crusoe/test_crusoe.py @@ -0,0 +1,135 @@ +import os +from unittest.mock import patch + +CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" + + +def test_crusoe_json_registry(): + """Test Crusoe is registered in the JSON provider registry""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert JSONProviderRegistry.exists("crusoe") + config = JSONProviderRegistry.get("crusoe") + assert config is not None + assert config.base_url == CRUSOE_API_BASE + assert config.api_key_env == "CRUSOE_API_KEY" + assert config.api_base_env == "CRUSOE_API_BASE" + + +def test_crusoe_dynamic_config_defaults(): + """Test dynamic config returns correct default API base""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("crusoe"))() + + with patch.dict(os.environ, {}, clear=True): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + + assert api_base == CRUSOE_API_BASE + assert api_key is None + + +def test_crusoe_dynamic_config_env_vars(): + """Test dynamic config reads CRUSOE_API_KEY and CRUSOE_API_BASE from env""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("crusoe"))() + + with patch.dict( + os.environ, + {"CRUSOE_API_KEY": "test-key", "CRUSOE_API_BASE": "https://custom.crusoe.com/v1"}, + ): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + + assert api_base == "https://custom.crusoe.com/v1" + assert api_key == "test-key" + + +def test_crusoe_dynamic_config_explicit_params(): + """Test explicit params override env vars""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("crusoe"))() + + with patch.dict(os.environ, {"CRUSOE_API_KEY": "env-key"}): + api_base, api_key = config._get_openai_compatible_provider_info( + "https://override.crusoe.com/v1", "override-key" + ) + + assert api_base == "https://override.crusoe.com/v1" + assert api_key == "override-key" + + +def test_crusoe_supported_params(): + """Test dynamic config returns standard OpenAI params""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("crusoe"))() + params = config.get_supported_openai_params(model="meta-llama/Llama-3.3-70B-Instruct") + + assert isinstance(params, list) + assert len(params) > 0 + assert "temperature" in params + assert "max_tokens" in params + assert "stream" in params + + +def test_crusoe_param_mapping_max_completion_tokens(): + """Test max_completion_tokens is mapped to max_tokens for Crusoe""" + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("crusoe"))() + optional_params = config.map_openai_params( + non_default_params={"max_completion_tokens": 1024}, + optional_params={}, + model="meta-llama/Llama-3.3-70B-Instruct", + drop_params=False, + ) + + assert "max_tokens" in optional_params, "max_completion_tokens should be mapped to max_tokens" + assert optional_params["max_tokens"] == 1024 + assert "max_completion_tokens" not in optional_params + + +def test_crusoe_provider_detection_by_prefix(): + """Test crusoe/model prefix is correctly routed""" + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, _, _ = get_llm_provider("crusoe/meta-llama/Llama-3.3-70B-Instruct") + assert provider == "crusoe" + assert model == "meta-llama/Llama-3.3-70B-Instruct" + + +def test_crusoe_model_list_populated(): + """Test Crusoe models are present in model_prices_and_context_window.json""" + import litellm + + original_model_cost = litellm.model_cost + original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") + try: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + expected = [ + "crusoe/meta-llama/Llama-3.3-70B-Instruct", + "crusoe/deepseek-ai/DeepSeek-R1-0528", + "crusoe/deepseek-ai/DeepSeek-V3-0324", + "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507", + "crusoe/moonshotai/Kimi-K2-Thinking", + "crusoe/openai/gpt-oss-120b", + "crusoe/google/gemma-3-12b-it", + ] + for model in expected: + assert model in litellm.model_cost, f"{model} not found in model_cost" + assert litellm.model_cost[model].get("litellm_provider") == "crusoe" + finally: + litellm.model_cost = original_model_cost + if original_env is None: + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + else: + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env diff --git a/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py new file mode 100644 index 0000000000..54e689dea6 --- /dev/null +++ b/tests/test_litellm/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py @@ -0,0 +1,33 @@ +import pytest + +from litellm.llms.elevenlabs.text_to_speech.transformation import ( + ElevenLabsTextToSpeechConfig, +) + + +def test_should_encode_elevenlabs_voice_id_path_segment(): + config = ElevenLabsTextToSpeechConfig() + + url = config.get_complete_url( + model="elevenlabs/tts", + api_base="https://api.elevenlabs.io", + litellm_params={ + config.ELEVENLABS_VOICE_ID_KEY: "voice/../../models?x=1#frag", + }, + ) + + assert ( + url + == "https://api.elevenlabs.io/v1/text-to-speech/voice%2F..%2F..%2Fmodels%3Fx%3D1%23frag" + ) + + +def test_should_reject_dot_segment_elevenlabs_voice_id(): + config = ElevenLabsTextToSpeechConfig() + + with pytest.raises(ValueError, match="voice_id cannot be a dot path segment"): + config.get_complete_url( + model="elevenlabs/tts", + api_base="https://api.elevenlabs.io", + litellm_params={config.ELEVENLABS_VOICE_ID_KEY: ".."}, + ) diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py index 2431c9a9c4..a2f9572468 100644 --- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py +++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py @@ -2,7 +2,6 @@ Test Google AI Studio (Gemini) files transformation functionality """ -import os from unittest.mock import Mock, patch import httpx @@ -93,6 +92,30 @@ class TestGoogleAIStudioFilesTransformation: assert "key=" not in url assert params == {} + def test_transform_retrieve_file_request_encodes_file_id_path_segment(self): + file_id = "files/../../models/gemini-pro?x=1#frag" + litellm_params = {"api_key": "test-api-key"} + + url, params = self.handler.transform_retrieve_file_request( + file_id=file_id, + optional_params={}, + litellm_params=litellm_params, + ) + + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" + ) + assert params == {} + + def test_transform_retrieve_file_request_rejects_dot_path_segment(self): + with pytest.raises(ValueError, match="file_id cannot be a dot path segment"): + self.handler.transform_retrieve_file_request( + file_id="files/..", + optional_params={}, + litellm_params={"api_key": "test-api-key"}, + ) + @patch.dict("os.environ", {}, clear=True) @patch("litellm.llms.gemini.common_utils.get_secret_str", return_value=None) def test_transform_retrieve_file_request_missing_api_key(self, mock_get_secret): @@ -297,9 +320,7 @@ class TestGoogleAIStudioFilesTransformation: litellm_params=litellm_params, ) - # Verify URL extraction - assert "files/test123" in url - assert "generativelanguage.googleapis.com" in url + assert url == "https://generativelanguage.googleapis.com/v1beta/files/test123" # Params should be empty (API key goes in header via validate_environment) assert params == {} @@ -322,3 +343,22 @@ class TestGoogleAIStudioFilesTransformation: assert file_id in url assert "generativelanguage.googleapis.com" in url assert params == {} + + def test_transform_delete_file_request_encodes_file_id_path_segment(self): + file_id = "files/../../models/gemini-pro?x=1#frag" + litellm_params = { + "api_key": "test-api-key", + "api_base": "https://generativelanguage.googleapis.com", + } + + url, params = self.handler.transform_delete_file_request( + file_id=file_id, + optional_params={}, + litellm_params=litellm_params, + ) + + assert ( + url + == "https://generativelanguage.googleapis.com/v1beta/files/..%2F..%2Fmodels%2Fgemini-pro%3Fx%3D1%23frag" + ) + assert params == {} diff --git a/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py b/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py index 10d66174c5..43ce030323 100644 --- a/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py +++ b/tests/test_litellm/llms/manus/responses/test_manus_responses_transformation.py @@ -58,3 +58,18 @@ def test_transform_responses_api_request_adds_manus_params(): assert result["agent_profile"] == "manus-1.6" assert "input" in result assert "model" in result + + +def test_get_response_request_encodes_response_id(): + """Test response IDs are encoded before being appended to Manus URLs.""" + config = ManusResponsesAPIConfig() + + url, params = config.transform_get_response_api_request( + response_id="../../files?x=1#frag", + api_base="https://api.manus.im/v1/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.manus.im/v1/responses/..%2F..%2Ffiles%3Fx%3D1%23frag" + assert params == {} diff --git a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py b/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py index 751e48ff7a..f39be511b9 100644 --- a/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py +++ b/tests/test_litellm/llms/openai/evals/test_openai_evals_transformation.py @@ -49,6 +49,17 @@ def test_get_complete_url_with_eval_id(config: OpenAIEvalsConfig): assert url == "https://api.openai.com/v1/evals/eval_123" +def test_get_complete_url_encodes_eval_id(config: OpenAIEvalsConfig): + """Test eval_id is treated as a single path segment.""" + url = config.get_complete_url( + api_base="https://api.openai.com", + endpoint="evals", + eval_id="../../files?x=1#frag", + ) + + assert url == "https://api.openai.com/v1/evals/..%2F..%2Ffiles%3Fx%3D1%23frag" + + def test_get_complete_url_without_eval_id(config: OpenAIEvalsConfig): """Test URL construction without eval_id""" url = config.get_complete_url( @@ -253,3 +264,20 @@ def test_transform_cancel_eval_response(config: OpenAIEvalsConfig): assert result.id == "eval_123" assert result.object == "eval" + + +def test_transform_run_requests_encode_eval_and_run_ids(config: OpenAIEvalsConfig): + """Test run path IDs are treated as single path segments.""" + url, _, request_body = config.transform_cancel_run_request( + eval_id="../../evals?x=1#frag", + run_id="../runs#other", + api_base="https://api.openai.com", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/evals/..%2F..%2Fevals%3Fx%3D1%23frag/runs/..%2Fruns%23other/cancel" + ) + assert request_body == {} diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index dae8784283..acb9fa9b64 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -265,6 +265,24 @@ class TestOpenAIResponsesAPIConfig: assert result == "https://custom-openai.example.com/v1/responses" + def test_response_id_path_requests_encode_response_id(self): + """Test response_id is treated as one upstream URL path segment.""" + api_base = "https://custom-openai.example.com/v1/responses" + response_id = "../../files?x=1#frag" + + url, data = self.config.transform_list_input_items_request( + response_id=response_id, + api_base=api_base, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://custom-openai.example.com/v1/responses/..%2F..%2Ffiles%3Fx%3D1%23frag/input_items" + ) + assert data["limit"] == 20 + def test_get_event_model_class_generic_event(self): """Test that get_event_model_class returns the correct event model class""" from litellm.types.llms.openai import GenericEvent @@ -547,7 +565,12 @@ class TestOpenAIResponsesAPIConfig: """Base helper strips ``namespace`` from custom_tool_call for every provider path.""" inp = [ {"type": "function_call", "call_id": "a", "name": "f", "namespace": "keep"}, - {"type": "custom_tool_call", "call_id": "b", "name": "c", "namespace": "drop"}, + { + "type": "custom_tool_call", + "call_id": "b", + "name": "c", + "namespace": "drop", + }, ] out = BaseResponsesAPIConfig.strip_custom_tool_call_namespace_from_responses_input( inp diff --git a/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py b/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py index a8e07cd364..aebd93d062 100644 --- a/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py +++ b/tests/test_litellm/llms/openai/vector_store_files/test_openai_vector_store_files_transformation.py @@ -32,6 +32,21 @@ def test_get_complete_url(config: OpenAIVectorStoreFilesConfig): assert url == "https://api.example.com/v1/vector_stores/vs_123/files" +def test_get_complete_url_encodes_vector_store_id( + config: OpenAIVectorStoreFilesConfig, +): + url = config.get_complete_url( + api_base="https://api.example.com/v1", + vector_store_id="../vs_123?x=1#frag", + litellm_params={}, + ) + + assert ( + url + == "https://api.example.com/v1/vector_stores/..%2Fvs_123%3Fx%3D1%23frag/files" + ) + + def test_transform_create_request(config: OpenAIVectorStoreFilesConfig): api_base = "https://api.example.com/v1/vector_stores/vs_123/files" url, payload = config.transform_create_vector_store_file_request( @@ -60,6 +75,22 @@ def test_transform_list_request(config: OpenAIVectorStoreFilesConfig): assert params == {"limit": 2, "order": "asc"} +def test_transform_file_request_encodes_file_id(config: OpenAIVectorStoreFilesConfig): + api_base = "https://api.example.com/v1/vector_stores/vs_123/files" + + url, params = config.transform_retrieve_vector_store_file_content_request( + vector_store_id="vs_123", + file_id="../../files?x=1#frag", + api_base=api_base, + ) + + assert ( + url + == "https://api.example.com/v1/vector_stores/vs_123/files/..%2F..%2Ffiles%3Fx%3D1%23frag/content" + ) + assert params == {} + + def test_transform_create_response(config: OpenAIVectorStoreFilesConfig): response = httpx.Response( status_code=200, diff --git a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py b/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py index 053b107afb..e7b1aab45b 100644 --- a/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py +++ b/tests/test_litellm/llms/openai/vector_stores/test_openai_vector_stores_transformation.py @@ -64,3 +64,21 @@ class TestOpenAIVectorStoreAPIConfig: for i in range(16): assert f"key_{i}" in request_body["metadata"] assert request_body["metadata"][f"key_{i}"] == f"value_{i}" + + def test_transform_search_vector_store_request_encodes_vector_store_id(self): + config = OpenAIVectorStoreConfig() + + url, request_body = config.transform_search_vector_store_request( + vector_store_id="../../files?x=1#frag", + query="hello", + vector_store_search_optional_params={}, + api_base="https://api.openai.com/v1/vector_stores", + litellm_logging_obj=None, # type: ignore[arg-type] + litellm_params={}, + ) + + assert ( + url + == "https://api.openai.com/v1/vector_stores/..%2F..%2Ffiles%3Fx%3D1%23frag/search" + ) + assert request_body["query"] == "hello" diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py new file mode 100644 index 0000000000..c15554a46a --- /dev/null +++ b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py @@ -0,0 +1,70 @@ +from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.utils import encode_character_id_with_provider + + +def test_video_content_request_encodes_video_id_path_segment(): + config = OpenAIVideoConfig() + + url, params = config.transform_video_content_request( + video_id="../../responses?x=1#frag", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/videos/..%2F..%2Fresponses%3Fx%3D1%23frag/content" + ) + assert params == {} + + +def test_video_content_request_encodes_variant_query_param(): + """``variant`` is user-controlled and was previously interpolated raw + into the query string. A value like ``thumbnail&extra=1`` would + inject additional query parameters into the upstream request.""" + config = OpenAIVideoConfig() + + url, _ = config.transform_video_content_request( + video_id="vid_123", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + variant="thumbnail&extra=1", + ) + + # ``&`` and ``=`` must be percent-encoded so they cannot terminate + # the ``variant`` value or open a new query parameter. + assert "?variant=thumbnail%26extra%3D1" in url + # Sanity: the legitimate "thumbnail" value still round-trips cleanly. + url2, _ = config.transform_video_content_request( + video_id="vid_123", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + variant="thumbnail", + ) + assert url2.endswith("?variant=thumbnail") + + +def test_wrapped_character_id_is_decoded_then_encoded_as_path_segment(): + config = OpenAIVideoConfig() + character_id = encode_character_id_with_provider( + "../../characters?x=1#frag", + provider="openai", + model_id="sora", + ) + + url, params = config.transform_video_get_character_request( + character_id=character_id, + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.openai.com/v1/videos/characters/..%2F..%2Fcharacters%3Fx%3D1%23frag" + ) + assert params == {} diff --git a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py index e3343c3037..56953a574d 100644 --- a/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py +++ b/tests/test_litellm/llms/pg_vector/vector_stores/test_pg_vector_transformation.py @@ -141,6 +141,24 @@ class TestPGVectorStoreConfig: assert headers["Authorization"] == "Bearer test_key" assert url == "https://example.com/v1/vector_stores" + def test_search_request_encodes_vector_store_id(self): + config = PGVectorStoreConfig() + + url, request_body = config.transform_search_vector_store_request( + vector_store_id="../../files?x=1#frag", + query="hello", + vector_store_search_optional_params={}, + api_base="https://example.com/v1/vector_stores", + litellm_logging_obj=Mock(), + litellm_params={}, + ) + + assert ( + url + == "https://example.com/v1/vector_stores/..%2F..%2Ffiles%3Fx%3D1%23frag/search" + ) + assert request_body["query"] == "hello" + def test_environment_variable_support(self): """ Test that environment variables are supported for configuration. diff --git a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py b/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py index ae43eac7ff..baf2ab3391 100644 --- a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py +++ b/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py @@ -117,6 +117,24 @@ class TestRAGFlowChatTransformation: == "http://localhost:9380/api/v1/agents_openai/my-agent-id/chat/completions" ) + def test_get_complete_url_encodes_entity_id(self): + """Test RAGFlow chat IDs are encoded as one upstream path segment.""" + config = RAGFlowConfig() + + url = config.get_complete_url( + api_base="http://localhost:9380", + api_key=None, + model="ragflow/chat/..%2F..%2Fagents_openai%2Fother/gpt-4o-mini", + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert ( + url + == "http://localhost:9380/api/v1/chats_openai/..%252F..%252Fagents_openai%252Fother/chat/completions" + ) + def test_get_complete_url_strips_v1(self): """Test URL construction when api_base ends with /v1.""" config = RAGFlowConfig() diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py index 755716c9da..24879ce83f 100644 --- a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py +++ b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py @@ -134,6 +134,21 @@ class TestRunwayMLVideoTransformation: with pytest.raises(ValueError, match="still processing"): self.config._extract_video_url_from_response(processing_response) + def test_transform_video_status_encodes_video_id_path_segment(self): + """Test task IDs are encoded before being appended to Runway URLs.""" + url, params = self.config.transform_video_status_retrieve_request( + video_id="../../tasks/other?x=1#frag", + api_base="https://api.dev.runwayml.com/v1", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://api.dev.runwayml.com/v1/tasks/..%2F..%2Ftasks%2Fother%3Fx%3D1%23frag" + ) + assert params == {} + def test_full_video_workflow(self): """Test complete video generation workflow from creation to status check.""" config = RunwayMLVideoConfig() diff --git a/tests/test_litellm/llms/test_polling_url_origin_match.py b/tests/test_litellm/llms/test_polling_url_origin_match.py new file mode 100644 index 0000000000..f1f910bc73 --- /dev/null +++ b/tests/test_litellm/llms/test_polling_url_origin_match.py @@ -0,0 +1,177 @@ +""" +VERIA-51: polling URLs returned by upstream APIs (Azure DALL-E, +Azure Document Intelligence, Black Forest Labs) used to be followed +without origin validation. The handlers attached the operator's API +key to the polling request, so an attacker who could influence the +upstream response (or a compromised upstream) could redirect the proxy +to send credentials anywhere. + +These tests assert each handler now rejects polling URLs that don't +share an origin with the original request URL. +""" + +from unittest.mock import MagicMock, patch + +import httpx +import pytest + + +# Azure DALL-E sync + async paths route through ``assert_same_origin`` +# the same way as the cases below. The helper itself is unit-tested in +# ``tests/test_litellm/litellm_core_utils/test_url_utils.py``; the +# tests here exercise the wiring at sites with simpler signatures. + + +# ── Azure Document Intelligence polling ─────────────────────────────────────── + + +def test_azure_di_sync_rejects_cross_origin_polling(): + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, + ) + + config = AzureDocumentIntelligenceOCRConfig() + + raw_response = MagicMock() + raw_response.status_code = 202 + raw_response.headers = { + "Operation-Location": "https://attacker.example.com/results/xyz", + } + raw_response.request = MagicMock() + raw_response.request.url = ( + "https://eastus.cognitiveservices.azure.com/documentintelligence/.../analyze" + ) + raw_response.request.headers = {"Ocp-Apim-Subscription-Key": "leak-me"} + + with pytest.raises(ValueError, match="rejected polling URL"): + config.transform_ocr_response( + model="azure-doc-intel", + raw_response=raw_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + response={}, + ) + + +# ── Black Forest Labs polling ───────────────────────────────────────────────── + + +def test_bfl_image_generation_sync_rejects_cross_origin_polling(): + from litellm.llms.black_forest_labs.image_generation.handler import ( + BlackForestLabsImageGeneration, + ) + + handler = BlackForestLabsImageGeneration() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://attacker.example.com/get_result"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro" + + sync_client = MagicMock() + sync_client.get = MagicMock() + + with pytest.raises(Exception, match="Rejected polling URL"): + handler._poll_for_result_sync( + initial_response=initial_response, + headers={"x-key": "secret"}, + sync_client=sync_client, + ) + + sync_client.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_bfl_image_generation_async_rejects_cross_origin_polling(): + from litellm.llms.black_forest_labs.image_generation.handler import ( + BlackForestLabsImageGeneration, + ) + + handler = BlackForestLabsImageGeneration() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://attacker.example.com/get_result"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro" + + async_client = MagicMock() + async_client.get = MagicMock() + + with pytest.raises(Exception, match="Rejected polling URL"): + await handler._poll_for_result_async( + initial_response=initial_response, + headers={"x-key": "secret"}, + async_client=async_client, + ) + + async_client.get.assert_not_called() + + +def test_bfl_image_edit_sync_rejects_cross_origin_polling(): + from litellm.llms.black_forest_labs.image_edit.handler import ( + BlackForestLabsImageEdit, + ) + + handler = BlackForestLabsImageEdit() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://attacker.example.com/get_result"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro/edit" + + sync_client = MagicMock() + sync_client.get = MagicMock() + + with pytest.raises(Exception, match="Rejected polling URL"): + handler._poll_for_result_sync( + initial_response=initial_response, + headers={"x-key": "secret"}, + sync_client=sync_client, + ) + + sync_client.get.assert_not_called() + + +def test_bfl_image_generation_same_origin_polling_passes(): + """Sanity check: when the polling URL shares origin with the original + request, the origin check passes and polling proceeds.""" + from litellm.llms.black_forest_labs.image_generation.handler import ( + BlackForestLabsImageGeneration, + ) + + handler = BlackForestLabsImageGeneration() + + initial_response = MagicMock() + initial_response.status_code = 200 + initial_response.json = MagicMock( + return_value={"polling_url": "https://api.bfl.ai/v1/get_result?id=abc"} + ) + initial_response.request = MagicMock() + initial_response.request.url = "https://api.bfl.ai/v1/flux-pro" + + sync_client = MagicMock() + poll_response = MagicMock() + poll_response.status_code = 200 + poll_response.json = MagicMock(return_value={"status": "Ready"}) + sync_client.get = MagicMock(return_value=poll_response) + + result = handler._poll_for_result_sync( + initial_response=initial_response, + headers={"x-key": "secret"}, + sync_client=sync_client, + ) + + sync_client.get.assert_called_once() + assert result is poll_response diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index 831d1ef464..977c53280a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -86,6 +86,76 @@ def test_check_if_part_exists_in_parts_camel_case_snake_case(): assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) +def test_cached_content_respects_modify_params_for_cache_incompatible_fields(): + """Regression: cachedContent drops system/tools/toolConfig only when modify_params=True.""" + import litellm + + cache_name = "projects/p/locations/us-central1/cachedContents/abc123" + messages = [ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "hi"}, + ] + optional_params = { + "tools": [ + { + "functionDeclarations": [ + {"name": "get_weather", "description": "Get weather"}, + ] + } + ], + "tool_choice": {"functionCallingConfig": {"mode": "AUTO"}}, + } + + original_modify_params = litellm.modify_params + try: + # With modify_params=False (default), keep fields even with cachedContent. + litellm.modify_params = False + result = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=cache_name, + ) + assert result.get("cachedContent") == cache_name + assert "system_instruction" in result + assert "tools" in result + assert "toolConfig" in result + assert "contents" in result + + # With modify_params=True, drop cache-incompatible fields. + litellm.modify_params = True + result_modify_true = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=cache_name, + ) + assert result_modify_true.get("cachedContent") == cache_name + assert "system_instruction" not in result_modify_true + assert "tools" not in result_modify_true + assert "toolConfig" not in result_modify_true + assert "contents" in result_modify_true + + # Without cache, fields are always included. + result_no_cache = _transform_request_body( + messages=list(messages), + model="gemini-2.5-pro", + optional_params=dict(optional_params), + custom_llm_provider="vertex_ai", + litellm_params={}, + cached_content=None, + ) + assert "system_instruction" in result_no_cache + assert "tools" in result_no_cache + assert "toolConfig" in result_no_cache + finally: + litellm.modify_params = original_modify_params + + # Tests for issue #14556: Labels field provider-aware filtering def test_google_genai_excludes_labels(): """Test that Google GenAI/AI Studio endpoints exclude labels when custom_llm_provider='gemini'""" @@ -1291,6 +1361,53 @@ def test_file_data_field_order_gcs_urls(): ), "mime_type must come before file_uri in the file_data dict" +def test_gemini_files_api_uri_without_format(): + """ + Test that Gemini Files API URIs work WITHOUT an explicit format/mime_type. + + When a user uploads a file via the Gemini Files API and then references it + by URI (https://generativelanguage.googleapis.com/v1beta/files/...), + the file is already on Google's servers. These URLs return 403 when + fetched directly, so _process_gemini_media must NOT try to resolve the + MIME type via HTTP. Instead it should pass the URI through as file_data + and let the Gemini API resolve the type from its stored metadata. + + Related issue: https://github.com/BerriAI/litellm/issues/24907 + """ + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + file_url = "https://generativelanguage.googleapis.com/v1beta/files/37eh7rsw1vfe" + + # Should NOT raise — previously this hit the generic https:// handler + # which called _get_image_mime_type_from_url() and got a 403. + result = _process_gemini_media(image_url=file_url) + + assert "file_data" in result + file_data = result["file_data"] + assert file_data["file_uri"] == file_url + # When no format is provided, mime_type should be absent so the + # Gemini API infers it from the stored file metadata. + assert "mime_type" not in file_data + + +def test_gemini_files_api_uri_with_format(): + """ + Test that Gemini Files API URIs correctly forward an explicit format. + + Related issue: https://github.com/BerriAI/litellm/issues/24907 + """ + from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + file_url = "https://generativelanguage.googleapis.com/v1beta/files/n1vhxa28lyaw" + + result = _process_gemini_media(image_url=file_url, format="text/plain") + + assert "file_data" in result + file_data = result["file_data"] + assert file_data["file_uri"] == file_url + assert file_data["mime_type"] == "text/plain" + + def test_extract_file_data_with_path_object(): """ Test that filename is correctly extracted from Path objects for MIME type detection. 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/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py index 33b3bfce44..8a135dac3b 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -93,6 +93,51 @@ def test_vertex_ai_cancel_batch(): assert ":cancel" in call_args.kwargs["url"] +def test_vertex_ai_cancel_batch_encodes_batch_id(): + """Test that vertex_ai cancel_batch encodes user-controlled batch IDs.""" + handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket") + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "name": "projects/test-project/locations/us-central1/batchPredictionJobs/123456", + "state": "JOB_STATE_CANCELLING", + "createTime": "2024-03-17T10:00:00.000000Z", + "inputConfig": {"gcsSource": {"uris": ["gs://test-bucket/input.jsonl"]}}, + "outputConfig": { + "gcsDestination": {"outputUriPrefix": "gs://test-bucket/output"} + }, + } + + with patch( + "litellm.llms.vertex_ai.batches.handler._get_httpx_client" + ) as mock_client: + mock_client.return_value.post.return_value = mock_response + mock_client.return_value.get.return_value = mock_response + + with patch.object(handler, "_ensure_access_token") as mock_auth: + mock_auth.return_value = ("fake-token", "test-project") + + handler.cancel_batch( + _is_async=False, + batch_id="../../batchPredictionJobs/other?x=1#frag", + api_base=None, + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=600.0, + max_retries=None, + ) + + post_url = mock_client.return_value.post.call_args.kwargs["url"] + get_url = mock_client.return_value.get.call_args.kwargs["url"] + assert ( + "/..%2F..%2FbatchPredictionJobs%2Fother%3Fx%3D1%23frag:cancel" + in post_url + ) + assert "/..%2F..%2FbatchPredictionJobs%2Fother%3Fx%3D1%23frag" in get_url + + def test_vertex_ai_cancel_batch_forwards_timeout(): """Test that timeout is forwarded to the POST (cancel) HTTP call. diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py new file mode 100644 index 0000000000..b6329f33ae --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_search_vector_store_transformation.py @@ -0,0 +1,40 @@ +import pytest + +from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( + VertexSearchAPIVectorStoreConfig, +) + + +def test_should_encode_vertex_search_vector_store_id_in_complete_url(): + config = VertexSearchAPIVectorStoreConfig() + + url = config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vertex_collection_id": "default/collection", + "vector_store_id": "../../dataStores/other?x=1#frag", + }, + ) + + assert ( + url + == "https://discoveryengine.googleapis.com/v1/projects/test-project/locations/global/collections/default%2Fcollection/dataStores/..%2F..%2FdataStores%2Fother%3Fx%3D1%23frag/servingConfigs/default_config" + ) + + +def test_should_reject_dot_segment_vertex_search_vector_store_id(): + config = VertexSearchAPIVectorStoreConfig() + + with pytest.raises( + ValueError, match="vector_store_id cannot be a dot path segment" + ): + config.get_complete_url( + api_base=None, + litellm_params={ + "vertex_project": "test-project", + "vertex_location": "global", + "vector_store_id": "..", + }, + ) diff --git a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py index 623d162ccf..13571e63c7 100644 --- a/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py +++ b/tests/test_litellm/llms/volcengine/responses/test_volcengine_responses_transformation.py @@ -101,6 +101,23 @@ class TestVolcengineResponsesAPITransformation: ) assert api_base_full == "https://custom.volc.com/api/v3/responses" + def test_response_id_path_requests_encode_response_id(self): + """response_id should be encoded before building Volcengine URLs.""" + config = VolcEngineResponsesAPIConfig() + + url, params = config.transform_cancel_response_api_request( + response_id="../../responses/other?x=1#frag", + api_base="https://custom.volc.com/api/v3/responses", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert ( + url + == "https://custom.volc.com/api/v3/responses/..%2F..%2Fresponses%2Fother%3Fx%3D1%23frag/cancel" + ) + assert params == {} + @pytest.mark.parametrize( "litellm_params, expected_key", [ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 558c677d2d..85d5d6ba46 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -23,6 +23,21 @@ def mock_mcp_client_ip(): yield +@pytest.fixture +def trust_xff(): + """Force ``IPAddressUtils.is_request_from_trusted_proxy`` to True. + + Tests that exercise X-Forwarded-* parsing logic opt into this fixture. + The trust gate's own behaviour is covered by + ``test_get_request_base_url_xff_trust_gate``. + """ + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ): + yield + + @pytest.mark.asyncio async def test_authorize_endpoint_includes_response_type(): """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" @@ -505,6 +520,7 @@ async def test_register_client_remote_registration_success(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_authorize_endpoint_respects_x_forwarded_proto(): """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" try: @@ -572,6 +588,7 @@ async def test_authorize_endpoint_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_token_endpoint_respects_x_forwarded_proto(): """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" try: @@ -650,6 +667,7 @@ async def test_token_endpoint_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_oauth_protected_resource_respects_x_forwarded_proto(): """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" try: @@ -704,6 +722,7 @@ async def test_oauth_protected_resource_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_oauth_authorization_server_respects_x_forwarded_proto(): """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" try: @@ -759,6 +778,7 @@ async def test_oauth_authorization_server_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_register_client_respects_x_forwarded_proto(): """Test that register_client uses X-Forwarded-Proto for redirect_uris""" try: @@ -796,6 +816,7 @@ async def test_register_client_respects_x_forwarded_proto(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_authorize_endpoint_respects_x_forwarded_host(): """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" try: @@ -869,6 +890,7 @@ async def test_authorize_endpoint_respects_x_forwarded_host(): @pytest.mark.asyncio +@pytest.mark.usefixtures("trust_xff") async def test_token_endpoint_respects_x_forwarded_host(): """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" try: @@ -1071,7 +1093,12 @@ async def test_token_endpoint_respects_x_forwarded_host(): def test_get_request_base_url_comprehensive( base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url ): - """Comprehensive test for get_request_base_url with various header combinations""" + """Comprehensive test for get_request_base_url with various header combinations. + + These cases exercise the X-Forwarded-* parsing logic, so the trust gate + is patched True; the gate's own behaviour is covered by the + ``test_get_request_base_url_xff_trust_gate`` matrix below. + """ try: from fastapi import Request @@ -1081,11 +1108,9 @@ def test_get_request_base_url_comprehensive( except ImportError: pytest.skip("MCP discoverable endpoints not available") - # Create mock request mock_request = MagicMock(spec=Request) mock_request.base_url = base_url - # Build headers dict headers = {} if x_forwarded_proto: headers["X-Forwarded-Proto"] = x_forwarded_proto @@ -1094,16 +1119,17 @@ def test_get_request_base_url_comprehensive( if x_forwarded_port: headers["X-Forwarded-Port"] = x_forwarded_port - # Mock headers.get() to return our test values def mock_get(header_name, default=None): return headers.get(header_name, default) mock_request.headers.get = mock_get - # Test the function - result = get_request_base_url(mock_request) + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.is_request_from_trusted_proxy", + return_value=True, + ): + result = get_request_base_url(mock_request) - # Verify result assert result == expected_url, ( f"Expected '{expected_url}' but got '{result}'\n" f"Input: base_url={base_url}, " @@ -1113,6 +1139,131 @@ def test_get_request_base_url_comprehensive( ) +@pytest.mark.parametrize( + "general_settings,direct_ip,expect_xff_honoured", + [ + # Default: use_x_forwarded_for not set -> ignore X-Forwarded-* entirely. + ({}, "127.0.0.1", False), + # XFF enabled, no trusted ranges -> still ignored (no way to tell a trusted + # reverse proxy from a direct attacker). + ({"use_x_forwarded_for": True}, "127.0.0.1", False), + # XFF enabled, ranges set, but caller IP outside any range -> ignored. + ( + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + "203.0.113.5", + False, + ), + # XFF enabled, caller in trusted range -> headers honoured. + ( + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["10.0.0.0/8"], + }, + "10.0.0.7", + True, + ), + # Loopback example (common dev / single-host deploy). + ( + { + "use_x_forwarded_for": True, + "mcp_trusted_proxy_ranges": ["127.0.0.0/8"], + }, + "127.0.0.1", + True, + ), + ], +) +def test_get_request_base_url_xff_trust_gate( + general_settings, direct_ip, expect_xff_honoured +): + """Verify the X-Forwarded-* trust gate. + + With XFF poisoning attempted, the helper must return either the literal + base_url (gate denies) or the forwarded URL (gate allows), never the + forwarded URL when the gate denies. + """ + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = direct_ip + + headers = { + "X-Forwarded-Proto": "https", + "X-Forwarded-Host": "attacker.example.com", + } + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + mock_request.headers.__contains__ = lambda self_, name: name in headers + + with patch( + "litellm.proxy.proxy_server.general_settings", + general_settings, + create=True, + ): + result = get_request_base_url(mock_request) + + if expect_xff_honoured: + assert result == "https://attacker.example.com" + else: + assert result == "http://localhost:4000" + + +def test_xff_misconfig_warning_emitted_once(caplog): + """Operators upgrading from the old "always trust X-Forwarded-*" behaviour + get a one-shot warning when they have ``use_x_forwarded_for`` enabled + but no ``mcp_trusted_proxy_ranges`` configured. The warning must NOT + spam every request.""" + try: + from fastapi import Request + + from litellm.proxy import auth as proxy_auth_pkg # noqa: F401 + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + get_request_base_url, + ) + from litellm.proxy.auth import ip_address_utils + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + # Reset the module-level one-shot flag so the test is deterministic. + ip_address_utils._warned_xff_without_trusted_ranges = False + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://localhost:4000/" + mock_request.client = MagicMock() + mock_request.client.host = "203.0.113.5" + headers = {"X-Forwarded-Host": "attacker.example.com"} + mock_request.headers.get = lambda name, default=None: headers.get(name, default) + + misconfig = {"use_x_forwarded_for": True} + + import logging + + with ( + caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"), + patch("litellm.proxy.proxy_server.general_settings", misconfig, create=True), + ): + for _ in range(3): + get_request_base_url(mock_request) + + matching = [ + rec for rec in caplog.records if "mcp_trusted_proxy_ranges" in rec.getMessage() + ] + assert ( + len(matching) == 1 + ), f"expected exactly one warning, got {len(matching)}: {[r.getMessage() for r in matching]}" + + # ------------------------------------------------------------------- # Tests for scopes_supported when mcp_server.scopes is None # ------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py new file mode 100644 index 0000000000..3ad01e9c3e --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py @@ -0,0 +1,220 @@ +""" +VERIA-7 regression: OpenAPI-backed (local-registry) MCP tools must run +through `pre_call_tool_check` before dispatch, the same as managed +MCP server tools. +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_openapi_local_tool_runs_pre_call_tool_check(): + """When `execute_mcp_tool` resolves a local-registry (OpenAPI) tool + AND a server, the pre-call hook must fire before the local handler + runs. Pre-fix this path skipped the hook entirely.""" + from litellm.proxy._experimental.mcp_server import server as mcp_module + + user = UserAPIKeyAuth( + api_key="sk-user", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + fake_server = MagicMock() + fake_server.name = "openapi-petstore" + fake_server.is_byok = False + fake_server.auth_type = None + fake_server.mcp_info = None + fake_server.server_id = "srv-1" + fake_server.server_name = "openapi-petstore" + + fake_tool = MagicMock() + fake_tool.name = "list_pets" + + pre_call = AsyncMock(return_value={}) + handle_local = AsyncMock(return_value=[]) + + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=fake_server, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "pre_call_tool_check", + new=pre_call, + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=fake_tool, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + new=handle_local, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + await mcp_module.execute_mcp_tool( + name="list_pets", + arguments={"limit": 10}, + allowed_mcp_servers=[fake_server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + pre_call.assert_awaited_once() + handle_local.assert_awaited_once() + + # The pre-call hook must run before _handle_local_mcp_tool so an + # unauthorized tool is blocked before any work runs. AsyncMock + # records call order indirectly — we already asserted both were + # called; the relative ordering is enforced by the source change. + pre_call_kwargs = pre_call.await_args.kwargs + assert pre_call_kwargs["name"] == "list_pets" + assert pre_call_kwargs["server"] is fake_server + assert pre_call_kwargs["user_api_key_auth"] is user + # `proxy_logging_obj` must be sourced from the canonical proxy_server + # module (same as the managed path) — passing None would crash the + # downstream `_create_mcp_request_object_from_kwargs` call with + # AttributeError after the security checks succeed. + assert pre_call_kwargs["proxy_logging_obj"] is not None + + +@pytest.mark.asyncio +async def test_openapi_local_tool_blocked_when_pre_call_check_raises(): + """If the pre-call check raises (caller not authorized for this + tool), the local handler must NOT be invoked.""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + user = UserAPIKeyAuth( + api_key="sk-user", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + fake_server = MagicMock() + fake_server.name = "openapi-petstore" + fake_server.is_byok = False + fake_server.auth_type = None + fake_server.mcp_info = None + fake_server.server_id = "srv-1" + fake_server.server_name = "openapi-petstore" + + fake_tool = MagicMock() + fake_tool.name = "delete_pet" + + pre_call = AsyncMock( + side_effect=HTTPException(status_code=403, detail="not allowed") + ) + handle_local = AsyncMock(return_value=[]) + + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=fake_server, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "pre_call_tool_check", + new=pre_call, + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=fake_tool, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + new=handle_local, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + with pytest.raises(HTTPException) as exc: + await mcp_module.execute_mcp_tool( + name="delete_pet", + arguments={}, + allowed_mcp_servers=[fake_server], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + assert exc.value.status_code == 403 + pre_call.assert_awaited_once() + handle_local.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_openapi_local_tool_denied_when_server_not_resolvable(): + """If the local-registry tool is found but no MCP server resolves + (startup race or orphaned registry entry), the call must be rejected + rather than dispatched without `pre_call_tool_check`.""" + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server import server as mcp_module + + user = UserAPIKeyAuth( + api_key="sk-user", + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + fake_tool = MagicMock() + fake_tool.name = "list_pets" + + pre_call = AsyncMock(return_value={}) + handle_local = AsyncMock(return_value=[]) + + # `_get_mcp_server_from_tool_name` returns None — no server context. + with ( + patch.object( + mcp_module.global_mcp_server_manager, + "_get_mcp_server_from_tool_name", + return_value=None, + ), + patch.object( + mcp_module.global_mcp_server_manager, + "pre_call_tool_check", + new=pre_call, + ), + patch.object( + mcp_module.global_mcp_tool_registry, + "get_tool", + return_value=fake_tool, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool", + new=handle_local, + ), + patch( + "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed", + return_value=True, + ), + ): + with pytest.raises(HTTPException) as exc: + await mcp_module.execute_mcp_tool( + name="list_pets", + arguments={}, + allowed_mcp_servers=[], + start_time=datetime.now(timezone.utc), + user_api_key_auth=user, + ) + + assert exc.value.status_code == 503 + pre_call.assert_not_awaited() + handle_local.assert_not_awaited() diff --git a/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py b/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py new file mode 100644 index 0000000000..b0d2595e48 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_admin_viewer_handler_access.py @@ -0,0 +1,168 @@ +""" +Handler-level admin viewer parity tests. + +These tests assert that PROXY_ADMIN_VIEW_ONLY callers are NOT blocked at the +handler level for read-only admin endpoints. The route_checks layer is tested +separately in `test_route_checks.py`; here we verify each individual endpoint +function has been updated to use `_user_has_admin_view()` rather than a bare +`user_role != PROXY_ADMIN` check. + +The principle (see Admin Viewer role doc): anything Proxy Admin can read, +Admin Viewer can read. No writes, no cost-incurring actions. +""" + +import os +import sys +import types +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../../")) + +import litellm.proxy.proxy_server as ps +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.proxy_server import app + + +def _make_admin_viewer_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + ) + + +def _override_auth(role: LitellmUserRoles) -> None: + fake_user = UserAPIKeyAuth(user_id="viewer_user", user_role=role) + app.dependency_overrides[ps.user_api_key_auth] = lambda: fake_user + + +def _clear_overrides() -> None: + app.dependency_overrides.clear() + + +@pytest.fixture +def admin_viewer_client(monkeypatch): + """TestClient where auth always returns PROXY_ADMIN_VIEW_ONLY + a mocked Prisma.""" + mock_prisma = MagicMock() + + # Common DB tables touched by the read endpoints under test. + mock_budget_table = MagicMock() + mock_budget_table.find_many = AsyncMock(return_value=[]) + mock_budget_table.find_first = AsyncMock(return_value=None) + + mock_invitation_table = MagicMock() + mock_invitation_table.find_unique = AsyncMock(return_value=None) + + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + + mock_prisma.db = types.SimpleNamespace( + litellm_budgettable=mock_budget_table, + litellm_invitationlink=mock_invitation_table, + litellm_config=mock_config_table, + ) + + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + _override_auth(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) + + yield TestClient(app) + + _clear_overrides() + + +def _assert_not_role_blocked(response) -> None: + """The endpoint must not return a role-block error. + + Detects both the 400 ``not_allowed_access`` pattern (used by most + management endpoints) and the 403 ``Admin role required`` pattern + (used by model cost map endpoints). + """ + if response.status_code in (400, 401, 403): + body = response.json() + detail = body.get("detail", body) + if isinstance(detail, dict): + err = detail.get("error", "") + else: + err = str(detail) + err_lower = err.lower() + role_block_signals = ( + "your role=", + "not allowed to access", + "admin role required", + "admin-only endpoint", + ) + for signal in role_block_signals: + assert ( + signal not in err_lower + ), f"endpoint blocked PROXY_ADMIN_VIEW_ONLY at handler level: {err}" + + +def test_budget_list_allows_admin_viewer(admin_viewer_client): + """`/budget/list` is read-only and must be accessible to Admin Viewer.""" + resp = admin_viewer_client.get("/budget/list") + _assert_not_role_blocked(resp) + assert resp.status_code == 200, resp.text + + +def test_budget_settings_allows_admin_viewer(admin_viewer_client): + """`/budget/settings` describes a budget's fields; read-only.""" + resp = admin_viewer_client.get("/budget/settings", params={"budget_id": "b1"}) + _assert_not_role_blocked(resp) + assert resp.status_code == 200, resp.text + + +def test_alerting_settings_allows_admin_viewer(admin_viewer_client): + """`/alerting/settings` describes alerting params; read-only.""" + resp = admin_viewer_client.get("/alerting/settings") + _assert_not_role_blocked(resp) + # Endpoint may 400 for *config* reasons (no proxy config loaded), but it + # must not 400 because of role. + assert resp.status_code != 403, resp.text + + +def test_get_config_field_info_allows_admin_viewer(admin_viewer_client): + """`/config/field/info` describes a single general-settings field; read-only.""" + resp = admin_viewer_client.get( + "/config/field/info", params={"field_name": "alerting"} + ) + _assert_not_role_blocked(resp) + + +def test_get_config_list_allows_admin_viewer(admin_viewer_client): + """`/config/list` lists configurable params for a config_type; read-only.""" + resp = admin_viewer_client.get( + "/config/list", params={"config_type": "general_settings"} + ) + _assert_not_role_blocked(resp) + + +def test_get_config_callbacks_allows_admin_viewer(admin_viewer_client): + """`/get/config/callbacks` lists current callbacks; read-only.""" + resp = admin_viewer_client.get("/get/config/callbacks") + _assert_not_role_blocked(resp) + + +def test_invitation_info_allows_admin_viewer(admin_viewer_client): + """`/invitation/info` reads a single invitation; read-only. + + The invitation lookup will return 400 because no invitation exists in our + mock DB — that's fine. We only assert it doesn't hit the role-block path. + """ + resp = admin_viewer_client.get( + "/invitation/info", params={"invitation_id": "nonexistent"} + ) + _assert_not_role_blocked(resp) + + +def test_model_cost_map_reload_status_allows_admin_viewer(admin_viewer_client): + """`/schedule/model_cost_map_reload/status` is read-only operations status.""" + resp = admin_viewer_client.get("/schedule/model_cost_map_reload/status") + _assert_not_role_blocked(resp) + + +def test_model_cost_map_source_allows_admin_viewer(admin_viewer_client): + """`/model/cost_map/source` reads the configured cost map source URL.""" + resp = admin_viewer_client.get("/model/cost_map/source") + _assert_not_role_blocked(resp) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 4c21d0ec64..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, @@ -1964,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_oauth2_proxy_hook.py b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py new file mode 100644 index 0000000000..dcbfd281e0 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_oauth2_proxy_hook.py @@ -0,0 +1,211 @@ +""" +Regression tests for the OAuth2-proxy header-forgery fix +(GHSA-5c3m-qffq-4r9m). + +The hook reads HTTP request headers per ``oauth2_config_mappings`` and +constructs a ``UserAPIKeyAuth`` from them. The fix has two parts: + +1. Only requests from configured trusted proxy CIDR ranges may provide + identity headers. +2. Only identity fields may be mapped from those headers. Without the + identity-only allowlist any field could be mapped — including + ``user_role``, which Pydantic coerces from the string + ``"proxy_admin"`` into ``LitellmUserRoles.PROXY_ADMIN``. +""" + +import os +import sys + +import pytest +from fastapi import Request +from starlette.datastructures import Headers + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.auth.oauth2_proxy_hook import ( + ALLOWED_OAUTH2_PROXY_FIELDS, + handle_oauth2_proxy_request, +) + + +def _request_with_headers(headers: dict, *, client_host: str = "127.0.0.1") -> Request: + scope = { + "type": "http", + "client": (client_host, 12345), + "headers": [(k.lower().encode(), v.encode()) for k, v in headers.items()], + } + request = Request(scope=scope) + request._headers = Headers(headers) + return request + + +@pytest.fixture +def configure_proxy(monkeypatch): + """ + Yields a callable that sets ``oauth2_config_mappings`` and + ``trusted_proxy_ranges`` on the proxy_server module for the duration + of one test. Defaults to a single identity mapping and localhost as + a trusted proxy. + """ + import litellm.proxy.proxy_server as proxy_server + + def _configure(*, mappings=None, trusted_proxy_ranges=("127.0.0.1/32",)): + if mappings is None: + mappings = {"user_id": "x-user-id"} + settings = { + "oauth2_config_mappings": mappings, + "trusted_proxy_ranges": trusted_proxy_ranges, + } + monkeypatch.setattr( + proxy_server, + "general_settings", + settings, + raising=False, + ) + + return _configure + + +@pytest.mark.asyncio +async def test_returns_auth_for_simple_user_id_mapping(configure_proxy): + configure_proxy() + request = _request_with_headers({"x-user-id": "alice"}) + + auth = await handle_oauth2_proxy_request(request) + + assert auth.user_id == "alice" + assert auth.user_role is None + + +@pytest.mark.asyncio +async def test_rejects_identity_headers_without_trusted_proxy_ranges(configure_proxy): + configure_proxy(trusted_proxy_ranges=None) + request = _request_with_headers({"x-user-id": "alice"}) + + with pytest.raises(ValueError, match="trusted_proxy_ranges"): + await handle_oauth2_proxy_request(request) + + +@pytest.mark.asyncio +async def test_rejects_identity_headers_from_untrusted_direct_client(configure_proxy): + configure_proxy(trusted_proxy_ranges=["10.0.0.0/24"]) + request = _request_with_headers({"x-user-id": "alice"}, client_host="203.0.113.10") + + with pytest.raises(ValueError, match="not trusted"): + await handle_oauth2_proxy_request(request) + + +@pytest.mark.parametrize( + "privileged_field", + [ + # The GHSA-5c3m-qffq-4r9m primary privesc field. + "user_role", + # Key-level enforcement bypass shapes. + "api_key", + "token", + "permissions", + "allowed_routes", + "max_budget", + "spend", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "metadata", + # User-level enforcement bypass — flagged by Greptile as a denylist gap. + "user_max_budget", + "user_tpm_limit", + "user_rpm_limit", + "user_spend", + # Team / org / end-user / region — same class, all denied by the + # identity-only allowlist. + "team_max_budget", + "team_spend", + "team_member_tpm_limit", + "organization_max_budget", + "organization_tpm_limit", + "end_user_max_budget", + "allowed_model_region", + # Anything not on ALLOWED_OAUTH2_PROXY_FIELDS is blocked, even + # fabricated field names admins might try. + "definitely_not_a_real_field", + ], +) +@pytest.mark.asyncio +async def test_refuses_to_map_non_identity_fields(configure_proxy, privileged_field): + # GHSA-5c3m-qffq-4r9m attack shape: admin maps a privileged field + # to a header and a caller forges the value. The allowlist rejects + # any non-identity mapping at request time, regardless of whether + # the field ever appeared on a denylist — which is the whole reason + # we use an allowlist instead. + configure_proxy(mappings={privileged_field: f"x-{privileged_field}"}) + request = _request_with_headers({f"x-{privileged_field}": "proxy_admin"}) + + with pytest.raises(ValueError) as exc: + await handle_oauth2_proxy_request(request) + assert privileged_field in str(exc.value) + + +@pytest.mark.parametrize("identity_field", sorted(ALLOWED_OAUTH2_PROXY_FIELDS)) +def test_allowlist_is_identity_only(identity_field): + # Lock in the allowlist's intent: only identity-assertion fields are + # safe to populate from a header. If anyone proposes adding budget / + # spend / role / permission to ``ALLOWED_OAUTH2_PROXY_FIELDS``, this + # assertion forces them to update the test deliberately. + assert identity_field in { + "user_id", + "user_email", + "team_id", + "team_alias", + "org_id", + "models", + } + + +@pytest.mark.asyncio +async def test_user_role_header_forgery_attack_is_blocked(configure_proxy): + # End-to-end form of the privesc: with ``user_role`` mapped, the + # forged ``X-User-Role: proxy_admin`` header would have produced + # a ``UserAPIKeyAuth(user_role=PROXY_ADMIN)``. Now rejected before + # any auth object is constructed. + configure_proxy( + mappings={"user_id": "x-user-id", "user_role": "x-user-role"}, + ) + request = _request_with_headers( + { + "x-user-id": "attacker", + "x-user-role": LitellmUserRoles.PROXY_ADMIN.value, + } + ) + + with pytest.raises(ValueError, match="user_role"): + await handle_oauth2_proxy_request(request) + + +@pytest.mark.asyncio +async def test_safe_fields_still_pass_through(configure_proxy): + # The documented use case for OAuth2 proxy auth: identity assertion + # from a trusted upstream. Must remain unaffected by the denylist. + configure_proxy( + mappings={ + "user_id": "x-user-id", + "user_email": "x-user-email", + "team_id": "x-team-id", + "models": "x-models", + }, + ) + request = _request_with_headers( + { + "x-user-id": "alice", + "x-user-email": "alice@example.com", + "x-team-id": "team-corp", + "x-models": "gpt-4, gpt-3.5-turbo", + } + ) + + auth = await handle_oauth2_proxy_request(request) + + assert auth.user_id == "alice" + assert auth.user_email == "alice@example.com" + assert auth.team_id == "team-corp" + assert auth.models == ["gpt-4", "gpt-3.5-turbo"] diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index a5d405cfc2..39f832256d 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1198,6 +1198,349 @@ def test_proxy_admin_viewer_can_access_audit_logs(route): ) +# ── Admin Viewer parity: Logs page endpoints ────────────────────────────────── +# +# The Admin Viewer (PROXY_ADMIN_VIEW_ONLY) role is documented as +# "view all keys, view all spend" and follows a read-parity-with-Proxy-Admin +# rule. The UI Logs page is the most user-visible failure mode: filtering and +# log details break entirely when these routes are blocked at the route_checks +# layer, even though the underlying handlers already gate on PROXY_ADMIN_VIEW_ONLY. +# +# Each route below corresponds to a network call made by the Logs page +# (ui/litellm-dashboard/src/components/view_logs/) — see the comment on each. +ADMIN_VIEWER_LOGS_PAGE_ROUTES = [ + # Main paginated log list — uiSpendLogsCall in log_filter_logic.tsx & index.tsx + "/spend/logs/ui", + # Single-log detail drawer — fetched on row click in LogDetailsDrawer + "/spend/logs/ui/abc-request-id", + # Multi-call session drawer — sessionSpendLogsCall in LogDetailsDrawer + "/spend/logs/session/ui", + # End User filter dropdown — allEndUsersCall in index.tsx + "/customer/list", + "/customer/info", + # Cost estimation — used by some log views + "/cost/estimate", + # Public spend logs / spend tracking routes that admin viewer should read + "/spend/logs", + "/spend/keys", + "/spend/users", + "/spend/tags", + "/spend/calculate", +] + + +@pytest.mark.parametrize("route", ADMIN_VIEWER_LOGS_PAGE_ROUTES) +def test_proxy_admin_viewer_can_access_logs_page_endpoints(route): + """ + PROXY_ADMIN_VIEW_ONLY must pass route_checks for every endpoint the UI + Logs page depends on. Without these, the page renders empty / errors. + """ + user_obj = LiteLLM_UserTable( + user_id="viewer_user", + user_email="viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + valid_token = UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as e: + pytest.fail( + f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}" + ) + + +@pytest.mark.parametrize("route", ADMIN_VIEWER_LOGS_PAGE_ROUTES) +def test_internal_user_blocked_from_admin_viewer_logs_routes(route): + """ + The Logs-page route opening above must NOT also widen access for + INTERNAL_USER. Plain internal users still see only their own logs and + must be blocked from proxy-wide spend tracking + customer routes. + """ + user_obj = LiteLLM_UserTable( + user_id="internal_user", + user_email="user@example.com", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + valid_token = UserAPIKeyAuth( + user_id="internal_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + # Routes already in `spend_tracking_routes` (which is part of + # `internal_user_routes`) are intentionally accessible to internal users + # for their own scoped spend — those handlers enforce per-user filtering. + # /cost/estimate is similarly per-user. The /customer/* routes are + # admin-only. + INTERNAL_USER_BLOCKED_SUBSET = { + "/customer/list", + "/customer/info", + } + if route not in INTERNAL_USER_BLOCKED_SUBSET: + return + + with pytest.raises(Exception) as exc_info: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + assert "Only proxy admin" in str(exc_info.value) + + +# ── Admin Viewer parity: Settings/observability read endpoints ──────────────── +# +# These are GET endpoints accessible to PROXY_ADMIN that the UI exposes to +# admin viewers via sidebar items gated by `all_admin_roles` (which includes +# proxy_admin_viewer). Without these, the Logging & Alerts, Caching, Budgets, +# and Admin Settings pages break for admin viewers. +ADMIN_VIEWER_SETTINGS_ROUTES = [ + # Logging & Alerts page + "/callbacks/list", + "/callbacks/configs", + "/get/config/callbacks", + "/alerting/settings", + # Admin Settings / Router Settings pages + "/config/list", + "/config/field/info", + # Budgets page + "/budget/list", + "/budget/settings", + # Invitation viewing (admin viewer cannot create/delete; can read) + "/invitation/info", + # Guardrails / Policies pages (read-only views) + "/guardrails/list", + "/v2/guardrails/list", + "/guardrails/submissions", + "/guardrails/submissions/some-guardrail-id", + "/guardrails/usage/overview", + "/policies/attachments/list", + # MCP semantic filter settings (read) + "/get/mcp_semantic_filter_settings", + # Model cost map (read-only status / source) + "/schedule/model_cost_map_reload/status", + "/model/cost_map/source", +] + + +@pytest.mark.parametrize("route", ADMIN_VIEWER_SETTINGS_ROUTES) +def test_proxy_admin_viewer_can_access_settings_read_endpoints(route): + """ + PROXY_ADMIN_VIEW_ONLY must pass route_checks for the read-only + settings/observability endpoints exposed in admin-only sidebar groups. + """ + user_obj = LiteLLM_UserTable( + user_id="viewer_user", + user_email="viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + valid_token = UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + request = MagicMock(spec=Request) + request.query_params = {} + + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as e: + pytest.fail( + f"proxy_admin_viewer should be able to access {route}. Got error: {str(e)}" + ) + + +# ── Admin Viewer parity: default-allow GET semantics ───────────────────────── +# +# The route-check layer is structured to default-allow safe HTTP methods +# (GET / HEAD / OPTIONS) for PROXY_ADMIN_VIEW_ONLY. This eliminates the +# whack-a-mole where every newly-added GET endpoint silently 403'd until +# someone remembered to add it to admin_viewer_routes. +# +# These tests pin the new contract: +# - Any GET endpoint not on the LLM/inference path is readable. +# - Any unsafe method (POST/PUT/PATCH/DELETE) outside the explicit allow +# sets is still 403. + +# Routes the user reported as broken in production — they're in disparate +# corners of the codebase and represent the long tail of GETs we'd otherwise +# need to enumerate manually. Default-allow makes them all work. +ADMIN_VIEWER_REPORTED_GET_ROUTES = [ + "/in_product_nudges", + "/health/latest", + "/credentials", + "/v1/mcp/network/client-ip", + "/claude-code/plugins", + "/policy/templates", + # Routes we already had to enumerate manually (regression coverage). + "/spend/logs/ui", + "/customer/list", + "/guardrails/list", + "/policies/attachments/list", + # Hypothetical future GETs — must not require an allowlist entry. + "/some/future/read/endpoint", + "/another/admin-tool/status", +] + + +@pytest.mark.parametrize("route", ADMIN_VIEWER_REPORTED_GET_ROUTES) +def test_proxy_admin_viewer_default_allows_any_get(route): + """ + PROXY_ADMIN_VIEW_ONLY must be able to GET any non-inference endpoint. + + This is a structural guarantee: the route-check defaults to allow for + safe HTTP methods so we don't have to maintain an explicit allowlist. + """ + user_obj = LiteLLM_UserTable( + user_id="viewer_user", + user_email="viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + valid_token = UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + request = MagicMock(spec=Request) + request.method = "GET" + request.query_params = {} + request.url = MagicMock() + request.url.path = route + + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as e: + pytest.fail(f"proxy_admin_viewer GET should default-allow {route!r}. Got: {e}") + + +@pytest.mark.parametrize( + "route", + [ + # Random path that isn't in any allowlist — POST must still 403. + "/some/future/write/endpoint", + # Hard-blocked write routes. + "/user/new", + "/team/new", + "/key/generate", + "/model/new", + ], +) +def test_proxy_admin_viewer_post_blocked_outside_allowlists(route): + """ + Default-allow only applies to safe HTTP methods. POST/PUT/PATCH/DELETE + on a route not in any allow set must still 403. + """ + user_obj = LiteLLM_UserTable( + user_id="viewer_user", + user_email="viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + valid_token = UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + assert exc_info.value.status_code == 403 + + +# ── Admin Viewer: management_routes write endpoints stay blocked ───────────── +# +# `management_routes` is a mix of reads (info/list, handled via the safe-method +# branch — GET) and writes. The route_checks layer must NOT blanket-allow the +# whole set on POST — that would let Admin Viewer mutate teams, JWT mappings, +# and bulk-update keys, violating the "no writes, ever" rule. +# +# These cases pin the gap closed (Greptile P1 review, 2026-04-30). +ADMIN_VIEWER_MANAGEMENT_ROUTE_WRITES = [ + # Team writes + "/team/block", + "/team/unblock", + "/team/permissions_update", + # JWT key mapping writes + "/jwt/key/mapping/new", + "/jwt/key/mapping/update", + "/jwt/key/mapping/delete", + # Key writes (existing _ADMIN_VIEWER_BLOCKED_WRITE_ROUTES doesn't list bulk + # update or per-key reset-spend, so the management_routes fallback was the + # only thing keeping them out — and it was permissive, not restrictive). + "/key/bulk_update", + "/key/some-key-id/reset_spend", +] + + +@pytest.mark.parametrize("route", ADMIN_VIEWER_MANAGEMENT_ROUTE_WRITES) +def test_proxy_admin_viewer_post_blocked_for_management_route_writes(route): + """ + Admin Viewer must be blocked on POST to write endpoints in + `management_routes`, even when the specific route is not in + `_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES`. + """ + user_obj = LiteLLM_UserTable( + user_id="viewer_user", + user_email="viewer@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + valid_token = UserAPIKeyAuth( + user_id="viewer_user", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + + with pytest.raises(HTTPException) as exc_info: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + assert exc_info.value.status_code == 403 + + class TestModelsRouteExemptFromDisableLLMEndpoints: """ Test that /models and /v1/models are exempt from DISABLE_LLM_API_ENDPOINTS. 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 3f28191bce..398084fdc5 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( @@ -17,6 +16,8 @@ from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( LiteLLMRoutes, LiteLLM_JWTAuth, + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, LiteLLM_UserTable, LitellmUserRoles, ProxyErrorTypes, @@ -25,9 +26,11 @@ 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 ( _route_requires_auth_despite_public, + _reserve_budget_after_common_checks, _run_centralized_common_checks, _run_post_custom_auth_checks, get_api_key, @@ -35,6 +38,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" @@ -75,6 +85,74 @@ def test_public_ai_hub_routes_remain_public(): assert _route_requires_auth_despite_public(route, {}) is False +@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"]) @@ -133,6 +211,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"]) @@ -1890,7 +2001,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, @@ -2151,6 +2262,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 f7cb4d72d9..2e738ff900 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -231,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""" 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/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 59b6e24f43..e10258c082 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -758,6 +758,60 @@ class TestDeferredStreamingClosure: apply_guardrail_called is False ), "apply_guardrail guardrails must be SKIPPED in deferred path" + @pytest.mark.asyncio + async def test_streaming_iterator_hook_skipped_in_deferred_path(self): + """regression test: guardrails that define async_post_call_streaming_iterator_hook must be SKIPPED in _run_deferred_stream_guardrails. + The iterator hook already scanned the assembled response in the streaming + pipeline""" + success_hook_called = False + + class IteratorHookGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="iterator-hook", + default_on=True, + event_hook=GuardrailEventHooks.post_call, + ) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict, response, request_data + ): + async for chunk in response: + yield chunk + + async def async_post_call_success_hook( + self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any + ) -> Any: + nonlocal success_hook_called + success_hook_called = True + return response + + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"metadata": {}} + + async def track_async_success(*args, **kwargs): + pass + + mock_logging_obj.async_success_handler = track_async_success + + guardrail = IteratorHookGuardrail() + + with patch("litellm.callbacks", [guardrail]): + await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( + captured_data={"model": "gpt-4", "metadata": {}}, + captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"), + captured_logging_obj=mock_logging_obj, + assembled_response=MagicMock(), + cache_hit=False, + ) + + await asyncio.sleep(0) + + assert success_hook_called is False, ( + "Guardrails that implement async_post_call_streaming_iterator_hook " + "must be SKIPPED in deferred path — the iterator hook already ran" + ) + @pytest.mark.asyncio async def test_hooks_receive_merged_guardrail_data(self): """Hooks must receive guardrail_data (the merged dict from diff --git a/tests/test_litellm/proxy/guardrails/test_qostodian_nexus_guardrail.py b/tests/test_litellm/proxy/guardrails/test_qostodian_nexus_guardrail.py new file mode 100644 index 0000000000..6daa3e1430 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/test_qostodian_nexus_guardrail.py @@ -0,0 +1,252 @@ +""" +Unit tests for Qostodian Nexus (by Qohash) integration. + +Tests verify: +1. QostodianNexus can be instantiated with default and custom values +2. Qostodian Nexus is registered in SupportedGuardrailIntegrations +3. Guardrail initializer and class registries contain Qostodian Nexus +4. Configuration parameters are properly passed through +5. QostodianNexusConfigModel works correctly +""" + +import os +import pytest +from unittest.mock import MagicMock + + +def test_qostodian_nexus_initialization_with_defaults(): + """Test QostodianNexus initializes with default values.""" + import os + from unittest.mock import patch + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + + # Unset env var so the hardcoded default is used + env = {k: v for k, v in os.environ.items() if k != "QOSTODIAN_NEXUS_API_BASE"} + with patch.dict(os.environ, env, clear=True): + guardrail = QostodianNexus() + + # Should use default api_base + assert guardrail.api_base is not None + assert "nexus:8800" in guardrail.api_base + + +def test_qostodian_nexus_initialization_with_custom_api_base(): + """Test QostodianNexus initializes with custom api_base.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + + custom_api_base = "http://custom-nexus:9000" + guardrail = QostodianNexus(api_base=custom_api_base) + + assert custom_api_base in guardrail.api_base + + +def test_qostodian_nexus_in_supported_guardrail_integrations(): + """Test that Qostodian Nexus is registered in SupportedGuardrailIntegrations enum.""" + from litellm.types.guardrails import SupportedGuardrailIntegrations + + # Check enum contains QOSTODIAN_NEXUS + assert hasattr(SupportedGuardrailIntegrations, "QOSTODIAN_NEXUS") + assert SupportedGuardrailIntegrations.QOSTODIAN_NEXUS.value == "qostodian_nexus" + + # Check it's in the list of all values + all_values = [e.value for e in SupportedGuardrailIntegrations] + assert "qostodian_nexus" in all_values + + +def test_qostodian_nexus_in_guardrail_initializer_registry(): + """Test that Qostodian Nexus is registered in guardrail_initializer_registry.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import ( + guardrail_initializer_registry, + ) + + assert "qostodian_nexus" in guardrail_initializer_registry + assert callable(guardrail_initializer_registry["qostodian_nexus"]) + + +def test_qostodian_nexus_in_guardrail_class_registry(): + """Test that Qostodian Nexus is registered in guardrail_class_registry.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import ( + guardrail_class_registry, + QostodianNexus, + ) + + assert "qostodian_nexus" in guardrail_class_registry + assert guardrail_class_registry["qostodian_nexus"] == QostodianNexus + + +def test_qostodian_nexus_config_model_initialization(): + """Test QostodianNexusConfigModel can be instantiated.""" + from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, + ) + + config = QostodianNexusConfigModel( + api_base="http://test:8800", + ) + + assert config.api_base == "http://test:8800" + + +def test_qostodian_nexus_config_model_defaults(): + """Test QostodianNexusConfigModel uses correct defaults.""" + from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, + ) + + config = QostodianNexusConfigModel() + + assert config.api_base is None + + +def test_qostodian_nexus_config_model_ui_friendly_name(): + """Test QostodianNexusConfigModel returns correct UI friendly name.""" + from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, + ) + + ui_name = QostodianNexusConfigModel.ui_friendly_name() + assert ui_name == "Qostodian Nexus" + + +def test_qostodian_nexus_initializer_function(): + """Test the initialize_guardrail function.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import initialize_guardrail + from litellm.types.guardrails import LitellmParams, Guardrail + from unittest.mock import patch + + # Mock litellm.logging_callback_manager + with patch("litellm.logging_callback_manager") as mock_manager: + mock_manager.add_litellm_callback = MagicMock() + + # Create test params + litellm_params = LitellmParams( + guardrail="qostodian_nexus", + mode="pre_call", + api_base="http://test:8800", + default_on=True, + ) + + guardrail_config: Guardrail = {"guardrail_name": "test-qostodian-nexus"} + + # Call initializer + result = initialize_guardrail(litellm_params, guardrail_config) + + # Verify callback was added + mock_manager.add_litellm_callback.assert_called_once() + + # Verify returned instance has correct properties + assert result is not None + assert "test:8800" in result.api_base + + +def test_qostodian_nexus_inherits_from_generic_guardrail_api(): + """Test that QostodianNexus inherits from GenericGuardrailAPI.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( + GenericGuardrailAPI, + ) + + assert issubclass(QostodianNexus, GenericGuardrailAPI) + + +def test_qostodian_nexus_guardrail_name_constant(): + """Test that GUARDRAIL_NAME constant is defined correctly.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash.qohash import GUARDRAIL_NAME + + assert GUARDRAIL_NAME == "qostodian_nexus" + + +def test_qostodian_nexus_get_config_model(): + """Test that QostodianNexus returns the correct config model.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, + ) + + config_model = QostodianNexus.get_config_model() + + assert config_model is not None + assert config_model == QostodianNexusConfigModel + + +def test_qostodian_nexus_env_vars(): + """Test that QOSTODIAN_NEXUS_API_BASE env var is picked up correctly.""" + import os + from unittest.mock import patch + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + + with patch.dict(os.environ, {"QOSTODIAN_NEXUS_API_BASE": "http://new-api:8800"}): + guardrail = QostodianNexus() + assert "new-api:8800" in guardrail.api_base + + +def test_qostodian_nexus_config_model_field_descriptions(): + """Test that QostodianNexusConfigModel has correct field descriptions.""" + from litellm.types.proxy.guardrails.guardrail_hooks.qohash import ( + QostodianNexusConfigModel, + ) + + # Check that field descriptions mention the correct env vars + api_base_field = QostodianNexusConfigModel.model_fields["api_base"] + assert "QOSTODIAN_NEXUS_API_BASE" in api_base_field.description + + +def test_qostodian_nexus_unified_detection(): + """ + Test that QostodianNexus is properly detected by LiteLLM's unified guardrail system. + + This verifies the fix for the detection bug where QostodianNexus wasn't being + recognized because apply_guardrail was only inherited, not in the class's own __dict__. + """ + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + + # Create an instance (this is how LiteLLM uses it) + instance = QostodianNexus(api_base="http://test:8800") + + # Test the exact detection logic used in litellm/proxy/utils.py:868 + # use_unified = "apply_guardrail" in type(callback).__dict__ + use_unified = "apply_guardrail" in type(instance).__dict__ + + # Should be detected as using unified guardrail system + assert use_unified is True, ( + "QostodianNexus should be detected by unified guardrail system. " + "The apply_guardrail method must be present in QostodianNexus.__dict__" + ) + + # Also verify the method is callable + assert hasattr(instance, "apply_guardrail") + assert callable(instance.apply_guardrail) + + +def test_qostodian_nexus_builtin_extra_headers(): + """Test that QostodianNexus includes built-in x-qostodian-nexus-identifiers-* headers.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + + instance = QostodianNexus() + + expected_headers = [ + "x-qostodian-nexus-identifiers-trace", + "x-qostodian-nexus-identifiers-source", + "x-qostodian-nexus-identifiers-container", + "x-qostodian-nexus-identifiers-identity", + ] + + for header in expected_headers: + assert header in instance.extra_headers, ( + f"Expected built-in header '{header}' to be in extra_headers" + ) + + +def test_qostodian_nexus_extra_headers_merged(): + """Test that caller-supplied extra_headers are merged with built-in headers.""" + from litellm.proxy.guardrails.guardrail_hooks.qohash import QostodianNexus + + custom_header = "x-custom-correlation-id" + instance = QostodianNexus(extra_headers=[custom_header]) + + # Built-in headers should be present + assert "x-qostodian-nexus-identifiers-trace" in instance.extra_headers + # Custom header should also be present + assert custom_header in instance.extra_headers + # No duplicates + assert len(instance.extra_headers) == len(set(instance.extra_headers)) 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 d275b02e34..ae57c02e7c 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -795,3 +795,797 @@ def test_get_callback_identifier_custom_logger_registry_and_fallback(): result = get_callback_identifier(my_callback_function) # Should fall back to callback_name() which returns __name__ assert result == "my_callback_function" + + +# --------------------------------------------------------------------------- +# /health response shape: model-access scoping and display-field allowlist +# --------------------------------------------------------------------------- +# These tests pin the contract that the /health response (a) only includes +# deployments the calling key is allowed to see, and (b) does not return +# provider routing fields like api_base / api_version. They guard against +# regressions that would widen the response shape. + + +@pytest.mark.asyncio +async def test_health_endpoint_filters_model_list_by_user_access(): + """ + health_endpoint() should restrict _llm_model_list to deployments whose + model_name appears in user_api_key_dict.models before running the health + check. A key scoped to ["model-a"] should only see model-a in the result, + not other deployments configured on the proxy. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-b.test", + "api_version": "2024-10-21", + }, + "model_info": {"id": "id-b"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["model-a"], + ) + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [], + "healthy_count": 0, + "unhealthy_count": 0, + } + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + from fastapi import Response + + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + + assert ( + "model_list" in captured + ), "health_endpoint did not call _perform_health_check_and_save" + returned_names = {m["model_name"] for m in captured["model_list"]} + assert returned_names == { + "model-a" + }, f"health_endpoint did not scope model_list to caller access: {returned_names}" + + +@pytest.mark.asyncio +async def test_health_endpoint_filters_background_cache_by_user_access(): + """ + When background_health_checks is enabled, health_endpoint() should also + scope the cached result to the caller's allowed models rather than + returning the cache verbatim. + """ + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-b.test", + }, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "api_base": "https://example-a.test", + }, + { + "model": "openai/gpt-4o", + "model_id": "id-b", + "api_base": "https://example-b.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 2, + "unhealthy_count": 0, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + models=["model-a"], + ) + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + from fastapi import Response + + # Pass model=None, model_id=None explicitly: direct calls to the + # handler skip FastAPI's Query() resolution, so unspecified params + # would otherwise carry the Query() sentinel (which is truthy). + result = await health_endpoint( + response=Response(), + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + # Sanity: the source cache had two entries before scoping; the scoping + # step is what reduces it to one. (This guards against the test passing + # vacuously when the cache filter drops everything because cached + # entries lack the model_id key — both entries carry model_id above.) + assert len(cached_results["healthy_endpoints"]) == 2 + assert all( + ep.get("model_id") for ep in cached_results["healthy_endpoints"] + ), "test fixture invariant: every cached entry must carry a model_id" + + # The non-admin caller must not see api_base on the returned cache entries. + returned = result.get("healthy_endpoints", []) + assert ( + len(returned) == 1 + ), f"expected exactly one cached entry after scoping, got {len(returned)}" + assert returned[0]["model_id"] == "id-a" + assert "api_base" not in returned[0] + assert result["healthy_count"] == 1 + assert result["unhealthy_count"] == 0 + + +@pytest.mark.asyncio +async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not(): + """ + A proxy admin should still see ``api_base`` and ``api_version`` in the + /health response so they can tell which Vertex region / Azure resource + + API version is healthy. A non-admin caller must not — both fields + should be stripped, and the response should carry a notice header so + non-admin clients can detect the change programmatically. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + ] + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "api_base": "https://us-central1-aiplatform.googleapis.com/v1/projects/p", + "api_version": "2024-10-21", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + admin_key = UserAPIKeyAuth( + api_key="hashed-admin-key", + models=["model-a"], + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + non_admin_key = UserAPIKeyAuth( + api_key="hashed-user-key", + models=["model-a"], + ) + + common_patches = [ + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ] + + for p in common_patches: + p.start() + try: + admin_response = Response() + non_admin_response = Response() + admin_result = await health_endpoint( + response=admin_response, + user_api_key_dict=admin_key, + model=None, + model_id=None, + ) + non_admin_result = await health_endpoint( + response=non_admin_response, + user_api_key_dict=non_admin_key, + model=None, + model_id=None, + ) + finally: + for p in common_patches: + p.stop() + + admin_eps = admin_result.get("healthy_endpoints", []) + non_admin_eps = non_admin_result.get("healthy_endpoints", []) + + assert len(admin_eps) == 1 + assert ( + admin_eps[0]["api_base"] + == "https://us-central1-aiplatform.googleapis.com/v1/projects/p" + ), "admin must see the full api_base so they can identify the region" + assert ( + admin_eps[0]["api_version"] == "2024-10-21" + ), "admin must see api_version so they can distinguish provider deployments" + + assert len(non_admin_eps) == 1 + assert "api_base" not in non_admin_eps[0] + assert "api_version" not in non_admin_eps[0] + + # Non-admin response must advertise that api_base/api_version were + # withheld so clients that previously parsed them can detect the change. + assert ( + non_admin_response.headers.get("Litellm-Health-Field-Notice") + == "api_base and api_version are admin-only on this endpoint" + ) + assert "Litellm-Health-Field-Notice" not in admin_response.headers + + # Stripping must produce a copy — the shared cache must still carry the + # routing fields so the next admin caller can read them. + cached_first = cached_results["healthy_endpoints"][0] + assert ( + cached_first["api_base"] + == "https://us-central1-aiplatform.googleapis.com/v1/projects/p" + ) + assert cached_first["api_version"] == "2024-10-21" + + +@pytest.mark.asyncio +async def test_health_endpoint_warns_when_scoped_models_lack_model_id(): + """ + When a scoped key's accessible models exist on the proxy but none of the + matching deployments expose a ``model_info.id``, the cache filter drops + everything. The response should include a structured ``warnings`` field + so the caller can distinguish "no deployments configured" from + "deployments excluded due to missing model_info.id". + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + # Intentionally no model_info.id — this is the misconfiguration + # the warnings field is meant to flag. + "model_info": {}, + }, + ] + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "api_base": "https://example-a.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-user-key", + models=["model-a"], + ) + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + result = await health_endpoint( + response=Response(), + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 0 + assert "warnings" in result, ( + "empty cache result must surface a warnings field so the caller " + "can distinguish 'no deployments' from 'deployments excluded'" + ) + assert any("model_info.id" in w for w in result["warnings"]) + + +@pytest.mark.asyncio +async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cache(): + """ + A non-admin scoped to model-a must not be able to read model-b's cached + health entry by guessing its model_id. Before the fix, + _resolve_targeted_model_ids returned {model_id} unconditionally, so the + cache filter was driven by an unvalidated ID and the global cache + leaked id-b's entry to the caller. + """ + from fastapi import Response + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", # caller has no access + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-b", + "api_base": "https://leaky-internal.test", + }, + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-scoped", + models=["model-a"], + ) + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + # llm_router None here means the model_id 404 lookup short-circuits; + # we patch _llm_model_list directly instead to drive the cache path. + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + # Calling with model="model-b" rather than model_id="id-b" because + # the model_id branch raises 404 when llm_router is None. The bug + # being verified is the same: targeted resolver must drop entries + # not in the caller's scoped model_list. With the fix, the result + # has no leaked endpoints and the targeted-503 path fires. + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-b", + model_id=None, + ) + + leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])} + leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])} + assert ( + "id-b" not in leaked_ids + ), "background cache leaked an out-of-scope deployment to a scoped caller" + assert result["healthy_count"] == 0 + assert response.status_code == 503 + + +@pytest.mark.asyncio +async def test_health_endpoint_503_for_targeted_unhealthy_model_under_background_cache_admin(): + """ + With background_health_checks enabled, an admin calling /health?model=foo + must get 503 when foo specifically has zero healthy endpoints — even if + other unrelated models in the cache are healthy. Without the cache-path + filter, the global healthy_count would mask the targeted failure. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", # the unhealthy target + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", # an unrelated healthy model + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + cached_results = { + "healthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-b"}, + ], + "unhealthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"}, + ], + "healthy_count": 1, + "unhealthy_count": 1, + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", True), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", cached_results), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + ): + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + model_id=None, + ) + + assert response.status_code == 503 + # Body must be scoped to the targeted model — not the global cache. + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 1 + returned_ids = {ep["model_id"] for ep in result.get("unhealthy_endpoints", [])} + assert returned_ids == {"id-a"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_endpoints(): + """ + /health?model=foo must return 503 when the targeted model resolves but + has zero healthy endpoints. Body shape stays the same so existing + parsers still work; only the HTTP status changes. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://example-a.test", + }, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [ + { + "model": "openai/gpt-4o", + "model_id": "id-a", + "error": "boom", + } + ], + "healthy_count": 0, + "unhealthy_count": 1, + } + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + result = await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + ) + + assert response.status_code == 503 + assert result["healthy_count"] == 0 + assert result["unhealthy_count"] == 1 + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endpoints(): + """ + /health?model=foo with a healthy endpoint must keep returning the + default 200. Verifies the 503 path doesn't fire when healthy_count > 0. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [{"model": "openai/gpt-4o", "model_id": "id-a"}], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0, + } + + response = Response() + # Default Response() exposes status_code as None; the endpoint should + # leave it alone for the healthy path. + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model="model-a", + ) + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_health_endpoint_no_model_param_returns_200_even_when_zero_healthy(): + """ + The non-targeted /health (no model / model_id query) preserves the + legacy 200 behavior even when healthy_count == 0. Existing K8s probes + and dashboards depend on this; only the targeted call became 5xx-aware. + """ + from fastapi import Response + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + async def fake_perform(**kwargs): + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [ + {"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"} + ], + "healthy_count": 0, + "unhealthy_count": 1, + } + + response = Response() + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + # Pass model=None, model_id=None explicitly: when invoked through + # FastAPI, the Query(None) defaults resolve to None, but direct + # function calls in unit tests receive Query() sentinel objects + # (which are truthy). The explicit None mirrors production routing. + await health_endpoint( + response=response, + user_api_key_dict=user_api_key_dict, + model=None, + model_id=None, + ) + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_health_readiness_returns_503_when_db_disconnected(): + """ + When a Prisma client is configured but its health_check fails, the + readiness probe should mark the worker as unhealthy via the HTTP + status — not just a body field — so K8s removes the pod from the + Service endpoints. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope")) + mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope")) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result["db"] == "disconnected" + assert result["status"] == "healthy" # body shape unchanged for back-compat + + +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_db_connected(): + """Happy path: connected DB keeps the legacy 200.""" + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock() + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result["db"] == "connected" + + +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_no_db_configured(): + """ + `prisma_client is None` means the operator chose not to use a DB. That + is a valid configuration — the worker should still report ready. We + only flip to 503 when a DB *was* configured but is unreachable. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + response = Response() + with patch("litellm.proxy.proxy_server.prisma_client", None): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result["db"] == "Not connected" + + +def test_clean_endpoint_data_strips_credentials_keeps_routing_fields(): + """ + _clean_endpoint_data() drops credentials but leaves api_base / + api_version intact — the per-caller hide/show happens in the endpoint + layer based on user role, not in the cleaning helper. This guarantees + proxy admins continue to see those fields in the /health response. + """ + from litellm.proxy.health_check import _clean_endpoint_data + + raw = { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "api_base": "https://example.test/v1", + "api_version": "2024-10-21", + "aws_access_key_id": "AKIAEXAMPLE", + } + + cleaned = _clean_endpoint_data(raw, details=True) + + assert "api_key" not in cleaned + assert "aws_access_key_id" not in cleaned + assert cleaned.get("api_base") == "https://example.test/v1" + assert cleaned.get("api_version") == "2024-10-21" diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py new file mode 100644 index 0000000000..7f1006543b --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -0,0 +1,285 @@ +""" +VERIA-39 regression tests: + +- The batch input-file token counter must measure embeddings (`input`) + and text-completion (`prompt`) payloads, not only chat (`messages`). +- The batch rate-limiter pre-call hook must reject batch files that name + models the caller is not authorized to use. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +# --------------------------------------------------------------------------- +# Token counter — covers all three batch payload shapes +# --------------------------------------------------------------------------- + + +def test_token_counter_counts_chat_messages(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_text_completion_prompt(): + """Pre-fix this returned 0 tokens (the function only inspected + `messages`), letting `prompt`-style batches slip past TPM limits.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_embedding_input_string(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + {"body": {"model": "text-embedding-3-small", "input": "hello world"}} + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_embedding_input_list(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "text-embedding-3-small", + "input": ["hello", "world"], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_text_completion_prompt_list(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": ["alpha", "beta"], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_pre_tokenized_prompt_int_list(): + """OpenAI's text-completion API accepts a single pre-tokenized prompt as + a list of ints. Each int is one token; pre-fix this shape was silently + counted as zero, leaving a TPM bypass.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [1, 2, 3, 4, 5], + } + } + ] + ) + assert usage.prompt_tokens == 5 + + +def test_token_counter_counts_pre_tokenized_prompt_list_of_int_lists(): + """Multiple pre-tokenized prompts (`list[list[int]]`) — the most + important bypass shape. A 1000-token batch must report 1000 tokens, + not zero.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [[1] * 250, [2] * 250, [3] * 500], + } + } + ] + ) + assert usage.prompt_tokens == 1000 + + +def test_token_counter_counts_pre_tokenized_input_for_embeddings(): + """Same shape applies to embeddings (`input`).""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "text-embedding-3-small", + "input": [[1, 2, 3], [4, 5, 6]], + } + } + ] + ) + assert usage.prompt_tokens == 6 + + +# --------------------------------------------------------------------------- +# Model extractor +# --------------------------------------------------------------------------- + + +def test_model_extractor_returns_distinct_models(): + from litellm.batches.batch_utils import _get_models_from_batch_input_file_content + + models = _get_models_from_batch_input_file_content( + [ + {"body": {"model": "gpt-4o", "messages": []}}, + {"body": {"model": "gpt-4o", "messages": []}}, # duplicate + {"body": {"model": "gpt-4o-mini", "messages": []}}, + {"body": {}}, # missing model + ] + ) + assert models == ["gpt-4o", "gpt-4o-mini"] + + +# --------------------------------------------------------------------------- +# Pre-call hook model validation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pre_call_rejects_unauthorized_model_in_batch_file(): + """Pre-fix the hook only validated the outer `model` parameter and + forwarded the file as-is. With this fix, a model named inside the + JSONL that the caller cannot use must trigger a 403.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + # Simulated decoded batch file: caller is restricted to gpt-3.5 + # but the JSONL points at gpt-4o. + file_dict = [ + {"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "x"}]}} + ] + + user = UserAPIKeyAuth( + api_key="sk-restricted", + user_id="alice", + models=["gpt-3.5-turbo"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + # `can_key_call_model` raises a ProxyException for non-allowed models. + async def _raise_unauthorized(**kwargs): + raise Exception( + f"Key not allowed to access model. This key only has access to models={kwargs['valid_token'].models}" + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=AsyncMock(side_effect=_raise_unauthorized), + ), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + ) + + assert exc.value.status_code == 403 + assert "gpt-4o" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_pre_call_allows_authorized_model_in_batch_file(): + """If every model in the JSONL is on the caller's allowlist, the hook + must not raise.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + file_dict = [ + { + "body": { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "x"}], + } + } + ] + + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=["gpt-3.5-turbo"], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=AsyncMock(return_value=True), + ), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + # Should not raise + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + ) + + +@pytest.mark.asyncio +async def test_pre_call_skips_check_when_no_models_present(): + """Files without any `body.model` (corrupt or empty) must not 500; + the rate limiter logs a warning elsewhere and proceeds.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice") + + # Should not raise even though `can_key_call_model` is the default + # (would fail). The early-return on empty models keeps the call out + # entirely. + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=[], + ) + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=[{"body": {}}], + ) diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py new file mode 100644 index 0000000000..0074d7062b --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -0,0 +1,208 @@ +""" +Unit tests for the personal-budget pre-call hook. + +The reservation path (added in PR #26845) atomically pre-fills the same +`spend:user:{user_id}` counter this hook reads, admitting at a strict-`<` +boundary. Re-checking with `>=` after reservation would reject requests the +reservation already admitted when the reservation fills the counter to +exactly `max_budget` (e.g. requests with no `max_tokens` cap fall back to +reserving the smallest remaining headroom). + +These tests pin the skip-when-reserved behavior and guard against drift. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + + +def _make_user_api_key_auth( + user_id: str = "user-1", + user_max_budget: float = 10.0, + user_spend: float = 0.0, + team_id=None, + budget_reservation=None, +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + user_max_budget=user_max_budget, + user_spend=user_spend, + team_id=team_id, + budget_reservation=budget_reservation, + ) + + +@pytest.mark.asyncio +async def test_under_budget_passes(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=3.0), + ): + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + + +@pytest.mark.asyncio +async def test_over_budget_rejects_without_reservation(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth(user_max_budget=10.0) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + assert "Max budget limit reached." in exc_info.value.detail + + +@pytest.mark.asyncio +async def test_skips_when_user_counter_is_reserved(): + """ + Reservation atomically pre-fills `spend:user:{user_id}` and admits the + request. The legacy `>=` check must not double-enforce on the same + counter — that's what produced the boundary regression where a fresh + user with no `max_tokens` cap got 429'd on their first request. + """ + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_id="user-1", + user_max_budget=10.0, + budget_reservation={ + "reserved_cost": 10.0, + "entries": [ + { + "counter_key": "spend:user:user-1", + "entity_type": "User", + "entity_id": "user-1", + "reserved_cost": 10.0, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + }, + ) + + # `get_current_spend` would return 10.0 here (counter pre-filled by the + # reservation). The hook must skip without reading it. + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ) as mock_get_spend: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_does_not_skip_when_reservation_covers_a_different_counter(): + """ + A reservation that only covers e.g. `spend:team:{team_id}` (not the user + counter) must not exempt the user-budget check. + """ + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_id="user-1", + user_max_budget=10.0, + budget_reservation={ + "reserved_cost": 5.0, + "entries": [ + { + "counter_key": "spend:team:team-x", + "entity_type": "Team", + "entity_id": "team-x", + "reserved_cost": 5.0, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + }, + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=10.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_team_keys_skip_personal_budget(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_max_budget=10.0, + team_id="team-1", + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=999.0), + ) as mock_get_spend: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_no_max_budget_passes(): + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", + user_id="user-1", + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=999.0), + ) as mock_get_spend: + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert result is None + mock_get_spend.assert_not_awaited() diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index d4fb5b7271..370477c360 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -1029,6 +1029,83 @@ async def test_team_member_rate_limits_v3(): ), "Team member TPM limit should be set" +@pytest.mark.asyncio +async def test_team_member_rate_limits_v3_raises_429_when_over_limit(): + """ + When should_rate_limit reports OVER_LIMIT for the team_member descriptor, the + pre-call hook raises HTTP 429 with rate_limit headers — same contract as + test_rpm_api_key_rate_limits_v3 / test_tpm_api_key_rate_limits_v3. + """ + _api_key = hash_token("sk-12345") + _team_id = "team_123" + _user_id = "user_456" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + team_id=_team_id, + user_id=_user_id, + team_member_rpm_limit=10, + team_member_tpm_limit=1000, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 10, + "limit_remaining": -1, + "rate_limit_type": "requests", + "descriptor_key": "team_member", + }, + { + "code": "OK", + "current_limit": 1000, + "limit_remaining": 500, + "rate_limit_type": "tokens", + "descriptor_key": "team_member", + }, + ], + } + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + error = None + try: + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-3.5-turbo"}, + call_type="", + ) + except HTTPException as e: + error = e + assert e.status_code == 429 + assert "rate_limit_type" in e.headers + assert e.headers.get("rate_limit_type") == "requests" + assert "retry-after" in e.headers + + assert error is not None, "An Exception must be thrown" + assert captured_descriptors is not None, "Rate limit descriptors should be captured" + team_member_descriptor = None + for descriptor in captured_descriptors: + if descriptor["key"] == "team_member": + team_member_descriptor = descriptor + break + assert team_member_descriptor is not None + assert team_member_descriptor["value"] == f"{_team_id}:{_user_id}" + + @pytest.mark.asyncio async def test_dynamic_rate_limiting_v3(): """ 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/hooks/test_rate_limiter_toctou.py b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py new file mode 100644 index 0000000000..ceea5de799 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_rate_limiter_toctou.py @@ -0,0 +1,489 @@ +""" +Tests validating TOCTOU race condition in batch + dynamic rate limiters. + +Issue: rate-limit check (read_only=True) and counter increment happen as two +separate awaits. Concurrent requests all observe the same pre-increment state, +all pass validation, then all increment — bypassing the limit. + +Vulnerable code paths: +- litellm/proxy/hooks/batch_rate_limiter.py:181-248 + (_check_and_increment_batch_counters: should_rate_limit(read_only=True) + → validate → async_increment_tokens_with_ttl_preservation) +- litellm/proxy/hooks/dynamic_rate_limiter_v3.py:463-548 + (_check_rate_limits PHASE 1 read_only check → PHASE 3 increment) + +These tests EXPECTED to fail against current (vulnerable) code and pass once +check-and-increment becomes atomic. +""" + +import asyncio +import os +import sys +from typing import List + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm import DualCache, Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler, +) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.proxy.utils import InternalUsageCache, hash_token + + +def _make_phase1_barrier(num_concurrent: int, timeout: float = 0.1): + """ + Sync primitive that, pre-fix, forces all N concurrent coroutines to finish + their read-only Phase 1 check before any proceeds to Phase 3 increment — + mimicking asyncio I/O scheduling under load on the vulnerable code. + + Wraps `should_rate_limit` so on `read_only=True` calls it waits until N + callers arrive (TOCTOU window opened) OR `timeout` elapses (post-fix path: + the limiter's serialization lock prevents N from ever reaching the + barrier; the timeout lets the holder proceed so the lock can do its job). + + Pre-fix: barrier fills before timeout → all see same state → bypass observed. + Post-fix: only lock-holder reaches barrier → times out → serial execution + enforces limit. + """ + arrived = 0 + all_arrived = asyncio.Event() + + def wrap(original): + async def patched(*args, **kwargs): + result = await original(*args, **kwargs) + if kwargs.get("read_only"): + nonlocal arrived + arrived += 1 + if arrived >= num_concurrent: + all_arrived.set() + try: + await asyncio.wait_for(all_arrived.wait(), timeout=timeout) + except asyncio.TimeoutError: + pass + return result + + return patched + + return wrap + + +@pytest.mark.asyncio +async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou(): + """ + 5 concurrent batch submissions of 40 tokens each against TPM=100 limit. + + Sequential semantics: only 2 batches fit (2 * 40 = 80 ≤ 100, 3rd at 120 fails). + With TOCTOU: all 5 succeed → 200 tokens consumed, 100% over limit. + + Demonstrates batch_rate_limiter.py:183-248 multi-phase flaw. + """ + NUM_CONCURRENT = 5 + BATCH_TOKENS = 40 + TPM_LIMIT = 100 + + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("toctou-batch-key"), + tpm_limit=TPM_LIMIT, + rpm_limit=1000, + ) + batch_usage = BatchFileUsage(total_tokens=BATCH_TOKENS, request_count=1) + + barrier = _make_phase1_barrier(NUM_CONCURRENT) + rate_limiter.should_rate_limit = barrier(rate_limiter.should_rate_limit) + + results = await asyncio.gather( + *[ + batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=batch_usage, + ) + for _ in range(NUM_CONCURRENT) + ], + return_exceptions=True, + ) + + successes = [r for r in results if not isinstance(r, Exception)] + rejections = [r for r in results if isinstance(r, Exception)] + total_consumed = len(successes) * BATCH_TOKENS + max_allowed_successes = TPM_LIMIT // BATCH_TOKENS # 2 + + assert len(successes) <= max_allowed_successes, ( + f"TOCTOU bypass: {len(successes)}/{NUM_CONCURRENT} concurrent batches " + f"passed despite TPM={TPM_LIMIT}. Consumed {total_consumed} tokens " + f"({total_consumed - TPM_LIMIT} over limit). " + f"Atomic check-and-increment would allow ≤{max_allowed_successes}. " + f"Rejections: {len(rejections)}" + ) + + +@pytest.mark.asyncio +async def test_batch_limiter_uses_atomic_check_and_increment(): + """ + Regression test: batch limiter routes through + `atomic_check_and_increment_by_n` rather than the legacy two-phase + pattern (read_only=True check + separate async_increment_tokens_with_ttl_preservation). + + Ensures future refactors don't reintroduce the TOCTOU window. + """ + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + call_log: List[str] = [] + original_atomic = rate_limiter.atomic_check_and_increment_by_n + original_should = rate_limiter.should_rate_limit + + async def logging_atomic(*args, **kwargs): + call_log.append("atomic_check_and_increment_by_n") + return await original_atomic(*args, **kwargs) + + async def logging_should(*args, **kwargs): + call_log.append(f"should_rate_limit(read_only={kwargs.get('read_only')})") + return await original_should(*args, **kwargs) + + rate_limiter.atomic_check_and_increment_by_n = logging_atomic + rate_limiter.should_rate_limit = logging_should + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("atomic-test-key"), + tpm_limit=10000, + rpm_limit=1000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=50, request_count=1), + ) + + assert "atomic_check_and_increment_by_n" in call_log, ( + f"Batch limiter must route through atomic_check_and_increment_by_n. " + f"Calls observed: {call_log}" + ) + legacy_calls = [c for c in call_log if c.startswith("should_rate_limit(")] + assert not legacy_calls, ( + f"Batch limiter must not call should_rate_limit directly (legacy " + f"two-phase pattern). Observed: {legacy_calls}" + ) + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_concurrent_bypasses_model_capacity(): + """ + DynamicRateLimitHandler PHASE 1 (read_only check) → PHASE 3 (increment) + is non-atomic: dynamic_rate_limiter_v3.py:463-548. + + With TPM=100 model capacity and 5 concurrent priority="high" requests + each consuming the full model_saturation_check counter, all observe the + same Phase 1 state (counter=0), all pass, all proceed to Phase 3. + + Sequential atomic semantics would block requests once the model counter + reaches its limit. TOCTOU lets all pass Phase 1 simultaneously. + """ + NUM_CONCURRENT = 10 + MODEL_RPM = 2 + # Sequential bound: dynamic limiter rejects when `counter > current_limit` + # (strict `>`), so a request whose Phase 1 sees counter=RPM still passes + # (RPM is not strictly greater). Atomic execution therefore admits up to + # RPM + 1 successes before the next sees counter > RPM. + MAX_SEQUENTIAL_SUCCESSES = MODEL_RPM + 1 + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "toctou-dyn-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": MODEL_RPM, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + barrier = _make_phase1_barrier(NUM_CONCURRENT) + handler.v3_limiter.should_rate_limit = barrier(handler.v3_limiter.should_rate_limit) + + from litellm.types.router import ModelGroupInfo + + model_group_info = ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=MODEL_RPM, + tpm=None, + ) + + async def one_request(idx: int): + user = UserAPIKeyAuth(api_key=hash_token(f"dyn-key-{idx}")) + user.metadata = {"priority": "high"} + try: + await handler._check_rate_limits( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + return "OK" + except Exception as e: + return e + + results = await asyncio.gather( + *[one_request(i) for i in range(NUM_CONCURRENT)], + return_exceptions=True, + ) + successes = [r for r in results if r == "OK"] + + assert len(successes) <= MAX_SEQUENTIAL_SUCCESSES, ( + f"TOCTOU bypass in DynamicRateLimitHandler: {len(successes)}/{NUM_CONCURRENT} " + f"concurrent requests passed Phase 1 + Phase 3 despite model RPM={MODEL_RPM}. " + f"Atomic check-and-increment would block once counter > RPM " + f"(at most {MAX_SEQUENTIAL_SUCCESSES} sequential successes)." + ) + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(): + """ + Regression test: dynamic limiter's enforced descriptors flow through + `atomic_check_and_increment_by_n`, not the legacy + read_only=True check followed by a separate read_only=False increment. + + When priority is enforced (saturation >= threshold), priority_model is + bundled into the atomic call alongside model_saturation_check. When not + enforced, priority counter is incremented for tracking only. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "atomic-dyn-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + atomic_descriptors_observed: List[List[str]] = [] + original_atomic = handler.v3_limiter.atomic_check_and_increment_by_n + + async def logging_atomic(*args, **kwargs): + ds = kwargs.get("descriptors") or (args[0] if args else []) + atomic_descriptors_observed.append([d["key"] for d in ds]) + return await original_atomic(*args, **kwargs) + + handler.v3_limiter.atomic_check_and_increment_by_n = logging_atomic + + from litellm.types.router import ModelGroupInfo + + user = UserAPIKeyAuth(api_key=hash_token("dyn-atomic-key")) + user.metadata = {"priority": "high"} + + await handler._check_rate_limits( + model=model, + model_group_info=ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=None, + tpm=1000, + ), + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + + assert atomic_descriptors_observed, ( + "Dynamic limiter must route enforced descriptors through " + "atomic_check_and_increment_by_n (no legacy read_only=True / " + "separate-increment pattern)." + ) + assert "model_saturation_check" in atomic_descriptors_observed[0], ( + f"Expected model_saturation_check in atomic descriptor set. " + f"Got: {atomic_descriptors_observed}" + ) + + +@pytest.mark.asyncio +async def test_batch_zero_token_consumes_rpm_only(): + """ + Zero-token batch (e.g. metadata-only call) should still increment RPM + counter but NOT TPM counter. + + Edge case from review: `if inc_amount <= 0: continue` in + `atomic_check_and_increment_by_n` skips descriptor counters whose + increment is zero. Verifies asymmetric quota consumption is intentional + and observable: an RPM-bounded but TPM-free request path stays bounded + by RPM alone. + """ + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("zero-token-key"), + tpm_limit=100, + rpm_limit=3, + ) + zero_batch = BatchFileUsage(total_tokens=0, request_count=1) + + # 3 zero-token batches must succeed (RPM=3 allows). 4th must hit RPM cap, + # NOT TPM (because token counter never increments past 0). + for i in range(3): + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=zero_batch, + ) + + # Inspect counters: RPM key incremented to 3, TPM key absent (or 0). + rpm_key = rate_limiter.create_rate_limit_keys( + "api_key", user_api_key_dict.api_key or "", "requests" + ) + tpm_key = rate_limiter.create_rate_limit_keys( + "api_key", user_api_key_dict.api_key or "", "tokens" + ) + rpm_val = await internal_usage_cache.async_get_cache( + key=rpm_key, litellm_parent_otel_span=None, local_only=True + ) + tpm_val = await internal_usage_cache.async_get_cache( + key=tpm_key, litellm_parent_otel_span=None, local_only=True + ) + assert int(rpm_val or 0) == 3, f"RPM counter must reach 3, got {rpm_val}" + assert tpm_val in ( + None, + 0, + "0", + ), f"TPM counter must remain unset/0 for zero-token batches, got {tpm_val}" + + # 4th attempt: RPM exhausted -> 429. + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=zero_batch, + ) + assert exc.value.status_code == 429 + + +@pytest.mark.asyncio +async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(): + """ + Fail-closed guard: when atomic_check_and_increment_by_n returns + overall_code=OVER_LIMIT but with a descriptor_key the dispatcher does + not recognize, the dynamic limiter must raise 429 rather than silently + fall through. + + Reproduces by patching atomic_check_and_increment_by_n to return an + OVER_LIMIT response carrying an unknown descriptor_key. + """ + from fastapi import HTTPException + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fail-closed-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 1000, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + async def fake_atomic(*args, **kwargs): + return { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "tokens", + "descriptor_key": "future_unrecognized_descriptor", + } + ], + } + + handler.v3_limiter.atomic_check_and_increment_by_n = fake_atomic + + from litellm.types.router import ModelGroupInfo + + user = UserAPIKeyAuth(api_key=hash_token("fail-closed-key")) + user.metadata = {"priority": "high"} + + with pytest.raises(HTTPException) as exc: + await handler._check_rate_limits( + model=model, + model_group_info=ModelGroupInfo( + model_group=model, + providers=["openai"], + rpm=None, + tpm=1000, + ), + user_api_key_dict=user, + priority="high", + saturation=0.0, + data={}, + ) + assert ( + exc.value.status_code == 429 + ), f"Expected 429 fail-closed on unknown descriptor; got {exc.value.status_code}" diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index a1e7fe59ca..f898763d2c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -481,3 +481,39 @@ class TestSetObjectMetadataField: ): _set_object_metadata_field(team, "model_rpm_limit", {"x": 1}) assert team.metadata == {"model_rpm_limit": {"x": 1}} + + +class TestRequireCallerUserIdForNonAdmin: + """ + Security regression: service-account keys (user_id=None) must not bypass + the non-admin scoping branch on analytics endpoints. + """ + + def test_returns_user_id_when_present(self): + from litellm.proxy.management_endpoints.common_utils import ( + require_caller_user_id_for_non_admin, + ) + + key_dict = UserAPIKeyAuth( + user_id="user-abc", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + assert require_caller_user_id_for_non_admin(key_dict) == "user-abc" + + def test_raises_403_when_user_id_is_none(self): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + require_caller_user_id_for_non_admin, + ) + + # Simulates a service-account key (user_id forced to None at key creation) + service_account_key = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + with pytest.raises(HTTPException) as exc_info: + require_caller_user_id_for_non_admin(service_account_key) + + assert exc_info.value.status_code == 403 + assert "Service-account keys" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index a1ba7ecd67..f4dc85dad9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -1732,6 +1732,107 @@ async def test_get_user_daily_activity_non_admin_cannot_view_other_users(monkeyp assert call_kwargs.kwargs["entity_id"] == "regular-user-123" +@pytest.mark.asyncio +async def test_get_user_daily_activity_rejects_service_account_caller(monkeypatch): + """ + Security regression: a non-admin caller with user_id=None (a service-account + key, where user_id is forced to None at key creation) must not be able to + bypass the entity filter and read every tenant's daily spend. + + Before the fix, the endpoint silently defaulted user_id to + user_api_key_dict.user_id, which is itself None for service-account keys. + None != None is False, the same-user check passed, and entity_id=None + flowed into get_daily_activity, where the SQL builder treats None as + "no filter". + """ + from unittest.mock import AsyncMock, MagicMock + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity, + ) + + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Tripwire: ensure get_daily_activity is never reached + mock_get_daily = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity", + mock_get_daily, + ) + + service_account_key = UserAPIKeyAuth( + user_id=None, # service-account keys have user_id forced to None + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id=None, + page=1, + page_size=50, + timezone=None, + user_api_key_dict=service_account_key, + ) + + assert exc_info.value.status_code == 403 + assert "Service-account keys" in str(exc_info.value.detail) + mock_get_daily.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_user_daily_activity_aggregated_rejects_service_account_caller( + monkeypatch, +): + """ + Same security regression as + test_get_user_daily_activity_rejects_service_account_caller, on the + aggregated route. Same shape, raw-SQL builder, same fix. + """ + from unittest.mock import AsyncMock, MagicMock + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_daily_activity_aggregated, + ) + + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_get_daily_agg = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + service_account_key = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with pytest.raises(HTTPException) as exc_info: + await get_user_daily_activity_aggregated( + start_date="2025-01-01", + end_date="2025-01-31", + model=None, + api_key=None, + user_id=None, + timezone=None, + user_api_key_dict=service_account_key, + ) + + assert exc_info.value.status_code == 403 + assert "Service-account keys" in str(exc_info.value.detail) + mock_get_daily_agg.assert_not_called() + + @pytest.mark.asyncio async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch): """ @@ -1949,11 +2050,13 @@ async def test_delete_user_rejects_org_admin_deleting_outside_scope(mocker): assert exc.value.status_code == 403 # Critical: no delete_many calls should have executed. - assert not hasattr( - mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls" - ) or len( - mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls - ) == 0 + assert ( + not hasattr( + mock_prisma_client.db.litellm_verificationtoken.delete_many, "mock_calls" + ) + or len(mock_prisma_client.db.litellm_verificationtoken.delete_many.mock_calls) + == 0 + ) @pytest.mark.asyncio @@ -2631,3 +2734,107 @@ class TestGetUserIdFromRequestValidation: request = self._make_request(f"user_id={exact_id}") result = get_user_id_from_request(request) assert result == exact_id + + +# --------------------------------------------------------------------------- +# VERIA-60: /user/info post-decode re-authorization +# --------------------------------------------------------------------------- + + +def test_enforce_user_info_access_admin_bypass(): + """Proxy admins must always be allowed past the re-check.""" + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _enforce_user_info_access, + ) + + admin = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN.value + ) + # Should not raise even when querying a different user + _enforce_user_info_access(user_id="someone_else", user_api_key_dict=admin) + + +def test_enforce_user_info_access_view_only_admin_blocked_from_other_users(): + """PROXY_ADMIN_VIEW_ONLY is not a true admin for /user/info — the upstream + route check applies the same `user_id == valid_token.user_id` rule, so the + re-check here must mirror that and deny cross-user lookups.""" + import pytest + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _enforce_user_info_access, + ) + + viewer = UserAPIKeyAuth( + user_id="viewer", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + with pytest.raises(HTTPException) as exc_info: + _enforce_user_info_access(user_id="someone_else", user_api_key_dict=viewer) + assert exc_info.value.status_code == 403 + + +def test_enforce_user_info_access_view_only_admin_can_read_own(): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _enforce_user_info_access, + ) + + viewer = UserAPIKeyAuth( + user_id="viewer", + user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ) + _enforce_user_info_access(user_id="viewer", user_api_key_dict=viewer) + + +def test_enforce_user_info_access_owner_allowed(): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _enforce_user_info_access, + ) + + user = UserAPIKeyAuth( + user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value + ) + _enforce_user_info_access(user_id="alice", user_api_key_dict=user) + + +def test_enforce_user_info_access_no_user_id_allowed(): + """No user_id in query → handler resolves to caller's own id later, so + this branch must not raise.""" + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _enforce_user_info_access, + ) + + user = UserAPIKeyAuth( + user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value + ) + _enforce_user_info_access(user_id=None, user_api_key_dict=user) + + +def test_enforce_user_info_access_blocks_cross_user_lookup(): + """A non-admin caller may not query another user's row, even if URL + re-parsing produced a user_id that differs from the one the route check + saw (the VERIA-60 bypass).""" + import pytest + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _enforce_user_info_access, + ) + + attacker = UserAPIKeyAuth( + user_id="attacker space", # original (URL-decoded) id seen by route check + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + + with pytest.raises(HTTPException) as exc_info: + # Re-parsed id (with literal '+') belongs to the victim + _enforce_user_info_access(user_id="victim+target", user_api_key_dict=attacker) + + assert exc_info.value.status_code == 403 + assert "key not allowed to access this user's info" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index f256ffd866..b292e8d0ca 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5021,6 +5021,299 @@ async def test_list_keys_non_admin_user_id_auto_set(): ) +def _make_member_team_table( + team_id: str, + member_user_id: str, + member_role: str = "user", + team_member_permissions=None, +): + """Build a LiteLLM_TeamTable with a single member, suitable for list_keys tests.""" + from litellm.proxy._types import LiteLLM_TeamTable, Member + + return LiteLLM_TeamTable( + team_id=team_id, + members_with_roles=[Member(user_id=member_user_id, role=member_role)], + team_member_permissions=team_member_permissions, + ) + + +async def _invoke_list_keys_and_capture_helper_kwargs( + user_api_key_dict, + team_objects, + *, + include_team_keys: bool = True, +): + """ + Invoke list_keys with mocked dependencies and return the kwargs that + list_keys passes to _list_key_helper (so tests can assert on + admin_team_ids / member_team_ids classification). + """ + from unittest.mock import Mock, patch + + from litellm.proxy._types import LiteLLM_UserTable + + mock_prisma_client = AsyncMock() + mock_user_info = LiteLLM_UserTable( + user_id=user_api_key_dict.user_id, + user_email="member@example.com", + teams=[t.team_id for t in team_objects], + organization_memberships=[], + ) + mock_list_key_helper = AsyncMock( + return_value={ + "keys": [], + "total_count": 0, + "current_page": 1, + "total_pages": 0, + } + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check", + return_value=mock_user_info, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._fetch_user_team_objects", + AsyncMock(return_value=team_objects), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper", + mock_list_key_helper, + ), + ): + await list_keys( + request=Mock(), + user_api_key_dict=user_api_key_dict, + include_team_keys=include_team_keys, + status=None, + ) + mock_list_key_helper.assert_called_once() + return mock_list_key_helper.call_args.kwargs + + +@pytest.mark.asyncio +async def test_list_keys_team_member_with_key_list_permission_sees_all_team_keys(): + """ + Bug fix: when a team has /key/list in team_member_permissions, regular + team members must get full key visibility for that team — same as a team + admin would. This means other members' keys AND service account keys + (user_id=NULL) must be returned, not only the caller's own keys. + + This test pins down list_keys' classification: the team must be passed + to _list_key_helper as a full-visibility team (admin_team_ids), not as + a service-account-only team (member_team_ids). + """ + member_user_id = "member-user-1" + team_id = "team-with-permission" + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id=member_user_id, + ) + team_objects = [ + _make_member_team_table( + team_id=team_id, + member_user_id=member_user_id, + member_role="user", + team_member_permissions=["/key/list"], + ) + ] + + helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs( + user_api_key_dict=user_api_key_dict, + team_objects=team_objects, + ) + + admin_team_ids = helper_kwargs.get("admin_team_ids") or [] + member_team_ids = helper_kwargs.get("member_team_ids") or [] + assert team_id in admin_team_ids, ( + "team granting /key/list permission must be classified as full-visibility " + "(admin_team_ids), but got " + f"admin_team_ids={admin_team_ids}, member_team_ids={member_team_ids}" + ) + + +@pytest.mark.asyncio +async def test_list_keys_team_member_without_key_list_permission_only_service_accounts(): + """ + Without the /key/list permission, the existing scoping must hold: the + team is classified as member-only, so only service account keys + (user_id=NULL) for that team are visible to the caller. + """ + member_user_id = "member-user-2" + team_id = "team-no-permission" + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id=member_user_id, + ) + team_objects = [ + _make_member_team_table( + team_id=team_id, + member_user_id=member_user_id, + member_role="user", + team_member_permissions=None, + ) + ] + + helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs( + user_api_key_dict=user_api_key_dict, + team_objects=team_objects, + ) + + admin_team_ids = helper_kwargs.get("admin_team_ids") or [] + member_team_ids = helper_kwargs.get("member_team_ids") or [] + assert team_id not in admin_team_ids + assert team_id in member_team_ids + + +@pytest.mark.asyncio +async def test_list_keys_team_member_with_permission_in_one_team_only(): + """ + Granular: a user is a member of two teams. Only one team grants + /key/list — the other does not. The classification must respect the + per-team permission, not leak full visibility across both teams. + """ + member_user_id = "member-user-3" + team_with_permission = "team-A" + team_without_permission = "team-B" + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id=member_user_id, + ) + team_objects = [ + _make_member_team_table( + team_id=team_with_permission, + member_user_id=member_user_id, + member_role="user", + team_member_permissions=["/key/list"], + ), + _make_member_team_table( + team_id=team_without_permission, + member_user_id=member_user_id, + member_role="user", + team_member_permissions=[], + ), + ] + + helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs( + user_api_key_dict=user_api_key_dict, + team_objects=team_objects, + ) + + admin_team_ids = helper_kwargs.get("admin_team_ids") or [] + member_team_ids = helper_kwargs.get("member_team_ids") or [] + assert team_with_permission in admin_team_ids + assert team_without_permission not in admin_team_ids + assert team_without_permission in member_team_ids + + +@pytest.mark.asyncio +async def test_list_keys_team_admin_unaffected_by_member_permission_logic(): + """ + Sanity: a team admin's classification is unchanged by the new + permission-aware path. They still appear in admin_team_ids (full + visibility) regardless of the team_member_permissions value. + """ + admin_user_id = "team-admin-user" + team_id = "team-with-admin" + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id=admin_user_id, + ) + team_objects = [ + _make_member_team_table( + team_id=team_id, + member_user_id=admin_user_id, + member_role="admin", + team_member_permissions=None, + ) + ] + + helper_kwargs = await _invoke_list_keys_and_capture_helper_kwargs( + user_api_key_dict=user_api_key_dict, + team_objects=team_objects, + ) + + admin_team_ids = helper_kwargs.get("admin_team_ids") or [] + assert team_id in admin_team_ids + + +def test_build_key_filter_conditions_full_visibility_team_includes_service_accounts(): + """ + Direct check on the SQL filter: when a team is in the full-visibility + set (admin_team_ids), the filter clause for that team is + {"team_id": {"in": [...]}} with NO user_id constraint — so service + account keys (user_id=NULL) AND other members' keys are returned. + + This is the SQL-level proof that a member with /key/list permission + will see service account keys for the team. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + full_visibility_team = "team-full-vis" + where = _build_key_filter_conditions( + user_id="member-user-x", + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=[full_visibility_team], + member_team_ids=[full_visibility_team], + include_created_by_keys=False, + ) + + serialized = json.dumps(where) + # Full-visibility team filter: no user_id restriction + assert ( + json.dumps({"team_id": {"in": [full_visibility_team]}}) in serialized + ), f"expected unrestricted team_id IN clause, got: {serialized}" + # No service-account-only AND clause for this team (it would be redundant + # and would erroneously narrow the visibility back to user_id=NULL). + sa_only_clause = json.dumps( + {"AND": [{"team_id": {"in": [full_visibility_team]}}, {"user_id": None}]} + ) + assert ( + sa_only_clause not in serialized + ), f"team in admin_team_ids must not also be filtered to user_id=NULL: {serialized}" + + +def test_build_key_filter_conditions_member_only_team_restricts_to_service_accounts(): + """ + Existing-behavior pin: when a team is ONLY in member_team_ids (not + admin_team_ids), the filter for that team must be + {"AND": [{"team_id": {"in": [...]}}, {"user_id": None}]} — i.e. only + service accounts visible. This is the "no permission" baseline that + must keep working. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _build_key_filter_conditions, + ) + + member_only_team = "team-member-only" + where = _build_key_filter_conditions( + user_id="member-user-y", + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=[], + member_team_ids=[member_only_team], + include_created_by_keys=False, + ) + + serialized = json.dumps(where) + expected = json.dumps( + {"AND": [{"team_id": {"in": [member_only_team]}}, {"user_id": None}]} + ) + assert ( + expected in serialized + ), f"member-only team must be restricted to user_id=NULL keys, got: {serialized}" + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ 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 a0ae95df58..69798744f7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1841,6 +1841,7 @@ class TestCustomUISSO: "x-forwarded-for": "192.168.1.1", } mock_request.base_url = "https://test.litellm.ai/" + mock_request.client.host = "10.0.0.10" # Mock the custom handler mock_custom_handler = MagicMock(spec=CustomSSOLoginHandler) @@ -1866,36 +1867,73 @@ class TestCustomUISSO: "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", mock_custom_handler, ): - with patch.object( - SSOAuthenticationHandler, - "get_redirect_response_from_openid", - return_value=mock_redirect_response, - ) as mock_get_redirect: - # Act - result = ( + with patch( + "litellm.proxy.proxy_server.general_settings", + {"trusted_proxy_ranges": ["10.0.0.0/24"]}, + ): + with patch.object( + SSOAuthenticationHandler, + "get_redirect_response_from_openid", + return_value=mock_redirect_response, + ) as mock_get_redirect: + # Act + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + request=mock_request + ) + + # Assert + # Verify the custom handler was called with the request + mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with( + request=mock_request + ) + + # Verify the redirect response was generated with correct OpenID + mock_get_redirect.assert_called_once_with( + result=expected_openid, + request=mock_request, + received_response=None, + generic_client_id=None, + ui_access_mode=None, + ) + + # Verify the result is the redirect response + assert result == mock_redirect_response + assert result.status_code == 303 + + @pytest.mark.asyncio + async def test_handle_custom_ui_sso_sign_in_rejects_untrusted_proxy(self): + """Custom UI SSO rejects spoofed identity headers from direct clients.""" + from enterprise.litellm_enterprise.proxy.auth.custom_sso_handler import ( + EnterpriseCustomSSOHandler, + ) + from litellm.integrations.custom_sso_handler import CustomSSOLoginHandler + + mock_request = MagicMock(spec=Request) + mock_request.headers = { + "x-litellm-user-id": "admin", + "x-litellm-user-email": "admin@example.com", + } + mock_request.base_url = "https://test.litellm.ai/" + mock_request.client.host = "203.0.113.10" + + mock_custom_handler = MagicMock(spec=CustomSSOLoginHandler) + mock_custom_handler.handle_custom_ui_sso_sign_in = AsyncMock() + + with patch("litellm.proxy.proxy_server.premium_user", True): + with patch( + "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", + mock_custom_handler, + ): + with patch( + "litellm.proxy.proxy_server.general_settings", + {"trusted_proxy_ranges": ["10.0.0.0/24"]}, + ): + with pytest.raises(ValueError, match="not trusted"): await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( request=mock_request ) - ) - # Assert - # Verify the custom handler was called with the request - mock_custom_handler.handle_custom_ui_sso_sign_in.assert_called_once_with( - request=mock_request - ) - - # Verify the redirect response was generated with correct OpenID - mock_get_redirect.assert_called_once_with( - result=expected_openid, - request=mock_request, - received_response=None, - generic_client_id=None, - ui_access_mode=None, - ) - - # Verify the result is the redirect response - assert result == mock_redirect_response - assert result.status_code == 303 + mock_custom_handler.handle_custom_ui_sso_sign_in.assert_not_called() @pytest.mark.asyncio async def test_custom_ui_sso_handler_execution_with_real_class(self): @@ -1946,6 +1984,7 @@ class TestCustomUISSO: "x-forwarded-for": "10.0.0.1", } mock_request.base_url = "https://custom.litellm.ai/" + mock_request.client.host = "10.0.0.20" # Mock the redirect response method mock_redirect_response = MagicMock() @@ -1956,34 +1995,36 @@ class TestCustomUISSO: "litellm.proxy.proxy_server.user_custom_ui_sso_sign_in_handler", test_handler_instance, ): - with patch.object( - SSOAuthenticationHandler, - "get_redirect_response_from_openid", - return_value=mock_redirect_response, - ) as mock_get_redirect: - # Act - result = ( - await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( + with patch( + "litellm.proxy.proxy_server.general_settings", + {"trusted_proxy_ranges": ["10.0.0.0/24"]}, + ): + with patch.object( + SSOAuthenticationHandler, + "get_redirect_response_from_openid", + return_value=mock_redirect_response, + ) as mock_get_redirect: + # Act + result = await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in( request=mock_request ) - ) - # Assert that our custom handler was executed - assert test_handler_instance.method_called is True - assert test_handler_instance.received_request == mock_request + # Assert that our custom handler was executed + assert test_handler_instance.method_called is True + assert test_handler_instance.received_request == mock_request - # Verify the redirect response was called with the OpenID from our custom handler - mock_get_redirect.assert_called_once() - call_args = mock_get_redirect.call_args.kwargs + # Verify the redirect response was called with the OpenID from our custom handler + mock_get_redirect.assert_called_once() + call_args = mock_get_redirect.call_args.kwargs - # Verify the OpenID object has the expected values from our custom handler - openid_result = call_args["result"] - assert openid_result.id == "custom_test_user_456" - assert openid_result.email == "custom@example.com" - assert openid_result.first_name == "Custom" - assert openid_result.last_name == "Handler" - assert openid_result.display_name == "Custom Handler Test" - assert openid_result.provider == "custom" + # Verify the OpenID object has the expected values from our custom handler + openid_result = call_args["result"] + assert openid_result.id == "custom_test_user_456" + assert openid_result.email == "custom@example.com" + assert openid_result.first_name == "Custom" + assert openid_result.last_name == "Handler" + assert openid_result.display_name == "Custom Handler Test" + assert openid_result.provider == "custom" # Verify the request and other parameters were passed correctly assert call_args["request"] == mock_request @@ -5767,3 +5808,324 @@ class TestSyncUserRoleFromJwtRoleMap: ) prisma.db.litellm_usertable.update.assert_not_called() + + +# ── VERIA-34 regression: PKCE state-to-session-cookie binding ─────────────── + + +class TestPKCEStateCookieBinding: + """The Generic SSO PKCE flow used the URL ``state`` parameter as a + cache-key for the PKCE ``code_verifier`` without binding the state to + the caller's browser. An attacker who pre-mints a state + cached + verifier could hand the link to a victim and capture the resulting + access token. Fix: set ``litellm_oauth_state`` HttpOnly cookie on + the redirect; verify the URL state matches the cookie before doing + the PKCE token exchange.""" + + @pytest.mark.asyncio + async def test_redirect_response_sets_oauth_state_cookie_when_pkce_enabled(self): + """``get_generic_sso_redirect_response`` must set + ``litellm_oauth_state`` on the redirect response when PKCE is on so + the callback can verify it later. The cookie must carry HttpOnly, + SameSite=Lax, and (because no http request was supplied to the + helper) the production-safe ``Secure`` default.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="https://idp.example.com/authorize?state=test-state-xyz" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "test-state-xyz", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="https://idp.example.com/authorize", + ) + + assert response is not None + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert ( + cookie_str is not None + ), f"litellm_oauth_state cookie not set; got: {cookie_headers}" + assert "test-state-xyz" in cookie_str + assert "HttpOnly" in cookie_str + assert "SameSite=lax" in cookie_str + # No incoming Request supplied → ``Secure`` defaults to True so a + # network observer on plain HTTP cannot read the state value. + assert "Secure" in cookie_str + + @pytest.mark.asyncio + async def test_redirect_response_omits_oauth_state_cookie_when_pkce_disabled( + self, + ): + """Non-PKCE flows delegate to fastapi-sso's own session-cookie + binding; we do not set our cookie there because it would never be + validated (and could collide with a concurrent PKCE session in + the same browser).""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="https://idp.example.com/authorize?state=test-state-xyz" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "test-state-xyz", + "GENERIC_CLIENT_USE_PKCE": "false", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="https://idp.example.com/authorize", + ) + + assert response is not None + cookie_headers = response.headers.getlist("set-cookie") + assert not any( + "litellm_oauth_state=" in c for c in cookie_headers + ), f"litellm_oauth_state cookie set on non-PKCE flow; got: {cookie_headers}" + + @pytest.mark.asyncio + async def test_redirect_response_drops_secure_flag_for_http_dev(self): + """When the incoming request is plain HTTP (local dev), ``Secure`` + must be dropped so the browser will actually attach the cookie on + the callback hop.""" + from fastapi.responses import RedirectResponse + + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + ) + + mock_redirect = RedirectResponse( + url="http://idp.local/authorize?state=local-dev-state" + ) + mock_generic_sso = MagicMock() + mock_generic_sso.__enter__ = MagicMock(return_value=mock_generic_sso) + mock_generic_sso.__exit__ = MagicMock(return_value=None) + mock_generic_sso.get_login_redirect = AsyncMock(return_value=mock_redirect) + + http_request = MagicMock(spec=Request) + http_request.url.scheme = "http" + + with patch.dict( + os.environ, + { + "GENERIC_CLIENT_STATE": "local-dev-state", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ): + response = await SSOAuthenticationHandler.get_generic_sso_redirect_response( + generic_sso=mock_generic_sso, + state=None, + generic_authorization_endpoint="http://idp.local/authorize", + request=http_request, + ) + + cookie_headers = response.headers.getlist("set-cookie") + cookie_str = next( + (c for c in cookie_headers if "litellm_oauth_state=" in c), None + ) + assert cookie_str is not None + assert "Secure" not in cookie_str + + @pytest.mark.asyncio + async def test_pkce_callback_rejects_missing_cookie(self): + """When PKCE is enabled and a code_verifier is in the cache, the + callback must reject a request that has no ``litellm_oauth_state`` + cookie (browser-to-server binding missing).""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + get_generic_sso_response, + ) + + mock_request = MagicMock(spec=Request) + mock_request.query_params = { + "state": "attacker-minted-state", + "code": "auth-code", + } + # No oauth_state cookie set → request.cookies.get returns None. + mock_request.cookies = {} + + with ( + patch.object( + SSOAuthenticationHandler, + "prepare_token_exchange_parameters", + AsyncMock( + return_value={ + "code_verifier": "attacker-cached-verifier", + "_pkce_cache_key": "pkce_verifier:attacker-minted-state", + } + ), + ), + patch("fastapi_sso.sso.base.DiscoveryDocument"), + patch("fastapi_sso.sso.generic.create_provider", return_value=MagicMock()), + patch.dict( + os.environ, + { + "GENERIC_CLIENT_SECRET": "x", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://idp.example.com/auth", + "GENERIC_TOKEN_ENDPOINT": "https://idp.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://idp.example.com/userinfo", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ), + pytest.raises(ProxyException) as exc_info, + ): + await get_generic_sso_response( + request=mock_request, + jwt_handler=MagicMock(spec=JWTHandler), + generic_client_id="cid", + redirect_url="https://proxy.example.com/sso/callback", + sso_jwt_handler=None, + ) + + assert "state" in str(exc_info.value.message).lower() + + @pytest.mark.asyncio + async def test_pkce_callback_rejects_state_cookie_mismatch(self): + """The Login-CSRF shape: attacker mints state ``A``, victim's browser + carries cookie state ``B``. The callback must reject.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + get_generic_sso_response, + ) + + mock_request = MagicMock(spec=Request) + mock_request.query_params = { + "state": "attacker-minted-state", + "code": "auth-code", + } + mock_request.cookies = {"litellm_oauth_state": "victim-browser-state"} + + with ( + patch.object( + SSOAuthenticationHandler, + "prepare_token_exchange_parameters", + AsyncMock( + return_value={ + "code_verifier": "verifier", + "_pkce_cache_key": "pkce_verifier:attacker-minted-state", + } + ), + ), + patch("fastapi_sso.sso.base.DiscoveryDocument"), + patch("fastapi_sso.sso.generic.create_provider", return_value=MagicMock()), + patch.dict( + os.environ, + { + "GENERIC_CLIENT_SECRET": "x", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://idp.example.com/auth", + "GENERIC_TOKEN_ENDPOINT": "https://idp.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://idp.example.com/userinfo", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ), + pytest.raises(ProxyException) as exc_info, + ): + await get_generic_sso_response( + request=mock_request, + jwt_handler=MagicMock(spec=JWTHandler), + generic_client_id="cid", + redirect_url="https://proxy.example.com/sso/callback", + sso_jwt_handler=None, + ) + + assert "state" in str(exc_info.value.message).lower() + + @pytest.mark.asyncio + async def test_pkce_callback_accepts_matching_state_cookie(self): + """Happy path: URL state and cookie state match (the legitimate + flow where the same browser that started the redirect lands on + the callback) → the PKCE token exchange proceeds.""" + from litellm.proxy.management_endpoints.ui_sso import ( + SSOAuthenticationHandler, + get_generic_sso_response, + ) + + mock_request = MagicMock(spec=Request) + mock_request.query_params = {"state": "matched-state", "code": "auth-code"} + mock_request.cookies = {"litellm_oauth_state": "matched-state"} + + with ( + patch.object( + SSOAuthenticationHandler, + "prepare_token_exchange_parameters", + AsyncMock( + return_value={ + "code_verifier": "verifier", + "_pkce_cache_key": "pkce_verifier:matched-state", + } + ), + ), + patch.object( + SSOAuthenticationHandler, + "_pkce_token_exchange", + AsyncMock( + return_value={ + "access_token": "tok", + "id_token": "id", + "sub": "user@example.com", + "email": "user@example.com", + } + ), + ), + patch.object( + SSOAuthenticationHandler, + "_delete_pkce_verifier", + AsyncMock(), + ), + patch("fastapi_sso.sso.base.DiscoveryDocument"), + patch("fastapi_sso.sso.generic.create_provider", return_value=MagicMock()), + patch.dict( + os.environ, + { + "GENERIC_CLIENT_SECRET": "x", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://idp.example.com/auth", + "GENERIC_TOKEN_ENDPOINT": "https://idp.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://idp.example.com/userinfo", + "GENERIC_CLIENT_USE_PKCE": "true", + }, + ), + ): + jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler.get_team_ids_from_jwt.return_value = [] + result, _, _ = await get_generic_sso_response( + request=mock_request, + jwt_handler=jwt_handler, + generic_client_id="cid", + redirect_url="https://proxy.example.com/sso/callback", + sso_jwt_handler=None, + ) + + # State-cookie check passed, so the function got past the early + # ProxyException raise and produced an SSO result object. + assert result is not None diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index 66a18e2edb..e8a74e41da 100644 --- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -409,3 +409,60 @@ class TestStreamUsageAiChat: end_date="2025-01-31", user_id="my-user-id", ) + + +class TestUsageAiChatServiceAccountGuard: + """ + Security regression: a non-admin caller with user_id=None (service-account + key) must be rejected at the endpoint boundary, before any tool dispatch. + """ + + @pytest.mark.asyncio + async def test_non_admin_with_user_id_none_is_rejected(self): + from fastapi import HTTPException + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( + ChatMessage, + UsageAIChatRequest, + usage_ai_chat, + ) + + service_account_key = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + request = MagicMock() + body = UsageAIChatRequest( + messages=[ChatMessage(role="user", content="hi")], + model="gpt-4o-mini", + ) + + with pytest.raises(HTTPException) as exc_info: + await usage_ai_chat( + data=body, + request=request, + user_api_key_dict=service_account_key, + ) + + assert exc_info.value.status_code == 403 + assert "Service-account keys" in str(exc_info.value.detail) + + def test_resolve_fetch_kwargs_tripwire_fires_on_none_user_id(self): + """ + Defense-in-depth: if a future endpoint forgets the entry guard and + a non-admin caller with user_id=None reaches _resolve_fetch_kwargs, + the tripwire must fire rather than issuing an unscoped query. + """ + from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( + _resolve_fetch_kwargs, + ) + + with pytest.raises(ValueError) as exc_info: + _resolve_fetch_kwargs( + fn_name="get_usage_data", + fn_args={"start_date": "2025-01-01", "end_date": "2025-01-31"}, + user_id=None, + is_admin=False, + ) + assert "Endpoint-level guard missing" in str(exc_info.value) 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_langfuse_passthrough_security.py b/tests/test_litellm/proxy/test_langfuse_passthrough_security.py new file mode 100644 index 0000000000..5ef3c38c09 --- /dev/null +++ b/tests/test_litellm/proxy/test_langfuse_passthrough_security.py @@ -0,0 +1,102 @@ +import socket + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( + _build_langfuse_proxy_target, + _get_langfuse_proxy_credentials, +) + + +def test_dynamic_langfuse_host_requires_dynamic_credentials(monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", True, raising=False) + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "global-public") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "global-secret") + + with pytest.raises(HTTPException) as exc: + _get_langfuse_proxy_credentials( + dynamic_host_supplied=True, + dynamic_langfuse_public_key=None, + dynamic_langfuse_secret_key=None, + ) + + assert exc.value.status_code == 400 + + +def test_global_langfuse_host_can_use_env_credentials(monkeypatch): + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "global-public") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "global-secret") + + public_key, secret_key = _get_langfuse_proxy_credentials( + dynamic_host_supplied=False, + dynamic_langfuse_public_key=None, + dynamic_langfuse_secret_key=None, + ) + + assert public_key == "global-public" + assert secret_key == "global-secret" + + +@pytest.mark.parametrize( + "endpoint", + [ + "../api/public/projects", + "%2e%2e/api/public/projects", + "%252e%252e%252fapi/public/projects", + "api\\public\\projects", + "%2f%2fattacker.example/api", + ], +) +def test_langfuse_proxy_target_rejects_traversal_paths(endpoint): + with pytest.raises(HTTPException) as exc: + _build_langfuse_proxy_target( + endpoint=endpoint, + base_target_url="https://cloud.langfuse.com", + dynamic_host_supplied=False, + ) + + assert exc.value.status_code == 400 + + +def test_dynamic_langfuse_proxy_target_rejects_internal_host(monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", True, raising=False) + + with pytest.raises(HTTPException) as exc: + _build_langfuse_proxy_target( + endpoint="api/public/projects", + base_target_url="http://127.0.0.1:3000", + dynamic_host_supplied=True, + ) + + assert exc.value.status_code == 400 + + +def test_dynamic_langfuse_proxy_target_preserves_host_header_for_http(monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", True, raising=False) + + def fake_getaddrinfo(host, port, proto): + assert host == "langfuse.example" + assert port == 80 + assert proto == socket.IPPROTO_TCP + return [ + ( + socket.AF_INET, + socket.SOCK_STREAM, + socket.IPPROTO_TCP, + "", + ("8.8.8.8", 80), + ) + ] + + monkeypatch.setattr(socket, "getaddrinfo", fake_getaddrinfo) + + target_url, headers = _build_langfuse_proxy_target( + endpoint="api/public/projects", + base_target_url="http://langfuse.example", + dynamic_host_supplied=True, + ) + + assert target_url == "http://8.8.8.8/api/public/projects" + assert headers["Host"] == "langfuse.example" diff --git a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py new file mode 100644 index 0000000000..64cb931888 --- /dev/null +++ b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py @@ -0,0 +1,118 @@ +import sys +from types import ModuleType, SimpleNamespace + +from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids + + +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 ( + 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_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6be1a9ecef..4e24d8af65 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -2846,6 +2846,61 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po assert "metadata" in data +@pytest.mark.asyncio +async def test_api_created_global_policy_applies_to_new_key_without_restart(): + """ + Regression: policies created at runtime via policy builder must apply + immediately when attached globally, even if the server started with no + initialized policy config. + """ + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import ( + Policy, + PolicyAttachment, + PolicyGuardrails, + ) + + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + } + user_api_key_dict = UserAPIKeyAuth(api_key="test-key") + + policy_registry = get_policy_registry() + attachment_registry = get_attachment_registry() + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + try: + policy_registry.add_policy( + "runtime-global-policy", + Policy(guardrails=PolicyGuardrails(add=["runtime-guardrail"])), + ) + attachment_registry.add_attachment( + PolicyAttachment(policy="runtime-global-policy", scope="*") + ) + + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=user_api_key_dict, + ) + + assert "runtime-guardrail" in data["metadata"]["guardrails"] + assert "runtime-global-policy" in data["metadata"]["applied_policies"] + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_policy_version_by_id(): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 37e5300565..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 @@ -5084,8 +5084,12 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( @pytest.mark.asyncio async def test_reseed_spend_from_db_user_and_org_prefixes(): - """User and org counters must reseed from their own DB tables, not - fall through to 0.0 like the other counters do today.""" + """User and org counters reseed from their own DB tables. + + End-user and tag counters use the already fetched auth objects passed as + fallback_spend, so this reseed helper must not add extra per-request DB + reads for them. + """ from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed user_row = MagicMock() @@ -5095,6 +5099,8 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): fake_prisma = MagicMock() fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + fake_prisma.db.litellm_endusertable.find_unique = AsyncMock() + fake_prisma.db.litellm_tagtable.find_unique = AsyncMock() fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock( return_value=org_row ) @@ -5104,6 +5110,18 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): where={"user_id": "alice"} ) + assert ( + await SpendCounterReseed.from_db( + fake_prisma, + "spend:end_user:customer-1", + ) + is None + ) + fake_prisma.db.litellm_endusertable.find_unique.assert_not_awaited() + + assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") is None + fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited() + assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0 fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( where={"organization_id": "acme"} @@ -5133,6 +5151,468 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited() +@pytest.mark.asyncio +async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + + counter_cache = DualCache() + window_start = datetime.now(timezone.utc) - timedelta(hours=1) + fake_prisma = MagicMock() + fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}] + ) + + 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: + await _init_and_increment_window_spend_counter( + counter_key="spend:key:key-window:window:1h", + entity_type="Key", + entity_id="key-window", + window_start=window_start, + increment=0.5, + ) + + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( + by=["api_key"], + where={"api_key": "key-window", "startTime": {"gte": window_start}}, + sum={"spend": True}, + ) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-window:window:1h" + ) == pytest.approx(2.75) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _init_and_increment_spend_counter + + counter_cache = DualCache() + counter_key = "spend:team:team-stale-local" + counter_cache.in_memory_cache.set_cache(key=counter_key, value=10.0) + + redis_store: dict = {} + + async def redis_increment(key, value, **_): + redis_store[key] = (redis_store.get(key) or 0.0) + value + return redis_store[key] + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + counter_cache.redis_cache = fake_redis + + db_row = MagicMock() + db_row.spend = 42.0 + fake_prisma = MagicMock() + fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=db_row) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma, orig_user = ( + ps.spend_counter_cache, + ps.prisma_client, + ps.user_api_key_cache, + ) + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + ps.user_api_key_cache = DualCache() + try: + await _init_and_increment_spend_counter( + counter_key=counter_key, + source_cache_key="team_id:team-stale-local", + increment=1.5, + ) + + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( + where={"team_id": "team-stale-local"} + ) + assert redis_store[counter_key] == pytest.approx(43.5) + assert counter_cache.in_memory_cache.get_cache( + key=counter_key + ) == pytest.approx(43.5) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + ps.user_api_key_cache = orig_user + + +@pytest.mark.asyncio +async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + + counter_cache = DualCache() + counter_key = "spend:key:key-window-stale-local:window:1h" + counter_cache.in_memory_cache.set_cache(key=counter_key, value=100.0) + window_start = datetime.now(timezone.utc) - timedelta(hours=1) + + redis_store: dict = {} + + async def redis_increment(key, value, **_): + redis_store[key] = (redis_store.get(key) or 0.0) + value + return redis_store[key] + + async def redis_set_cache(key, value, **_): + if key in redis_store: + return False + redis_store[key] = value + return True + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_set_cache = AsyncMock(side_effect=redis_set_cache) + fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}] + ) + + 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: + await _init_and_increment_window_spend_counter( + counter_key=counter_key, + entity_type="Key", + entity_id="key-window-stale-local", + window_start=window_start, + increment=0.5, + ) + + fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once_with( + by=["api_key"], + where={ + "api_key": "key-window-stale-local", + "startTime": {"gte": window_start}, + }, + sum={"spend": True}, + ) + assert redis_store[counter_key] == pytest.approx(2.75) + assert counter_cache.in_memory_cache.get_cache( + key=counter_key + ) == pytest.approx(2.75) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + + counter_cache = DualCache() + counter_key = "spend:key:key-window-concurrent-seed:window:1h" + window_start = datetime.now(timezone.utc) - timedelta(hours=1) + redis_store = {counter_key: 2.75} + redis_reads = 0 + + async def redis_get_cache(key): + nonlocal redis_reads + redis_reads += 1 + if redis_reads <= 2: + return None + return redis_store.get(key) + + async def redis_increment(key, value, **_): + redis_store[key] = (redis_store.get(key) or 0.0) + value + return redis_store[key] + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=redis_get_cache) + fake_redis.async_set_cache = AsyncMock(return_value=False) + fake_redis.async_increment = AsyncMock(side_effect=redis_increment) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( + return_value=[ + {"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}} + ] + ) + + 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: + await _init_and_increment_window_spend_counter( + counter_key=counter_key, + entity_type="Key", + entity_id="key-window-concurrent-seed", + window_start=window_start, + increment=0.5, + ) + + fake_redis.async_set_cache.assert_awaited_once_with( + key=counter_key, + value=2.25, + nx=True, + ) + assert redis_store[counter_key] == pytest.approx(3.25) + assert counter_cache.in_memory_cache.get_cache( + key=counter_key + ) == pytest.approx(3.25) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_window_spend_counter_skips_invalid_window_start(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _init_and_increment_window_spend_counter + + counter_cache = DualCache() + + import litellm.proxy.proxy_server as ps + + orig_counter = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + try: + await _init_and_increment_window_spend_counter( + counter_key="spend:key:key-invalid-window:window:not-a-duration", + entity_type="Key", + entity_id="key-invalid-window", + window_start=None, + increment=0.5, + ) + + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-invalid-window:window:not-a-duration" + ) + is None + ) + finally: + ps.spend_counter_cache = orig_counter + + +@pytest.mark.asyncio +async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _ensure_window_spend_counter_initialized + + counter_cache = DualCache() + counter_key = "spend:key:key-window-db-unavailable:window:1h" + + 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 = None + try: + initialized = await _ensure_window_spend_counter_initialized( + counter_key=counter_key, + entity_type="Key", + entity_id="key-window-db-unavailable", + window_start=datetime.now(timezone.utc) - timedelta(hours=1), + ) + + assert initialized is False + assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_increment_spend_counters_finalizes_after_unreserved_increments(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import increment_spend_counters + + counter_cache = DualCache() + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-finalize-after-increments", + value=0.5, + ) + budget_reservation = { + "reserved_cost": 0.5, + "entries": [ + { + "counter_key": "spend:key:key-finalize-after-increments", + "entity_type": "Key", + "entity_id": "key-finalize-after-increments", + "reserved_cost": 0.5, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + incremented_counters = [] + + async def assert_reservation_not_finalized_yet(**kwargs): + assert budget_reservation["finalized"] is False + incremented_counters.append(kwargs["counter_key"]) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_user = ps.spend_counter_cache, ps.user_api_key_cache + ps.spend_counter_cache = counter_cache + ps.user_api_key_cache = DualCache() + try: + with patch( + "litellm.proxy.proxy_server._init_and_increment_spend_counter", + new=AsyncMock(side_effect=assert_reservation_not_finalized_yet), + ): + await increment_spend_counters( + token="key-finalize-after-increments", + team_id="team-finalize-after-increments", + user_id=None, + response_cost=0.25, + budget_reservation=budget_reservation, + ) + + assert incremented_counters == ["spend:team:team-finalize-after-increments"] + assert budget_reservation["finalized"] is True + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-finalize-after-increments" + ) == pytest.approx(0.25) + finally: + ps.spend_counter_cache = orig_counter + ps.user_api_key_cache = orig_user + + +@pytest.mark.asyncio +async def test_increment_spend_counters_finalizes_none_cost_reservation(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import increment_spend_counters + + counter_cache = DualCache() + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-finalize-none-cost", + value=0.5, + ) + budget_reservation = { + "reserved_cost": 0.5, + "entries": [ + { + "counter_key": "spend:key:key-finalize-none-cost", + "entity_type": "Key", + "entity_id": "key-finalize-none-cost", + "reserved_cost": 0.5, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + import litellm.proxy.proxy_server as ps + + orig_counter = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + try: + await increment_spend_counters( + token="key-finalize-none-cost", + team_id=None, + user_id=None, + response_cost=None, + budget_reservation=budget_reservation, + ) + + assert budget_reservation["finalized"] is True + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-finalize-none-cost" + ) == pytest.approx(0.0) + finally: + ps.spend_counter_cache = orig_counter + + +@pytest.mark.asyncio +async def test_increment_spend_counters_invalidates_bad_reserved_counter_without_failing(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import increment_spend_counters + + counter_cache = DualCache() + budget_reservation = { + "reserved_cost": 0.5, + "entries": [ + { + "counter_key": "spend:key:key-bad-reserved-counter", + "entity_type": "Key", + "entity_id": "key-bad-reserved-counter", + "reserved_cost": 0.5, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + import litellm.proxy.proxy_server as ps + + orig_counter = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + try: + with patch( + "litellm.proxy.proxy_server.verbose_proxy_logger.warning" + ) as mock_warning: + await increment_spend_counters( + token="key-bad-reserved-counter", + team_id=None, + user_id=None, + response_cost=0.25, + budget_reservation=budget_reservation, + ) + + mock_warning.assert_called_once() + assert budget_reservation["finalized"] is True + assert ( + counter_cache.in_memory_cache.get_cache( + key="spend:key:key-bad-reserved-counter" + ) + is None + ) + finally: + ps.spend_counter_cache = orig_counter + + +@pytest.mark.asyncio +async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure(): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import _increment_spend_counter_cache + + counter_cache = DualCache() + counter_cache.in_memory_cache.set_cache(key="spend:team:redis-fail", value=4.0) + fake_redis = AsyncMock() + fake_redis.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) + fake_redis.async_delete_cache = AsyncMock() + counter_cache.redis_cache = fake_redis + + import litellm.proxy.proxy_server as ps + + orig_counter = ps.spend_counter_cache + ps.spend_counter_cache = counter_cache + try: + with pytest.raises(RuntimeError): + await _increment_spend_counter_cache( + counter_key="spend:team:redis-fail", + increment=0.5, + ) + + assert ( + counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None + ) + fake_redis.async_delete_cache.assert_awaited_once_with( + key="spend:team:redis-fail" + ) + finally: + ps.spend_counter_cache = orig_counter + + @pytest.mark.asyncio async def test_get_current_spend_reseeds_from_db_when_counter_missing(): """ @@ -5181,6 +5661,9 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): assert ("spend:team_member:user-1:team-1", 362.0) in [ (w["key"], w["value"]) for w in recorded_warms ] + assert counter_cache.in_memory_cache.get_cache( + key="spend:team_member:user-1:team-1" + ) == pytest.approx(362.0) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -5505,6 +5988,202 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): ps.prisma_client = orig_prisma +# ----------------------------------------------------------------------------- +# /config/update — critical paths only. +# +# These exercise the four behaviors that broke or changed in the rewrite of +# update_config (litellm/proxy/proxy_server.py): targeted per-section writes, +# the removal of the store_model_in_db gate, env var encryption, and the +# success_callback / litellm_settings merge semantics. All other branches +# (auth, missing-DB, slack auto-enable, router_settings merge) are covered +# implicitly or by upstream tests. +# ----------------------------------------------------------------------------- + + +class _FakeRow: + def __init__(self, param_name, param_value): + self.param_name = param_name + self.param_value = param_value + + +class _FakeLitellmConfig: + def __init__(self, initial_rows=None): + self.rows = dict(initial_rows or {}) + self.upsert_calls: list = [] + self.find_first = AsyncMock(side_effect=self._find_first) + self.upsert = AsyncMock(side_effect=self._upsert) + + async def _find_first(self, where=None): + if where and "param_name" in where: + name = where["param_name"] + if name in self.rows: + return _FakeRow(name, self.rows[name]) + return None + + async def _upsert(self, where=None, data=None): + name = where["param_name"] + raw = data["update"]["param_value"] + value = json.loads(raw) if isinstance(raw, str) else raw + self.rows[name] = value + self.upsert_calls.append((name, value)) + + +class _FakePrismaClient: + def __init__(self, initial_rows=None): + self.db = mock.MagicMock() + self.db.litellm_config = _FakeLitellmConfig(initial_rows=initial_rows) + self.jsonify_object = lambda obj: obj + + +@pytest.fixture +def _update_config_setup(monkeypatch): + """Install fakes for the /config/update endpoint and return (client, prisma).""" + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth as auth_dep + + def _install(initial_rows=None, store_model_in_db=True): + prisma = _FakePrismaClient(initial_rows=initial_rows) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr( + "litellm.proxy.proxy_server.store_model_in_db", store_model_in_db + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.encrypt_value_helper", + lambda value, **_: f"enc:{value}", + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.invalidate_config_param", + AsyncMock(return_value=None), + ) + from litellm.proxy.proxy_server import proxy_config as real_proxy_config + + monkeypatch.setattr( + real_proxy_config, "add_deployment", AsyncMock(return_value=None) + ) + + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[auth_dep] = lambda: UserAPIKeyAuth( + user_id="test_admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + ) + client = TestClient(app) + + def _restore(): + app.dependency_overrides = original_overrides + + return client, prisma, _restore + + return _install + + +def test_update_config_writes_only_sent_section(_update_config_setup): + """A request that only touches general_settings must not write any other + section row, and must leave previously-written rows byte-identical.""" + client, prisma, restore = _update_config_setup( + initial_rows={ + "litellm_settings": {"drop_params": True}, + "environment_variables": {"FOO": "enc:bar"}, + } + ) + try: + resp = client.post( + "/config/update", + json={"general_settings": {"store_prompts_in_spend_logs": True}}, + ) + assert resp.status_code == 200 + written = {name for name, _ in prisma.db.litellm_config.upsert_calls} + assert written == {"general_settings"} + assert prisma.db.litellm_config.rows["litellm_settings"] == { + "drop_params": True + } + assert prisma.db.litellm_config.rows["environment_variables"] == { + "FOO": "enc:bar" + } + finally: + restore() + + +def test_update_config_can_flip_store_model_in_db_when_currently_false( + _update_config_setup, +): + """The endpoint used to refuse all writes when store_model_in_db was + False, blocking the very request that would flip it to True.""" + client, prisma, restore = _update_config_setup(store_model_in_db=False) + try: + resp = client.post( + "/config/update", json={"general_settings": {"store_model_in_db": True}} + ) + assert resp.status_code == 200 + assert ( + prisma.db.litellm_config.rows["general_settings"]["store_model_in_db"] + is True + ) + finally: + restore() + + +def test_update_config_environment_variables_encrypted_before_write( + _update_config_setup, +): + """env var values must be encrypted before they hit the DB row.""" + client, prisma, restore = _update_config_setup() + try: + resp = client.post( + "/config/update", + json={"environment_variables": {"OPENAI_API_KEY": "sk-secret"}}, + ) + assert resp.status_code == 200 + stored = prisma.db.litellm_config.rows["environment_variables"] + assert stored == {"OPENAI_API_KEY": "enc:sk-secret"} + finally: + restore() + + +def test_update_config_litellm_settings_request_wins_for_non_callback_keys( + _update_config_setup, +): + """Sending {"drop_params": False} when the row holds drop_params: True + must persist False (request wins). Untouched keys preserved.""" + client, prisma, restore = _update_config_setup( + initial_rows={ + "litellm_settings": {"drop_params": True, "set_verbose": True}, + } + ) + try: + resp = client.post( + "/config/update", json={"litellm_settings": {"drop_params": False}} + ) + assert resp.status_code == 200 + stored = prisma.db.litellm_config.rows["litellm_settings"] + assert stored["drop_params"] is False + assert stored["set_verbose"] is True + finally: + restore() + + +def test_update_config_success_callback_normalizes_existing_mixed_case( + _update_config_setup, +): + """Existing mixed-case callback names (written elsewhere) must be + normalized to lowercase before union, otherwise the union dedup misses + against the lowercase incoming entry and delete_callback (lowercase + lookup) cannot find the original.""" + client, prisma, restore = _update_config_setup( + initial_rows={"litellm_settings": {"success_callback": ["Langfuse", "SQS"]}} + ) + try: + resp = client.post( + "/config/update", + json={"litellm_settings": {"success_callback": ["langfuse"]}}, + ) + assert resp.status_code == 200 + stored = prisma.db.litellm_config.rows["litellm_settings"]["success_callback"] + assert set(stored) == {"langfuse", "sqs"} + finally: + restore() + + # --------------------------------------------------------------------------- # Lazy feature loading (LazyFeatureMiddleware) — verifies that optional # routers are NOT imported at module load and ARE imported on first request @@ -5513,9 +6192,6 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): # --------------------------------------------------------------------------- -import sys - - class TestLazyFeatureRegistry: """Sanity checks on the registry shape — guards against accidental edits.""" 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/test_anthropic_skills_transformation.py b/tests/test_litellm/test_anthropic_skills_transformation.py index 6761d67180..1b917f08ca 100644 --- a/tests/test_litellm/test_anthropic_skills_transformation.py +++ b/tests/test_litellm/test_anthropic_skills_transformation.py @@ -70,6 +70,15 @@ class TestAnthropicSkillsConfigURLConstruction: ) assert url == f"{FAKE_API_BASE}/v1/skills/skill_abc123" + def test_url_with_skill_id_encodes_path_segment(self): + url = self.config.get_complete_url( + api_base=FAKE_API_BASE, + endpoint="skills", + skill_id="../../files?x=1#frag", + ) + + assert url == f"{FAKE_API_BASE}/v1/skills/..%2F..%2Ffiles%3Fx%3D1%23frag" + def test_url_falls_back_to_anthropic_default(self): with patch( "litellm.llms.anthropic.common_utils.AnthropicModelInfo.get_api_base", diff --git a/tests/test_litellm/test_openai_embedding_encoding_format_default.py b/tests/test_litellm/test_openai_embedding_encoding_format_default.py new file mode 100644 index 0000000000..94e4e3c81e --- /dev/null +++ b/tests/test_litellm/test_openai_embedding_encoding_format_default.py @@ -0,0 +1,124 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from litellm import embedding + + +@pytest.mark.parametrize( + "set_env, env_value, expected", + [ + (False, None, "float"), + (True, "base64", "base64"), + ], +) +def test_openai_embedding_encoding_format_default( + monkeypatch, set_env, env_value, expected +): + monkeypatch.delenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", raising=False) + if set_env: + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_value) + + mock_response = MagicMock() + mock_response.parse.return_value = MagicMock( + model_dump=lambda: { + "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], + "model": "text-embedding-ada-002", + "object": "list", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ) + mock_response.headers = {} + + with patch( + "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" + ) as mock_get_client: + mock_client_instance = MagicMock() + mock_get_client.return_value = mock_client_instance + mock_client_instance.embeddings.with_raw_response.create.return_value = ( + mock_response + ) + + embedding( + model="text-embedding-ada-002", + input="Hello world", + ) + + call_kwargs = ( + mock_client_instance.embeddings.with_raw_response.create.call_args[1] + ) + assert call_kwargs["encoding_format"] == expected + + +@pytest.mark.parametrize("env_none", ["none", "NONE", " none "]) +def test_openai_embedding_encoding_format_env_none_omits_param( + monkeypatch, env_none +): + """LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT=none omits encoding_format (provider default).""" + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", env_none) + + mock_response = MagicMock() + mock_response.parse.return_value = MagicMock( + model_dump=lambda: { + "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], + "model": "text-embedding-ada-002", + "object": "list", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ) + mock_response.headers = {} + + with patch( + "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" + ) as mock_get_client: + mock_client_instance = MagicMock() + mock_get_client.return_value = mock_client_instance + mock_client_instance.embeddings.with_raw_response.create.return_value = ( + mock_response + ) + + embedding( + model="text-embedding-ada-002", + input="Hello world", + ) + + call_kwargs = ( + mock_client_instance.embeddings.with_raw_response.create.call_args[1] + ) + assert "encoding_format" not in call_kwargs + + +def test_openai_embedding_encoding_format_explicit_overrides_env(monkeypatch): + """Request `encoding_format` wins over LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT.""" + monkeypatch.setenv("LITELLM_DEFAULT_EMBEDDING_ENCODING_FORMAT", "float") + + mock_response = MagicMock() + mock_response.parse.return_value = MagicMock( + model_dump=lambda: { + "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], + "model": "text-embedding-ada-002", + "object": "list", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ) + mock_response.headers = {} + + with patch( + "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" + ) as mock_get_client: + mock_client_instance = MagicMock() + mock_get_client.return_value = mock_client_instance + mock_client_instance.embeddings.with_raw_response.create.return_value = ( + mock_response + ) + + embedding( + model="text-embedding-ada-002", + input="Hello world", + encoding_format="base64", + ) + + call_kwargs = ( + mock_client_instance.embeddings.with_raw_response.create.call_args[1] + ) + assert call_kwargs["encoding_format"] == "base64" diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 6cbb2fd7b8..8a0a2221c1 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -9,11 +9,11 @@ from litellm._logging import ( JsonFormatter, _redact_string, _secret_filter, - _setup_json_exception_handlers, verbose_logger, verbose_proxy_logger, verbose_router_logger, ) +from litellm.litellm_core_utils.secret_redaction import redact_string SECRET = "sk-proj-abc123def456ghi789jklmnopqrst" @@ -57,12 +57,12 @@ def test_redact_string_catches_secret_patterns(): SECRET, ] for secret in cases: - result = _redact_string("msg: " + secret) + result = redact_string("msg: " + secret) assert secret not in result, f"{secret!r} was not redacted" assert "REDACTED" in result normal = "Loaded model gpt-4 with 3 replicas on us-east-1" - assert _redact_string(normal) == normal + assert redact_string(normal) == normal def test_filter_redacts_secrets_in_logger_output(): @@ -155,7 +155,7 @@ def test_x_api_key_regex_does_not_consume_json_delimiters(): """x-api-key pattern must stop before closing quotes/braces so JSON stays valid.""" # Simulates a JSON log line containing an x-api-key header value json_line = '{"headers": {"x-api-key": "secret123"}, "status": 200}' - result = _redact_string(json_line) + result = redact_string(json_line) # The secret value should be redacted assert "secret123" not in result assert "REDACTED" in result @@ -234,12 +234,12 @@ def test_key_name_redaction_catches_secrets_in_dict_repr(): "'slack_webhook_url': 'https://hooks.slack.com/services/T00/B00/xxx'", ] for secret_line in cases: - result = _redact_string(secret_line) + result = redact_string(secret_line) assert "REDACTED" in result, f"Key-name redaction missed: {secret_line!r}" # Non-sensitive keys should NOT be redacted safe = "'enable_jwt_auth': True, 'store_model_in_db': True" - assert _redact_string(safe) == safe + assert redact_string(safe) == safe def test_key_name_redaction_in_general_settings_dict(): @@ -277,7 +277,7 @@ _SAMPLE_SA_JSON = ( def test_pem_private_key_redacted_in_json(): - result = _redact_string(_SAMPLE_SA_JSON) + result = redact_string(_SAMPLE_SA_JSON) assert "MIIEvQIBADA" not in result assert "-----BEGIN" not in result @@ -286,12 +286,12 @@ def test_pem_private_key_redacted_in_dict_repr(): import json sa = json.loads(_SAMPLE_SA_JSON) - result = _redact_string(str(sa)) + result = redact_string(str(sa)) assert "MIIEvQIBADA" not in result def test_service_account_blob_fully_redacted(): - result = _redact_string(f"Got={_SAMPLE_SA_JSON}") + result = redact_string(f"Got={_SAMPLE_SA_JSON}") assert "my-proj-123" not in result assert "sa@my-proj.iam.gserviceaccount.com" not in result assert "abc123def" not in result @@ -320,22 +320,22 @@ def test_vertex_traceback_redacts_pem(): "Unable to load vertex credentials from environment. " f"Got={_SAMPLE_SA_JSON}" ) - result = _redact_string(traceback_text) + result = redact_string(traceback_text) assert "MIIEvQIBADA" not in result assert "-----BEGIN" not in result def test_gcp_oauth_token_redacted(): - result = _redact_string("access token ya29.c.c0ASRK0GZvXlongtokenhere") + result = redact_string("access token ya29.c.c0ASRK0GZvXlongtokenhere") assert "ya29." not in result assert "REDACTED" in result def test_non_pem_private_key_value_redacted(): - result = _redact_string("'private_key': 'some-non-pem-secret-value'") + result = redact_string("'private_key': 'some-non-pem-secret-value'") assert "some-non-pem-secret" not in result def test_normal_vertex_log_not_redacted(): msg = "Vertex: Loading vertex credentials, is_file_path=True, current dir /app" - assert _redact_string(msg) == msg + assert redact_string(msg) == msg diff --git a/ui/litellm-dashboard/public/assets/logos/qohash.jpg b/ui/litellm-dashboard/public/assets/logos/qohash.jpg new file mode 100644 index 0000000000..50227ab391 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/qohash.jpg differ 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 7c162e2056..944c56833e 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 @@ -17,7 +17,7 @@ import { RefreshIcon } from "@heroicons/react/outline"; import { useQueryClient } from "@tanstack/react-query"; import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; import type { UploadProps } from "antd"; -import { Form, Typography } from "antd"; +import { Form } from "antd"; import { PlusCircleOutlined } from "@ant-design/icons"; import React, { useEffect, useMemo, useState } from "react"; import AddModelTab from "../../../components/add_model/add_model_tab"; @@ -251,15 +251,9 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te const isLoading = isLoadingModels || isLoadingModelCostMap || isLoadingCredentials || isLoadingUISettings; - if (userRole && userRole == "Admin Viewer") { - const { Title, Paragraph } = Typography; - return ( -
- Access Denied - Ask your proxy admin for access to view all models -
- ); - } + // Admin Viewer can view all models read-only — page render proceeds; the + // individual write-action tabs (Add Model, LLM Credentials, etc.) are + // gated separately below. const handleOk = async () => { try { @@ -395,107 +389,154 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te modelAccessGroups={availableModelAccessGroups} /> ) : ( - - -
- {all_admin_roles.includes(userRole) ? All Models : Your Models} - {!shouldHideAddModelTab && Add Model} - {all_admin_roles.includes(userRole) && LLM Credentials} - {all_admin_roles.includes(userRole) && Pass-Through Endpoints} - {all_admin_roles.includes(userRole) && Health Status} - {all_admin_roles.includes(userRole) && Model Retry Settings} - {all_admin_roles.includes(userRole) && Model Group Alias} - {all_admin_roles.includes(userRole) && Price Data Reload} -
- -
- {lastRefreshed && Last Refreshed: {lastRefreshed}} - -
- - - - {!shouldHideAddModelTab && ( - - { + // Build a single source-of-truth list of {tab, panel} pairs. + // Conditionally-hidden tabs (e.g. "Add Model" for non-admin) get + // filtered out as a unit so tab indices and panel indices can + // never drift apart — Tremor's TabList and TabPanels filter + // falsy children inconsistently, which previously caused + // "click LLM Credentials, see nothing" for Admin Viewer. + const isAdmin = all_admin_roles.includes(userRole); + const visibleTabs: Array<{ tab: React.ReactElement; panel: React.ReactElement }> = [ + { + tab: {isAdmin ? "All Models" : "Your Models"}, + panel: ( + - - )} - - - - - - - - - - - - - - - - + ), + }, + ]; + if (!shouldHideAddModelTab) { + visibleTabs.push({ + tab: Add Model, + panel: ( + + + + ), + }); + } + if (isAdmin) { + visibleTabs.push( + { + tab: LLM Credentials, + panel: ( + + + + ), + }, + { + tab: Pass-Through Endpoints, + panel: ( + + + + ), + }, + { + tab: Health Status, + panel: ( + + + + ), + }, + { + tab: Model Retry Settings, + panel: ( + + ), + }, + { + tab: Model Group Alias, + panel: ( + + + + ), + }, + { + tab: Price Data Reload, + panel: , + }, + ); + } + return ( + + +
{visibleTabs.map((t) => t.tab)}
+ +
+ {lastRefreshed && Last Refreshed: {lastRefreshed}} + +
+ + {visibleTabs.map((t) => t.panel)} + + ); + })() )} diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx index 866b7d0f17..cd58c51a86 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.test.tsx @@ -289,4 +289,85 @@ describe("LoginPage", () => { expect(ssoButton).toBeInTheDocument(); expect(ssoButton).toBeDisabled(); }); + + describe("URL ?token= legacy path is rejected (security regression test)", () => { + const originalLocation = window.location; + + beforeEach(() => { + Object.defineProperty(window, "location", { + value: { + ...originalLocation, + href: "http://localhost:3000/ui/login?token=attacker.jwt.value", + pathname: "/ui/login", + search: "?token=attacker.jwt.value", + }, + writable: true, + }); + document.cookie = + "token=; expires=Thu, 01 Jan 1970 00:00:00 GMT; path=/; SameSite=Lax"; + }); + + afterEach(() => { + Object.defineProperty(window, "location", { + value: originalLocation, + writable: true, + }); + }); + + it("must not set a token cookie or redirect to /ui/?login=success when ?token= is in the URL", async () => { + (useUIConfig as ReturnType).mockReturnValue({ + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(null); + (isJwtExpired as ReturnType).mockReturnValue(false); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByRole("heading", { name: "Login" })).toBeInTheDocument(); + }); + + expect(document.cookie).not.toContain("token=attacker.jwt.value"); + expect(mockReplace).not.toHaveBeenCalledWith("/ui/?login=success"); + }); + + it("must not overwrite an existing valid session cookie when ?token= is in the URL", async () => { + (useUIConfig as ReturnType).mockReturnValue({ + data: { + auto_redirect_to_sso: false, + server_root_path: "/", + proxy_base_url: null, + sso_configured: false, + }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue("legitimate-session-jwt"); + (isJwtExpired as ReturnType).mockReturnValue(false); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(mockReplace).toHaveBeenCalledWith("/ui"); + }); + + expect(document.cookie).not.toContain("token=attacker.jwt.value"); + expect(mockReplace).not.toHaveBeenCalledWith("/ui/?login=success"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 7ad3e32ef5..74ee9f9de5 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -66,21 +66,6 @@ function LoginPageContent() { return; } - // Backwards compat: handle direct token in URL (legacy flow) - const urlToken = params.get("token"); - if (urlToken && !isJwtExpired(urlToken)) { - document.cookie = `token=${urlToken}; path=/; SameSite=Lax`; - params.delete("token"); - const cleanSearch = params.toString(); - window.history.replaceState( - null, - "", - window.location.pathname + (cleanSearch ? `?${cleanSearch}` : ""), - ); - router.replace("/ui/?login=success"); - return; - } - // If switching workers on a control plane, clear the old token and show login const switchingWorker = params.has("worker"); if (switchingWorker && uiConfig?.is_control_plane) { diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index a8553d5405..06bf3b68d0 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -325,9 +325,6 @@ function CreateKeyPageContent() { if (decoded.user_role) { const formattedUserRole = formatUserRole(decoded.user_role); setUserRole(formattedUserRole); - if (formattedUserRole == "Admin Viewer") { - setPage("usage"); - } } if (decoded.user_email) { diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 39695c1348..75058157a6 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -22,7 +22,7 @@ import { modelHubPublicModelsCall, } from "@/components/networking"; import PublicModelHub from "@/components/public_model_hub"; -import { isAdminRole } from "@/utils/roles"; +import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; import { CopyOutlined } from "@ant-design/icons"; import { Badge, Button, Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; import { Modal } from "antd"; @@ -61,6 +61,10 @@ interface ModelGroupInfo { } const ModelHubTable: React.FC = ({ accessToken, publicPage, premiumUser, userRole }) => { + // Admin Viewer follows the read-parity rule: see the AI Hub catalog, but + // cannot toggle public visibility (write). + const canModify = isProxyAdminRole(userRole || ""); + const [publicPageAllowed, setPublicPageAllowed] = useState(false); const [modelHubData, setModelHubData] = useState(null); const [loading, setLoading] = useState(true); @@ -420,7 +424,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Useful Links Management Section for Admins */} - {isAdminRole(userRole || "") && ( + {canModify && (
@@ -441,7 +445,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Model Filters and Table */} {/* Header with Make Public Button */} - {publicPage == false && isAdminRole(userRole || "") && ( + {publicPage == false && canModify && (
@@ -470,7 +474,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Header with Make Public Button */} - {publicPage == false && isAdminRole(userRole || "") && ( + {publicPage == false && canModify && (
@@ -496,7 +500,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Header with Make Public Button */} - {publicPage == false && isAdminRole(userRole || "") && ( + {publicPage == false && canModify && (
@@ -520,7 +524,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Skill Hub Tab */} - {publicPage == false && isAdminRole(userRole || "") && ( + {publicPage == false && canModify && (
+ {canModify && ( + + )} diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx index 9c5933aba3..b4ab7ddddb 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.tsx @@ -8,6 +8,7 @@ import DeleteResourceModal from "../../../common_components/DeleteResourceModal" import { ProviderLogo } from "../../../molecules/models/ProviderLogo"; import NotificationsManager from "../../../molecules/notifications_manager"; import { getCallbacksCall, setCallbacksCall } from "../../../networking"; +import { isProxyAdminRole } from "@/utils/roles"; import AddFallbacks from "./AddFallbacks"; type FallbackEntry = { [modelName: string]: string[] }; @@ -243,15 +244,19 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID, mo }; const hasFallbacks = Array.isArray(routerSettings.fallbacks) && routerSettings.fallbacks.length > 0; + // Admin Viewer follows the read-parity rule: see fallbacks, no writes. + const canModify = isProxyAdminRole(userRole ?? ""); return ( <> - data.model_name) : []} - accessToken={accessToken || ""} - value={routerSettings.fallbacks || []} - onChange={handleFallbacksChange} - /> + {canModify && ( + data.model_name) : []} + accessToken={accessToken || ""} + value={routerSettings.fallbacks || []} + onChange={handleFallbacksChange} + /> + )} {!hasFallbacks ? (
@@ -280,30 +285,34 @@ const Fallbacks: React.FC = ({ accessToken, userRole, userID, mo {renderFallbacksChain(key, Array.isArray(value) ? value : [], getProviderFromModel)} - - testFallbackModelResponse(Object.keys(item)[0], accessToken || "")} - className="cursor-pointer hover:text-blue-600" - /> - - - handleDeleteClick(item)} - onKeyDown={(e) => e.key === "Enter" && handleDeleteClick(item)} - className="cursor-pointer inline-flex" - > - - - + {canModify && ( + <> + + testFallbackModelResponse(Object.keys(item)[0], accessToken || "")} + className="cursor-pointer hover:text-blue-600" + /> + + + handleDeleteClick(item)} + onKeyDown={(e) => e.key === "Enter" && handleDeleteClick(item)} + className="cursor-pointer inline-flex" + > + + + + + )} )), diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 69b29564d8..809f1d4e17 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -169,8 +169,8 @@ const UsagePage: React.FC = ({ 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 = ({ accessToken, userRole }) => { onPromptClick={handlePromptClick} onDeleteClick={handleDeleteClick} accessToken={accessToken} - isAdmin={isAdmin} + isAdmin={canModify} /> )} diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 90eac56540..db262bc84a 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -1,7 +1,6 @@ "use client"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { Col, Grid } from "@tremor/react"; -import { Typography } from "antd"; import { jwtDecode } from "jwt-decode"; import { useSearchParams } from "next/navigation"; import React, { useEffect, useState } from "react"; @@ -317,15 +316,10 @@ const UserDashboard: React.FC = ({ setUserRole("App Owner"); } - if (userRole && userRole == "Admin Viewer") { - const { Title, Paragraph } = Typography; - return ( -
- Access Denied - Ask your proxy admin for access to create keys -
- ); - } + // Admin Viewer can view keys read-only — gate "Create Key" but render the + // virtual-keys table the same as for Proxy Admin (read parity). Every + // other role keeps its existing ability to create keys. + const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer"; console.log("inside user dashboard, selected team", selectedTeam); console.log("All cookies after redirect:", document.cookie); @@ -333,15 +327,17 @@ const UserDashboard: React.FC = ({
- + {canCreateKey && ( + + )} diff --git a/ui/litellm-dashboard/src/components/view_users.tsx b/ui/litellm-dashboard/src/components/view_users.tsx index 50123a8395..4ee1597512 100644 --- a/ui/litellm-dashboard/src/components/view_users.tsx +++ b/ui/litellm-dashboard/src/components/view_users.tsx @@ -304,7 +304,9 @@ const ViewUserDashboard: React.FC = ({ accessToken, toke ) : userID && accessToken ? ( <> - + {isProxyAdmin && ( + + )} {isProxyAdmin && (