diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 1a2c81d1f92..2dc85fce05b 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -47,6 +47,10 @@ After: the same request comes back with real token counts, so the dashboard show +## Affected release + + + ## Linear ticket @@ -154,3 +158,4 @@ Example checklists: ## Final Attestation - [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR + diff --git a/.github/scripts/auto_update_price_and_context_window_file.py b/.github/scripts/auto_update_price_and_context_window_file.py index 461d8d347d9..4c9a05bd4f0 100644 --- a/.github/scripts/auto_update_price_and_context_window_file.py +++ b/.github/scripts/auto_update_price_and_context_window_file.py @@ -1,6 +1,8 @@ import asyncio import aiohttp import json +import math +from typing import Any # Asynchronously fetch data from a given URL async def fetch_data(url): @@ -21,11 +23,157 @@ async def fetch_data(url): print("Error fetching data from URL:", e) return None + +FRIENDLI_API_URL = "https://api.friendli.ai/serverless/v1/models" +FRIENDLI_PROVIDER = "friendliai" + +INHERITABLE_BASE_KEYS = ( + "supports_pdf_input", + "supports_assistant_prefill", + "supports_adaptive_thinking", + "supports_output_config", +) + +REASONING_EFFORT_LEVEL_ORDER = ("none", "minimal", "low", "medium", "high", "xhigh", "max") + + +def _find_base_model_entry(base_model: str, local_data: dict) -> str | None: + if not base_model: + return None + bm_tail = base_model.split("/")[-1].lower() + if base_model in local_data: + return base_model + for key in local_data: + if key.startswith("sample_spec") or key == "fallback_generalizations": + continue + if key.split("/")[-1].lower() == bm_tail: + return key + return None + + +def _reasoning_effort_levels(reasoning_options: list) -> list: + offered = { + val + for opt in reasoning_options or [] + if opt.get("type") == "effort" + for val in opt.get("values", []) + } + return [level for level in REASONING_EFFORT_LEVEL_ORDER if level in offered] + + +def _valid_token_price(value: object) -> bool: + try: + price = float(value) # pyright: ignore[reportArgumentType] # non-numeric values are rejected via the except + except (TypeError, ValueError): + return False + return math.isfinite(price) and price >= 0 + + +def _has_valid_token_prices(pricing: dict | None) -> bool: + prices = pricing or {} + return _valid_token_price(prices.get("input")) and _valid_token_price(prices.get("output")) + + +def _pricing(pricing: dict) -> dict: + out: dict[str, Any] = {} + if not pricing: + return out + if "input" in pricing: + out["input_cost_per_token"] = float(pricing["input"]) + if "output" in pricing: + out["output_cost_per_token"] = float(pricing["output"]) + if "input_cache_read" in pricing and pricing["input_cache_read"] is not None: + out["cache_read_input_token_cost"] = float(pricing["input_cache_read"]) + return out + + +def _modality_flags(input_mods: list) -> dict: + mods = input_mods or [] + has_image = "image" in mods + return { + "supports_vision": has_image, + "supports_image_input": has_image, + "supports_video_input": "video" in mods, + } + + +def transform_friendli_data(data: list, local_data: dict) -> dict: + transformed: dict[str, dict] = {} + if not data: + return transformed + for model in data: + # An unpriced row must never wholesale-replace an already priced local entry: + # missing prices cost-calculate as zero, silently zeroing tracked spend + if not _has_valid_token_prices(model.get("pricing")): + continue + model_id = model["id"] + base_model = model.get("base_model") or "" + entry: dict[str, Any] = { + "litellm_provider": FRIENDLI_PROVIDER, + } + + base_key = _find_base_model_entry(base_model, local_data) + if base_key: + base_entry = local_data[base_key] + for k in INHERITABLE_BASE_KEYS: + if k in base_entry: + entry[k] = base_entry[k] + + ctx = model.get("context_length") + if ctx is not None: + entry["max_input_tokens"] = int(ctx) + max_out = model.get("max_completion_tokens") + if max_out is not None: + entry["max_output_tokens"] = int(max_out) + entry["max_tokens"] = int(max_out) + + pricing = _pricing(model.get("pricing", {})) + entry.update(pricing) + entry["supports_prompt_caching"] = "cache_read_input_token_cost" in pricing + + reasoning = model.get("reasoning") is True + entry["supports_reasoning"] = reasoning + if reasoning: + entry["reasoning_effort_levels"] = _reasoning_effort_levels( + model.get("reasoning_options", []) + ) + + func = model.get("functionality", {}) + entry["supports_function_calling"] = func.get("tool_call") is True + entry["supports_parallel_function_calling"] = func.get("parallel_tool_call") is True + is_struct = func.get("structured_output") is True + entry["supports_response_schema"] = is_struct + entry["supports_native_structured_output"] = is_struct + entry["supports_system_messages"] = func.get("system_messages") is True + entry["supports_tool_choice"] = func.get("tool_choice") is True + + entry.update(_modality_flags(model.get("input_modalities", []))) + + entry["mode"] = model.get("mode", "chat") + + desc = model.get("description") + if desc: + entry["comment"] = desc + + dep = model.get("deprecation_date") + if dep: + entry["deprecation_date"] = dep.split("T")[0] + + entry["source"] = FRIENDLI_API_URL + + transformed[f"{FRIENDLI_PROVIDER}/{model_id}"] = entry + return transformed + # Synchronize local data with remote data -def sync_local_data_with_remote(local_data, remote_data): +def sync_local_data_with_remote(local_data, remote_data, replace_keys=frozenset()): # Update existing keys in local_data with values from remote_data + # (replace_keys entries are swapped wholesale so a field the remote catalog + # dropped, e.g. cache pricing, cannot survive as a stale value) for key in (set(local_data) & set(remote_data)): - local_data[key].update(remote_data[key]) + if key in replace_keys: + local_data[key] = remote_data[key] + else: + local_data[key].update(remote_data[key]) # Add new keys from remote_data to local_data for key in (set(remote_data) - set(local_data)): @@ -46,6 +194,8 @@ def write_to_file(file_path, data): # Update the existing models and add the missing models for OpenRouter def transform_openrouter_data(data): transformed = {} + if not data: + return transformed for row in data: # Add the fields 'max_tokens' and 'input_cost_per_token' obj = { @@ -84,7 +234,14 @@ def transform_openrouter_data(data): # Update the existing models and add the missing models for Vercel AI Gateway def transform_vercel_ai_gateway_data(data): transformed = {} + if not data: + return transformed for row in data: + # Rows without token pricing or token limits (video/embedding models) previously KeyError'd the whole sync + if any(row.get(k) is None for k in ("context_window", "max_tokens")) or any( + row.get("pricing", {}).get(k) is None for k in ("input", "output") + ): + continue obj = { "max_tokens": row["context_window"], "input_cost_per_token": float(row["pricing"]["input"]), @@ -143,13 +300,16 @@ def main(): vercel_data = asyncio.run(fetch_data(vercel_ai_gateway_url)) # Transform the fetched Vercel AI Gateway data vercel_data = transform_vercel_ai_gateway_data(vercel_data) + + friendli_data = asyncio.run(fetch_data(FRIENDLI_API_URL)) + friendli_data = transform_friendli_data(friendli_data, local_data) # Combine both datasets - all_remote_data = {**openrouter_data, **vercel_data} + all_remote_data = {**openrouter_data, **vercel_data, **friendli_data} # If both local and openrouter data are available, synchronize and save if local_data and all_remote_data: - sync_local_data_with_remote(local_data, all_remote_data) + sync_local_data_with_remote(local_data, all_remote_data, replace_keys=frozenset(friendli_data)) write_to_file(local_file_path, local_data) else: print("Failed to fetch model data from either local file or URL.") diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index c7e1b94a2ef..4899b87da7a 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -27,10 +27,13 @@ import litellm from litellm import Router, verbose_logger from litellm._uuid import uuid from litellm.caching.caching import DualCache +from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_metadata, ) +from openai.types.file_deleted import FileDeleted + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, @@ -48,7 +51,6 @@ from litellm.proxy._types import ( from litellm.proxy.openai_files_endpoints.common_utils import ( BATCH_CREATE_HIDDEN_PARAM, FILE_LIST_CONTINUATION_CHUNK_SIZE, - MAX_FILE_LIST_LIMIT, _is_base64_encoded_unified_file_id, apply_unified_file_ids, decode_model_from_file_id, @@ -1787,7 +1789,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span: Optional[Span], llm_router: Router, **data: Dict, - ) -> OpenAIFileObject: + ) -> FileDeleted: # Check if file deletion should be blocked due to batch references await self._check_file_deletion_allowed(file_id) @@ -1795,7 +1797,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # file_id = convert_b64_uid_to_unified_uid(file_id) model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) - delete_response = None specific_model_file_id_mapping = model_file_id_mapping.get(file_id) if specific_model_file_id_mapping: # Remove conflicting keys from data to avoid duplicate keyword arguments @@ -1810,23 +1811,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): else {} ), } - delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data) + await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data) - stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) + await self.delete_unified_file_id(file_id, litellm_parent_otel_span) - # Record successful deletion metric only on actual success - if stored_file_object or delete_response: - prom_logger = self._get_prometheus_logger() - if prom_logger: - prom_logger.record_managed_file_deleted(result="success") - - if stored_file_object: - return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id}) - elif delete_response: - delete_response.id = file_id - return delete_response - else: - raise Exception(f"LiteLLM Managed File object with id={file_id} not found") + prom_logger = self._get_prometheus_logger() + if prom_logger: + prom_logger.record_managed_file_deleted(result="success") + return FileDeleted(id=file_id, object="file", deleted=True) async def afile_content( self, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql new file mode 100644 index 00000000000..3bf6b819715 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120000_add_litellm_call_id_spend_logs/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "litellm_call_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql new file mode 100644 index 00000000000..62ad5c42ba7 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql @@ -0,0 +1,12 @@ +-- CreateIndex (CONCURRENTLY) +-- +-- Disclaimer: +-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a +-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction. +-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is +-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated. +-- - Do not edit this file after it has been applied to any database: Prisma checksums +-- migrations; add a new migration instead. +-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration +-- without IF NOT EXISTS if you must support older versions). +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7d521d54791..dd7967aafe3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -659,12 +659,14 @@ model LiteLLM_SpendLogs { mcp_namespaced_tool_name String? agent_id String? proxy_server_request Json? @default("{}") + litellm_call_id String? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([litellm_call_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/constants.py b/litellm/constants.py index be5a4e4b076..5751e6e46af 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -53,6 +53,7 @@ S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64 S3_PREFIX_DIGEST_CHARS: Final = 16 # s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 +MAX_FILE_LIST_LIMIT: Final = 10000 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) budget_reservation_disabled_info_emitted = False @@ -143,6 +144,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3) ) MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) +MAX_LITELLM_CALL_ID_LENGTH: Final = 256 MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048 DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000 diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 22bdb016dc1..440d97d13be 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -481,7 +481,8 @@ def cost_per_token( else: model_with_provider = f"{custom_llm_provider}/{model}" if region_name is not None: - model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{model}" + bare_model: Final = model[len(_prov_prefix) :] if model_is_str and model.startswith(_prov_prefix) else model + model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{bare_model}" if model_with_provider_and_region in model_cost_ref: # use region based pricing, if it's available model_with_provider = model_with_provider_and_region else: @@ -778,6 +779,7 @@ def _select_model_name_for_cost_calc( custom_pricing: bool | None = None, custom_llm_provider: str | None = None, router_model_id: str | None = None, + region_name: str | None = None, ) -> str | None: """ 1. If custom pricing is true, return received model name @@ -799,8 +801,8 @@ def _select_model_name_for_cost_calc( provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model") explicit_pricing: Final = custom_pricing is True or base_model is not None priced_from_response: Final = provider_response_model is not None or completion_response_model is not None - region_name: Final = ( - _get_hidden_str_for_cost_calc(hidden_params, "region_name") + priced_region: Final = ( + _get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name if not explicit_pricing and priced_from_response else None ) @@ -837,8 +839,10 @@ def _select_model_name_for_cost_calc( and custom_llm_provider is not None and not _model_contains_known_llm_provider(return_model) ): # add provider prefix if not already present, to match model_cost - provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}" - return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", region_name) + provider_prefix: Final = ( + custom_llm_provider if priced_region is None else f"{custom_llm_provider}/{priced_region}" + ) + return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", priced_region) return return_model @@ -1300,6 +1304,7 @@ def completion_cost( service_tier = _normalize_service_tier(service_tier) + explicit_pricing: Final = custom_pricing is True or base_model is not None selected_model: Final = _select_model_name_for_cost_calc( model=model, completion_response=completion_response, @@ -1307,6 +1312,7 @@ def completion_cost( custom_pricing=custom_pricing, base_model=base_model, router_model_id=router_model_id, + region_name=region_name, ) potential_model_names: Final = [ @@ -1651,7 +1657,7 @@ def completion_cost( completion_tokens=completion_tokens or 0, custom_llm_provider=custom_llm_provider, response_time_ms=total_time, - region_name=region_name, + region_name=None if explicit_pricing else region_name, custom_cost_per_second=custom_cost_per_second, custom_cost_per_token=custom_cost_per_token, prompt_characters=prompt_characters, @@ -1861,6 +1867,7 @@ def response_cost_calculator( data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ### VERTEX LOCATION ### vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") + region_name: str | None = None, ) -> float: """ Returns @@ -1894,6 +1901,7 @@ def response_cost_calculator( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + region_name=region_name, ) return response_cost except Exception as e: diff --git a/litellm/files/main.py b/litellm/files/main.py index 218518eb3cd..1d5da29fe6f 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -682,6 +682,10 @@ def file_list( ) if provider_config is not None: litellm_params_dict: Final = get_litellm_params(**kwargs) + add_trusted_model_credentials_to_litellm_params( + litellm_params_dict=litellm_params_dict, + kwargs=kwargs, + ) litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index ce51fb19970..02681d8b499 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -601,12 +601,15 @@ def _get_openai_compatible_provider_info( dynamic_api_key, ) = litellm.GroqChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "bedrock_mantle": + from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix + ( api_base, dynamic_api_key, ) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info( api_base, api_key, litellm_params=litellm_params, model=model ) + model = split_mantle_region_prefix(model)[1] # rebind-ok: the prefix is routing only, not a Mantle model id elif custom_llm_provider == "nvidia_nim": # nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 api_base = api_base or get_secret("NVIDIA_NIM_API_BASE") or "https://integrate.api.nvidia.com/v1" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2d3a99abe81..5fe94320491 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -454,6 +454,17 @@ def _resolve_vertex_location_for_cost( return VertexBase.get_vertex_region(configured_location, model) +def _resolve_mantle_region_for_cost( + custom_llm_provider: str | None, + litellm_params: Mapping[str, object] | None, +) -> str | None: + if custom_llm_provider != "bedrock_mantle": + return None + from litellm.llms.bedrock_mantle.common_utils import resolve_mantle_region + + return resolve_mantle_region(litellm_params or MappingProxyType({})) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -1768,6 +1779,10 @@ class Logging(LiteLLMLoggingBaseClass): optional_params=self.optional_params, model=litellm_model_name or self.model, ), + "region_name": _resolve_mantle_region_for_cost( + custom_llm_provider=self.model_call_details.get("custom_llm_provider", None), + litellm_params=self.model_call_details.get("litellm_params"), + ), } except Exception as e: # error creating kwargs for cost calculation debug_info = StandardLoggingModelCostFailureDebugInformation( diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 9158ff4569f..e7179aad25b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -19,6 +19,8 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.exceptions import MidStreamFallbackError +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic import ( AppliedEdit, CompactionBlock, @@ -58,6 +60,25 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: return value if value else () +def _error_status_and_message(exc: Exception) -> tuple[int, str]: + if isinstance(exc, (BaseLLMException, MidStreamFallbackError)): + return exc.status_code, exc.message + return 500, str(exc) or "Upstream stream ended before completion" + + +def _mid_stream_error_sse_event(exc: Exception) -> bytes: + from litellm.anthropic_interface.exceptions.exception_mapping_utils import ( + AnthropicExceptionMapping, + ) + + status_code, message = _error_status_and_message(exc) + error_response = AnthropicExceptionMapping.transform_to_anthropic_error( + status_code=status_code, + raw_message=message, + ) + return f"event: error\ndata: {json.dumps(error_response)}\n\n".encode() + + def _delta_payload_field(delta_type: StreamingContentBlockDeltaType) -> str: match delta_type: case "text_delta": @@ -990,14 +1011,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): Async version of anthropic_sse_wrapper. Convert AnthropicStreamWrapper dict chunks to Server-Sent Events format. """ - async for chunk in self: - if isinstance(chunk, dict): - event_type: str = str(chunk.get("type", "message")) - payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" - yield payload.encode() - else: - # For non-dict chunks, forward the original value unchanged - yield chunk + try: + async for chunk in self: + if isinstance(chunk, dict): + event_type: str = str(chunk.get("type", "message")) + payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n" + yield payload.encode() + else: + yield chunk + except Exception as e: # noqa: BLE001 # boundary before the socket: any upstream failure becomes an Anthropic error event + verbose_logger.exception("Anthropic Adapter - mid-stream error, emitting Anthropic error event: %s", e) + yield _mid_stream_error_sse_event(e) def _increment_content_block_index(self): self.current_content_block_index += 1 diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 7a7088c2fb5..6d16a1cea69 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Any, Union import httpx @@ -160,6 +160,15 @@ class BaseFilesConfig(BaseConfig): ) -> tuple[str, dict]: """Transform file list request into provider-specific format.""" + def transform_list_files_next_request( + self, + raw_response: httpx.Response, + optional_params: Mapping[str, object], + litellm_params: dict, # mutable-ok: carries provider stashes from the request transform to the response one + ) -> tuple[str, dict[str, str]] | None: + """Request for the page after `raw_response`, or None once the listing is complete.""" + return None + @abstractmethod def transform_list_files_response( self, @@ -258,7 +267,7 @@ class BaseFileEndpoints(ABC): litellm_parent_otel_span: Span | None, llm_router: Router, **data: dict, - ) -> OpenAIFileObject: + ) -> FileDeleted: pass @abstractmethod diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 9875ac2b9c3..6e2b0c12090 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1,14 +1,18 @@ import base64 import json import os +import posixpath import time +import xml.etree.ElementTree as ET from collections.abc import Iterable, Mapping, MutableMapping, Sequence from contextlib import suppress +from dataclasses import dataclass +from datetime import datetime from functools import cache from itertools import chain from types import MappingProxyType from typing import Any, Final, Literal, TypeAlias, TypedDict -from urllib.parse import unquote +from urllib.parse import quote, unquote, urlencode import httpx from httpx import Headers, Response @@ -23,6 +27,7 @@ from litellm.files.utils import FilesAPIUtils from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, + BEDROCK_MANAGED_S3_OUTPUT_PREFIX, BEDROCK_MANAGED_S3_PREFIXES, BEDROCK_MANAGED_S3_UPLOAD_PREFIX, build_managed_cloud_object_name, @@ -62,6 +67,10 @@ from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resol S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" +LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose" + +LIST_FILES_LOCATION_PARAM: Final = "_s3_list_files_location" + class _S3DeleteContext(BaseModel): file_id: str = Field(min_length=1) @@ -152,6 +161,13 @@ class _BedrockS3RequestParams(BaseModel): s3_endpoint_url: str | None = None +@dataclass(frozen=True, slots=True) +class _S3RequestTarget: + endpoint_url: str + aws_region_name: str + request_params: _BedrockS3RequestParams + + class _TrustedS3ModelCredentials(BaseModel): """The S3 buckets the server trusts file ids against, from the deployment snapshot.""" @@ -248,6 +264,128 @@ def _validate_file_id_against_configured_buckets( return validate_against(configured_bucket_names[-1]) +_REJECTED_FILE_ID_REQUEST_URL: Final = "https://litellm.ai" + + +def _rejected_file_id(reason: ValueError) -> BedrockError: + message: Final = str(reason) + return BedrockError( + status_code=400, + message=message, + response=httpx.Response( + status_code=400, + text=message, + request=httpx.Request(method="GET", url=_REJECTED_FILE_ID_REQUEST_URL), + ), + ) + + +def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: + configured_bucket_names: Final = get_configured_s3_bucket_names(litellm_params) + allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params) + try: + return _validate_file_id_against_configured_buckets( + s3_uri=extract_s3_uri_from_file_id(file_id), + configured_bucket_names=configured_bucket_names, + allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, + ) + except ValueError as reason: + raise _rejected_file_id(reason) from reason + + +_ANY_MANAGED_LISTING_PREFIX: Final = os.path.commonprefix(BEDROCK_MANAGED_S3_PREFIXES) +_MANAGED_LISTING_PREFIX_BY_PURPOSE: Final = MappingProxyType( + { + "batch": os.path.commonprefix((BEDROCK_MANAGED_S3_BATCH_PREFIX, BEDROCK_MANAGED_S3_UPLOAD_PREFIX)), + "batch_output": BEDROCK_MANAGED_S3_OUTPUT_PREFIX, + } +) + + +_EMPTY_LISTING_QUERY: Final = (("list-type", "2"), ("max-keys", "0")) + + +def _managed_listing_prefix(configured_prefix: str, purpose: str | None) -> str | None: + managed_prefix: Final = _MANAGED_LISTING_PREFIX_BY_PURPOSE.get(purpose) if purpose else _ANY_MANAGED_LISTING_PREFIX + if managed_prefix is None: + return None + return f"{configured_prefix}/{managed_prefix}" if configured_prefix else managed_prefix + + +def _listing_query(configured_prefix: str, purpose: str | None) -> tuple[tuple[str, str], ...]: + listing_prefix: Final = _managed_listing_prefix(configured_prefix, purpose) + if listing_prefix is None: + return _EMPTY_LISTING_QUERY + return (("list-type", "2"), ("prefix", listing_prefix)) + + +def _requested_listing_purpose(litellm_params: Mapping[str, object]) -> str | None: + requested_purpose: Final = litellm_params.get(LIST_FILES_PURPOSE_PARAM) + return requested_purpose if isinstance(requested_purpose, str) else None + + +def _walked_listing_purpose(litellm_params: Mapping[str, object]) -> str | None: + walked_purpose: Final = litellm_params.get(LIST_FILES_LOCATION_PARAM) + return walked_purpose if isinstance(walked_purpose, str) else _requested_listing_purpose(litellm_params) + + +def _output_location_still_unlisted(litellm_params: Mapping[str, object]) -> bool: + if _walked_listing_purpose(litellm_params) is not None: + return False + return _listing_bucket_name(litellm_params, "batch_output") != _listing_bucket_name(litellm_params, None) + + +def _listing_bucket_name(litellm_params: Mapping[str, object], purpose: str | None) -> str: + if purpose != "batch_output": + return get_configured_s3_bucket_name(litellm_params) + trusted: Final = _trusted_s3_model_credentials(litellm_params) + return ( + trusted.s3_output_bucket_name + or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") + or get_configured_s3_bucket_name(litellm_params) + ) + + +def _listed_object_created_at(entry: ET.Element) -> int: + last_modified: Final = entry.findtext("{*}LastModified") + if not last_modified: + return 0 + return int(datetime.fromisoformat(last_modified.replace("Z", "+00:00")).timestamp()) + + +def _listed_managed_file( + entry: ET.Element, + bucket_name: str, + configured_bucket_name: str, + allow_legacy_cloud_file_ids: bool, +) -> OpenAIFileObject | None: + object_key: Final = entry.findtext("{*}Key") + if not object_key: + return None + file_id: Final = f"s3://{bucket_name}/{object_key}" + try: + validate_managed_cloud_file_id( + file_id=file_id, + scheme="s3://", + configured_bucket_name=configured_bucket_name, + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, + ) + except ValueError: + return None + _, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name) + relative_key: Final = object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key + return OpenAIFileObject( + id=file_id, + bytes=int(entry.findtext("{*}Size") or 0), + created_at=_listed_object_created_at(entry), + filename=posixpath.basename(object_key), + object="file", + purpose="batch_output" if relative_key.startswith(BEDROCK_MANAGED_S3_OUTPUT_PREFIX) else "batch", + status="uploaded", + ) + + def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int: """ S3 answers PutObject with an empty body, so the stored object size comes from the @@ -1213,18 +1351,86 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): def transform_list_files_request( self, purpose: str | None, - optional_params: dict, - litellm_params: dict, - ) -> tuple[str, dict]: - raise NotImplementedError("BedrockFilesConfig does not support file listing") + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + litellm_params[LIST_FILES_PURPOSE_PARAM] = purpose # rebind-ok: handed to the response transform + litellm_params[LIST_FILES_LOCATION_PARAM] = purpose # rebind-ok: names the location the next page walks + return self._signed_listing_request(purpose, optional_params, litellm_params, continuation_token=None) + + def transform_list_files_next_request( + self, + raw_response: httpx.Response, + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]] | None: + if raw_response.status_code >= 400: + return None + continuation_token: Final = ET.fromstring(raw_response.content).findtext("{*}NextContinuationToken") + if continuation_token: + return self._signed_listing_request( + _walked_listing_purpose(litellm_params), optional_params, litellm_params, continuation_token + ) + if not _output_location_still_unlisted(litellm_params): + return None + litellm_params[LIST_FILES_LOCATION_PARAM] = "batch_output" # rebind-ok: the input location is fully listed + return self._signed_listing_request("batch_output", optional_params, litellm_params, continuation_token=None) + + def _signed_listing_request( + self, + purpose: str | None, + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + continuation_token: str | None, + ) -> tuple[str, dict[str, str]]: + bucket_name, configured_prefix = split_configured_cloud_bucket_name( + _listing_bucket_name(litellm_params, purpose) + ) + target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params) + url: Final = f"{target.endpoint_url}/{bucket_name}/" + listing_query: Final = _listing_query(configured_prefix, purpose) + continuation_query: Final = (("continuation-token", continuation_token),) if continuation_token else () + query: Final[dict[str, str]] = dict( # mutable-ok: the base files contract returns the query as a dict + listing_query + continuation_query + ) + signed_headers: Final = self._sign_s3_request_without_body( + method="GET", + api_base=f"{url}?{urlencode(query, quote_via=quote, safe='')}", + aws_region_name=target.aws_region_name, + request_params=target.request_params, + ) + litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment + return url, query def transform_list_files_response( self, raw_response: httpx.Response, logging_obj: LiteLLMLoggingObj, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> list[OpenAIFileObject]: - raise NotImplementedError("BedrockFilesConfig does not support file listing") + if raw_response.status_code >= 400: + raise BedrockError( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + response=raw_response, + ) + purpose: Final = _requested_listing_purpose(litellm_params) + configured_bucket_name: Final = _listing_bucket_name(litellm_params, _walked_listing_purpose(litellm_params)) + allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params) + listing: Final = ET.fromstring(raw_response.content) + bucket_name: Final = ( + listing.findtext("{*}Name") or split_configured_cloud_bucket_name(configured_bucket_name)[0] + ) + listed_files: Final = ( + _listed_managed_file(entry, bucket_name, configured_bucket_name, allow_legacy_cloud_file_ids) + for entry in listing.iterfind("{*}Contents") + ) + return [ # mutable-ok: the base files contract returns a list + listed_file + for listed_file in listed_files + if listed_file is not None and (purpose is None or listed_file.purpose == purpose) + ] def transform_file_content_request( self, @@ -1255,39 +1461,54 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: MutableMapping[str, object], ) -> tuple[str, dict[str, str]]: - s3_uri: Final = extract_s3_uri_from_file_id(file_id) - bucket_name, object_key = _validate_file_id_against_configured_buckets( - s3_uri=s3_uri, - configured_bucket_names=get_configured_s3_bucket_names(litellm_params), - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), + bucket_name, object_key = _resolve_managed_s3_object(file_id=file_id, litellm_params=litellm_params) + target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params) + url: Final = f"{target.endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" + signed_headers: Final = self._sign_s3_request_without_body( + method=method, + api_base=url, + aws_region_name=target.aws_region_name, + request_params=target.request_params, ) + litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment + return url, {} # mutable-ok: the base files contract returns the query as a dict - request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params}) - + def _s3_request_target( + self, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + ) -> _S3RequestTarget: + """ + The shared files handler passes optional_params={}, so AWS credentials and + region arrive via litellm_params here (unlike the upload path). + s3_region_name wins over aws_region_name, same priority as get_complete_file_url. + """ + request_params: Final = _BedrockS3RequestParams.model_validate( + MappingProxyType({**litellm_params, **optional_params}) + ) region_preference: Final = request_params.s3_region_name or request_params.aws_region_name - region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference} - aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="") - - s3_endpoint_url: Final = ( + aws_region_name: Final = self._get_aws_region_name( + optional_params={"aws_region_name": region_preference}, # mutable-ok: BaseAWSLLM takes a dict + model="", + ) + endpoint_url: Final = ( request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" ).rstrip("/") - url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" - - litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body( - api_base=url, - aws_region_name=aws_region_name, - request_params=request_params, - method=method, + return _S3RequestTarget( + endpoint_url=endpoint_url, aws_region_name=aws_region_name, request_params=request_params ) - return url, {} def _sign_s3_request_without_body( self, + method: Literal["GET", "DELETE"], api_base: str, aws_region_name: str, request_params: _BedrockS3RequestParams, - method: Literal["GET", "DELETE"] = "GET", - ) -> dict[str, str]: + ) -> Mapping[str, str]: + """ + SigV4-sign a bodiless S3 request (GetObject, DeleteObject, ListObjectsV2), + mirroring `_sign_s3_request` (PUT). + """ try: import hashlib @@ -1313,11 +1534,11 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped method=method, url=api_base, - headers={"x-amz-content-sha256": empty_body_hash}, + headers={"x-amz-content-sha256": empty_body_hash}, # mutable-ok: botocore AWSRequest takes a dict ) auth: Final = S3SigV4Auth(credentials, "s3", aws_region_name) # any-ok: botocore untyped auth.add_auth(aws_request) # any-ok: botocore request mutation is untyped - return dict(aws_request.headers) # any-ok: botocore headers are untyped + return MappingProxyType(dict(aws_request.headers)) # any-ok: botocore headers are untyped def transform_file_content_response( self, @@ -1330,6 +1551,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): status_code=raw_response.status_code, message=raw_response.text, headers=raw_response.headers, + response=raw_response, ) return HttpxBinaryResponseContent(response=raw_response) diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index d91157c3d10..a1153dffc93 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -29,7 +29,7 @@ from litellm.types.router import GenericLiteLLMParams from ...base_llm.chat.transformation import BaseLLMException from ...bedrock.common_utils import BedrockError from ...openai_like.chat.transformation import OpenAILikeChatConfig -from ..common_utils import mantle_base_segment +from ..common_utils import mantle_base_segment, split_mantle_region_prefix class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): @@ -61,8 +61,10 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): litellm_params: GenericLiteLLMParams | None = None, model: str | None = None, ) -> tuple[str | None, str | None]: + prefix_region, base_model = split_mantle_region_prefix(model) if model else (None, None) region: Final = ( (litellm_params.aws_region_name if litellm_params else None) + or prefix_region or get_secret_str("BEDROCK_MANTLE_REGION") or get_secret_str("AWS_REGION_NAME") or get_secret_str("AWS_REGION") @@ -75,7 +77,7 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): api_base = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") - or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(model, litellm.model_cost)}" + or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(base_model, litellm.model_cost)}" ) dynamic_api_key: Final = self._resolve_bearer_token(api_key) return api_base, dynamic_api_key diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index ac94d9cb922..9892ef6224e 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -24,9 +24,11 @@ from botocore.exceptions import ( ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS +from litellm.llms.bedrock.common_utils import AmazonBedrockGlobalConfig from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" +BEDROCK_REGIONS: Final = frozenset(AmazonBedrockGlobalConfig().get_all_regions()) # Standard Mantle host: https://bedrock-mantle..api.aws (group 1 = region). MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws(?=/|$)", re.IGNORECASE) @@ -36,6 +38,13 @@ def resolve_mantle_bearer_token(api_key: str | None) -> str | None: return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") +def split_mantle_region_prefix(model: str) -> tuple[str | None, str]: + head, sep, tail = model.partition("/") + if sep and head in BEDROCK_REGIONS: + return head, tail + return None, model + + def resolve_mantle_region(params: Mapping[str, object]) -> str: region: Final = params.get("aws_region_name") if isinstance(region, str) and region: @@ -130,7 +139,7 @@ def mantle_supports_responses(model: str | None, model_cost: dict) -> bool: gpt-oss substring), so a substring gate would be wrong. A model absent from model_cost simply has no signal and returns False (chat-completions emulation). """ - entry: Final = model_cost.get(f"bedrock_mantle/{model}", {}) + entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {} if "/v1/responses" in (entry.get("supported_endpoints") or []): return True return entry.get("mode") == "responses" @@ -147,5 +156,5 @@ def mantle_base_segment(model: str | None, model_cost: dict) -> str: the base for the model's whole OpenAI-compatible surface, so both the chat and responses configs derive from it -- there is no separate model-name rule. """ - entry: Final = model_cost.get(f"bedrock_mantle/{model}", {}) + entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {} return "openai/v1" if entry.get("use_openai_responses_path") is True else "v1" diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7109e6942d1..2fe4130a310 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -18,7 +18,7 @@ import litellm.types import litellm.types.utils 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.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, @@ -4981,15 +4981,16 @@ class BaseLLMHTTPHandler: ) try: - response: Final = sync_httpx_client.get(url=url, headers=headers, params=params) + response: Final = sync_httpx_client.get(url=url, headers=headers, params=params, timeout=timeout) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) - return provider_config.transform_list_files_response( - raw_response=response, - logging_obj=logging_obj, - litellm_params=litellm_params, + files_per_page: Final = self._files_per_listing_page( + response, provider_config, logging_obj, litellm_params, headers, sync_httpx_client, timeout ) + return [ # mutable-ok: the files contract returns the listing as a list + listed_file for page_files in files_per_page for listed_file in page_files + ] async def async_list_files( self, @@ -5037,16 +5038,101 @@ class BaseLLMHTTPHandler: ) try: - response: Final = await async_httpx_client.get(url=url, headers=headers, params=params) + response: Final = await async_httpx_client.get(url=url, headers=headers, params=params, timeout=timeout) except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) - return provider_config.transform_list_files_response( - raw_response=response, - logging_obj=logging_obj, + files_per_page: Final = self._files_per_async_listing_page( + response, provider_config, logging_obj, litellm_params, headers, async_httpx_client, timeout + ) + return [ # mutable-ok: the files contract returns the listing as a list + listed_file async for page_files in files_per_page for listed_file in page_files + ] + + def _files_per_listing_page( + self, + first_page: httpx.Response, + provider_config: BaseFilesConfig, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict + client: HTTPHandler, + timeout: float | httpx.Timeout | None, + ) -> Iterator[list[OpenAIFileObject]]: # mutable-ok: each page arrives as the list the files contract returns + latest_page = first_page # rebind-ok: advances one page per loop turn + listed_count = 0 # rebind-ok: grows per page so the listing stops at MAX_FILE_LIST_LIMIT, OpenAI's ceiling + while True: + page_files = provider_config.transform_list_files_response( + raw_response=latest_page, logging_obj=logging_obj, litellm_params=litellm_params + ) + yield page_files[: MAX_FILE_LIST_LIMIT - listed_count] + listed_count += len(page_files) + next_request = self._next_listing_request(latest_page, provider_config, litellm_params, listed_count) + if next_request is None: + return + url, params = next_request + next_headers = self._next_listing_page_headers(provider_config, headers, litellm_params) + try: + latest_page = client.get(url=url, headers=next_headers, params=params, timeout=timeout) + except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the first page's fetch + raise self._handle_error(e=e, provider_config=provider_config) + + async def _files_per_async_listing_page( + self, + first_page: httpx.Response, + provider_config: BaseFilesConfig, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict + client: AsyncHTTPHandler, + timeout: float | httpx.Timeout | None, + ) -> AsyncIterator[list[OpenAIFileObject]]: # mutable-ok: each page arrives as the list the files contract returns + latest_page = first_page # rebind-ok: advances one page per loop turn + listed_count = 0 # rebind-ok: grows per page so the listing stops at MAX_FILE_LIST_LIMIT, OpenAI's ceiling + while True: + page_files = provider_config.transform_list_files_response( + raw_response=latest_page, logging_obj=logging_obj, litellm_params=litellm_params + ) + yield page_files[: MAX_FILE_LIST_LIMIT - listed_count] + listed_count += len(page_files) + next_request = self._next_listing_request(latest_page, provider_config, litellm_params, listed_count) + if next_request is None: + return + url, params = next_request + next_headers = self._next_listing_page_headers(provider_config, headers, litellm_params) + try: + latest_page = await client.get(url=url, headers=next_headers, params=params, timeout=timeout) + except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the first page's fetch + raise self._handle_error(e=e, provider_config=provider_config) + + def _next_listing_page_headers( + self, + provider_config: BaseFilesConfig, + headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + ) -> dict: # mutable-ok: validate_environment returns the header dict the files contract types + return provider_config.validate_environment( + api_key=litellm_params.get("api_key"), + headers=headers, + model="", + messages=[], + optional_params={}, litellm_params=litellm_params, ) + def _next_listing_request( + self, + latest_page: httpx.Response, + provider_config: BaseFilesConfig, + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + listed_count: int, + ) -> tuple[str, dict[str, str]] | None: # mutable-ok: the base files contract returns the query as a dict + if listed_count >= MAX_FILE_LIST_LIMIT: + return None + return provider_config.transform_list_files_next_request( + raw_response=latest_page, optional_params={}, litellm_params=litellm_params + ) + def retrieve_file_content( self, file_content_request: "FileContentRequest", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 06c7a6aa46e..2220d0e1fe5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23083,58 +23083,206 @@ }, "friendliai/zai-org/GLM-5.3-Flash": { "litellm_provider": "friendliai", - "supports_reasoning": true, - "supports_function_calling": true, "max_input_tokens": 1048576, - "max_tokens": 1048576, "max_output_tokens": 1048576, + "max_tokens": 1048576, "input_cost_per_token": 1.5e-07, "output_cost_per_token": 5e-07, "cache_read_input_token_cost": 3e-08, "supports_prompt_caching": true, + "supports_reasoning": true, "reasoning_effort_levels": [ "low", "high", "max" ], + "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, - "mode": "chat", - "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", - "source": "https://api.friendli.ai/serverless/v1/models", "supports_vision": true, "supports_image_input": true, - "supports_video_input": true + "supports_video_input": true, + "mode": "chat", + "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", + "source": "https://api.friendli.ai/serverless/v1/models" }, "friendliai/zai-org/GLM-5.3": { "litellm_provider": "friendliai", - "supports_reasoning": true, - "supports_function_calling": true, "max_input_tokens": 1048576, - "max_tokens": 1048576, "max_output_tokens": 1048576, + "max_tokens": 1048576, "input_cost_per_token": 1.26e-06, "output_cost_per_token": 3.96e-06, "cache_read_input_token_cost": 2.34e-07, "supports_prompt_caching": true, + "supports_reasoning": true, "reasoning_effort_levels": [ "low", "high", "max" ], + "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, "mode": "chat", "comment": "Flagship GLM model for long-horizon coding, agents, and complex project delivery", - "source": "https://api.friendli.ai/serverless/v1/models", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/google/gemma-4-31B-it": { + "litellm_provider": "friendliai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4e-07, + "supports_prompt_caching": false, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_image_input": true, + "supports_video_input": false, + "mode": "chat", + "comment": "Largest Gemma 4 instruction model for open, self-hosted chat and reasoning", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/zai-org/GLM-5.2": { + "litellm_provider": "friendliai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [ + "high", + "max" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, "supports_vision": false, - "supports_image_input": false + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/LGAI-EXAONE/K-EXAONE-2.0-750B-A37B": { + "litellm_provider": "friendliai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1.2e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Frontier-scale multilingual language model developed by LG AI Research", + "deprecation_date": "2026-09-06", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/deepseek-ai/DeepSeek-V3.2": { + "litellm_provider": "friendliai", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 2.5e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "DeepSeek chat model for instruction following, coding, and analysis", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "friendliai", + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "max_tokens": 196608, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-08, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/zai-org/GLM-5.1": { + "litellm_provider": "friendliai", + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "max_tokens": 202752, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Strong GLM coding model for agentic engineering, terminals, and repository generation", + "source": "https://api.friendli.ai/serverless/v1/models" }, "ft:babbage-002": { "deprecation_date": "2026-10-23", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ae6c042ab3a..b315a2beac9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3882,6 +3882,7 @@ class SpendLogsPayload(TypedDict): session_id: str | None request_duration_ms: int | None status: Literal["success", "failure"] + litellm_call_id: ReadOnly[str | None] class SpanAttributes(str, enum.Enum): @@ -4580,7 +4581,6 @@ class UserManagementEndpointParamDocStringEnums(str, enum.Enum): ) metadata_doc_str = """Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }""" max_parallel_requests_doc_str = """Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.""" - soft_budget_doc_str = """Optional[float] - Get alerts when user crosses given budget, doesn't block requests.""" model_max_budget_doc_str = """Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)""" model_rpm_limit_doc_str = """Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" model_tpm_limit_doc_str = """Optional[float] - Model-specific tpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)""" diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 9c175242a9a..576585ee9a3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -72,6 +72,7 @@ from litellm.proxy.auth.budget_throttle import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, @@ -80,6 +81,7 @@ from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, + NO_TEAM_MEMBERSHIP_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, end_user_cache_key, @@ -2164,10 +2166,10 @@ async def get_team_membership( _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) # check if in cache - cached_membership_obj: Final = await user_api_key_cache.async_get_cache( - key=_key, - model_type=LiteLLM_TeamMembership, - ) + cached: Final[object] = await user_api_key_cache.async_get_cache(key=_key) + if cached == NO_TEAM_MEMBERSHIP_SENTINEL: + return None + cached_membership_obj: Final = CacheCodec.deserialize(cached, model_type=LiteLLM_TeamMembership) if cached_membership_obj is not None: return cached_membership_obj @@ -2179,6 +2181,11 @@ async def get_team_membership( ) if response is None: + await user_api_key_cache.async_set_cache( + key=_key, + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ttl=get_management_object_ttl(user_api_key_cache), + ) return None _response: Final = LiteLLM_TeamMembership.model_validate(response.dict()) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 69091ee8344..4304542fc83 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -52,6 +52,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import can_team_access_model +from litellm.proxy.auth.resolvers.grants import GrantResolver, UserLookup, canonical_user_id from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.common_utils.user_api_key_cache import ( @@ -1656,9 +1657,7 @@ class JWTAuthManager: ``get_user_object`` resolved a legacy row with a different ``user_id``, use that row's id; otherwise keep the claim. GH #26789. """ - if user_object is not None and user_object.user_id: - return user_object.user_id - return user_id + return canonical_user_id(user_id=user_id, user_object=user_object) @staticmethod async def get_objects( @@ -1725,22 +1724,23 @@ class JWTAuthManager: code=403, ) - user_object: LiteLLM_UserTable | None = None - if user_id: - user_object = ( - await get_user_object( - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email), - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - user_email=user_email, - sso_user_id=user_id, - ) - if user_id - else None - ) + user_object, team_membership_object, effective_user_id = await GrantResolver( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + load_user=get_user_object, + load_team=get_team_object, + load_membership=get_team_membership, + ).resolve_identity( + UserLookup( + user_id=user_id, + user_email=user_email, + sso_user_id=user_id, + upsert=jwt_handler.is_upsert_user_id(valid_user_email=valid_user_email), + ), + team_id=team_id, + ) end_user_object: LiteLLM_EndUserTable | None = None if end_user_id: @@ -1757,37 +1757,12 @@ class JWTAuthManager: else None ) - # Rebind to resolved DB user_id for team_membership + auth_builder (GH #26789). - effective_user_id: Final = JWTAuthManager._canonical_user_id_from_db(user_id=user_id, user_object=user_object) - if effective_user_id != user_id: - verbose_proxy_logger.debug( - "JWT Auth: rebinding user_id %r -> DB user_id %r (email/sso match)", - user_id, - effective_user_id, - ) - user_id = effective_user_id - - team_membership_object: LiteLLM_TeamMembership | None = None - if user_id and team_id: - team_membership_object = ( - await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - if user_id and team_id - else None - ) - return ( user_object, org_object, end_user_object, team_membership_object, - user_id, + effective_user_id, ) @staticmethod diff --git a/litellm/proxy/auth/resolvers/grants.py b/litellm/proxy/auth/resolvers/grants.py new file mode 100644 index 00000000000..eb39d2a6812 --- /dev/null +++ b/litellm/proxy/auth/resolvers/grants.py @@ -0,0 +1,267 @@ +"""Load a caller's user row, team row, and team membership from the database and validate them together. + +The virtual-key path reads these off the combined-view SQL join. Every other credential (an IdP JWT, a +``lite login`` session token) carries only identifiers, or a snapshot of grants taken when it was minted, so +it has to read the live rows on each request. Both of those paths resolve the same rows with the same +membership rule, and ``GrantResolver`` is the one place that rule lives. +""" + +from __future__ import annotations + +from collections.abc import Coroutine, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, NoReturn, Protocol, TypeAlias + +from fastapi import HTTPException, status +from pydantic import BaseModel, ValidationError +from pydantic.main import IncEx +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + ProxyErrorTypes, + ProxyException, +) +from litellm.proxy.auth.auth_checks import ( + TeamNotFoundError, + UserNotFoundError, + get_team_membership, + get_team_object, + get_user_object, +) + +if TYPE_CHECKING: + from litellm.proxy._types import Span + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class UserLoader(Protocol): + def __call__( + self, + *, + user_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + user_id_upsert: bool, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, + sso_user_id: str | None, + user_email: str | None, + ) -> Coroutine[object, object, LiteLLM_UserTable | None]: ... + + +class TeamLoader(Protocol): + def __call__( + self, + *, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, + ) -> Coroutine[object, object, LiteLLM_TeamTableCachedObj]: ... + + +class MembershipLoader(Protocol): + def __call__( + self, + *, + user_id: str, + team_id: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, + ) -> Coroutine[object, object, LiteLLM_TeamMembership | None]: ... + + +@dataclass(frozen=True, slots=True) +class UserLookup: + """The user a credential names, plus the hints ``get_user_object`` may fall back to when the id alone + matches no row.""" + + user_id: str | None + user_email: str | None = None + sso_user_id: str | None = None + upsert: bool = False + + +@dataclass(frozen=True, slots=True) +class ResolvedGrants: + """The live rows behind a credential. ``effective_user_id`` is the DB row's id when a fuzzy match found a + legacy row under a different id (GH #26789), otherwise the id the credential named.""" + + user_object: LiteLLM_UserTable | None + team_object: LiteLLM_TeamTableCachedObj | None + team_membership: LiteLLM_TeamMembership | None + effective_user_id: str | None + + +@dataclass(frozen=True, slots=True) +class UserGone: + user_id: str + + +@dataclass(frozen=True, slots=True) +class TeamGone: + team_id: str + + +@dataclass(frozen=True, slots=True) +class NotAMember: + user_id: str + team_id: str + + +@dataclass(frozen=True, slots=True) +class LookupDegraded: + """A row could not be read for a reason that says nothing about the caller: the database is down or a + loader failed. The caller decides whether a grant it already holds may stand in.""" + + error: Exception + + +GrantDenial: TypeAlias = UserGone | TeamGone | NotAMember +GrantOutcome: TypeAlias = ResolvedGrants | GrantDenial | LookupDegraded + + +_MODELS_COLUMN: Final[Mapping[str, IncEx | bool]] = MappingProxyType({"models": True}) + + +class _UserModelColumn(BaseModel): + """``LiteLLM_UserTable.models`` is a bare ``list``; re-read it with the shape a token's ``models`` takes.""" + + models: tuple[str, ...] = () + + +def user_models(user_object: LiteLLM_UserTable) -> tuple[str, ...]: + try: + return _UserModelColumn.model_validate(user_object.model_dump(include=_MODELS_COLUMN)).models + except ValidationError: + return () + + +def canonical_user_id(user_id: str | None, user_object: LiteLLM_UserTable | None) -> str | None: + if user_object is not None and user_object.user_id: + return user_object.user_id + return user_id + + +def raise_public(denial: GrantDenial) -> NoReturn: + match denial: + case UserGone(user_id=user_id): + raise ProxyException( + message=f"Authentication Error, user '{user_id}' no longer exists.", + type=ProxyErrorTypes.auth_error, + param="user_id", + code=status.HTTP_401_UNAUTHORIZED, + ) + case TeamGone(team_id=team_id): + raise TeamNotFoundError(team_id=team_id) + case NotAMember(team_id=team_id): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Team '{team_id}' is not in your team memberships.", + ) + case _: + assert_never(denial) + + +class GrantResolver: + """Reads the user, membership, and team rows for a credential through injected loaders. + + The loaders default to the shared ``auth_checks`` readers. A caller passes its own module's names for them + so the reads stay interceptable where that module's callers already intercept them. ``resolve_identity`` + is the JWT half: user and membership only, since the JWT builder selects the team itself and lets loader + errors surface as they are. ``resolve`` also reads the team row and applies the membership rule, which is + what a credential carrying a grant snapshot needs to refresh it. + """ + + def __init__( + self, + prisma_client: PrismaClient | None, + cache: UserApiKeyCache, + *, + parent_otel_span: Span | None = None, + proxy_logging_obj: ProxyLogging | None = None, + load_user: UserLoader = get_user_object, + load_team: TeamLoader = get_team_object, + load_membership: MembershipLoader = get_team_membership, + ) -> None: + self._prisma = prisma_client + self._cache = cache + self._parent_otel_span = parent_otel_span + self._proxy_logging_obj = proxy_logging_obj + self._load_user = load_user + self._load_team = load_team + self._load_membership = load_membership + + async def resolve_identity( + self, lookup: UserLookup, team_id: str | None + ) -> tuple[LiteLLM_UserTable | None, LiteLLM_TeamMembership | None, str | None]: + user_object: Final = await self._user(lookup) if lookup.user_id else None + effective_user_id: Final = canonical_user_id(lookup.user_id, user_object) + if effective_user_id != lookup.user_id: + verbose_proxy_logger.debug( + "Auth: rebinding user_id %r -> DB user_id %r (email/sso match)", + lookup.user_id, + effective_user_id, + ) + membership: Final = ( + await self._membership(user_id=effective_user_id, team_id=team_id) + if effective_user_id and team_id + else None + ) + return user_object, membership, effective_user_id + + async def resolve(self, lookup: UserLookup, team_id: str | None) -> GrantOutcome: + try: + user_object, membership, effective_user_id = await self.resolve_identity(lookup, team_id) + except UserNotFoundError: + return UserGone(user_id=lookup.user_id or "") + except Exception as error: + return LookupDegraded(error=error) + if team_id is None: + return ResolvedGrants(user_object, None, membership, effective_user_id) + if user_object is not None and team_id not in user_object.teams: + return NotAMember(user_id=user_object.user_id, team_id=team_id) + try: + team_object: Final = await self._load_team( + team_id=team_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) + except TeamNotFoundError: + return TeamGone(team_id=team_id) + except Exception as error: + return LookupDegraded(error=error) + return ResolvedGrants(user_object, team_object, membership, effective_user_id) + + async def _user(self, lookup: UserLookup) -> LiteLLM_UserTable | None: + return await self._load_user( + user_id=lookup.user_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + user_id_upsert=lookup.upsert, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + user_email=lookup.user_email, + sso_user_id=lookup.sso_user_id, + ) + + async def _membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None: + return await self._load_membership( + user_id=user_id, + team_id=team_id, + prisma_client=self._prisma, + user_api_key_cache=self._cache, + parent_otel_span=self._parent_otel_span, + proxy_logging_obj=self._proxy_logging_obj, + ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9828311112e..687f36bbe8b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -56,6 +56,7 @@ from litellm.proxy.auth.auth_checks import ( get_jwt_key_mapping_object, get_object_permission, get_project_object, + get_team_membership, get_team_object, get_user_object, is_valid_fallback_model, @@ -82,6 +83,14 @@ from litellm.proxy.auth.network import TrustedProxyConfig, resolve_network_conte from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.resolvers import CredentialRef, Principal +from litellm.proxy.auth.resolvers.grants import ( + GrantResolver, + LookupDegraded, + ResolvedGrants, + UserLookup, + raise_public, + user_models, +) from litellm.proxy.auth.resolvers.store import IdentityStore from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.team_grants import team_grants @@ -1230,6 +1239,52 @@ async def _record_unparsable_body_failure( verbose_proxy_logger.exception("Failed to log the request rejected for an unparsable body: %s", e) +async def _refresh_session_token_grants( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging, +) -> UserAPIKeyAuth: + """Rebuild a ``lite login`` session token's grants from the live user and team rows. + + The blob only proves who logged in and which team they picked. Team models, aliases, the user's own model + list, and their role are re-read every request, so a `/team/update` or a demotion shows up without a + re-login, and a user removed from the team or deleted outright is refused. When a row cannot be read for + a reason unrelated to the caller, the minted grants stand in exactly as they did before this refresh. + """ + outcome: Final = await GrantResolver( + prisma_client, + user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + load_user=get_user_object, + load_team=get_team_object, + load_membership=get_team_membership, + ).resolve(UserLookup(user_id=valid_token.user_id), team_id=valid_token.team_id) + match outcome: + case ResolvedGrants( + user_object=LiteLLM_UserTable() as user_object, team_object=team_object, team_membership=team_membership + ): + return UserAPIKeyAuth.model_validate( + MappingProxyType( + { + **valid_token.model_dump(exclude_none=True), + **team_grants(team_object, team_membership, user_object.user_id), + "user_role": _get_user_role(user_object), + "models": () if team_object is not None else user_models(user_object), + } + ) + ) + case ResolvedGrants(): + return valid_token + case LookupDegraded(error=error): + verbose_proxy_logger.debug("Session token grants not refreshed, keeping minted grants: %s", error) + return valid_token + case _: + raise_public(outcome) + + async def _resolve_object_permission_for_unresolvable_team( object_permission_id: str | None, prisma_client: PrismaClient | None, @@ -1774,6 +1829,15 @@ async def _user_api_key_auth_builder( ): valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(api_key) + if valid_token is not None and valid_token.is_session_token and prisma_client is not None: + valid_token = await _refresh_session_token_grants( # rebind-ok: later checks read this name + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if ( valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e6ed60ba177..fc603701e44 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -24,6 +24,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, LITELLM_DETAILED_TIMING, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, + MAX_LITELLM_CALL_ID_LENGTH, MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, NON_INFERENCE_CALL_TYPES, RETURN_RAW_MODEL_NAME_METADATA_KEY, @@ -217,6 +218,12 @@ def _withheld_provider_output(response: object) -> bool: return getattr(response, "has_buffered_provider_output", False) is True +def resolve_litellm_call_id(client_call_id: str | None) -> str: + if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH: + return client_call_id + return str(uuid.uuid4()) + + def _should_return_raw_model_name(request_data: dict[str, object]) -> bool: return any( isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True @@ -1938,7 +1945,7 @@ class ProxyBaseLLMRequestProcessing: if alias_target is not None: self.data["model"] = alias_target - self.data["litellm_call_id"] = request.headers.get("x-litellm-call-id", str(uuid.uuid4())) + self.data["litellm_call_id"] = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) DDSpanTagger.tag_call_id(self.data.get("litellm_call_id")) DDSpanTagger.tag_request( user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index cb72088ee4a..1c7a379897f 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -336,6 +336,14 @@ def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: return f"team_membership:{user_id}:{team_id}" +#: Cached under ``team_membership_reservation_cache_key`` when a member has no ``LiteLLM_TeamMembership`` +#: row, so a session-token member without a per-member budget costs no DB read per request. Lives beside +#: the key builder because it is part of the same cache protocol: every reader of the key must know that +#: a plain string here means "no row", distinct from a serialized membership. The two budget readers +#: already treat a non-model value as "no row", so they need no change to stay correct. +NO_TEAM_MEMBERSHIP_SENTINEL: Final = "__no_team_membership__" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index b9964c0e342..db6ec754c6e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -443,6 +443,11 @@ async def health_services_endpoint( } return pointfive_health if service == "webhook": + if not _is_proxy_admin(user_api_key_dict): + webhook_non_admin_detail: Final[_ServiceTestErrorDetail] = { + "error": "Only proxy admins can trigger the webhook test alert." + } + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=webhook_non_admin_detail) user_info: Final = CallInfo( token=user_api_key_dict.token or "", spend=1, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 000b7f874ee..e3efda507f6 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -478,7 +478,6 @@ async def new_user( - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -1651,7 +1650,6 @@ async def user_update( - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e9e37540dd8..0b7f69bbb7f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -14,7 +14,7 @@ import copy import json import math import traceback -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timezone from types import MappingProxyType @@ -1903,6 +1903,39 @@ def validate_team_org_change( return True +def _member_user_ids(members_with_roles: Sequence[dict[str, object]]) -> tuple[str, ...]: + """Extract the string ``user_id`` of each team member, dropping rows without one. + + ``members_with_roles`` is a Prisma-deserialized JSON column, so its ``user_id`` is typed + ``object``; the ``isinstance`` narrows it to the ``str`` ``invalidate_team_member_spend_state`` needs. + """ + return tuple(user_id for member in members_with_roles if isinstance((user_id := member.get("user_id")), str)) + + +async def _evict_created_membership_caches( + user_ids: Iterable[str], + team_id: str, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Evict the ``get_team_membership`` negative-cache sentinel for members whose row was just created. + + A session-token request caches ``NO_TEAM_MEMBERSHIP_SENTINEL`` for a member with no + ``LiteLLM_TeamMembership`` row. When a create path (``/team/member_add`` or the ``/team/update`` + budget backfill) later writes that row with a per-member budget, the stale sentinel keeps the + member's budget unenforced until the membership cache TTL expires, so it must be evicted here. + """ + await asyncio.gather( + *( + invalidate_team_member_spend_state( + user_id=user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + ) + for user_id in user_ids + ) + ) + + @router.post("/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @management_endpoint_wrapper async def update_team( @@ -2239,6 +2272,11 @@ async def update_team( team_member_budget_id=_backfill_budget_id, prisma_client=prisma_client, ) + await _evict_created_membership_caches( + user_ids=_member_user_ids(existing_team_row.members_with_roles), + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) elif _team_member_fields_in_request: updated_kv = await TeamMemberBudgetHandler.clear_team_member_budget_fields( team_table=existing_team_row, @@ -3191,6 +3229,12 @@ async def team_member_add( litellm_proxy_admin_name=litellm_proxy_admin_name, ) + await _evict_created_membership_caches( + user_ids=(tm.user_id for tm in updated_team_memberships), + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) + _emit_team_members_metric(complete_team_data) await _create_team_member_add_audit_logs( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b1f282a0978..38a907892b4 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -16,6 +16,7 @@ from typing import ( ) from litellm.batches.batch_utils import batch_cost_is_final +from litellm.constants import MAX_FILE_LIST_LIMIT from litellm.proxy._types import ProxyException from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -34,8 +35,6 @@ if TYPE_CHECKING: from litellm.types.utils import LiteLLMBatch -MAX_FILE_LIST_LIMIT: Final = 10000 - FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500 BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create" diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index c315d30b8f3..07cdc33e306 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( ) from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.llms.base_llm.managed_resources.isolation import build_list_page from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -90,6 +91,7 @@ from litellm.router import Router from litellm.types.llms.openai import ( CREATE_FILE_REQUESTS_PURPOSE, FileExpiresAfter, + FileListPage, OpenAIFileObject, OpenAIFilesPurpose, ) @@ -97,6 +99,7 @@ from litellm.types.llms.openai import ( router: Final = APIRouter() _MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None) +_LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject]) class UploadedFileInfo(TypedDict): @@ -1287,6 +1290,11 @@ async def delete_file( user_api_key_dict=user_api_key_dict, managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), ) + if is_managed_cloud_storage_uri(file_id) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Raw cloud storage file ids can only be deleted by a proxy admin key. Use the LiteLLM managed file id returned when the file was created.", + ) custom_llm_provider: Final = ( provider @@ -1446,6 +1454,12 @@ async def delete_file( ) +def _as_file_list_page(response: object) -> object: + if not isinstance(response, list): + return response + return FileListPage(**build_list_page(_LISTED_FILES_ADAPTER.validate_python(response))) + + @router.get( "/{provider}/v1/files", dependencies=[Depends(user_api_key_auth)], @@ -1524,7 +1538,7 @@ async def list_files( if should_route and credentials is not None: # Use model-based routing with credentials from config - prepare_data_with_credentials(data=data, credentials=credentials) + prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True) response = await litellm.afile_list( custom_llm_provider=credentials["custom_llm_provider"], purpose=purpose, @@ -1550,7 +1564,7 @@ async def list_files( model_id=target_model_names_list[0], operation_context="file list", ) - prepare_data_with_credentials(data=data, credentials=credentials) + prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True) response = await litellm.afile_list( custom_llm_provider=credentials["custom_llm_provider"], purpose=purpose, @@ -1592,6 +1606,7 @@ async def list_files( status_code=500, detail="Either 'provider' or 'target_model_names' must be provided e.g. `?target_model_names=gpt-4o`", ) + response = _as_file_list_page(response) # rebind-ok: each dispatch branch above binds response ## POST CALL HOOKS ### _response: Final = await proxy_logging_obj.post_call_success_hook( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7d521d54791..dd7967aafe3 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -659,12 +659,14 @@ model LiteLLM_SpendLogs { mcp_namespaced_tool_name String? agent_id String? proxy_server_request Json? @default("{}") + litellm_call_id String? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([litellm_call_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8cfb6354dd0..0ec2788fbaa 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3,6 +3,7 @@ import collections import json import os from collections.abc import Mapping, Sequence +from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from itertools import groupby from types import MappingProxyType @@ -17,6 +18,7 @@ from typing import ( TypeAlias, TypedDict, TypeVar, + cast, # noqa: TID251 # custom-logger and cold-storage payloads are untyped JSON ) import fastapi @@ -75,6 +77,7 @@ _SPEND_LOG_LIST_COLUMNS: Final = """ cache_hit, cache_key, request_tags, team_id, organization_id, end_user, requester_ip_address, session_id, status, mcp_namespaced_tool_name, agent_id, + litellm_call_id, COALESCE(request_duration_ms, (EXTRACT(EPOCH FROM ("endTime" - "startTime")) * 1000)::INTEGER) AS request_duration_ms """ @@ -91,9 +94,9 @@ class _SupportsModelDump(Protocol): def model_dump(self) -> Mapping[str, object]: ... -class _SpendLogOwnershipRow(Protocol): - user: str | None - team_id: str | None +class _SpendLogOwnerRow(TypedDict): + user: ReadOnly[str | None] + team_id: ReadOnly[str | None] class _ActivityRow(TypedDict): @@ -330,12 +333,37 @@ async def _find_spend_logs( return rows -async def _find_spend_log_row(prisma_client: PrismaClient, request_id: str) -> _SpendLogOwnershipRow | None: - """Read the single spend log row identified by ``request_id``.""" - return await _spend_logs_table(prisma_client).find_unique( - where={"request_id": request_id}, - include=None, - ) +class _RequestIdEquals(TypedDict): + request_id: ReadOnly[str] + + +class _LitellmCallIdEquals(TypedDict): + litellm_call_id: ReadOnly[str] + + +def _request_id_or_call_id_clause(request_id: str) -> tuple[_RequestIdEquals, _LitellmCallIdEquals]: + request_id_clause: Final[_RequestIdEquals] = {"request_id": request_id} + call_id_clause: Final[_LitellmCallIdEquals] = {"litellm_call_id": request_id} + return (request_id_clause, call_id_clause) + + +async def _find_spend_log_owners(prisma_client: PrismaClient, request_id: str) -> Sequence[_SpendLogOwnerRow]: + """Read the distinct ``(user, team_id)`` owner pairs across every spend log row + identified by ``request_id`` or ``litellm_call_id``. + + ``litellm_call_id`` is populated from the client-settable ``x-litellm-call-id`` + request header, so it is not guaranteed unique to one tenant: any number of rows + can match one id. The read is uncapped because a flood of another tenant's rows + carrying the caller's id could otherwise push the caller's own owner pair past a + row-sample cap and lock them out of their own lookup. + """ + sql_query: Final = """ + SELECT DISTINCT "user", team_id + FROM "LiteLLM_SpendLogs" + WHERE request_id = $1 OR litellm_call_id = $1 + """ + owners: Final[Sequence[_SpendLogOwnerRow] | None] = await _query_raw_or_none(prisma_client, sql_query, request_id) + return owners if owners is not None else () async def _count_spend_logs(prisma_client: PrismaClient, where: Mapping[str, object]) -> int: @@ -2599,10 +2627,11 @@ async def ui_view_spend_logs( if max_spend is not None: where_conditions["spend"]["lte"] = max_spend # A request_id lookup drops the date window, so a non-admin could otherwise - # reach any single row by id; require they own it, mirroring the detail - # endpoint. That ownership check fully authorizes the one row, so the - # general scoping below is skipped for id lookups. Scoped to the UI route - # so the public v2 contract is unchanged. + # reach any single row by id; require they own one of the matches, mirroring + # the detail endpoint, and keep the general scoping below so a colliding + # foreign row is filtered out rather than served or allowed to deny the + # caller their own row. Scoped to the UI route so the public v2 contract is + # unchanged. if request_id is not None and not is_v2 and not is_admin_view: await _assert_user_can_view_request_id( prisma_client=prisma_client, @@ -2610,10 +2639,9 @@ async def ui_view_spend_logs( request_id=request_id, ) user_scope_applies: Final = ( - not is_request_id_lookup - and not is_admin_view + not is_admin_view and team_id is None - and _can_user_view_spend_log(user_api_key_dict=user_api_key_dict) + and (is_request_id_lookup or _can_user_view_spend_log(user_api_key_dict=user_api_key_dict)) ) permitted_team_ids: Final = ( await _get_permitted_team_ids_for_spend_logs_or_empty( @@ -2626,7 +2654,7 @@ async def ui_view_spend_logs( explicit_user_requires_caller_scope: Final = ( user_scope_applies and not permitted_team_ids and user_id is not None ) - if not is_request_id_lookup and not is_admin_view: + if not is_admin_view: if team_id is not None: can_view_team: Final = await _can_team_member_view_log( prisma_client=prisma_client, @@ -2696,7 +2724,6 @@ async def ui_view_spend_logs( ("team_id", "team_id"), ('"user"', "user"), ("api_key", "api_key"), - ("request_id", "request_id"), ("model", "model"), ("model_id", "model_id"), ("model_group", "model_group"), @@ -2708,6 +2735,13 @@ async def ui_view_spend_logs( sql_params.append(val) p += 1 + request_id_filter: Final = where_conditions.get("request_id") + exact_request_id_first: Final = f"(request_id = ${p}) DESC, " if isinstance(request_id_filter, str) else "" + if isinstance(request_id_filter, str): + sql_conditions.append(f"(request_id = ${p} OR litellm_call_id = ${p})") + sql_params.append(request_id_filter) + p += 1 + # Multi-team OR filter: (user = $X OR team_id = ANY($Y)) if permitted_team_ids: or_clause: Final = f'("user" = ${p} OR team_id = ANY(${p + 1}::text[]))' @@ -2837,7 +2871,7 @@ async def ui_view_spend_logs( WHERE {joined_conditions} ORDER BY {_SESSION_GROUP_KEY_SQL}, call_type IN {_MCP_CALL_TYPES_SQL}, "startTime" DESC ) AS session_representatives - ORDER BY {_order_expr} {_sql_dir}{_nulls_clause}, request_id + ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause}, request_id LIMIT ${p} OFFSET ${p + 1} """ if session_grouping @@ -2846,7 +2880,7 @@ async def ui_view_spend_logs( {_SPEND_LOG_LIST_COLUMNS} FROM "LiteLLM_SpendLogs" WHERE {joined_conditions} - ORDER BY {_order_expr} {_sql_dir}{_nulls_clause} + ORDER BY {exact_request_id_first}{_order_expr} {_sql_dir}{_nulls_clause} LIMIT ${p} OFFSET ${p + 1} """ ) @@ -2854,6 +2888,14 @@ async def ui_view_spend_logs( data: Final = await prisma_client.db.query_raw(sql_query, *sql_params) + if request_id is not None and not is_v2 and not is_admin_view: + await _assert_user_owns_fetched_spend_rows( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + rows=data, + request_id=request_id, + ) + _hydrate_spend_log_metadata(data) # Calculate total pages @@ -3104,7 +3146,7 @@ def _hydrate_spend_log_metadata(rows: Sequence[Mapping[str, object]]) -> None: def _cold_storage_object_key_from_metadata( - metadata: str | dict | None, + metadata: str | Mapping[str, object] | None, ) -> str | None: if isinstance(metadata, str): try: @@ -3209,7 +3251,8 @@ async def ui_view_request_response_for_request_id( """ from litellm.proxy.proxy_server import prisma_client - if not _is_admin_view_safe(user_api_key_dict=user_api_key_dict): + caller_is_admin: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict) + if not caller_is_admin: if prisma_client is None: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, @@ -3234,38 +3277,45 @@ async def ui_view_request_response_for_request_id( if end_date is not None: end_date_obj = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S").replace(tzinfo=timezone.utc) + spend_log_row: Final = ( + None + if prisma_client is None + else await _resolve_spend_log_payload_row( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + request_id=request_id, + caller_is_admin=caller_is_admin, + ) + ) + stored_request_id: Final = _stored_request_id(spend_log_row, request_id) + for custom_logger in custom_loggers: payload = await custom_logger.get_request_response_payload( - request_id=request_id, + request_id=stored_request_id, start_time_utc=start_date_obj, end_time_utc=end_date_obj, ) if payload is not None: + if not caller_is_admin and prisma_client is not None: + await _assert_user_owns_cold_storage_payload( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + payload=cast(Mapping[str, object], payload), # cast-ok: custom-logger payload is untyped + request_id=request_id, + ) return payload + if spend_log_row is None: + return None + # Fallback: the list endpoint omits the heavy columns for performance, so # serve them here. When prompts were offloaded to cold storage the DB holds # only placeholders, so _resolve_request_response_payload fetches the real # payload from the configured cold storage backend by object key. - if prisma_client is not None: - from litellm.proxy.spend_tracking.cold_storage_handler import ( - ColdStorageHandler, - ) + from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler - sql_query: Final = """ - SELECT messages, response, proxy_server_request, metadata - FROM "LiteLLM_SpendLogs" - WHERE request_id = $1 - LIMIT 1 - """ - db_result: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none( - prisma_client, sql_query, request_id - ) - if db_result and len(db_result) > 0: - resolved = await _resolve_request_response_payload(db_result[0], cold_storage_handler=ColdStorageHandler()) - return resolved._asdict() - - return None + resolved: Final = await _resolve_request_response_payload(spend_log_row, cold_storage_handler=ColdStorageHandler()) + return resolved._asdict() @router.get( @@ -3391,7 +3441,7 @@ async def view_spend_logs( if api_key is not None and isinstance(api_key, str): filter_query["api_key"] = summary_api_key if request_id is not None and isinstance(request_id, str): - filter_query["request_id"] = request_id + filter_query["OR"] = _request_id_or_call_id_clause(request_id) if user_id is not None and isinstance(user_id, str): filter_query["user"] = user_id @@ -3439,7 +3489,7 @@ async def view_spend_logs( return [*summary_items, *padding] else: - scoped_filter: Final[dict[str, str]] = {} + scoped_filter: Final[dict[str, object]] = {} if api_key is not None and isinstance(api_key, str): if api_key.startswith("sk-"): hashed_token = prisma_client.hash_token(token=api_key) @@ -3447,7 +3497,7 @@ async def view_spend_logs( hashed_token = api_key scoped_filter["api_key"] = hashed_token if request_id is not None and isinstance(request_id, str): - scoped_filter["request_id"] = request_id + scoped_filter["OR"] = _request_id_or_call_id_clause(request_id) if user_id is not None and isinstance(user_id, str): scoped_filter["user"] = user_id @@ -4676,39 +4726,190 @@ def _can_user_view_spend_log(user_api_key_dict: UserAPIKeyAuth) -> bool: ) +async def _user_can_view_spend_log_owner( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + owner_user: str | None, + owner_team_id: str | None, +) -> bool: + if owner_user is not None and owner_user == user_api_key_dict.user_id: + return True + if owner_team_id: + return await _can_team_member_view_log( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + team_id=owner_team_id, + ) + return False + + +def _spend_log_forbidden(request_id: str) -> HTTPException: + return HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={"error": f"Not authorized to view spend log for request_id={request_id}"}, + ) + + async def _assert_user_can_view_request_id( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, request_id: str, ) -> None: """ - Verify the requesting non-admin user is allowed to view this spend-log row. - Allowed when the log belongs to the user directly, or to one of their - permitted teams (admin or ``/spend/logs`` permission). - Raises HTTP 403 if not, including when no spend-log row exists for the - request_id (e.g. it was pruned by retention), so a missing row can't be - used to read a payload out of cold storage via the detail endpoint. + Verify the requesting non-admin user is allowed to view at least one spend-log + row identified by ``request_id`` or ``litellm_call_id``. The latter is + client-settable, so an id lookup can match rows across different tenants; the + data queries scope a non-admin's results to rows they own directly or via a + permitted team, so a colliding foreign row can neither be served nor deny the + caller their own. Raises HTTP 403 when none of the matching rows is theirs to + view, including when no row exists at all (e.g. it was pruned by retention), + so a missing row can't be used to read a payload out of cold storage via the + detail endpoint. """ - row: Final = await _find_spend_log_row(prisma_client, request_id) + owners: Final = await _find_spend_log_owners(prisma_client, request_id) + for owner in owners: + if await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner["user"], owner["team_id"]): + return + raise _spend_log_forbidden(request_id) - if row is not None and row.user is not None and row.user == user_api_key_dict.user_id: - return - if row is not None and row.team_id: - can_view: Final = await _can_team_member_view_log( +@dataclass(frozen=True, slots=True) +class _SpendLogViewer: + user_id: str | None + team_ids: tuple[str, ...] + + +async def _spend_log_viewer(prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth) -> _SpendLogViewer: + return _SpendLogViewer( + user_id=user_api_key_dict.user_id, + team_ids=await _get_permitted_team_ids_for_spend_logs_or_empty( prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, - team_id=row.team_id, - ) - if can_view: - return - - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": f"Not authorized to view spend log for request_id={request_id}"}, + ), ) +def _viewer_scope_clause(viewer: _SpendLogViewer | None) -> tuple[str, tuple[object, ...]]: + match viewer: + case None: + return ("", ()) + case _SpendLogViewer(user_id=user_id, team_ids=()): + return (' AND "user" = $2', (user_id,)) + case _SpendLogViewer(user_id=user_id, team_ids=team_ids): + return (' AND ("user" = $2 OR team_id = ANY($3::text[]))', (user_id, team_ids)) + + +def _spend_log_payload_query(request_id: str, viewer: _SpendLogViewer | None) -> tuple[str, tuple[object, ...]]: + """ + Fetch the one row an id lookup resolves to, preferring the exact ``request_id`` + match over rows that merely carry the id as their client-set ``litellm_call_id``. + A non-admin viewer only ever gets rows they own or rows of a team they may view. + """ + scope, scope_params = _viewer_scope_clause(viewer) + return ( + f""" + SELECT request_id, messages, response, proxy_server_request, metadata, "user", team_id + FROM "LiteLLM_SpendLogs" + WHERE (request_id = $1 OR litellm_call_id = $1){scope} + ORDER BY (request_id = $1) DESC + LIMIT 1 + """, + (request_id, *scope_params), + ) + + +async def _resolve_spend_log_payload_row( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + request_id: str, + caller_is_admin: bool, +) -> Mapping[str, object] | None: + """ + Resolve an id lookup to the caller's own spend-log row before any payload + store is consulted. Cold storage is keyed by the provider ``request_id``, so + asking it for the raw lookup id could hand back another tenant's payload when + that id is only the caller's ``litellm_call_id``; the row's stored + ``request_id`` is the key that names the caller's own request. + """ + viewer: Final = None if caller_is_admin else await _spend_log_viewer(prisma_client, user_api_key_dict) + sql_query, sql_params = _spend_log_payload_query(request_id, viewer) + rows: Final[Sequence[Mapping[str, object]] | None] = await _query_raw_or_none(prisma_client, sql_query, *sql_params) + if not rows: + return None + if not caller_is_admin: + await _assert_user_owns_fetched_spend_rows( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + rows=rows, + request_id=request_id, + ) + return rows[0] + + +def _stored_request_id(row: Mapping[str, object] | None, lookup_id: str) -> str: + stored: Final = None if row is None else row.get("request_id") + return stored if isinstance(stored, str) else lookup_id + + +def _fetched_row_owner(row: Mapping[str, object]) -> tuple[str | None, str | None]: + user: Final = row.get("user") + team_id: Final = row.get("team_id") + return ( + user if isinstance(user, str) else None, + team_id if isinstance(team_id, str) else None, + ) + + +async def _assert_user_owns_fetched_spend_rows( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + rows: Sequence[Mapping[str, object]], + request_id: str, +) -> None: + """ + Re-verify ownership on the rows an id lookup actually fetched. + ``_assert_user_can_view_request_id`` and the data query read the table at + different moments, so a foreign row inserted between them could otherwise be + returned even though the pre-check passed. Checking the fetched rows + themselves means no interleaving can return another tenant's row. + """ + for user, team_id in frozenset(_fetched_row_owner(row) for row in rows): + if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, user, team_id): + raise _spend_log_forbidden(request_id) + + +def _cold_storage_payload_owner(payload: Mapping[str, object]) -> tuple[str | None, str | None]: + metadata: Final = payload.get("metadata") + if not isinstance(metadata, Mapping): + return (None, None) + owner: Final = cast(Mapping[str, object], metadata) # cast-ok: cold-storage JSON is untyped + user: Final = owner.get("user_api_key_user_id") + team_id: Final = owner.get("user_api_key_team_id") + return ( + user if isinstance(user, str) else None, + team_id if isinstance(team_id, str) else None, + ) + + +async def _assert_user_owns_cold_storage_payload( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + payload: Mapping[str, object], + request_id: str, +) -> None: + """ + Authorize a cold-storage payload against the owner recorded inside it. + The custom logger reads the payload straight from cold storage, written + independently of the spend-log table and able to outlive its row, so a + request_id lookup could otherwise hand back another tenant's stored payload + when no row exists for the pre-check to catch. Verifying the payload's own + owner closes that gap, and a payload that records no owner fails closed. + """ + owner_user, owner_team_id = _cold_storage_payload_owner(payload) + if not await _user_can_view_spend_log_owner(prisma_client, user_api_key_dict, owner_user, owner_team_id): + raise _spend_log_forbidden(request_id) + + async def _get_permitted_team_ids_for_spend_logs( prisma_client: PrismaClient, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3fd20bb7e81..3ef21996b9c 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -621,6 +621,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs status=_get_status_for_spend_log( metadata=metadata, ), + litellm_call_id=litellm_call_id, ) verbose_proxy_logger.debug( diff --git a/litellm/types/router.py b/litellm/types/router.py index fc09c40fe08..c7363502017 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -268,6 +268,12 @@ class CredentialLiteLLMParams(BaseModel): # callers see it, breaking Azure deployments configured with # ``azure_ad_token`` instead of a static ``api_key`` (#30235). azure_ad_token: str | None = None + tenant_id: str | None = None + client_id: str | None = None + client_secret: str | None = None + azure_scope: str | None = None + azure_username: str | None = None + azure_password: str | None = None ## VERTEX AI ## vertex_project: str | None = None vertex_location: str | None = None diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 06c7a6aa46e..2220d0e1fe5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23083,58 +23083,206 @@ }, "friendliai/zai-org/GLM-5.3-Flash": { "litellm_provider": "friendliai", - "supports_reasoning": true, - "supports_function_calling": true, "max_input_tokens": 1048576, - "max_tokens": 1048576, "max_output_tokens": 1048576, + "max_tokens": 1048576, "input_cost_per_token": 1.5e-07, "output_cost_per_token": 5e-07, "cache_read_input_token_cost": 3e-08, "supports_prompt_caching": true, + "supports_reasoning": true, "reasoning_effort_levels": [ "low", "high", "max" ], + "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, - "mode": "chat", - "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", - "source": "https://api.friendli.ai/serverless/v1/models", "supports_vision": true, "supports_image_input": true, - "supports_video_input": true + "supports_video_input": true, + "mode": "chat", + "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", + "source": "https://api.friendli.ai/serverless/v1/models" }, "friendliai/zai-org/GLM-5.3": { "litellm_provider": "friendliai", - "supports_reasoning": true, - "supports_function_calling": true, "max_input_tokens": 1048576, - "max_tokens": 1048576, "max_output_tokens": 1048576, + "max_tokens": 1048576, "input_cost_per_token": 1.26e-06, "output_cost_per_token": 3.96e-06, "cache_read_input_token_cost": 2.34e-07, "supports_prompt_caching": true, + "supports_reasoning": true, "reasoning_effort_levels": [ "low", "high", "max" ], + "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, "mode": "chat", "comment": "Flagship GLM model for long-horizon coding, agents, and complex project delivery", - "source": "https://api.friendli.ai/serverless/v1/models", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/google/gemma-4-31B-it": { + "litellm_provider": "friendliai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4e-07, + "supports_prompt_caching": false, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_image_input": true, + "supports_video_input": false, + "mode": "chat", + "comment": "Largest Gemma 4 instruction model for open, self-hosted chat and reasoning", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/zai-org/GLM-5.2": { + "litellm_provider": "friendliai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [ + "high", + "max" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, "supports_vision": false, - "supports_image_input": false + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/LGAI-EXAONE/K-EXAONE-2.0-750B-A37B": { + "litellm_provider": "friendliai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.4e-06, + "cache_read_input_token_cost": 1.2e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Frontier-scale multilingual language model developed by LG AI Research", + "deprecation_date": "2026-09-06", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/deepseek-ai/DeepSeek-V3.2": { + "litellm_provider": "friendliai", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "input_cost_per_token": 5e-07, + "output_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 2.5e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "DeepSeek chat model for instruction following, coding, and analysis", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "friendliai", + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "max_tokens": 196608, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-08, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", + "source": "https://api.friendli.ai/serverless/v1/models" + }, + "friendliai/zai-org/GLM-5.1": { + "litellm_provider": "friendliai", + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "max_tokens": 202752, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, + "supports_prompt_caching": true, + "supports_reasoning": true, + "reasoning_effort_levels": [], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "supports_image_input": false, + "supports_video_input": false, + "mode": "chat", + "comment": "Strong GLM coding model for agentic engineering, terminals, and repository generation", + "source": "https://api.friendli.ai/serverless/v1/models" }, "ft:babbage-002": { "deprecation_date": "2026-10-23", diff --git a/schema.prisma b/schema.prisma index 7d521d54791..dd7967aafe3 100644 --- a/schema.prisma +++ b/schema.prisma @@ -659,12 +659,14 @@ model LiteLLM_SpendLogs { mcp_namespaced_tool_name String? agent_id String? proxy_server_request Json? @default("{}") + litellm_call_id String? created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@index([startTime]) @@index([startTime, request_id]) @@index([end_user]) @@index([session_id]) + @@index([litellm_call_id]) } model LiteLLM_BudgetWindowSpend { diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index a58c13d6a1c..0541ce25d4b 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -69,7 +69,7 @@ Each suite provides its own `client` fixture (see `llm_translation/passthrough_c Request and response bodies are typed pydantic models in `models.py`; only the fields a test reads are modelled, and nothing passes raw dicts. Outcomes come back as a `Result[R]` tagged union (`Success`, `NetworkError`, `UnauthorizedError`, `RateLimitedError`, `ValidationError`, `UnknownApiError`). Handle them with `match`, or call `unwrap(...)` when a non-success should fail the test. The harness hard-fails and never skips: a test marked `e2e` fails when no proxy answers its liveness probe, and once a request reaches the proxy any wrong behavior is likewise a hard failure, so a missing proxy turns the run red instead of being mistaken for a pass -Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure coverage of the harness itself carries no marker and runs regardless. Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache +Mark live tests with `@pytest.mark.e2e` (on the class or the module). Use `scoped_key` for a fresh all-models key that auto-deletes, `resources` when you need to create and tear down more than a key, and `unique_marker()` from `e2e_config` to keep prompts, tags, and customer ids from colliding across concurrent runs and the shared response cache ## Record and replay fixtures @@ -221,7 +221,7 @@ other... ``` ## Hard Rules -- no monkeypatching or mock tests, and never substitute a unit test for e2e feature coverage: a product feature is proven end to end against a live proxy, not with a unit test. if a contributor asks you to write an end to end test, do NOT stage a unit test of the feature with it; if you find a product gap, call it out in the PR description. tests that cover the harness itself are the exception and are allowed (for example `coverage_registry/test_collector.py`, which unit-tests the coverage collector): they carry no `e2e` marker, exercise harness plumbing rather than a product feature, and run whether or not a proxy is up +- no unit tests of any kind under `tests/e2e`. a product feature is proven end to end against a live proxy, never with a unit test, and the harness itself is not unit-tested here either. no monkeypatching or mock tests. if a contributor asks you to write an end to end test, do NOT stage a unit test with it; if you find a product gap, call it out in the PR description - use model management endpoints to create new models for a test. this could be in a conftest / inline for each test. ask the user what they want. diff --git a/tests/logging_callback_tests/test_gcs_pub_sub.py b/tests/logging_callback_tests/test_gcs_pub_sub.py index 10957fa2f92..1f1ca8960f6 100644 --- a/tests/logging_callback_tests/test_gcs_pub_sub.py +++ b/tests/logging_callback_tests/test_gcs_pub_sub.py @@ -28,6 +28,7 @@ verbose_logger.setLevel(logging.DEBUG) ignored_keys = [ "request_id", + "litellm_call_id", "metadata.litellm_call_id", "session_id", "startTime", diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 48fceb50403..cb08e00ff65 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -1193,7 +1193,7 @@ async def test_afile_delete_returns_managed_id_for_stored_provider_output(): assert response.id == unified_file_id assert response.object == "file" - assert response.filename == stored_file.filename + assert response.deleted is True assert stored_file.id == provider_file_id router.afile_delete.assert_awaited_once_with(model="model-123", file_id=provider_file_id) table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id}) @@ -1730,3 +1730,100 @@ async def test_batch_retrieve_hook_does_not_claim_attribution(): managed_files.store_unified_object_id.assert_awaited_once() assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False + + +@pytest.mark.asyncio +async def test_afile_delete_passes_trusted_model_credentials_to_router(): + """ + afile_delete must hand the deployment's credential snapshot to the router + call, since Bedrock validates the s3:// file id against the bucket in it. + """ + from types import MappingProxyType + + managed_files = _make_managed_files_instance() + unified_file_id = "unified-file-id" + s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl" + managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}}) + managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id)) + + mock_router = MagicMock() + mock_router.get_deployment_credentials_with_provider = MagicMock( + return_value={ + "custom_llm_provider": "bedrock", + "s3_bucket_name": "my-bucket", + "aws_region_name": "us-west-2", + } + ) + mock_router.afile_delete = AsyncMock(return_value=MagicMock()) + + await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=mock_router, + ) + + call_kwargs = mock_router.afile_delete.call_args.kwargs + assert call_kwargs["model"] == "model-123" + assert call_kwargs["file_id"] == s3_uri + trusted_credentials = call_kwargs["_litellm_internal_model_credentials"] + assert isinstance(trusted_credentials, MappingProxyType) + assert trusted_credentials["s3_bucket_name"] == "my-bucket" + + +@pytest.mark.asyncio +async def test_afile_delete_bedrock_unified_id_end_to_end(monkeypatch): + """ + Proxy repro for deleting a Bedrock batch input file by unified id: the + s3:// object must be removed via a SigV4-signed S3 DELETE using the + deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env). + + Regression test for "BedrockFilesConfig does not support file deletion" + raised on this path. + """ + import httpx + import respx + + import litellm + from litellm import Router + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + router = Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + "model_info": {"id": "model-123"}, + } + ] + ) + + managed_files = _make_managed_files_instance() + unified_file_id = "unified-file-id" + s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl" + managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}}) + managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id)) + + expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/job-123/input.jsonl" + with respx.mock: + route = respx.delete(expected_url).mock(return_value=httpx.Response(204)) + + response = await managed_files.afile_delete( + file_id=unified_file_id, + litellm_parent_otel_span=None, + llm_router=router, + ) + + assert route.called + assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert response.id == unified_file_id + assert response.model_dump() == {"id": unified_file_id, "object": "file", "deleted": True} + managed_files.delete_unified_file_id.assert_awaited_once_with(unified_file_id, None) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 2f6339dcdbb..451b740dd71 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5890,6 +5890,81 @@ def test_resolve_vertex_location_for_cost_default_region(monkeypatch): assert _resolve("vertex_ai", None, None, "gemini-3.5-flash") == "us-central1" +def test_resolve_mantle_region_for_cost(monkeypatch): + """Bedrock Mantle requests resolve the served region the way dispatch does (explicit + aws_region_name, then the api_base host, then the default); other providers get None.""" + from litellm.litellm_core_utils.litellm_logging import _resolve_mantle_region_for_cost + + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION_NAME", "AWS_REGION"): + monkeypatch.delenv(var, raising=False) + + assert _resolve_mantle_region_for_cost("bedrock", {"aws_region_name": "us-gov-west-1"}) is None + assert _resolve_mantle_region_for_cost(None, {"aws_region_name": "us-gov-west-1"}) is None + assert _resolve_mantle_region_for_cost("bedrock_mantle", {"aws_region_name": "us-gov-west-1"}) == "us-gov-west-1" + assert ( + _resolve_mantle_region_for_cost( + "bedrock_mantle", + {"api_base": "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"}, + ) + == "us-gov-west-1" + ) + assert _resolve_mantle_region_for_cost("bedrock_mantle", None) == "us-east-1" + + +def test_response_cost_calculator_prices_mantle_calls_on_the_served_region(monkeypatch): + """ + Mantle responses carry no region of their own (the OpenAI-compatible transform rebuilds the + response, and streams never had one), so the logging layer must price them from the region + the deployment was served in: an explicit aws_region_name or the api_base host, both of which + must select the GovCloud row over the commercial one. + """ + from datetime import datetime + + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url="")) + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION_NAME", "AWS_REGION"): + monkeypatch.delenv(var, raising=False) + + def cost_with(litellm_params): + logging_obj = LitellmLogging( + model="xai.grok-4.3", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="mantle-region", + function_id="f", + ) + logging_obj.update_environment_variables( + model="xai.grok-4.3", + user="", + optional_params={}, + litellm_params=litellm_params, + custom_llm_provider="bedrock_mantle", + ) + response = ModelResponse( + id="resp-1", + model="xai.grok-4.3", + choices=[{"message": {"role": "assistant", "content": "hello"}, "index": 0, "finish_reason": "stop"}], + usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + ) + return logging_obj._response_cost_calculator(result=response) + + commercial = litellm.model_cost["bedrock_mantle/xai.grok-4.3"] + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + expected_commercial = 38 * commercial["input_cost_per_token"] + 20 * commercial["output_cost_per_token"] + expected_gov = 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + assert expected_gov != expected_commercial + + assert cost_with({"api_base": ""}) == pytest.approx(expected_commercial) + assert cost_with({"aws_region_name": "us-gov-west-1"}) == pytest.approx(expected_gov) + assert cost_with( + {"api_base": "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"} + ) == pytest.approx(expected_gov) + + def test_response_cost_calculator_prices_proxy_vertex_calls_on_the_configured_location(monkeypatch): """ Proxy-shaped logging objects (created before the router picks a deployment) carry the diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py new file mode 100644 index 00000000000..45ec18733f7 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_mid_stream_error.py @@ -0,0 +1,142 @@ +""" +Regression tests for the ``/v1/messages`` async adapter dropping the socket on a +mid-stream provider error. + +When a non-Anthropic model (e.g. Bedrock Converse) is served through +``/v1/messages``, the proxy hands Starlette the async SSE iterator directly. If +the upstream provider stream raises while being pulled (Bedrock raises +``BedrockError`` when a ConverseStream ends without a terminal ``messageStop`` +event, common on cross-region inference profiles), the exception escaped the +request handler's try/except and tore down the connection. Clients like Claude +Code then showed a bare "Connection closed mid-response". + +The async SSE wrapper must instead surface the failure as a well-formed +Anthropic ``error`` event so the stream stays valid and the client can retry. +""" + +import json +import os +import sys +from typing import List, Optional +from unittest.mock import MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.exceptions import MidStreamFallbackError +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, + _mid_stream_error_sse_event, +) +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.utils import Delta, StreamingChoices + + +def _make_chunk(delta: Delta, finish_reason: Optional[str] = None) -> MagicMock: + chunk = MagicMock() + chunk.choices = [ + StreamingChoices(finish_reason=finish_reason, index=0, delta=delta, logprobs=None) + ] + chunk.usage = None + chunk._hidden_params = {} + return chunk + + +class _AsyncStreamThenRaise: + """Yields the given chunks, then raises ``exc`` (mimics a provider stream + that terminates mid-response).""" + + def __init__(self, items: List[MagicMock], exc: BaseException): + self._it = iter(items) + self._exc = exc + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self._it) + except StopIteration: + raise self._exc + + +def _parse_sse(raw: bytes) -> tuple[str, dict]: + text = raw.decode() + event_line, data_line = text.strip().split("\n", 1) + return event_line.removeprefix("event: "), json.loads(data_line.removeprefix("data: ")) + + +async def _drain_sse(wrapper: AnthropicStreamWrapper) -> List[bytes]: + return [event async for event in wrapper.async_anthropic_sse_wrapper()] + + +@pytest.mark.asyncio +async def test_mid_stream_bedrock_error_becomes_anthropic_error_event(): + """A ``BedrockError`` raised after partial content must be surfaced as a + terminal Anthropic ``error`` event, not propagated (which drops the socket + and yields "Connection closed mid-response").""" + chunks = [_make_chunk(Delta(content="Creating a file"))] + bedrock_err = BedrockError( + status_code=500, + message="Bedrock ConverseStream ended without a terminal 'messageStop' event", + ) + wrapper = AnthropicStreamWrapper( + completion_stream=_AsyncStreamThenRaise(chunks, bedrock_err), + model="bedrock-converse-sonnet-4-6", + ) + + events = await _drain_sse(wrapper) + + parsed = [_parse_sse(e) for e in events] + event_types = [name for name, _ in parsed] + assert "message_start" in event_types + assert event_types[-1] == "error" + _, error_payload = parsed[-1] + assert error_payload["type"] == "error" + assert error_payload["error"]["type"] == "api_error" + assert "messageStop" in error_payload["error"]["message"] + + +@pytest.mark.asyncio +async def test_mid_stream_error_does_not_raise_out_of_wrapper(): + """The async wrapper must fully drain without letting the upstream exception + escape — escaping is exactly what tore down the connection before the fix.""" + wrapper = AnthropicStreamWrapper( + completion_stream=_AsyncStreamThenRaise([], BedrockError(status_code=500, message="boom")), + model="claude-x", + ) + events = await _drain_sse(wrapper) + assert _parse_sse(events[-1])[0] == "error" + + +@pytest.mark.parametrize( + "status_code, expected_type", + [(500, "api_error"), (529, "overloaded_error"), (429, "rate_limit_error")], +) +def test_error_event_maps_status_code_to_anthropic_type(status_code, expected_type): + raw = _mid_stream_error_sse_event(BedrockError(status_code=status_code, message="upstream failed")) + name, payload = _parse_sse(raw) + assert name == "error" + assert payload["error"]["type"] == expected_type + assert payload["error"]["message"] == "upstream failed" + + +def test_error_event_defaults_to_500_when_status_missing(): + raw = _mid_stream_error_sse_event(ValueError("no status here")) + _, payload = _parse_sse(raw) + assert payload["error"]["type"] == "api_error" + assert payload["error"]["message"] == "no status here" + + +def test_error_event_preserves_midstream_fallback_error(): + exc = MidStreamFallbackError( + message="BedrockException - internalServerException", + model="bedrock-converse-sonnet-4-6", + llm_provider="bedrock", + original_exception=BedrockError(status_code=500, message="internalServerException"), + ) + name, payload = _parse_sse(_mid_stream_error_sse_event(exc)) + assert name == "error" + assert payload["error"]["type"] == "api_error" + assert "internalServerException" in payload["error"]["message"] diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 3b01a4f2054..c609455f3d8 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1948,11 +1948,13 @@ class TestBedrockFileDeletion: def test_delete_rejects_untrusted_objects_before_signing( self, file_id: str, message: str, monkeypatch: pytest.MonkeyPatch ) -> None: + from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.files.transformation import BedrockFilesConfig monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") - with pytest.raises(ValueError, match=message): + with pytest.raises(BedrockError, match=message) as rejection: BedrockFilesConfig().transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params={}) + assert rejection.value.status_code == 400 class TestBedrockFileContentTransformation: @@ -2030,11 +2032,12 @@ class TestBedrockFileContentTransformation: assert url == self.EXPECTED_URL def test_transform_file_content_request_rejects_foreign_bucket(self, monkeypatch): + from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.files.transformation import BedrockFilesConfig monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") - with pytest.raises(ValueError, match="configured storage bucket"): + with pytest.raises(BedrockError, match="configured storage bucket") as rejection: BedrockFilesConfig().transform_file_content_request( file_content_request={ "file_id": "s3://other-bucket/litellm-batch-outputs/job/x.jsonl.out" @@ -2043,18 +2046,25 @@ class TestBedrockFileContentTransformation: litellm_params=self._litellm_params(), ) + assert rejection.value.status_code == 400 + + def test_transform_file_content_request_rejects_unmanaged_key(self, monkeypatch): + from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.files.transformation import BedrockFilesConfig monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") - with pytest.raises(ValueError, match="LiteLLM-managed"): + with pytest.raises(BedrockError, match="LiteLLM-managed") as rejection: BedrockFilesConfig().transform_file_content_request( file_content_request={"file_id": "s3://my-bucket/private/x.jsonl"}, optional_params={}, litellm_params=self._litellm_params(), ) + assert rejection.value.status_code == 400 + + def test_extract_s3_uri_rejects_non_managed_file_id(self): """A file id that is neither an s3:// URI nor a unified id must be rejected.""" from litellm.llms.bedrock.files.transformation import ( @@ -2183,12 +2193,13 @@ class TestBedrockFileContentTransformation: def test_rejects_bucket_outside_input_and_output(self, monkeypatch): """A file id whose bucket is neither the input nor the output bucket is still rejected (SSRF / bucket-confusion guard).""" + from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.files.transformation import BedrockFilesConfig monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) - with pytest.raises(ValueError, match="configured storage bucket"): + with pytest.raises(BedrockError, match="configured storage bucket") as rejection: BedrockFilesConfig().transform_file_content_request( file_content_request={ "file_id": "s3://other-bucket/litellm-batch-outputs/job/x.jsonl.out" @@ -2199,6 +2210,9 @@ class TestBedrockFileContentTransformation: ), ) + assert rejection.value.status_code == 400 + + def test_sign_request_without_botocore_raises_helpful_error(self, monkeypatch): """A missing botocore must surface an actionable 'install boto3' error rather than a raw import failure.""" @@ -2605,6 +2619,7 @@ def test_sign_s3_request_without_body_assumes_role_with_external_id(monkeypatch) with patch.object(boto3, "client", return_value=FakeSTSClient()): signed_headers = BedrockFilesConfig()._sign_s3_request_without_body( + method="GET", api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", aws_region_name="us-east-1", request_params=request_params, @@ -2612,3 +2627,1058 @@ def test_sign_s3_request_without_body_assumes_role_with_external_id(monkeypatch) authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] assert "ASIAFILESGETROLE" in authorization + + +def _s3_signature_for(method: str, url: str, headers: Mapping[str, str]) -> str: + sent = {name.lower(): value for name, value in headers.items()} + signed_names = sent["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";") + request = AWSRequest( + method=method, + url=url, + headers={name: sent[name] for name in signed_names if name in sent}, + ) + request.context["timestamp"] = sent["x-amz-date"] + signer = S3SigV4Auth(Credentials("AKIAEXAMPLE", "secret"), "s3", "us-west-2") + return signer.signature(signer.string_to_sign(request, signer.canonical_request(request)), request) + + +def _sent_signature(headers: Mapping[str, str]) -> str: + authorization = {name.lower(): value for name, value in headers.items()}["authorization"] + return authorization.split("Signature=")[1].strip() + + +def _bedrock_s3_params() -> dict: + return { + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + } + + +def _trusted_bucket_snapshot(**deployment_litellm_params) -> dict: + from types import MappingProxyType + + from litellm.types.router import CredentialLiteLLMParams + + snapshot = CredentialLiteLLMParams(**deployment_litellm_params).model_dump(exclude_none=True) + return {**_bedrock_s3_params(), "_litellm_internal_model_credentials": MappingProxyType(snapshot)} + + +class TestBedrockFileDeletionTransformation: + """SigV4-signed S3 DeleteObject for LiteLLM-managed Bedrock batch files.""" + + S3_URI = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl" + EXPECTED_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/job-123/input.jsonl" + + def test_transform_delete_file_request_signs_s3_delete(self, monkeypatch): + import hashlib + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + + url, params = BedrockFilesConfig().transform_delete_file_request( + file_id=self.S3_URI, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.EXPECTED_URL + assert params == {} + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + lowered = {name.lower(): value for name, value in signed_headers.items()} + assert lowered["x-amz-content-sha256"] == hashlib.sha256(b"").hexdigest() + assert "/us-west-2/s3/aws4_request" in lowered["authorization"] + assert _sent_signature(signed_headers) == _s3_signature_for("DELETE", url, signed_headers) + + def test_transform_delete_file_request_decodes_unified_file_id(self, monkeypatch): + import base64 + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.utils import SpecialEnums + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "unified-id", "", self.S3_URI, "model-id" + ) + encoded_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") + litellm_params = _bedrock_s3_params() + + url, _ = BedrockFilesConfig().transform_delete_file_request( + file_id=encoded_file_id, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.EXPECTED_URL + + def test_transform_delete_file_request_uses_trusted_snapshot_bucket(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + + url, _ = BedrockFilesConfig().transform_delete_file_request( + file_id=self.S3_URI, + optional_params={}, + litellm_params=_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + assert url == self.EXPECTED_URL + + def test_transform_delete_file_request_rejects_foreign_bucket(self, monkeypatch): + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with pytest.raises(BedrockError, match="configured storage bucket") as rejection: + BedrockFilesConfig().transform_delete_file_request( + file_id="s3://other-bucket/litellm-bedrock-files/job-123/input.jsonl", + optional_params={}, + litellm_params=_bedrock_s3_params(), + ) + + assert rejection.value.status_code == 400 + + + def test_transform_delete_file_request_rejects_unmanaged_key(self, monkeypatch): + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with pytest.raises(BedrockError, match="LiteLLM-managed") as rejection: + BedrockFilesConfig().transform_delete_file_request( + file_id="s3://my-bucket/private/x.jsonl", + optional_params={}, + litellm_params=_bedrock_s3_params(), + ) + + assert rejection.value.status_code == 400 + + + def test_transform_delete_file_response_echoes_the_deleted_id(self): + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + deleted = BedrockFilesConfig().transform_delete_file_response( + raw_response=httpx.Response(204), + logging_obj=MagicMock(model_call_details={"additional_args": {"file_id": self.S3_URI}}), + litellm_params={}, + ) + + assert deleted.id == self.S3_URI + assert deleted.deleted is True + assert deleted.object == "file" + + def test_transform_delete_file_response_raises_on_s3_error(self): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + with pytest.raises(BedrockError) as excinfo: + BedrockFilesConfig().transform_delete_file_response( + raw_response=httpx.Response(403, text="AccessDenied"), + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert excinfo.value.status_code == 403 + + def test_file_delete_end_to_end_sends_signed_delete(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock: + route = respx.delete(self.EXPECTED_URL).mock(return_value=httpx.Response(204)) + + response = litellm.file_delete( + file_id=self.S3_URI, + custom_llm_provider="bedrock", + **_bedrock_s3_params(), + ) + + assert route.called + request = route.calls[0].request + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert _sent_signature(request.headers) == _s3_signature_for("DELETE", str(request.url), request.headers) + assert response.id == self.S3_URI + assert response.deleted is True + + @pytest.mark.asyncio + async def test_afile_delete_end_to_end_sends_signed_delete(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock: + route = respx.delete(self.EXPECTED_URL).mock(return_value=httpx.Response(204)) + + response = await litellm.afile_delete( + file_id=self.S3_URI, + custom_llm_provider="bedrock", + **_bedrock_s3_params(), + ) + + assert route.called + request = route.calls[0].request + assert _sent_signature(request.headers) == _s3_signature_for("DELETE", str(request.url), request.headers) + assert response.id == self.S3_URI + assert response.deleted is True + + def test_file_delete_end_to_end_answers_400_for_a_foreign_bucket(self, monkeypatch): + import respx + + import litellm + from litellm.llms.bedrock.common_utils import BedrockError + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock, pytest.raises(BedrockError) as rejection: + litellm.file_delete( + file_id="s3://other-bucket/litellm-bedrock-files/job-123/input.jsonl", + custom_llm_provider="bedrock", + **_bedrock_s3_params(), + ) + + assert rejection.value.status_code == 400 + assert "configured storage bucket" in rejection.value.message + + def test_file_delete_end_to_end_answers_400_for_a_non_managed_id(self, monkeypatch): + import respx + + import litellm + from litellm.llms.bedrock.common_utils import BedrockError + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock, pytest.raises(BedrockError) as rejection: + litellm.file_delete(file_id="file-1234567890", custom_llm_provider="bedrock", **_bedrock_s3_params()) + + assert rejection.value.status_code == 400 + assert "managed LiteLLM S3 file id" in rejection.value.message + + def test_file_delete_end_to_end_surfaces_the_s3_error_body(self, monkeypatch): + import httpx + import respx + + import litellm + from litellm.llms.bedrock.common_utils import BedrockError + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock: + respx.delete(self.EXPECTED_URL).mock( + return_value=httpx.Response(403, text="AccessDenied") + ) + with pytest.raises(BedrockError) as denied: + litellm.file_delete(file_id=self.S3_URI, custom_llm_provider="bedrock", **_bedrock_s3_params()) + + assert denied.value.status_code == 403 + assert "AccessDenied" in denied.value.message + + +class TestBedrockFileListTransformation: + """SigV4-signed S3 ListObjectsV2 over the LiteLLM-managed key prefixes.""" + + BUCKET_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/" + MANAGED_QUERY = {"list-type": "2", "prefix": "litellm-b"} + BATCH_QUERY = {"list-type": "2", "prefix": "litellm-bedrock-files"} + OUTPUT_QUERY = {"list-type": "2", "prefix": "litellm-batch-outputs/"} + OUTPUT_BUCKET_URL = "https://s3.us-west-2.amazonaws.com/my-output-bucket/" + OUTPUT_BUCKET_LISTING = b""" + + my-output-bucket + litellm-batch-outputs/ + + litellm-batch-outputs/job-9/input.jsonl.out + 2026-09-04T08:00:00.000Z + 70 + +""" + OUTPUT_BUCKET_ID = "s3://my-output-bucket/litellm-batch-outputs/job-9/input.jsonl.out" + LISTING = b""" + + my-bucket + litellm-b + 4 + false + + litellm-bedrock-files-model-abc.jsonl + 2026-09-01T10:00:00.000Z + 120 + + + litellm-bedrock-files/job-123/input.jsonl + 2026-09-02T11:30:00.000Z + 340 + + + litellm-batch-outputs/job-123/input.jsonl.out + 2026-09-03T12:45:00.000Z + 560 + + + litellm-bogus/other.jsonl + 2026-09-03T12:45:00.000Z + 1 + +""" + BATCH_IDS = ( + "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl", + "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl", + ) + OUTPUT_ID = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + CONTINUATION_TOKEN = "1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM=" + FIRST_PAGE = b""" + + my-bucket + litellm-bedrock-files + 1 + 1 + true + 1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM= + + litellm-bedrock-files/job-1/input.jsonl + 2026-09-01T10:00:00.000Z + 10 + +""" + LAST_PAGE = b""" + + my-bucket + litellm-bedrock-files + 1 + 1 + false + 1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM= + + litellm-bedrock-files/job-2/input.jsonl + 2026-09-02T10:00:00.000Z + 20 + +""" + PAGED_IDS = ( + "s3://my-bucket/litellm-bedrock-files/job-1/input.jsonl", + "s3://my-bucket/litellm-bedrock-files/job-2/input.jsonl", + ) + + def test_transform_list_files_request_signs_managed_prefix_listing(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import ( + LIST_FILES_PURPOSE_PARAM, + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + + url, params = BedrockFilesConfig().transform_list_files_request( + purpose="batch", + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.BUCKET_URL + assert params == self.BATCH_QUERY + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + assert _sent_signature(signed_headers) == _s3_signature_for( + "GET", f"{url}?list-type=2&prefix=litellm-bedrock-files", signed_headers + ) + assert litellm_params[LIST_FILES_PURPOSE_PARAM] == "batch" + + def test_transform_list_files_request_scopes_to_configured_prefix(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket/LLM AI Projects") + litellm_params = _bedrock_s3_params() + + url, params = BedrockFilesConfig().transform_list_files_request( + purpose=None, + optional_params={}, + litellm_params=litellm_params, + ) + + assert url == self.BUCKET_URL + assert params == {"list-type": "2", "prefix": "LLM AI Projects/litellm-b"} + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + assert _sent_signature(signed_headers) == _s3_signature_for( + "GET", f"{url}?list-type=2&prefix=LLM%20AI%20Projects%2Flitellm-b", signed_headers + ) + + def test_transform_list_files_request_uses_trusted_snapshot_bucket(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + + url, params = BedrockFilesConfig().transform_list_files_request( + purpose=None, + optional_params={}, + litellm_params=_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + assert url == self.BUCKET_URL + assert params == self.MANAGED_QUERY + + def _list_response(self, purpose: str | None, listing: bytes | None = None, status_code: int = 200): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + LIST_FILES_PURPOSE_PARAM, + BedrockFilesConfig, + ) + + return BedrockFilesConfig().transform_list_files_response( + raw_response=httpx.Response(status_code, content=listing if listing is not None else self.LISTING), + logging_obj=MagicMock(), + litellm_params={**_bedrock_s3_params(), LIST_FILES_PURPOSE_PARAM: purpose}, + ) + + def test_transform_list_files_response_maps_managed_objects(self, monkeypatch): + from datetime import datetime, timezone + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + files = self._list_response(purpose=None) + + assert [file.id for file in files] == [*self.BATCH_IDS, self.OUTPUT_ID] + assert [file.purpose for file in files] == ["batch", "batch", "batch_output"] + assert [file.bytes for file in files] == [120, 340, 560] + assert [file.filename for file in files] == [ + "litellm-bedrock-files-model-abc.jsonl", + "input.jsonl", + "input.jsonl.out", + ] + assert files[1].created_at == int(datetime(2026, 9, 2, 11, 30, tzinfo=timezone.utc).timestamp()) + assert {file.object for file in files} == {"file"} + assert {file.status for file in files} == {"uploaded"} + + def test_transform_list_files_response_filters_by_purpose(self, monkeypatch): + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + assert [file.id for file in self._list_response(purpose="batch")] == list(self.BATCH_IDS) + assert [file.id for file in self._list_response(purpose="batch_output")] == [self.OUTPUT_ID] + + def test_transform_list_files_response_scopes_to_configured_prefix(self, monkeypatch): + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket/team-a") + listing = b""" + + my-bucket + team-a/litellm-bedrock-files/job-1/input.jsonl10 + team-a/litellm-batch-outputs/job-1/input.jsonl.out20 + litellm-bedrock-files/job-2/input.jsonl30 +""" + + files = self._list_response(purpose=None, listing=listing) + + assert [(file.id, file.purpose) for file in files] == [ + ("s3://my-bucket/team-a/litellm-bedrock-files/job-1/input.jsonl", "batch"), + ("s3://my-bucket/team-a/litellm-batch-outputs/job-1/input.jsonl.out", "batch_output"), + ] + + def test_transform_list_files_response_raises_on_s3_error(self, monkeypatch): + from litellm.llms.bedrock.common_utils import BedrockError + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with pytest.raises(BedrockError) as excinfo: + self._list_response(purpose=None, listing=b"AccessDenied", status_code=403) + + assert excinfo.value.status_code == 403 + + def test_file_list_end_to_end_sends_signed_listing(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock: + route = respx.get(self.BUCKET_URL, params__contains=self.BATCH_QUERY).mock( + return_value=httpx.Response(200, content=self.LISTING) + ) + + files = litellm.file_list(custom_llm_provider="bedrock", purpose="batch", **_bedrock_s3_params()) + + assert route.called + request = route.calls[0].request + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert _sent_signature(request.headers) == _s3_signature_for("GET", str(request.url), request.headers) + assert [file.id for file in files] == list(self.BATCH_IDS) + + def test_file_list_uses_trusted_snapshot_bucket_without_env(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + + with respx.mock: + route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock( + return_value=httpx.Response(200, content=self.LISTING) + ) + + files = litellm.file_list( + custom_llm_provider="bedrock", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + assert route.called + assert [file.id for file in files] == [*self.BATCH_IDS, self.OUTPUT_ID] + + @pytest.mark.asyncio + async def test_afile_list_end_to_end_sends_signed_listing(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock: + route = respx.get(self.BUCKET_URL, params__contains=self.OUTPUT_QUERY).mock( + return_value=httpx.Response(200, content=self.LISTING) + ) + + files = await litellm.afile_list( + custom_llm_provider="bedrock", purpose="batch_output", **_bedrock_s3_params() + ) + + assert route.called + request = route.calls[0].request + assert _sent_signature(request.headers) == _s3_signature_for("GET", str(request.url), request.headers) + assert [file.id for file in files] == [self.OUTPUT_ID] + + def test_transform_list_files_request_narrows_prefix_to_requested_purpose(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + batch_url, batch_params = BedrockFilesConfig().transform_list_files_request( + purpose="batch", optional_params={}, litellm_params=_bedrock_s3_params() + ) + output_url, output_params = BedrockFilesConfig().transform_list_files_request( + purpose="batch_output", optional_params={}, litellm_params=_bedrock_s3_params() + ) + + assert (batch_url, batch_params) == (self.BUCKET_URL, self.BATCH_QUERY) + assert (output_url, output_params) == (self.BUCKET_URL, self.OUTPUT_QUERY) + + def test_transform_list_files_request_lists_configured_output_bucket_for_batch_output(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + litellm_params = _trusted_bucket_snapshot( + s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket/team-a" + ) + + url, params = BedrockFilesConfig().transform_list_files_request( + purpose="batch_output", optional_params={}, litellm_params=litellm_params + ) + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + input_url, input_params = BedrockFilesConfig().transform_list_files_request( + purpose="batch", optional_params={}, litellm_params=dict(litellm_params) + ) + + assert url == self.OUTPUT_BUCKET_URL + assert params == {"list-type": "2", "prefix": "team-a/litellm-batch-outputs/"} + assert _sent_signature(signed_headers) == _s3_signature_for( + "GET", f"{url}?list-type=2&prefix=team-a%2Flitellm-batch-outputs%2F", signed_headers + ) + assert (input_url, input_params) == (self.BUCKET_URL, self.BATCH_QUERY) + + def test_transform_list_files_request_reads_output_bucket_from_env(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + monkeypatch.setenv("AWS_S3_OUTPUT_BUCKET_NAME", "my-output-bucket") + + url, params = BedrockFilesConfig().transform_list_files_request( + purpose="batch_output", optional_params={}, litellm_params=_bedrock_s3_params() + ) + + assert (url, params) == (self.OUTPUT_BUCKET_URL, self.OUTPUT_QUERY) + + EMPTY_LISTING = b""" + + my-bucket + + 0 + 0 + false +""" + NO_KEYS_QUERY = {"list-type": "2", "max-keys": "0"} + + def test_transform_list_files_request_asks_for_no_keys_when_bedrock_never_stores_the_purpose(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + config = BedrockFilesConfig() + + url, params = config.transform_list_files_request( + purpose="user_data", optional_params={}, litellm_params=litellm_params + ) + next_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.EMPTY_LISTING), + optional_params={}, + litellm_params=litellm_params, + ) + + assert (url, params) == (self.BUCKET_URL, self.NO_KEYS_QUERY) + assert next_request is None + + def test_file_list_never_walks_the_bucket_for_a_purpose_bedrock_never_stores(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock: + route = respx.get(self.BUCKET_URL, params__contains=self.NO_KEYS_QUERY).mock( + return_value=httpx.Response(200, content=self.EMPTY_LISTING) + ) + + files = litellm.file_list(custom_llm_provider="bedrock", purpose="user_data", **_bedrock_s3_params()) + + assert route.call_count == 1 + assert "prefix" not in route.calls[0].request.url.params + assert files == [] + + def test_transform_list_files_request_lists_the_output_bucket_without_an_input_bucket(self, monkeypatch): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + litellm_params = _trusted_bucket_snapshot(s3_output_bucket_name="my-output-bucket") + + url, params = BedrockFilesConfig().transform_list_files_request( + purpose="batch_output", optional_params={}, litellm_params=litellm_params + ) + + assert (url, params) == (self.OUTPUT_BUCKET_URL, self.OUTPUT_QUERY) + with pytest.raises(ValueError, match="s3_bucket_name"): + BedrockFilesConfig().transform_list_files_request( + purpose="batch", optional_params={}, litellm_params=dict(litellm_params) + ) + + def test_transform_list_files_response_accepts_output_bucket_objects(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + LIST_FILES_PURPOSE_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + files = BedrockFilesConfig().transform_list_files_response( + raw_response=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING), + logging_obj=MagicMock(), + litellm_params={ + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"), + LIST_FILES_PURPOSE_PARAM: "batch_output", + }, + ) + + assert [(file.id, file.purpose, file.bytes) for file in files] == [(self.OUTPUT_BUCKET_ID, "batch_output", 70)] + + def test_file_list_batch_output_end_to_end_lists_output_bucket(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + route = respx.get(self.OUTPUT_BUCKET_URL, params__contains=self.OUTPUT_QUERY).mock( + return_value=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING) + ) + + files = litellm.file_list( + custom_llm_provider="bedrock", + purpose="batch_output", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"), + ) + + assert route.called + request = route.calls[0].request + assert _sent_signature(request.headers) == _s3_signature_for("GET", str(request.url), request.headers) + assert [file.id for file in files] == [self.OUTPUT_BUCKET_ID] + + def test_file_list_without_purpose_also_walks_a_separate_output_bucket(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + input_route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock( + return_value=httpx.Response(200, content=self.LISTING) + ) + output_route = respx.get(self.OUTPUT_BUCKET_URL, params__contains=self.OUTPUT_QUERY).mock( + return_value=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING) + ) + + files = litellm.file_list( + custom_llm_provider="bedrock", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"), + ) + + assert (input_route.call_count, output_route.call_count) == (1, 1) + output_request = output_route.calls[0].request + assert _sent_signature(output_request.headers) == _s3_signature_for( + "GET", str(output_request.url), output_request.headers + ) + assert [file.id for file in files] == [*self.BATCH_IDS, self.OUTPUT_ID, self.OUTPUT_BUCKET_ID] + + def test_file_list_without_purpose_walks_the_output_bucket_after_the_last_input_page(self, monkeypatch): + import httpx + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + respx.get(self.BUCKET_URL, params__contains={"continuation-token": self.CONTINUATION_TOKEN}).mock( + return_value=httpx.Response(200, content=self.LAST_PAGE) + ) + respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock( + return_value=httpx.Response(200, content=self.FIRST_PAGE) + ) + respx.get(self.OUTPUT_BUCKET_URL, params__contains=self.OUTPUT_QUERY).mock( + return_value=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING) + ) + + files = litellm.file_list( + custom_llm_provider="bedrock", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"), + ) + requested_urls = [str(call.request.url) for call in respx.calls] + + assert requested_urls == [ + f"{self.BUCKET_URL}?list-type=2&prefix=litellm-b", + f"{self.BUCKET_URL}?list-type=2&prefix=litellm-b" + "&continuation-token=1ueGcxLPRx1Tr%2FXYExHnhbYLgveDs2J%2Fwm36Hy4vbOwM%3D", + f"{self.OUTPUT_BUCKET_URL}?list-type=2&prefix=litellm-batch-outputs%2F", + ] + assert [file.id for file in files] == [*self.PAGED_IDS, self.OUTPUT_BUCKET_ID] + + @pytest.mark.parametrize( + ("purpose", "bucket_snapshot"), + [ + pytest.param(None, {"s3_bucket_name": "my-bucket"}, id="outputs-share-the-input-bucket"), + pytest.param( + "batch", + {"s3_bucket_name": "my-bucket", "s3_output_bucket_name": "my-output-bucket"}, + id="input-purpose-requested", + ), + ], + ) + def test_file_list_leaves_the_output_bucket_alone_unless_an_unfiltered_list_needs_it( + self, monkeypatch, purpose, bucket_snapshot + ): + import httpx + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + input_route = respx.get(self.BUCKET_URL).mock(return_value=httpx.Response(200, content=self.LISTING)) + output_route = respx.get(self.OUTPUT_BUCKET_URL).mock( + return_value=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING) + ) + + files = litellm.file_list( + custom_llm_provider="bedrock", purpose=purpose, **_trusted_bucket_snapshot(**bucket_snapshot) + ) + + assert (input_route.call_count, output_route.call_count) == (1, 0) + assert [file.id for file in files] == [*self.BATCH_IDS, *(() if purpose else (self.OUTPUT_ID,))] + + def test_transform_list_files_next_request_walks_an_output_prefix_inside_the_input_bucket(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + litellm_params = _trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-bucket/out") + config = BedrockFilesConfig() + config.transform_list_files_request(purpose=None, optional_params={}, litellm_params=litellm_params) + litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM) + + output_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.LISTING), optional_params={}, litellm_params=litellm_params + ) + after_output_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.LISTING), optional_params={}, litellm_params=litellm_params + ) + + assert output_request == (self.BUCKET_URL, {"list-type": "2", "prefix": "out/litellm-batch-outputs/"}) + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + assert _sent_signature(signed_headers) == _s3_signature_for( + "GET", f"{self.BUCKET_URL}?list-type=2&prefix=out%2Flitellm-batch-outputs%2F", signed_headers + ) + assert after_output_request is None + + def test_transform_list_files_next_request_signs_the_continuation_page(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + config = BedrockFilesConfig() + config.transform_list_files_request(purpose="batch", optional_params={}, litellm_params=litellm_params) + first_signature = _sent_signature(litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM)) + + next_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.FIRST_PAGE), + optional_params={}, + litellm_params=litellm_params, + ) + + assert next_request == (self.BUCKET_URL, {**self.BATCH_QUERY, "continuation-token": self.CONTINUATION_TOKEN}) + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + signed_url = ( + f"{self.BUCKET_URL}?list-type=2&prefix=litellm-bedrock-files" + "&continuation-token=1ueGcxLPRx1Tr%2FXYExHnhbYLgveDs2J%2Fwm36Hy4vbOwM%3D" + ) + assert _sent_signature(signed_headers) == _s3_signature_for("GET", signed_url, signed_headers) + assert _sent_signature(signed_headers) != first_signature + + @pytest.mark.parametrize( + ("status_code", "content"), + [ + pytest.param(200, LAST_PAGE, id="last-page"), + pytest.param(403, b"AccessDenied", id="error-page"), + ], + ) + def test_transform_list_files_next_request_stops_after_the_last_page(self, monkeypatch, status_code, content): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + + next_request = BedrockFilesConfig().transform_list_files_next_request( + raw_response=httpx.Response(status_code, content=content), + optional_params={}, + litellm_params=litellm_params, + ) + + assert next_request is None + assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params + + def _mock_paged_listing(self, respx_module): + import httpx + + last_page = respx_module.get( + self.BUCKET_URL, params__contains={"continuation-token": self.CONTINUATION_TOKEN} + ).mock(return_value=httpx.Response(200, content=self.LAST_PAGE)) + first_page = respx_module.get(self.BUCKET_URL, params__contains=self.BATCH_QUERY).mock( + return_value=httpx.Response(200, content=self.FIRST_PAGE) + ) + return first_page, last_page + + def _assert_paged_listing(self, first_page, last_page, files): + assert (first_page.call_count, last_page.call_count) == (1, 1) + read_timeouts = [call.request.extensions["timeout"]["read"] for call in (*first_page.calls, *last_page.calls)] + assert read_timeouts == [12.0, 12.0] + assert "continuation-token" not in str(first_page.calls[0].request.url) + last_request = last_page.calls[0].request + assert _sent_signature(last_request.headers) == _s3_signature_for( + "GET", str(last_request.url), last_request.headers + ) + assert [file.id for file in files] == list(self.PAGED_IDS) + + def test_file_list_follows_continuation_tokens_across_pages(self, monkeypatch): + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + first_page, last_page = self._mock_paged_listing(respx) + files = litellm.file_list( + custom_llm_provider="bedrock", + purpose="batch", + timeout=12, + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + self._assert_paged_listing(first_page, last_page, files) + + @pytest.mark.asyncio + async def test_afile_list_follows_continuation_tokens_across_pages(self, monkeypatch): + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock: + first_page, last_page = self._mock_paged_listing(respx) + files = await litellm.afile_list( + custom_llm_provider="bedrock", + purpose="batch", + timeout=12, + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + self._assert_paged_listing(first_page, last_page, files) + + OVERSIZED_PAGE_SIZE = 3000 + OVERSIZED_PAGE_COUNT = 6 + + def _oversized_listing_page(self, page_index: int) -> bytes: + contents = "".join( + f"litellm-bedrock-files/page-{page_index}/obj-{index}.jsonl" + "2026-09-01T10:00:00.000Z1" + for index in range(self.OVERSIZED_PAGE_SIZE) + ) + continuation = ( + f"truepage-{page_index + 1}" + if page_index < self.OVERSIZED_PAGE_COUNT - 1 + else "false" + ) + return ( + '' + '' + f"{continuation}{contents}" + ).encode() + + def _mock_oversized_listing(self, respx_module): + import httpx + + def page_for(request): + token = request.url.params.get("continuation-token", "page-0") + return httpx.Response(200, content=self._oversized_listing_page(int(token.removeprefix("page-")))) + + return respx_module.get(self.BUCKET_URL, params__contains=self.BATCH_QUERY).mock(side_effect=page_for) + + def _assert_capped_listing(self, route, files): + from litellm.constants import MAX_FILE_LIST_LIMIT + + pages_needed = -(-MAX_FILE_LIST_LIMIT // self.OVERSIZED_PAGE_SIZE) + last_index = MAX_FILE_LIST_LIMIT - (pages_needed - 1) * self.OVERSIZED_PAGE_SIZE - 1 + assert pages_needed < self.OVERSIZED_PAGE_COUNT + assert route.call_count == pages_needed + assert len(files) == MAX_FILE_LIST_LIMIT + assert files[-1].id == f"s3://my-bucket/litellm-bedrock-files/page-{pages_needed - 1}/obj-{last_index}.jsonl" + + def test_file_list_stops_at_the_openai_listing_ceiling(self, monkeypatch): + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + route = self._mock_oversized_listing(respx) + files = litellm.file_list( + custom_llm_provider="bedrock", + purpose="batch", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + self._assert_capped_listing(route, files) + + @pytest.mark.asyncio + async def test_afile_list_stops_at_the_openai_listing_ceiling(self, monkeypatch): + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock: + route = self._mock_oversized_listing(respx) + files = await litellm.afile_list( + custom_llm_provider="bedrock", + purpose="batch", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + self._assert_capped_listing(route, files) + + def test_file_list_end_to_end_surfaces_the_s3_error_body(self, monkeypatch): + import httpx + import respx + + import litellm + from litellm.llms.bedrock.common_utils import BedrockError + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + + with respx.mock: + respx.get(self.BUCKET_URL, params__contains=self.BATCH_QUERY).mock( + return_value=httpx.Response(403, text="AccessDenied") + ) + with pytest.raises(BedrockError) as denied: + litellm.file_list(custom_llm_provider="bedrock", purpose="batch", **_bedrock_s3_params()) + + assert denied.value.status_code == 403 + assert "AccessDenied" in denied.value.message diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index e83a844c87e..c948dfb3553 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -7,7 +7,7 @@ API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.ht import json import asyncio -from unittest.mock import patch +from unittest.mock import Mock, patch import httpx @@ -151,6 +151,27 @@ class TestBedrockMantleConfig: # /openai/v1 base per the AWS model card. assert api_base == "https://bedrock-mantle.us-east-2.api.aws/openai/v1" + def test_region_prefixed_model_routes_to_that_region(self, monkeypatch, local_cost_map): + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(var, raising=False) + cfg = BedrockMantleChatConfig() + api_base, _ = cfg._get_openai_compatible_provider_info(None, None, model="us-gov-west-1/xai.grok-4.3") + assert api_base == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1" + + def test_aws_region_name_param_beats_model_region_prefix(self, monkeypatch, local_cost_map): + from litellm.types.router import GenericLiteLLMParams + + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(var, raising=False) + cfg = BedrockMantleChatConfig() + api_base, _ = cfg._get_openai_compatible_provider_info( + None, + None, + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1"), + model="us-gov-west-1/xai.grok-4.3", + ) + assert api_base == "https://bedrock-mantle.us-east-1.api.aws/openai/v1" + def test_default_api_base_fallback_to_us_east_1(self, monkeypatch): monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) @@ -686,6 +707,120 @@ class TestBedrockMantleProviderResolution: assert model == "openai.gpt-oss-20b" + def test_get_llm_provider_strips_region_prefix(self, monkeypatch, local_cost_map): + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(var, raising=False) + model, provider, _, api_base = litellm.get_llm_provider("bedrock_mantle/us-gov-west-1/xai.grok-4.3") + assert provider == "bedrock_mantle" + assert model == "xai.grok-4.3" + assert api_base == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1" + + def test_completion_region_prefixed_model_sends_bare_model_to_that_region(self, monkeypatch, local_cost_map): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + for var in ( + "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_BASE", + "BEDROCK_MANTLE_REGION", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_PROFILE", + ): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0") + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "xai.grok-4.3", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + }, + request=request, + ) + + handler = Mock(side_effect=respond) + response = litellm.completion( + model="bedrock_mantle/us-gov-west-1/xai.grok-4.3", + messages=[{"role": "user", "content": "hello"}], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + sent = handler.call_args.args[0] + assert str(sent.url) == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions" + assert json.loads(sent.content)["model"] == "xai.grok-4.3" + assert "/us-gov-west-1/bedrock/aws4_request" in sent.headers["Authorization"] + assert response._hidden_params["response_cost"] == pytest.approx( + 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + ) + + def test_responses_region_prefixed_model_prices_from_that_region_over_env_region(self, monkeypatch, local_cost_map): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + for var in ( + "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_BASE", + "BEDROCK_MANTLE_REGION", + "AWS_REGION", + "AWS_PROFILE", + ): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0") + + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "resp_test", + "object": "response", + "created_at": 1733529600, + "status": "completed", + "model": "xai.grok-4.3", + "output": [ + { + "type": "message", + "id": "msg_test", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "usage": {"input_tokens": 38, "output_tokens": 20, "total_tokens": 58}, + }, + request=request, + ) + + handler = Mock(side_effect=respond) + response = litellm.responses( + model="bedrock_mantle/us-gov-west-1/xai.grok-4.3", + input="hello", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + sent = handler.call_args.args[0] + assert str(sent.url) == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/responses" + assert json.loads(sent.content)["model"] == "xai.grok-4.3" + assert "/us-gov-west-1/bedrock/aws4_request" in sent.headers["Authorization"] + assert response._hidden_params["response_cost"] == pytest.approx( + 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + ) + + class TestBedrockMantlePricing: """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing.""" diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d9a065db3cb..ea2c925212d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -6390,6 +6390,106 @@ async def test_get_team_membership_db_fetch_returns_validated_membership(): assert result.spend == 1.5 +@pytest.mark.asyncio +async def test_get_team_membership_negative_caches_a_missing_row(): + """ + Regression (LIT-7358): a member with no LiteLLM_TeamMembership row is the common lite-login case, + and the session-token refresh reads this loader on every request. Before the fix a missing row + returned None without caching, so every request re-queried the DB. The miss must be cached so the + second request serves from cache and never touches the DB. + """ + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, + team_membership_reservation_cache_key, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + + cache = UserApiKeyCache() + + first = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + second = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + assert first is None + assert second is None + mock_prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + cached = await cache.async_get_cache( + key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1") + ) + assert cached == NO_TEAM_MEMBERSHIP_SENTINEL + + +@pytest.mark.asyncio +async def test_get_team_membership_reads_sentinel_as_no_membership_not_a_model(): + """ + The negative-cache sentinel is a plain string sharing the key a serialized membership uses. + A pre-seeded sentinel must read back as None (no DB read), never be mistaken for a membership. + """ + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, + team_membership_reservation_cache_key, + ) + + cache = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1"), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + + result = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + + assert result is None + mock_prisma_client.db.litellm_teammembership.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_evicts_the_negative_cache_sentinel(): + """ + A member who later gains a per-member budget writes a membership row and calls + invalidate_team_member_spend_state. That must drop a cached "no membership" sentinel so the next + request re-reads the DB and honors the new budget instead of serving the stale miss until TTL. + """ + from litellm.proxy.auth.auth_checks import get_team_membership, invalidate_team_member_spend_state + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + cache = UserApiKeyCache() + membership_row = MagicMock() + membership_row.dict = lambda: {"user_id": "u-1", "team_id": "t-1", "spend": 0.0} + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(side_effect=[None, membership_row]) + + before = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + assert before is None + + await invalidate_team_member_spend_state(user_id="u-1", team_id="t-1", user_api_key_cache=cache) + assert ( + await cache.async_get_cache(key=team_membership_reservation_cache_key(user_id="u-1", team_id="t-1")) + is None + ) + + after = await get_team_membership( + user_id="u-1", team_id="t-1", prisma_client=mock_prisma_client, user_api_key_cache=cache + ) + assert after is not None + assert after.user_id == "u-1" + assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 + + @pytest.mark.asyncio async def test_get_access_object_db_fetch_returns_validated_access_group(): from litellm.proxy._types import LiteLLM_AccessGroupTable diff --git a/tests/test_litellm/proxy/auth/test_resolvers_grants.py b/tests/test_litellm/proxy/auth/test_resolvers_grants.py new file mode 100644 index 00000000000..3f6d943bf98 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_resolvers_grants.py @@ -0,0 +1,201 @@ +from fastapi import HTTPException +import pytest + +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, + ProxyException, +) +from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError +from litellm.proxy.auth.resolvers.grants import ( + GrantResolver, + LookupDegraded, + NotAMember, + ResolvedGrants, + TeamGone, + UserGone, + UserLookup, + raise_public, + user_models, +) + +USER_ID = "user-1" +TEAM_ID = "team-1" + + +class _Loaders: + """Fake row readers standing in for the ``auth_checks`` loaders, recording every call they receive.""" + + def __init__(self, *, user=None, team=None, membership=None, user_error=None, team_error=None): + self._user = user + self._team = team + self._membership = membership + self._user_error = user_error + self._team_error = team_error + self.user_calls = [] + self.team_calls = [] + self.membership_calls = [] + + async def load_user(self, **kwargs): + self.user_calls.append(kwargs) + if self._user_error is not None: + raise self._user_error + return self._user + + async def load_team(self, **kwargs): + self.team_calls.append(kwargs) + if self._team_error is not None: + raise self._team_error + return self._team + + async def load_membership(self, **kwargs): + self.membership_calls.append(kwargs) + return self._membership + + def resolver(self) -> GrantResolver: + return GrantResolver( + object(), + object(), + load_user=self.load_user, + load_team=self.load_team, + load_membership=self.load_membership, + ) + + +def _user(teams=(TEAM_ID,), user_id=USER_ID) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=user_id, user_role="internal_user", teams=list(teams), models=["gpt-5.5"]) + + +def _team(models=("gpt-5.5",)) -> LiteLLM_TeamTableCachedObj: + return LiteLLM_TeamTableCachedObj(team_id=TEAM_ID, team_alias="alias", models=list(models)) + + +async def test_resolve_returns_live_rows_for_a_member(): + membership = LiteLLM_TeamMembership(user_id=USER_ID, team_id=TEAM_ID, spend=1.5) + loaders = _Loaders(user=_user(), team=_team(models=("new-a", "new-b")), membership=membership) + + outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert outcome == ResolvedGrants( + user_object=_user(), + team_object=_team(models=("new-a", "new-b")), + team_membership=membership, + effective_user_id=USER_ID, + ) + assert loaders.team_calls[0]["team_id"] == TEAM_ID + assert loaders.membership_calls[0]["user_id"] == USER_ID + assert loaders.membership_calls[0]["team_id"] == TEAM_ID + + +async def test_resolve_denies_a_user_removed_from_the_team_without_reading_the_team(): + loaders = _Loaders(user=_user(teams=("other-team",)), team=_team()) + + outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert outcome == NotAMember(user_id=USER_ID, team_id=TEAM_ID) + assert loaders.team_calls == [] + + +async def test_resolve_reports_a_deleted_user(): + loaders = _Loaders(user_error=UserNotFoundError(user_id=USER_ID), team=_team()) + + outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert outcome == UserGone(user_id=USER_ID) + assert loaders.team_calls == [] + + +async def test_resolve_reports_a_deleted_team(): + loaders = _Loaders(user=_user(), team_error=TeamNotFoundError(team_id=TEAM_ID)) + + outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert outcome == TeamGone(team_id=TEAM_ID) + + +@pytest.mark.parametrize( + "loaders", + [ + _Loaders(user_error=Exception("No db connected")), + _Loaders(user=_user(), team_error=HTTPException(status_code=500, detail="db timeout")), + ], + ids=["user-read-failed", "team-read-failed"], +) +async def test_resolve_marks_an_unreadable_row_as_degraded_not_denied(loaders): + outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert isinstance(outcome, LookupDegraded) + + +async def test_resolve_without_a_team_skips_team_and_membership_reads(): + loaders = _Loaders(user=_user(teams=())) + + outcome = await loaders.resolver().resolve(UserLookup(user_id=USER_ID), team_id=None) + + assert outcome == ResolvedGrants( + user_object=_user(teams=()), team_object=None, team_membership=None, effective_user_id=USER_ID + ) + assert loaders.team_calls == [] + assert loaders.membership_calls == [] + + +async def test_resolve_identity_reads_membership_under_the_matched_rows_id(): + legacy_uuid = "bb8ab11f-09aa-47ae-b063-6e80506ac3bc" + loaders = _Loaders(user=_user(user_id=legacy_uuid)) + + user_object, _membership, effective_user_id = await loaders.resolver().resolve_identity( + UserLookup(user_id="matt@example.com", user_email="matt@example.com", sso_user_id="matt@example.com"), + team_id=TEAM_ID, + ) + + assert user_object is not None and user_object.user_id == legacy_uuid + assert effective_user_id == legacy_uuid + assert loaders.membership_calls[0]["user_id"] == legacy_uuid + assert loaders.user_calls[0]["user_email"] == "matt@example.com" + + +async def test_resolve_identity_without_a_user_id_reads_nothing(): + loaders = _Loaders(user=_user()) + + outcome = await loaders.resolver().resolve_identity(UserLookup(user_id=None), team_id=TEAM_ID) + + assert outcome == (None, None, None) + assert loaders.user_calls == [] + assert loaders.membership_calls == [] + + +async def test_resolve_identity_lets_loader_errors_surface(): + loaders = _Loaders(user_error=UserNotFoundError(user_id=USER_ID)) + + with pytest.raises(UserNotFoundError): + await loaders.resolver().resolve_identity(UserLookup(user_id=USER_ID), team_id=None) + + +def test_raise_public_maps_a_deleted_user_to_401(): + with pytest.raises(ProxyException) as exc_info: + raise_public(UserGone(user_id=USER_ID)) + assert exc_info.value.code == "401" + assert USER_ID in exc_info.value.message + + +def test_raise_public_maps_a_removed_member_to_403(): + with pytest.raises(HTTPException) as exc_info: + raise_public(NotAMember(user_id=USER_ID, team_id=TEAM_ID)) + assert exc_info.value.status_code == 403 + assert TEAM_ID in str(exc_info.value.detail) + + +def test_raise_public_maps_a_deleted_team_to_404(): + with pytest.raises(TeamNotFoundError) as exc_info: + raise_public(TeamGone(team_id=TEAM_ID)) + assert exc_info.value.status_code == 404 + + +@pytest.mark.parametrize( + ("stored", "expected"), + [(["gpt-5.5", "claude-opus-5"], ("gpt-5.5", "claude-opus-5")), ([], ()), ([{"not": "a model"}], ())], + ids=["models", "empty", "unusable-column"], +) +def test_user_models_reads_the_column_as_a_tuple_of_names(stored, expected): + assert user_models(LiteLLM_UserTable(user_id=USER_ID, models=stored)) == expected 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 0cdbcde6abc..fded86d43af 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 @@ -32,7 +32,7 @@ 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.auth_checks import TeamNotFoundError, UserNotFoundError, get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, @@ -6266,6 +6266,192 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert result.is_session_token is True +SESSION_TEAM_ID = "team-abc" +SESSION_USER_ID = "member-1" + + +def _mint_session_token( + monkeypatch, + *, + role=LitellmUserRoles.INTERNAL_USER, + team_id=SESSION_TEAM_ID, + team_models=("stale-model",), + models=(), +): + """Mint a ``lite login`` token carrying the grants as they were at login time.""" + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test") + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + user_info = LiteLLM_UserTable( + user_id=SESSION_USER_ID, user_email="user@example.com", user_role=role.value, models=list(models) + ) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token( + user_info, team_id=team_id, team_alias="stale-alias", team_models=list(team_models) + ) + + +def _session_user_row(*, teams=(SESSION_TEAM_ID,), role=LitellmUserRoles.INTERNAL_USER, models=()): + return LiteLLM_UserTable(user_id=SESSION_USER_ID, user_role=role.value, teams=list(teams), models=list(models)) + + +async def _authenticate_session_token_against_db( + cli_token, *, user_row=None, team_row=None, membership_row=None, user_error=None, team_error=None +): + """Drive the real builder for a session token with the DB row readers replaced by the given rows or + errors. Returns the ``_return_user_api_key_auth_obj`` mock so the caller can read the token it was + handed; a denial surfaces as the exception the builder raises.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + attrs = _proxy_attrs_for_db_lookup() + attrs["prisma_client"].db.litellm_teammembership.find_first = AsyncMock(return_value=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + assemble = AsyncMock(return_value=UserAPIKeyAuth(user_id=SESSION_USER_ID, is_session_token=True)) + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( # test-quality-ok: the builder has no injection seam for its assembler yet + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", assemble + ), + patch( # test-quality-ok: the builder reads its DB row loaders off module globals + "litellm.proxy.auth.user_api_key_auth.get_user_object", + AsyncMock(return_value=user_row, side_effect=user_error), + ), + patch( # test-quality-ok: the builder reads its DB row loaders off module globals + "litellm.proxy.auth.user_api_key_auth.get_team_object", + AsyncMock(return_value=team_row, side_effect=team_error), + ), + patch( # test-quality-ok: the builder reads its DB row loaders off module globals + "litellm.proxy.auth.user_api_key_auth.get_team_membership", + AsyncMock(return_value=membership_row), + ), + ): + await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {cli_token}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + return assemble + + +@pytest.mark.asyncio +async def test_session_token_reads_team_grants_from_the_live_team_row(monkeypatch): + """LIT-7358: a lite login token snapshots the team's models at login, so adding a model to the team did + nothing for that CLI until the user logged in again. The team row has to be re-read on every request.""" + from litellm.models.team import LiteLLM_ModelTable + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj + + cli_token = _mint_session_token(monkeypatch, team_models=("stale-model",)) + live_team = LiteLLM_TeamTableCachedObj( + team_id=SESSION_TEAM_ID, + team_alias="renamed-team", + models=["gpt-5.5", "claude-opus-5"], + litellm_model_table=LiteLLM_ModelTable(model_aliases={"fast": "gpt-5.5"}, created_by="a", updated_by="a"), + ) + membership = LiteLLM_TeamMembership(user_id=SESSION_USER_ID, team_id=SESSION_TEAM_ID, spend=2.5) + + assemble = await _authenticate_session_token_against_db( + cli_token, user_row=_session_user_row(), team_row=live_team, membership_row=membership + ) + + token = assemble.call_args.kwargs["valid_token_dict"] + assert token["team_models"] == ["gpt-5.5", "claude-opus-5"] + assert token["team_alias"] == "renamed-team" + assert token["team_model_aliases"] == {"fast": "gpt-5.5"} + assert token["team_member_spend"] == 2.5 + assert token["is_session_token"] is True + + +@pytest.mark.asyncio +async def test_session_token_without_a_team_reads_models_from_the_live_user_row(monkeypatch): + cli_token = _mint_session_token(monkeypatch, team_id=None, team_models=(), models=("stale-model",)) + + assemble = await _authenticate_session_token_against_db( + cli_token, user_row=_session_user_row(teams=(), models=("gpt-5.5",)) + ) + + assert assemble.call_args.kwargs["valid_token_dict"]["models"] == ["gpt-5.5"] + + +@pytest.mark.asyncio +async def test_demoted_admin_session_token_loses_admin_on_the_next_request(monkeypatch): + """The role baked into the token used to send a former admin down the admin early return forever.""" + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + cli_token = _mint_session_token(monkeypatch, role=LitellmUserRoles.PROXY_ADMIN) + + assemble = await _authenticate_session_token_against_db( + cli_token, + user_row=_session_user_row(role=LitellmUserRoles.INTERNAL_USER), + team_row=LiteLLM_TeamTableCachedObj(team_id=SESSION_TEAM_ID, models=["gpt-5.5"]), + ) + + assemble.assert_awaited_once() + assert assemble.call_args.kwargs["valid_token_dict"]["user_role"] == LitellmUserRoles.INTERNAL_USER + + +@pytest.mark.asyncio +async def test_session_token_is_refused_once_the_user_leaves_the_team(monkeypatch): + cli_token = _mint_session_token(monkeypatch) + + with pytest.raises(ProxyException) as exc_info: + await _authenticate_session_token_against_db(cli_token, user_row=_session_user_row(teams=("other-team",))) + + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + assert SESSION_TEAM_ID in exc_info.value.message + + +@pytest.mark.asyncio +async def test_session_token_is_refused_once_the_user_is_deleted(monkeypatch): + cli_token = _mint_session_token(monkeypatch) + + with pytest.raises(ProxyException) as exc_info: + await _authenticate_session_token_against_db(cli_token, user_error=UserNotFoundError(user_id=SESSION_USER_ID)) + + assert exc_info.value.code == str(status.HTTP_401_UNAUTHORIZED) + assert exc_info.value.type == ProxyErrorTypes.auth_error + + +@pytest.mark.asyncio +async def test_session_token_is_refused_once_the_team_is_deleted(monkeypatch): + cli_token = _mint_session_token(monkeypatch) + + with pytest.raises(ProxyException) as exc_info: + await _authenticate_session_token_against_db( + cli_token, user_row=_session_user_row(), team_error=TeamNotFoundError(team_id=SESSION_TEAM_ID) + ) + + assert exc_info.value.code == str(status.HTTP_404_NOT_FOUND) + + +@pytest.mark.asyncio +async def test_session_token_keeps_minted_grants_when_the_team_row_cannot_be_read(monkeypatch): + """A DB hiccup says nothing about the caller, so the grants minted at login stand for that request.""" + from fastapi import HTTPException + + cli_token = _mint_session_token(monkeypatch, team_models=("stale-model",)) + + assemble = await _authenticate_session_token_against_db( + cli_token, user_row=_session_user_row(), team_error=HTTPException(status_code=500, detail="db timeout") + ) + + token = assemble.call_args.kwargs["valid_token_dict"] + assert token["team_models"] == ["stale-model"] + assert token["team_alias"] == "stale-alias" + + @pytest.mark.asyncio async def test_cli_session_token_authenticates_when_jwt_auth_enabled_without_license(monkeypatch): """A lite login token is an encrypted (non-JWT) session blob. With 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 1ab50cc30de..761cd0685f2 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1195,6 +1195,61 @@ async def test_health_services_endpoint_newrelic_allows_proxy_admin(admin_role): mock_instance.async_health_check.assert_awaited_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "role", + [ + None, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + LitellmUserRoles.TEAM, + LitellmUserRoles.CUSTOMER, + ], +) +async def test_health_services_endpoint_webhook_blocks_non_admin(role): + """ + /health/services?service=webhook fires a real budget_crossed alert for the + caller's user_id and writes the same dedup cache entry the auth-time user + budget alert uses, so a non-admin could suppress their own real alert for + the cache TTL. Only proxy admins may trigger it. + """ + mock_proxy_logging = MagicMock() + mock_proxy_logging.budget_alerts = AsyncMock() + user_api_key_dict = UserAPIKeyAuth(token="non-admin-token", user_id="non-admin-user", user_role=role) + + with patch( # test-quality-ok: endpoint reads proxy_server module globals, same pattern as sibling tests + "litellm.proxy.proxy_server.proxy_logging_obj", + mock_proxy_logging, + ): + with pytest.raises(ProxyException) as exc_info: + await health_services_endpoint(user_api_key_dict=user_api_key_dict, service="webhook") + + assert str(exc_info.value.code) == "403" + mock_proxy_logging.budget_alerts.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "admin_role", + [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY], +) +async def test_health_services_endpoint_webhook_allows_proxy_admin(admin_role): + mock_proxy_logging = MagicMock() + mock_proxy_logging.budget_alerts = AsyncMock() + user_api_key_dict = UserAPIKeyAuth(token="admin-token", user_id="admin-user", user_role=admin_role) + + with patch( # test-quality-ok: endpoint reads proxy_server module globals, same pattern as sibling tests + "litellm.proxy.proxy_server.proxy_logging_obj", + mock_proxy_logging, + ): + await health_services_endpoint(user_api_key_dict=user_api_key_dict, service="webhook") + + mock_proxy_logging.budget_alerts.assert_awaited_once() + sent = mock_proxy_logging.budget_alerts.await_args.kwargs + assert sent["type"] == "user_budget" + assert sent["user_info"].user_id == "admin-user" + + @pytest.fixture(scope="function") def proxy_client(monkeypatch): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 2fb496d6231..0e1831614ac 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -13941,6 +13941,52 @@ async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent( assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5 +@pytest.mark.asyncio +async def test_evict_created_membership_caches_drops_the_negative_sentinel(): + """ + Regression: a membership-create path (/team/member_add, the /team/update budget backfill) must + evict any cached "no membership" sentinel a prior session-token read left, so a per-member budget + attached at create time is enforced on the next request instead of after the membership cache TTL. + Uses a real cache so the assertion is that the sentinel is actually gone, not that a mock was called. + """ + from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, + UserApiKeyCache, + team_membership_reservation_cache_key, + ) + from litellm.proxy.management_endpoints.team_endpoints import _evict_created_membership_caches + + cache = UserApiKeyCache() + kept_key = team_membership_reservation_cache_key(user_id="carol", team_id="team-eviction") + evicted_key = team_membership_reservation_cache_key(user_id="bob", team_id="team-eviction") + await cache.async_set_cache(key=kept_key, value=NO_TEAM_MEMBERSHIP_SENTINEL) + await cache.async_set_cache(key=evicted_key, value=NO_TEAM_MEMBERSHIP_SENTINEL) + + await _evict_created_membership_caches(user_ids=("bob",), team_id="team-eviction", user_api_key_cache=cache) + + assert await cache.async_get_cache(key=evicted_key) is None + assert await cache.async_get_cache(key=kept_key) == NO_TEAM_MEMBERSHIP_SENTINEL + + +def test_member_user_ids_keeps_only_string_user_ids(): + """ + The /team/update backfill feeds Prisma-deserialized member dicts here; a row can be missing + user_id or carry a non-string value. Only real string ids may reach invalidate_team_member_spend_state, + so those get eviction and the malformed rows are dropped rather than crashing the update. + """ + from litellm.proxy.management_endpoints.team_endpoints import _member_user_ids + + members = [ + {"user_id": "alice", "role": "admin"}, + {"role": "user"}, + {"user_id": None, "role": "user"}, + {"user_id": 123, "role": "user"}, + {"user_id": "bob", "role": "user"}, + ] + + assert _member_user_ids(members) == ("alice", "bob") + + def _team_spend_by_user_team(team_id: str, team_alias: str, member: Member, permissions: list[str]) -> MagicMock: team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = team_id diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 5f1e7e1fe0c..b696d9ebe5f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -2441,6 +2441,63 @@ def test_list_files_resolves_wildcard_deployment_credentials( proxy_logging_obj.post_call_failure_hook.assert_not_called() +def test_list_files_by_model_returns_an_openai_page_for_a_provider_listing( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + from litellm.types.llms.openai import FileListPage + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=None) + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + listed_files = [ + OpenAIFileObject( + id=f"file-{index}", + bytes=index, + created_at=index, + filename=f"{index}.jsonl", + object="file", + purpose="batch", + status="uploaded", + ) + for index in (1, 2) + ] + + async def _mock_afile_list(**kwargs): + return list(listed_files) + + monkeypatch.setattr(litellm, "afile_list", _mock_afile_list) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + + try: + response = client.get( + "/v1/files?target_model_names=gpt-3.5-turbo", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + body = response.json() + assert body["object"] == "list" + assert [listed["id"] for listed in body["data"]] == ["file-1", "file-2"] + assert (body["first_id"], body["last_id"], body["has_more"]) == ("file-1", "file-2", False) + hook_response = proxy_logging_obj.post_call_success_hook.call_args.kwargs["response"] + assert isinstance(hook_response, FileListPage) + assert [listed.id for listed in hook_response.data] == ["file-1", "file-2"] + + def test_list_files_model_routing_does_not_forward_custom_llm_provider_twice( mocker: MockerFixture, monkeypatch, llm_router: Router ): @@ -4668,6 +4725,249 @@ def test_create_file_path_traversal_filename_rejected_before_forwarding(monkeypa assert forwarded_calls == [] +def test_list_files_target_model_names_passes_trusted_bedrock_credentials( + mocker: MockerFixture, monkeypatch +): + """ + GET /v1/files?target_model_names= must hand the deployment's + immutable credential snapshot to litellm.afile_list, since Bedrock resolves + the S3 bucket to list from that snapshot rather than from request params. + """ + from types import MappingProxyType + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + bedrock_router = Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + }, + ] + ) + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=[]) + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_list(**kwargs): + captured_kwargs.update(kwargs) + return [] + + monkeypatch.setattr(litellm, "afile_list", _mock_afile_list) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + + try: + response = client.get( + "/v1/files?target_model_names=bedrock-claude&purpose=batch", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs["custom_llm_provider"] == "bedrock" + assert captured_kwargs["purpose"] == "batch" + trusted_credentials = captured_kwargs["_litellm_internal_model_credentials"] + assert isinstance(trusted_credentials, MappingProxyType) + assert trusted_credentials["s3_bucket_name"] == "my-bucket" + proxy_logging_obj.post_call_failure_hook.assert_not_called() + + +def test_delete_file_answers_400_for_an_id_outside_the_configured_bucket(mocker: MockerFixture, monkeypatch): + from urllib.parse import quote + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + bedrock_router = Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + }, + ] + ) + + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=[]) + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + foreign_file_id: Final = quote("s3://other-bucket/litellm-bedrock-files/job-123/input.jsonl", safe="") + + try: + with respx.mock: + response = client.delete( + f"/v1/files/{foreign_file_id}?model=bedrock-claude", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 400, response.text + assert "configured storage bucket" in response.json()["error"]["message"] + + +def _cloud_files_router() -> Router: + return Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "secret", + "aws_region_name": "us-west-2", + "s3_bucket_name": "my-bucket", + }, + }, + { + "model_name": "vertex-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-3.8-flash", + "vertex_project": "my-project", + "vertex_location": "us-central1", + "gcs_bucket_name": "my-gcs-bucket", + }, + }, + ] + ) + + +RAW_S3_FILE_ID: Final = "s3://my-bucket/litellm-batch-outputs/job-123/abc/input.jsonl.out" +RAW_GCS_FILE_ID: Final = "gs://my-gcs-bucket/litellm-vertex-files/publishers/google/models/gemini-3.8-flash/abc123" + + +@pytest.mark.parametrize( + ("route_prefix", "raw_file_id", "model_name"), + ( + ("/bedrock/v1/files", RAW_S3_FILE_ID, "bedrock-claude"), + ("/v1/files", RAW_S3_FILE_ID, "bedrock-claude"), + ("/files", RAW_S3_FILE_ID, "bedrock-claude"), + ("/vertex_ai/v1/files", RAW_GCS_FILE_ID, "vertex-gemini"), + ), +) +def test_delete_file_answers_403_for_a_raw_cloud_id_from_a_non_admin_key( + mocker: MockerFixture, monkeypatch, route_prefix: str, raw_file_id: str, model_name: str +): + from urllib.parse import quote + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + bedrock_router = _cloud_files_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + afile_delete = mocker.AsyncMock() + monkeypatch.setattr(litellm, "afile_delete", afile_delete) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + models=["bedrock-claude", "vertex-gemini"], + ) + + try: + response = client.delete( + f"{route_prefix}/{quote(raw_file_id, safe='')}?model={model_name}", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 403, response.text + assert "proxy admin" in response.json()["error"]["message"] + afile_delete.assert_not_called() + + +def test_delete_file_forwards_a_raw_cloud_id_from_a_proxy_admin_key(mocker: MockerFixture, monkeypatch): + from urllib.parse import quote + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + bedrock_router = _cloud_files_router() + proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router) + proxy_logging_obj.update_request_status = mocker.AsyncMock() + proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock() + + captured_kwargs: dict = {} + + async def _mock_afile_delete(**kwargs): + captured_kwargs.update(kwargs) + return OpenAIFileObject( + id=RAW_S3_FILE_ID, + object="file", + bytes=2, + created_at=1234567890, + filename="input.jsonl.out", + purpose="batch_output", + status="processed", + ) + + monkeypatch.setattr(litellm, "afile_delete", _mock_afile_delete) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="test-user", + ) + + try: + response = client.delete( + f"/bedrock/v1/files/{quote(RAW_S3_FILE_ID, safe='')}?model=bedrock-claude", + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + assert captured_kwargs.get("file_id") == RAW_S3_FILE_ID + assert captured_kwargs.get("custom_llm_provider") == "bedrock" + proxy_logging_obj.post_call_failure_hook.assert_not_called() + + def _setup_managed_file_route_answering_404( mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router ) -> None: diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 6e43ac4a12b..8283ee8395a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -124,7 +124,10 @@ def _reconstruct_ui_where_from_sql(sql_query, params): sess = re.fullmatch(r"session_id LIKE \$(\d+)", cond) status = re.fullmatch(r"status = \$(\d+)", cond) api_key_not_in = re.fullmatch(r"api_key NOT IN \(\$(\d+), \$(\d+)\)", cond) - if gte: + req_or_call = re.fullmatch(r"\(request_id = \$(\d+) OR litellm_call_id = \$\1\)", cond) + if req_or_call: + where["request_id_or_call_id"] = params[int(req_or_call.group(1)) - 1] + elif gte: date_bounds["gte"] = _iso(params[int(gte.group(1)) - 1]) elif lte: date_bounds["lte"] = _iso(params[int(lte.group(1)) - 1]) @@ -213,6 +216,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No query_observer(sql_query, params) if "mcp_tool_call_count" in sql_query: return [] + if 'SELECT DISTINCT "user", team_id' in sql_query: + return _emulate_spend_log_owner_lookup(mock_spend_logs, sql_query, params) filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) total = len(filtered) if "COUNT(*)" in sql_query: @@ -220,7 +225,13 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return [{"total_count": min(total, cap_plus_one)}] page_size = params[-2] if len(params) >= 2 else 50 skip = params[-1] if len(params) >= 1 else 0 - return [row for row in filtered[skip : skip + page_size]] + exact_first = re.search(r"ORDER BY \(request_id = \$(\d+)\) DESC", sql_query) + ordered = ( + sorted(filtered, key=lambda row: row["request_id"] == params[int(exact_first.group(1)) - 1], reverse=True) + if exact_first + else filtered + ) + return [row for row in ordered[skip : skip + page_size]] class MockPrismaClient: def __init__(self): @@ -424,37 +435,173 @@ def test_can_user_view_spend_log_false_for_other_roles(): assert spend_management_endpoints._can_user_view_spend_log(auth) is False +def _emulate_spend_log_owner_lookup(rows, sql_query, params): + """Emulate the ownership lookup SQL over an in-memory spend-log corpus, + honoring DISTINCT and any literal LIMIT the query carries so a capped or + non-distinct query produces the truncated result it would in Postgres.""" + lookup_id = params[0] + matches = [ + {"user": row.get("user"), "team_id": row.get("team_id")} + for row in rows + if lookup_id in (row.get("request_id"), row.get("litellm_call_id")) + ] + if "DISTINCT" in sql_query: + deduped = [] + for match in matches: + if match not in deduped: + deduped.append(match) + matches = deduped + limit = re.search(r"LIMIT\s+(\d+)", sql_query, re.IGNORECASE) + if limit is not None: + matches = matches[: int(limit.group(1))] + return matches + + +def _make_owner_lookup_prisma(rows): + class MockDB: + async def query_raw(self, sql_query, *params): + return _emulate_spend_log_owner_lookup(rows, sql_query, params) + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + return MockPrisma() + + @pytest.mark.asyncio async def test_assert_user_can_view_request_id_rejects_both_users_none(): """ API keys with user_id=None must not be treated as owning a log whose user field is None (avoid None == None bypass). """ - - class MockRow: - user = None - team_id = None - - class MockSpendLogs: - async def find_unique(self, where, include=None): - return MockRow() - - class MockDB: - def __init__(self): - self.litellm_spendlogs = MockSpendLogs() - - class MockPrisma: - def __init__(self): - self.db = MockDB() + prisma = _make_owner_lookup_prisma( + [{"request_id": "req-none-user", "litellm_call_id": None, "user": None, "team_id": None}] + ) auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=None) with pytest.raises(HTTPException) as exc_info: await spend_management_endpoints._assert_user_can_view_request_id( - MockPrisma(), auth, "req-none-user" + prisma, auth, "req-none-user" ) assert exc_info.value.status_code == 403 +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_rejects_when_no_match_is_owned(): + """An id whose every matching row belongs to other tenants is refused outright, + so the relaxed date window of an id lookup cannot reach a foreign row.""" + prisma = _make_owner_lookup_prisma( + [ + {"request_id": "foreign-request", "litellm_call_id": "shared-id", "user": "tenant_a", "team_id": None}, + {"request_id": "shared-id", "litellm_call_id": "other-call-id", "user": "tenant_b", "team_id": None}, + ] + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller") + with pytest.raises(HTTPException) as exc_info: + await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "shared-id") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_allows_owner_despite_foreign_collision(): + """ + litellm_call_id comes from the client-settable x-litellm-call-id header, so + another tenant can mint a row whose call id equals the caller's request_id. + That collision must not lock the caller out of their own row: the pre-check + passes once one match is theirs, and the scoped data queries keep the foreign + row out of the result. Regression for the every-match-must-be-owned rule that + let any tenant deny another's lookup by reusing their id. + """ + prisma = _make_owner_lookup_prisma( + [ + { + "request_id": "attacker-own-request", + "litellm_call_id": "victim-request-id", + "user": "attacker", + "team_id": None, + }, + { + "request_id": "victim-request-id", + "litellm_call_id": "victim-call-id", + "user": "victim", + "team_id": None, + }, + ] + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="victim") + result = await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "victim-request-id") + + assert result is None + + +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_finds_owner_past_any_row_cap(): + """ + An attacker can mint hundreds of rows carrying the victim's request_id as + their litellm_call_id, so a capped or sampled ownership read could exhaust + its cap on attacker-owned rows and never see the victim's own row, locking + the victim out of their lookup. The ownership read must consider every + matching row's owner no matter how many rows match. Regression for the + find_many(take=100) sample the first fix used. + """ + rows = [ + { + "request_id": f"attacker-request-{i}", + "litellm_call_id": "victim-request-id", + "user": "attacker", + "team_id": None, + } + for i in range(150) + ] + rows.append( + { + "request_id": "victim-request-id", + "litellm_call_id": "victim-call-id", + "user": "victim", + "team_id": None, + } + ) + prisma = _make_owner_lookup_prisma(rows) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="victim") + result = await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "victim-request-id") + + assert result is None + + +@pytest.mark.asyncio +async def test_assert_user_can_view_request_id_allows_when_every_match_is_owned(): + """The same ambiguous id matching more than one row is fine when every match + belongs to the caller (e.g. two of the caller's own requests happen to share + a request_id/litellm_call_id pairing); only a foreign match should block it.""" + prisma = _make_owner_lookup_prisma( + [ + { + "request_id": "shared-request-id", + "litellm_call_id": "caller-call-a", + "user": "caller", + "team_id": None, + }, + { + "request_id": "caller-request-b", + "litellm_call_id": "shared-request-id", + "user": "caller", + "team_id": None, + }, + ] + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller") + result = await spend_management_endpoints._assert_user_can_view_request_id( + prisma, auth, "shared-request-id" + ) + + assert result is None + + @pytest.mark.asyncio async def test_assert_user_can_view_request_id_rejects_missing_row(): """ @@ -462,24 +609,11 @@ async def test_assert_user_can_view_request_id_rejects_missing_row(): authorize reading the payload from cold storage; a missing row is not the same as an owned row. """ - - class MockSpendLogs: - async def find_unique(self, where, include=None): - return None - - class MockDB: - def __init__(self): - self.litellm_spendlogs = MockSpendLogs() - - class MockPrisma: - def __init__(self): - self.db = MockDB() + prisma = _make_owner_lookup_prisma([]) auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1") with pytest.raises(HTTPException) as exc_info: - await spend_management_endpoints._assert_user_can_view_request_id( - MockPrisma(), auth, "req-missing-row" - ) + await spend_management_endpoints._assert_user_can_view_request_id(prisma, auth, "req-missing-row") assert exc_info.value.status_code == 403 @@ -507,6 +641,7 @@ def test_ui_view_request_response_forbids_non_admin_without_db(client, monkeypat ignored_keys = [ "request_id", + "litellm_call_id", "metadata.litellm_call_id", "session_id", "startTime", @@ -2216,7 +2351,10 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window( def filter_fn(where): captured["where"] = where rows = _filter_logs_by_date_range(mock_spend_logs, where) - if where.get("request_id"): + rid_either = where.get("request_id_or_call_id") + if rid_either: + rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))] + elif where.get("request_id"): rows = [r for r in rows if r["request_id"] == where["request_id"]] return rows @@ -2246,9 +2384,82 @@ async def test_ui_view_spend_logs_request_id_lookup_ignores_date_window( data = response.json() assert data["total"] == 1 assert data["data"][0]["request_id"] == "req-old" - # Query dropped the time window and scoped solely by the primary key. + # Query dropped the time window and scoped solely by the id lookup. assert "startTime" not in captured["where"] - assert captured["where"]["request_id"] == "req-old" + assert captured["where"]["request_id_or_call_id"] == "req-old" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_lookup_matches_litellm_call_id( + client, monkeypatch +): + """ + LIT-6302: success rows are keyed by the upstream provider response id, so a + lookup with the x-litellm-call-id response header value found nothing. The id + lookup now matches request_id OR litellm_call_id, resolving the header value. + """ + today = datetime.datetime.now(timezone.utc) + mock_spend_logs = [ + { + "id": "log_provider_keyed", + "request_id": "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm", + "litellm_call_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.05, + "startTime": today.isoformat(), + "model": "gpt-4", + }, + { + "id": "log_other", + "request_id": "chatcmpl-other", + "litellm_call_id": "11111111-2222-3333-4444-555555555555", + "api_key": "sk-test-key", + "user": "test_user_1", + "team_id": "team1", + "spend": 0.01, + "startTime": today.isoformat(), + "model": "gpt-4", + }, + ] + + def filter_fn(where): + rid_either = where.get("request_id_or_call_id") + if rid_either: + return [ + r + for r in mock_spend_logs + if rid_either in (r["request_id"], r.get("litellm_call_id")) + ] + if where.get("request_id"): + return [ + r for r in mock_spend_logs if r["request_id"] == where["request_id"] + ] + return list(mock_spend_logs) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.prisma_client", + make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "b980eea9-5cd9-4099-93cd-8291e46c76fd"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200 + data = response.json() + assert data["total"] == 1 + assert ( + data["data"][0]["request_id"] == "chatcmpl-9ZKMURhVYSi9D6r6PJ9vLcayIK0Vm" + ) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -2303,23 +2514,18 @@ async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatc """A non-admin looking up a request_id they do not own is rejected (403), so the relaxed date window cannot read another tenant's log by id.""" - class _ForeignRow: - user = "other_user" - team_id = None + prisma = _make_owner_lookup_prisma( + [ + { + "request_id": "foreign-req", + "litellm_call_id": None, + "user": "other_user", + "team_id": None, + } + ] + ) - class _SpendLogs: - async def find_unique(self, where, include=None): - return _ForeignRow() - - class _DB: - def __init__(self): - self.litellm_spendlogs = _SpendLogs() - - class _Prisma: - def __init__(self): - self.db = _DB() - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _Prisma()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" ) @@ -2335,13 +2541,437 @@ async def test_ui_view_spend_logs_request_id_blocks_non_owner(client, monkeypatc @pytest.mark.asyncio -async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( +async def test_ui_view_spend_logs_request_id_collision_serves_only_callers_rows(client, monkeypatch): + """Two tenants share one id: the attacker minted a row whose client-set + litellm_call_id equals the victim's request_id. Each side's lookup of that id + returns only their own row, so the collision neither leaks the other tenant's + row nor denies the victim theirs (Veria: identifier collision could deny access).""" + now_iso = datetime.datetime.now(timezone.utc).isoformat() + corpus = [ + { + "id": "log_attacker", + "request_id": "attacker-req", + "litellm_call_id": "victim-req", + "api_key": "sk-attacker-key", + "user": "attacker_user", + "team_id": None, + "spend": 0.05, + "startTime": now_iso, + "model": "gpt-4", + }, + { + "id": "log_victim", + "request_id": "victim-req", + "litellm_call_id": "victim-call-id", + "api_key": "sk-victim-key", + "user": "victim_user", + "team_id": None, + "spend": 0.07, + "startTime": now_iso, + "model": "gpt-4", + }, + ] + + def filter_fn(where): + rid_either = where.get("request_id_or_call_id") + rows = [r for r in corpus if rid_either in (r["request_id"], r["litellm_call_id"])] + return [r for r in rows if where.get("user") is None or r["user"] == where["user"]] + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", make_ui_spend_logs_mock_prisma(corpus, filter_fn)) + try: + for caller, own_request_id, other in ( + ("victim_user", "victim-req", "attacker_user"), + ("attacker_user", "attacker-req", "victim_user"), + ): + app.dependency_overrides[ps.user_api_key_auth] = lambda caller=caller: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller + ) + response = client.get( + "/spend/logs/ui", + params={"request_id": "victim-req"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["total"] == 1 + assert data["data"][0]["request_id"] == own_request_id + assert other not in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_id_lookup_scopes_every_non_admin_role(client, monkeypatch): + """An org admin reaches /spend/logs/ui without the internal-user row scope. An + id lookup still fetches only rows they own, so another tenant's row carrying + that id as its client-set litellm_call_id neither leaks nor turns the + org admin's own lookup into a 403 (Bugbot: non-internal id lookup 403s on collision).""" + now_iso = datetime.datetime.now(timezone.utc).isoformat() + corpus = [ + { + "id": "log_attacker", + "request_id": "attacker-req", + "litellm_call_id": "victim-req", + "api_key": "sk-attacker-key", + "user": "attacker_user", + "team_id": None, + "spend": 0.05, + "startTime": now_iso, + "model": "gpt-4", + }, + { + "id": "log_victim", + "request_id": "victim-req", + "litellm_call_id": "victim-call-id", + "api_key": "sk-victim-key", + "user": "victim_user", + "team_id": None, + "spend": 0.07, + "startTime": now_iso, + "model": "gpt-4", + }, + ] + + def filter_fn(where): + rid_either = where.get("request_id_or_call_id") + rows = [r for r in corpus if rid_either in (r["request_id"], r["litellm_call_id"])] + return [r for r in rows if where.get("user") is None or r["user"] == where["user"]] + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", make_ui_spend_logs_mock_prisma(corpus, filter_fn)) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.ORG_ADMIN, user_id="victim_user" + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "victim-req"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["total"] == 1 + assert [row["request_id"] for row in data["data"]] == ["victim-req"] + assert "attacker_user" not in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_rejects_foreign_row_inserted_after_owner_check(client, monkeypatch): + """The SQL scope keeps foreign rows out of an id lookup; this backstop covers a + row the scope did not filter (the mock ignores it on purpose). The rows actually + fetched are ownership-checked again, so the lookup answers 403 instead of serving + the other tenant's row.""" + now_iso = datetime.datetime.now(timezone.utc).isoformat() + owned_row = { + "id": "log_owned", + "request_id": "attacker-req", + "litellm_call_id": "shared-id", + "api_key": "sk-test-key", + "user": "user_1", + "team_id": None, + "spend": 0.05, + "startTime": now_iso, + "model": "gpt-4", + } + foreign_row = { + "id": "log_foreign", + "request_id": "shared-id", + "litellm_call_id": None, + "api_key": "sk-victim-key", + "user": "victim_user", + "team_id": None, + "spend": 0.07, + "startTime": now_iso, + "model": "gpt-4", + } + + mock_prisma = make_ui_spend_logs_mock_prisma([owned_row], lambda where: [owned_row, foreign_row]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "shared-id"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert "victim_user" not in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +def _make_payload_lookup_prisma(rows): + """Emulate the detail endpoint's SQL over an in-memory corpus: the owner + pre-check, the caller scope on ``"user"`` and permitted teams, and the + exact-request_id-first ordering with LIMIT 1.""" + + class MockDB: + async def query_raw(self, sql_query, *params): + if 'SELECT DISTINCT "user", team_id' in sql_query: + return _emulate_spend_log_owner_lookup(rows, sql_query, params) + lookup_id = params[0] + matches = [r for r in rows if lookup_id in (r["request_id"], r["litellm_call_id"])] + if '"user" = $2' in sql_query: + team_ids = params[2] if "ANY($3::text[])" in sql_query else () + matches = [r for r in matches if r["user"] == params[1] or r["team_id"] in team_ids] + if "ORDER BY (request_id = $1) DESC" in sql_query: + matches = sorted(matches, key=lambda r: r["request_id"] == lookup_id, reverse=True) + return matches[:1] + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + return MockPrisma() + + +def _payload_row(request_id, litellm_call_id, user, prompt): + return { + "request_id": request_id, + "litellm_call_id": litellm_call_id, + "messages": [{"role": "user", "content": prompt}], + "response": {"id": request_id}, + "proxy_server_request": None, + "metadata": None, + "user": user, + "team_id": None, + } + + +@pytest.mark.asyncio +async def test_ui_view_request_response_collision_serves_callers_own_row(client, monkeypatch): + """The attacker's row carries the victim's request_id as its client-set call id + and was written first. Each tenant's detail lookup of that id serves only their + own payload, and an admin's lookup resolves the exact request_id match rather + than whichever colliding row the database happens to return first.""" + prisma = _make_payload_lookup_prisma( + [ + _payload_row("attacker-req", "victim-req", "attacker_user", "attacker prompt"), + _payload_row("victim-req", "victim-call-id", "victim_user", "victim prompt"), + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + try: + for role, user_id, own_prompt, other_prompt in ( + (LitellmUserRoles.INTERNAL_USER, "victim_user", "victim prompt", "attacker prompt"), + (LitellmUserRoles.INTERNAL_USER, "attacker_user", "attacker prompt", "victim prompt"), + (LitellmUserRoles.PROXY_ADMIN, "admin", "victim prompt", "attacker prompt"), + ): + app.dependency_overrides[ps.user_api_key_auth] = lambda role=role, user_id=user_id: UserAPIKeyAuth( + user_role=role, user_id=user_id + ) + response = client.get("/spend/logs/ui/victim-req", headers={"Authorization": "Bearer sk-test"}) + assert response.status_code == 200, response.text + assert own_prompt in response.text + assert other_prompt not in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_request_response_rejects_foreign_row_inserted_after_owner_check(client, monkeypatch): + """Backstop behind the SQL scope on the detail endpoint (the mock ignores the + scope on purpose): the payload row fetched by id is itself ownership-checked, so + a foreign row the scope did not filter cannot have its payload served.""" + + class MockDB: + async def query_raw(self, sql_query, *params): + if 'SELECT DISTINCT "user", team_id' in sql_query: + return [{"user": "user_1", "team_id": None}] + return [ + { + "messages": [{"role": "user", "content": "victim prompt"}], + "response": {"id": "resp-1"}, + "proxy_server_request": None, + "metadata": None, + "user": "victim_user", + "team_id": None, + } + ] + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrisma()) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui/shared-id", + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert "victim prompt" not in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_request_response_custom_logger_denies_foreign_payload_owner(client, monkeypatch): + """The custom-logger payload comes straight from cold storage, written independently + of the spend-log table and able to outlive its row. When an id lookup matches no row, + the DB owner pre-check has nothing to verify, so the payload is authorized against the + owner recorded inside it. A foreign tenant's stored payload is denied even though no + spend-log row exists for the pre-check to catch.""" + + class MockDB: + async def query_raw(self, sql_query, *params): + return [] + + class MockPrisma: + def __init__(self): + self.db = MockDB() + + class ColdStorageLogger: + async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc): + return { + "messages": [{"role": "user", "content": "victim prompt"}], + "response": {"id": "r"}, + "metadata": {"user_api_key_user_id": "victim_user", "user_api_key_team_id": None}, + } + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrisma()) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [ColdStorageLogger()], + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user_1" + ) + try: + response = client.get( + "/spend/logs/ui/shared-id", + params={"start_date": "2026-01-01 00:00:00"}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 403 + assert "victim prompt" not in response.text + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_request_response_custom_logger_is_keyed_by_callers_own_request_id(client, monkeypatch): + """Cold storage is keyed by the provider request_id. The caller's row carries the + lookup id only as its client-set litellm_call_id while another tenant's row owns + that id as its request_id. The custom logger is asked for the caller's own stored + request_id, so the caller gets their payload rather than a 403 from the foreign + payload's owner check, and the foreign payload is never fetched.""" + prisma = _make_payload_lookup_prisma( + [ + _payload_row("shared-id", "other-call-id", "other_user", "other tenant prompt"), + _payload_row("caller-req", "shared-id", "caller_user", "caller prompt"), + ] + ) + cold_storage = { + "shared-id": { + "messages": [{"role": "user", "content": "other tenant prompt"}], + "response": {"id": "shared-id"}, + "metadata": {"user_api_key_user_id": "other_user", "user_api_key_team_id": None}, + }, + "caller-req": { + "messages": [{"role": "user", "content": "caller prompt"}], + "response": {"id": "caller-req"}, + "metadata": {"user_api_key_user_id": "caller_user", "user_api_key_team_id": None}, + }, + } + requested_ids = [] + + class ColdStorageLogger: + async def get_request_response_payload(self, request_id, start_time_utc, end_time_utc): + requested_ids.append(request_id) + return cold_storage.get(request_id) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr( + litellm.logging_callback_manager, + "get_active_additional_logging_utils_from_custom_logger", + lambda: [ColdStorageLogger()], + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="caller_user" + ) + try: + response = client.get("/spend/logs/ui/shared-id", headers={"Authorization": "Bearer sk-test"}) + assert response.status_code == 200, response.text + assert "caller prompt" in response.text + assert "other tenant prompt" not in response.text + assert requested_ids == ["caller-req"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group_by_session", [False, True]) +async def test_ui_view_spend_logs_id_lookup_lists_exact_request_id_row_first(client, monkeypatch, group_by_session): + """The dashboard's deep link fetches a single row for ``?log_id=``. When a newer + row carries that id as its client-set litellm_call_id, the row whose request_id + is the id still comes first, so the link opens the request it names. The + session-grouped page orders its representatives the same way.""" + today = datetime.datetime.now(timezone.utc) + corpus = [ + { + "id": "log_colliding", + "request_id": "colliding-req", + "litellm_call_id": "victim-req", + "api_key": "sk-test-key", + "user": "other_user", + "team_id": None, + "spend": 0.01, + "startTime": today.isoformat(), + "model": "gpt-4", + }, + { + "id": "log_victim", + "request_id": "victim-req", + "litellm_call_id": "victim-call-id", + "api_key": "sk-test-key", + "user": "victim_user", + "team_id": None, + "spend": 0.02, + "startTime": (today - datetime.timedelta(minutes=5)).isoformat(), + "model": "gpt-4", + }, + ] + + def filter_fn(where): + rows = _filter_logs_by_date_range(corpus, where) + rid_either = where.get("request_id_or_call_id") + if rid_either: + return [r for r in rows if rid_either in (r["request_id"], r["litellm_call_id"])] + return rows + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", make_ui_spend_logs_mock_prisma(corpus, filter_fn)) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin" + ) + try: + response = client.get( + "/spend/logs/ui", + params={"request_id": "victim-req", "page_size": 1, "group_by_session": str(group_by_session).lower()}, + headers={"Authorization": "Bearer sk-test"}, + ) + assert response.status_code == 200, response.text + data = response.json() + assert data["total"] == 2 + assert [row["request_id"] for row in data["data"]] == ["victim-req"] + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.asyncio +async def test_ui_view_spend_logs_request_id_owner_lookup_drops_window_keeps_scope( client, monkeypatch ): - """A non-admin owner looking up their own request_id resolves across all time. - The ownership check authorizes the single row, so the query drops both the date - window and the general user/team scoping and filters by the primary key alone; - without that skip an internal user would have a `user`/`OR` clause added.""" + """A non-admin owner looking up their own request_id resolves across all time: + the query drops the date window the dashboard sends, while the caller's own-user + scope stays on the id lookup so a colliding foreign row can never be served.""" today = datetime.datetime.now(timezone.utc) mock_spend_logs = [ { @@ -2361,20 +2991,14 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( def filter_fn(where): captured["where"] = where rows = _filter_logs_by_date_range(mock_spend_logs, where) - if where.get("request_id"): + rid_either = where.get("request_id_or_call_id") + if rid_either: + rows = [r for r in rows if rid_either in (r["request_id"], r.get("litellm_call_id"))] + elif where.get("request_id"): rows = [r for r in rows if r["request_id"] == where["request_id"]] return rows mock_prisma = make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn) - - class _OwnedRow: - user = "user_1" - team_id = "team1" - - async def _find_unique(where, include=None): - return _OwnedRow() - - mock_prisma.db.find_unique = _find_unique monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) # A 5-day window that EXCLUDES the 90-day-old log, as the dashboard sends. @@ -2399,9 +3023,8 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only( assert data["total"] == 1 assert data["data"][0]["request_id"] == "req-old" assert "startTime" not in captured["where"] - assert captured["where"]["request_id"] == "req-old" - assert "user" not in captured["where"] - assert "OR" not in captured["where"] + assert captured["where"]["request_id_or_call_id"] == "req-old" + assert captured["where"]["user"] == "user_1" finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -5095,7 +5718,7 @@ async def test_view_spend_logs_internal_user_combines_user_with_request_id( where = mock_client.db.captured_where assert where is not None assert where["user"] == "internal-user-2" - assert where["request_id"] == "req-abc" + assert where["OR"] == ({"request_id": "req-abc"}, {"litellm_call_id": "req-abc"}) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -5122,7 +5745,7 @@ async def test_view_spend_logs_non_date_range_combines_user_with_request_id( where = mock_client.db.captured_where assert where is not None assert where["user"] == "internal-user-3" - assert where["request_id"] == "req-xyz" + assert where["OR"] == ({"request_id": "req-xyz"}, {"litellm_call_id": "req-xyz"}) finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -7052,25 +7675,19 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess def _fake_prisma_with_owned_spend_log(owner_user_id, messages_json, response_json): - class _Row: - user = owner_user_id - team_id = None - - class _SpendLogs: - async def find_unique(self, where, include=None): - return _Row() - class _DB: - def __init__(self): - self.litellm_spendlogs = _SpendLogs() - - async def query_raw(self, _sql, *_args): + async def query_raw(self, sql, *_args): + if 'SELECT DISTINCT "user", team_id' in sql: + return [{"user": owner_user_id, "team_id": None}] return [ { + "request_id": "req-owned-by-user-a", "messages": messages_json, "response": response_json, "proxy_server_request": "{}", "metadata": "{}", + "user": owner_user_id, + "team_id": None, } ] @@ -7127,8 +7744,8 @@ def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monk """ A different internal_user requesting someone else's row is forbidden; guards against _assert_user_can_view_request_id being skipped in the - detail-drawer handler. Also proves the handler stops before it ever asks - a custom logger or the DB for the payload. + detail-drawer handler. Also proves the handler stops at the owner lookup, + before it ever asks a custom logger or the DB for the payload. """ messages_json = json.dumps([{"role": "user", "content": "hi"}]) response_json = json.dumps({"choices": [{"message": {"content": "hello"}}]}) @@ -7160,7 +7777,7 @@ def test_ui_view_request_response_internal_user_non_owner_forbidden(client, monk ) assert response.status_code == 403 assert custom_logger.requested_ids == [] - assert query_raw_calls == [] + assert [args[0] for args, _kwargs in query_raw_calls if "messages" in args[0]] == [] finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) @@ -7171,24 +7788,27 @@ def test_ui_view_request_response_internal_user_missing_row_forbidden(client, mo request_id with no spend-log row (e.g. pruned by retention) must be denied before the handler ever consults a custom logger, otherwise a non-admin who guesses/obtains a request_id could read another tenant's - payload out of cold storage. Fails if `if row is None: return` is - reintroduced. + payload out of cold storage, and an existing payload that is not theirs + would still confirm the id exists. Even a payload recorded as the caller's + own is never fetched once the row is gone. Fails if an empty owner lookup + is allowed to fall through to the loggers. """ - class _SpendLogs: - async def find_unique(self, where, include=None): - return None - class _DB: - def __init__(self): - self.litellm_spendlogs = _SpendLogs() + async def query_raw(self, _sql, *_args): + return [] from types import SimpleNamespace fake_prisma = SimpleNamespace(db=_DB()) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", fake_prisma) - custom_logger = _RecordingAdditionalLoggingUtils({"messages": "should-not-be-returned"}) + custom_logger = _RecordingAdditionalLoggingUtils( + { + "messages": "should-not-be-returned", + "metadata": {"user_api_key_user_id": "user_a", "user_api_key_team_id": None}, + } + ) monkeypatch.setattr( litellm.logging_callback_manager, "get_active_additional_logging_utils_from_custom_logger", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e9bbaf0c96e..a72b4e28143 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1329,6 +1329,33 @@ def test_get_logging_payload_includes_agent_id_from_kwargs(): assert payload["agent_id"] == test_agent_id, f"Expected agent_id '{test_agent_id}', got '{payload.get('agent_id')}'" +def test_get_logging_payload_populates_litellm_call_id_alongside_provider_request_id(): + """ + LIT-6302: request_id stays the provider response id, so clients holding the + x-litellm-call-id header value could never find their row. The payload now + also carries litellm_call_id as its own column for lookups by either id. + """ + call_id = "b980eea9-5cd9-4099-93cd-8291e46c76fd" + + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_call_id": call_id, + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse( + id="chatcmpl-provider-id", + choices=[], + usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + assert payload["request_id"] == "chatcmpl-provider-id" + assert payload["litellm_call_id"] == call_id + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index efbb5eedad4..69e89d1c604 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -13,7 +13,7 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY +from litellm.constants import MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( @@ -30,6 +30,7 @@ from litellm.proxy.common_request_processing import ( _has_attribute_error_in_chain, _is_azure_model_router_request, open_sse_before_first_byte, + resolve_litellm_call_id, ttft_keepalive_interval, _override_openai_response_model, _parse_event_data_for_error, @@ -8060,6 +8061,19 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_ assert (records[0].exc_info is not None) is expect_traceback +class TestResolveLitellmCallId: + def test_client_call_id_within_the_bound_is_kept(self): + assert resolve_litellm_call_id("req-abc-123") == "req-abc-123" + at_bound: Final = "y" * MAX_LITELLM_CALL_ID_LENGTH + assert resolve_litellm_call_id(at_bound) == at_bound + + @pytest.mark.parametrize("client_call_id", [None, "", "x" * (MAX_LITELLM_CALL_ID_LENGTH + 1), "z" * 3000]) + def test_missing_empty_or_oversized_client_call_id_gets_a_generated_uuid(self, client_call_id): + resolved: Final = resolve_litellm_call_id(client_call_id) + assert resolved != client_call_id + assert uuid.UUID(resolved).version == 4 + + class _FailureHookRecorder: """Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it.""" diff --git a/tests/test_litellm/test_auto_update_price_and_context_window_file.py b/tests/test_litellm/test_auto_update_price_and_context_window_file.py new file mode 100644 index 00000000000..435747e9a09 --- /dev/null +++ b/tests/test_litellm/test_auto_update_price_and_context_window_file.py @@ -0,0 +1,222 @@ +"""Unit tests for the Friendli transform in +`.github/scripts/auto_update_price_and_context_window_file.py`.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / ".github" + / "scripts" + / "auto_update_price_and_context_window_file.py" +) + + +@pytest.fixture(scope="module") +def sync_module(): + spec = importlib.util.spec_from_file_location( + "auto_update_price_and_context_window_file", SCRIPT_PATH + ) + assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" + module = importlib.util.module_from_spec(spec) + sys.modules["auto_update_price_and_context_window_file"] = module + spec.loader.exec_module(module) + return module + + +def _reasoning_model(**overrides: object) -> dict: + model = { + "id": "zai-org/GLM-Test", + "base_model": "zhipuai/glm-test", + "context_length": 1048576, + "max_completion_tokens": 131072, + "pricing": {"input": "0.00000015", "output": "0.0000005", "input_cache_read": "0.00000003"}, + "reasoning": True, + "reasoning_options": [{"type": "effort", "values": ["max", "high", "low"]}], + "functionality": { + "tool_call": True, + "parallel_tool_call": True, + "structured_output": True, + "system_messages": True, + "tool_choice": True, + }, + "input_modalities": ["text", "image", "video"], + "mode": "chat", + } + model.update(overrides) + return model + + +def test_transform_emits_declared_effort_levels_in_canonical_order(sync_module): + entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ + "friendliai/zai-org/GLM-Test" + ] + assert entry["supports_reasoning"] is True + assert entry["reasoning_effort_levels"] == ["low", "high", "max"] + assert not any(k.endswith("_reasoning_effort") for k in entry) + + +def test_transform_reasoning_model_without_effort_options_declares_empty_levels(sync_module): + model = _reasoning_model(reasoning_options=[{"type": "budget_tokens", "values": []}]) + entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] + assert entry["reasoning_effort_levels"] == [] + + +def test_transform_non_reasoning_model_declares_no_levels(sync_module): + model = _reasoning_model(reasoning=False, reasoning_options=[]) + entry = sync_module.transform_friendli_data([model], {})["friendliai/zai-org/GLM-Test"] + assert entry["supports_reasoning"] is False + assert "reasoning_effort_levels" not in entry + + +def test_transform_max_tokens_mirrors_output_cap_not_context(sync_module): + entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ + "friendliai/zai-org/GLM-Test" + ] + assert entry["max_input_tokens"] == 1048576 + assert entry["max_output_tokens"] == 131072 + assert entry["max_tokens"] == entry["max_output_tokens"] + + +def test_transform_prompt_caching_follows_cache_pricing(sync_module): + cached = sync_module.transform_friendli_data([_reasoning_model()], {})[ + "friendliai/zai-org/GLM-Test" + ] + assert cached["supports_prompt_caching"] is True + assert cached["cache_read_input_token_cost"] == 3e-08 + + uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) + uncached = sync_module.transform_friendli_data([uncached_model], {})[ + "friendliai/zai-org/GLM-Test" + ] + assert uncached["supports_prompt_caching"] is False + assert "cache_read_input_token_cost" not in uncached + + +def test_transform_modalities_set_vision_image_and_video_flags(sync_module): + entry = sync_module.transform_friendli_data([_reasoning_model()], {})[ + "friendliai/zai-org/GLM-Test" + ] + assert entry["supports_vision"] is True + assert entry["supports_image_input"] is True + assert entry["supports_video_input"] is True + + text_only = _reasoning_model(input_modalities=["text"]) + entry_text = sync_module.transform_friendli_data([text_only], {})[ + "friendliai/zai-org/GLM-Test" + ] + assert entry_text["supports_vision"] is False + assert entry_text["supports_image_input"] is False + assert entry_text["supports_video_input"] is False + + +def test_transform_skips_rows_without_valid_token_prices_so_priced_local_entries_survive(sync_module): + local = { + "friendliai/zai-org/GLM-Test": { + "litellm_provider": "friendliai", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + } + } + unpriced_rows = [ + _reasoning_model(pricing={}), + _reasoning_model(pricing=None), + _reasoning_model(pricing={"input": "0.00000015"}), + _reasoning_model(pricing={"output": "0.0000005"}), + _reasoning_model(pricing={"input": "not-a-number", "output": "0.0000005"}), + _reasoning_model(pricing={"input": "-0.00000015", "output": "0.0000005"}), + _reasoning_model(pricing={"input": "inf", "output": "0.0000005"}), + _reasoning_model(pricing={"input": "nan", "output": "0.0000005"}), + ] + remote = sync_module.transform_friendli_data(unpriced_rows, local) + assert remote == {} + sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) + assert local["friendliai/zai-org/GLM-Test"]["input_cost_per_token"] == 1.5e-07 + assert local["friendliai/zai-org/GLM-Test"]["output_cost_per_token"] == 5e-07 + + +def test_transform_keeps_zero_priced_rows(sync_module): + free_model = _reasoning_model(pricing={"input": "0", "output": "0"}) + entry = sync_module.transform_friendli_data([free_model], {})["friendliai/zai-org/GLM-Test"] + assert entry["input_cost_per_token"] == 0.0 + assert entry["output_cost_per_token"] == 0.0 + + +def test_transforms_survive_failed_fetch(sync_module): + assert sync_module.transform_friendli_data(None, {}) == {} + assert sync_module.transform_friendli_data([], {}) == {} + assert sync_module.transform_openrouter_data(None) == {} + assert sync_module.transform_vercel_ai_gateway_data(None) == {} + + +def test_vercel_transform_skips_rows_without_token_pricing_or_limits(sync_module): + rows = [ + { + "id": "wan-video", + "pricing": {"video_duration_pricing": [{"resolution": "720p", "cost_per_second": "0.1"}]}, + }, + { + "id": "qwen3-embedding", + "context_window": 32768, + "max_tokens": 32768, + "pricing": {"input": "0.00000001"}, + }, + { + "id": "no-limits-chat", + "pricing": {"input": "0.000001", "output": "0.000002"}, + }, + { + "id": "good-chat", + "context_window": 128000, + "max_tokens": 8192, + "pricing": {"input": "0.000001", "output": "0.000002"}, + }, + ] + transformed = sync_module.transform_vercel_ai_gateway_data(rows) + assert list(transformed) == ["vercel_ai_gateway/good-chat"] + assert transformed["vercel_ai_gateway/good-chat"]["input_cost_per_token"] == 1e-06 + assert transformed["vercel_ai_gateway/good-chat"]["output_cost_per_token"] == 2e-06 + + +def test_sync_replaces_friendli_entries_so_dropped_cache_pricing_does_not_survive(sync_module): + local = { + "friendliai/zai-org/GLM-Test": { + "litellm_provider": "friendliai", + "cache_read_input_token_cost": 3e-08, + "supports_prompt_caching": True, + } + } + uncached_model = _reasoning_model(pricing={"input": "0.00000014", "output": "0.0000004"}) + remote = sync_module.transform_friendli_data([uncached_model], local) + sync_module.sync_local_data_with_remote(local, remote, replace_keys=frozenset(remote)) + synced = local["friendliai/zai-org/GLM-Test"] + assert "cache_read_input_token_cost" not in synced + assert synced["supports_prompt_caching"] is False + + +def test_sync_still_merges_entries_outside_replace_keys(sync_module): + local = {"openrouter/some-model": {"input_cost_per_token": 1e-06, "supports_vision": True}} + remote = {"openrouter/some-model": {"input_cost_per_token": 2e-06}} + sync_module.sync_local_data_with_remote(local, remote) + assert local["openrouter/some-model"] == {"input_cost_per_token": 2e-06, "supports_vision": True} + + +def test_transform_inherits_allowlisted_keys_from_base_model_entry(sync_module): + local = { + "zhipuai/glm-test": { + "supports_pdf_input": True, + "supports_assistant_prefill": True, + "input_cost_per_token": 9e-06, + } + } + entry = sync_module.transform_friendli_data([_reasoning_model()], local)[ + "friendliai/zai-org/GLM-Test" + ] + assert entry["supports_pdf_input"] is True + assert entry["supports_assistant_prefill"] is True + assert entry["input_cost_per_token"] == 1.5e-07 diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 3ef768790f8..f2659cee3fd 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4313,6 +4313,59 @@ def test_select_model_name_applies_region_to_private_provider_response_model(_lo assert selected == "bedrock/us-east-1/anthropic.claude-v2:1" +def test_completion_cost_region_name_prices_mantle_on_the_regional_row(_local_model_cost_map): + """completion_cost(region_name=...) must price a Bedrock Mantle call from the + bedrock_mantle// row when one exists, for the bare and the provider-prefixed + model alike, and keep the flat row for regions without their own row.""" + + response = litellm.ModelResponse( + id="x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="xai.grok-4.3", + usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + ) + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + flat = litellm.model_cost["bedrock_mantle/xai.grok-4.3"] + expected_gov = 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + expected_flat = 38 * flat["input_cost_per_token"] + 20 * flat["output_cost_per_token"] + assert expected_gov != expected_flat + + for model in ("xai.grok-4.3", "bedrock_mantle/xai.grok-4.3"): + assert litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock_mantle", + region_name="us-gov-west-1", + ) == pytest.approx(expected_gov) + assert litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock_mantle", + region_name="eu-west-1", + ) == pytest.approx(expected_flat) + assert litellm.completion_cost( + completion_response=response, model="xai.grok-4.3", custom_llm_provider="bedrock_mantle" + ) == pytest.approx(expected_flat) + + +def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_model_cost_map): + """A provider-prefixed model must still find its bedrock_mantle// row instead of + composing the region key with the provider segment twice.""" + + prompt_cost, completion_cost = litellm.cost_per_token( + model="bedrock_mantle/xai.grok-4.3", + prompt_tokens=38, + completion_tokens=20, + custom_llm_provider="bedrock_mantle", + region_name="us-gov-west-1", + ) + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + + assert prompt_cost + completion_cost == pytest.approx( + 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + ) + + def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map): """An explicit base_model keeps pricing on that model's own key even when the request carries a region with different regional rates, so the private provider model never widens region pricing.""" @@ -4329,6 +4382,29 @@ def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map assert selected == "bedrock/moonshotai.kimi-k2.5" +def test_completion_cost_base_model_ignores_regional_row(_local_model_cost_map): + """A deployment with base_model set is priced from that model's own row even when the response + carries a region whose regional row charges different rates.""" + + response = litellm.ModelResponse( + id="x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="my-bedrock-deployment", + usage={"prompt_tokens": 1000, "completion_tokens": 0, "total_tokens": 1000}, + ) + response._hidden_params = {"custom_llm_provider": "bedrock", "region_name": "eu-central-1"} + flat = litellm.model_cost["anthropic.claude-instant-v1"] + regional = litellm.model_cost["bedrock/eu-central-1/anthropic.claude-instant-v1"] + assert flat["input_cost_per_token"] != regional["input_cost_per_token"] + + assert litellm.completion_cost( + completion_response=response, + model="my-bedrock-deployment", + custom_llm_provider="bedrock", + base_model="anthropic.claude-instant-v1", + ) == pytest.approx(1000 * flat["input_cost_per_token"]) + + def test_completion_cost_nonzero_for_slash_alias_model_name(_local_model_cost_map): """End-to-end cost through a "/"-containing alias must price above zero (#38069).""" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f5e9b2091a0..b8a0d70f5bc 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5625,6 +5625,41 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): assert credentials.get(key) == value, key +def test_get_deployment_credentials_preserves_azure_entra_id_params(): + entra_params = { + "tenant_id": "deployment-tenant", + "client_id": "deployment-client", + "client_secret": "deployment-client-secret", + "azure_scope": "https://cognitiveservices.azure.us/.default", + "azure_username": "deployment-user", + "azure_password": "deployment-password", + } + router = litellm.Router( + model_list=[ + { + "model_name": "azure-entra-model", + "litellm_params": { + "model": "azure/gpt-5.4", + "api_base": "https://example.openai.azure.com/", + "api_version": "2024-10-21", + **entra_params, + }, + "model_info": {"id": "azure-entra-model-id"}, + } + ], + ) + + credentials = router.get_deployment_credentials(model_id="azure-entra-model-id") + credentials_with_provider = router.get_deployment_credentials_with_provider(model_id="azure-entra-model-id") + + assert credentials is not None + assert credentials_with_provider is not None + assert "api_key" not in credentials + for key, value in entra_params.items(): + assert credentials.get(key) == value, key + assert credentials_with_provider.get(key) == value, key + + def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict: return { "model_name": f"model_name_team-1_{model_id}", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index b4e300d9fc3..ccb9f90f9a3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -48,7 +48,7 @@ const AllModelsTab = ({ setSelectedTeamId, }: AllModelsTabProps) => { const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap(); - const { accessToken, userId, userRole } = useAuthorized(); + const { accessToken, userId, userRole, isViewOnly } = useAuthorized(); const { data: teams, isLoading: isLoadingTeams } = useTeams(); const queryClient = useQueryClient(); @@ -281,6 +281,7 @@ const AllModelsTab = ({ availableModelAccessGroups={availableModelAccessGroups} userRole={userRole} userID={userId} + isViewOnly={isViewOnly} onModelIdClick={setSelectedModelId} onTeamIdClick={setSelectedTeamId} onDeleteClick={handleDeleteClick} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx index 8ba71e82d48..726070c4bb6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.test.tsx @@ -59,6 +59,7 @@ const baseProps = { availableModelAccessGroups: ["sales-team"], userRole: "Admin", userID: "alice", + isViewOnly: false, onModelIdClick: vi.fn(), onTeamIdClick: vi.fn(), onDeleteClick: vi.fn(), @@ -254,6 +255,17 @@ describe("AllModelsTable", () => { expect(onTogglePauseClick).not.toHaveBeenCalled(); }); + it("does not let a view-only admin toggle a model", async () => { + const user = userEvent.setup(); + const onTogglePauseClick = vi.fn(); + render(); + + const toggle = screen.getByTestId("model-pause-toggle-model-1"); + expect(toggle).toHaveAttribute("data-disabled"); + await user.click(toggle); + expect(onTogglePauseClick).not.toHaveBeenCalled(); + }); + it("does not let anyone toggle a config model", async () => { const user = userEvent.setup(); const onTogglePauseClick = vi.fn(); @@ -309,6 +321,17 @@ describe("AllModelsTable", () => { expect(onDeleteClick).not.toHaveBeenCalled(); }); + it("blocks a view-only admin from deleting a DB model they created", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + render(); + + const deleteButton = screen.getByTestId("model-delete-model-1"); + expect(deleteButton).toBeDisabled(); + await user.click(deleteButton); + expect(onDeleteClick).not.toHaveBeenCalled(); + }); + it("blocks deleting a config model", async () => { const user = userEvent.setup(); const onDeleteClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx index 8482d0832c3..f46130d2386 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTable.tsx @@ -73,6 +73,7 @@ interface AllModelsTableProps { availableModelAccessGroups: string[]; userRole: string; userID: string; + isViewOnly: boolean; onModelIdClick: (modelId: string) => void; onTeamIdClick: (teamId: string) => void; onDeleteClick: (modelId: string) => void; @@ -120,6 +121,7 @@ export function AllModelsTable({ availableModelAccessGroups, userRole, userID, + isViewOnly, onModelIdClick, onTeamIdClick, onDeleteClick, @@ -132,6 +134,7 @@ export function AllModelsTable({ const columnDeps = { userRole, userID, + isViewOnly, onModelIdClick, onTeamIdClick, onDeleteClick, @@ -139,7 +142,7 @@ export function AllModelsTable({ pausingModelId, }; return getModelsTableColumns(columnDeps); - }, [userRole, userID, onModelIdClick, onTeamIdClick, onDeleteClick, onTogglePauseClick, pausingModelId]); + }, [userRole, userID, isViewOnly, onModelIdClick, onTeamIdClick, onDeleteClick, onTogglePauseClick, pausingModelId]); const modelGroupOptions = useMemo( () => [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx index f3ae687447e..0cc1207e547 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/ModelsTableColumns.tsx @@ -247,6 +247,7 @@ interface ModelRowActionsProps { model: ModelData; userRole: string; userID: string; + isViewOnly: boolean; isPausing: boolean; onDeleteClick?: (modelId: string) => void; onTogglePauseClick?: (modelId: string, blocked: boolean) => void | Promise; @@ -256,14 +257,15 @@ function ModelRowActions({ model, userRole, userID, + isViewOnly, isPausing, onDeleteClick, onTogglePauseClick, }: ModelRowActionsProps) { const modelId = model.model_info?.id; const isConfigModel = !model.model_info?.db_model; - const isAdmin = userRole === "Admin"; - const canEditModel = isAdmin || model.model_info?.created_by === userID; + const isAdmin = userRole === "Admin" && !isViewOnly; + const canEditModel = !isViewOnly && (isAdmin || model.model_info?.created_by === userID); const isBlocked = model.model_info?.blocked === true; const isPauseToggleable = !isConfigModel && isAdmin && Boolean(onTogglePauseClick); @@ -340,6 +342,7 @@ function ModelRowActions({ export interface ModelsTableColumnDeps { userRole: string; userID: string; + isViewOnly: boolean; onModelIdClick: (modelId: string) => void; onTeamIdClick: (teamId: string) => void; onDeleteClick?: (modelId: string) => void; @@ -350,6 +353,7 @@ export interface ModelsTableColumnDeps { export const getModelsTableColumns = ({ userRole, userID, + isViewOnly, onModelIdClick, onTeamIdClick, onDeleteClick, @@ -479,6 +483,7 @@ export const getModelsTableColumns = ({ model={row.original} userRole={userRole} userID={userID} + isViewOnly={isViewOnly} isPausing={pausingModelId === row.original.model_info?.id} onDeleteClick={onDeleteClick} onTogglePauseClick={onTogglePauseClick} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 84a05113177..105f6ff3043 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -14,6 +14,7 @@ vi.mock("./panels/HealthStatusPanel", () => ({ default: () =>
({ default: () =>
})); vi.mock("./panels/ModelGroupAliasPanel", () => ({ default: () =>
})); vi.mock("./panels/PriceDataPanel", () => ({ default: () =>
})); +vi.mock("./panels/AccessGroupBudgetsPanel", () => ({ default: () =>
})); const detailState = { modelId: null as string | null, teamId: null as string | null }; vi.mock("./detailNavigation", () => ({ @@ -25,7 +26,11 @@ vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, })); vi.mock("@/components/team/TeamInfo", () => ({ - default: ({ teamId }: { teamId: string }) =>
team:{teamId}
, + default: ({ teamId, is_team_admin }: { teamId: string; is_team_admin: boolean }) => ( +
+ team:{teamId} +
+ ), })); const mockUseAuthorized = vi.fn(); @@ -95,10 +100,19 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByRole("tab", { name: "All Models" })).not.toBeInTheDocument(); }); - it("renders the team detail overlay from the ?team drill-in", () => { + it("renders the team detail overlay from the ?team drill-in with admin edit rights", () => { detailState.teamId = "team-9"; renderPage(); expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9"); + expect(screen.getByTestId("team-info")).toHaveAttribute("data-team-admin", "true"); + }); + + it("opens the ?team drill-in without edit rights for a view-only admin", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + detailState.teamId = "team-9"; + renderPage(); + expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9"); + expect(screen.getByTestId("team-info")).toHaveAttribute("data-team-admin", "false"); }); it("hides admin-only tabs for a non-admin user", () => { @@ -108,6 +122,35 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); }); + it("keeps the full admin tab order for a real admin", () => { + renderPage(); + expect(screen.getAllByRole("tab").map((tab) => tab.textContent)).toEqual([ + "All Models", + "Add Model", + "Auto-Routers Beta", + "LLM Credentials", + "Pass-Through Endpoints", + "Health Status", + "Model Retry Settings", + "Model Group Alias", + "Model Access Group Budgets Beta", + "Price Data Reload", + ]); + }); + + it("hides the admin write-form tabs from a view-only admin, keeping the read views", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + renderPage(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "Health Status" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "LLM Credentials" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Pass-Through Endpoints" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Model Retry Settings" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Model Group Alias" })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: /Model Access Group Budgets/ })).not.toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Price Data Reload" })).not.toBeInTheDocument(); + }); + // POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one. it("hides the Add Model tab for a view-only admin session", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 34c9d87004e..4d6a90fc56e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -106,19 +106,16 @@ export default function ModelsAndEndpointsPage() { "", ...(canCreate ? (["add"] as const) : []), ...(isAdmin || canCreate ? (["auto-routers"] as const) : []), - ...(isAdmin - ? ([ - "llm-credentials", - "pass-through", - "health", - "retry-settings", - "model-group-alias", - "access-group-budgets", - "price-data", - ] as const) + // effectiveSessionRole reports proxy_admin_viewer as "Admin", so isAdmin alone would show a + // viewer these write-only panels; only the raw-role isViewOnly separates them. Health Status + // stays: it is the bucket's one read view, and viewers keep read parity with admins. + ...(isAdmin && !isViewOnly ? (["llm-credentials", "pass-through"] as const) : []), + ...(isAdmin ? (["health"] as const) : []), + ...(isAdmin && !isViewOnly + ? (["retry-settings", "model-group-alias", "access-group-budgets", "price-data"] as const) : []), ], - [canCreate, isAdmin], + [canCreate, isAdmin, isViewOnly], ); const allModelsLabel = isAdmin ? "All Models" : "Your Models"; @@ -148,7 +145,7 @@ export default function ModelsAndEndpointsPage() { teamId={teamId} onClose={close} accessToken={accessToken} - is_team_admin={userRole === "Admin"} + is_team_admin={userRole === "Admin" && !isViewOnly} is_proxy_admin={userRole === "Proxy Admin"} userModels={allModelsOnProxy} editTeam={false} diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx new file mode 100644 index 00000000000..ab91e10c2fd --- /dev/null +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.test.tsx @@ -0,0 +1,97 @@ +import userEvent from "@testing-library/user-event"; +import React from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; +import type { LogEntry as SpendLogEntry } from "@/components/view_logs/columns"; +import { LogViewer } from "./LogViewer"; + +vi.mock("@/components/networking", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, uiSpendLogsCall: vi.fn() }; +}); + +vi.mock("@/components/view_logs/LogDetailsDrawer", () => ({ + LogDetailsDrawer: function LogDetailsDrawerMock({ + open, + logEntry, + }: { + open: boolean; + logEntry?: { request_id: string } | null; + }) { + return ( +
+ {open ? "open" : "closed"} +
+ ); + }, +})); + +import { uiSpendLogsCall } from "@/components/networking"; + +const spendLog = (overrides: Partial): SpendLogEntry => ({ + request_id: "req-1", + api_key: "key-1", + team_id: "team-1", + model: "gpt-4o", + model_id: "model-1", + call_type: "acompletion", + spend: 0.01, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + startTime: "2026-09-02T09:50:13Z", + endTime: "2026-09-02T09:50:14Z", + cache_hit: "false", + messages: [], + response: {}, + ...overrides, +}); + +const guardrailLog = { + id: "provider-victim", + timestamp: "2026-09-02 09:50:13", + action: "passed" as const, + input_snippet: "victim prompt", +}; + +describe("GuardrailsMonitor LogViewer drawer", () => { + beforeEach(() => { + vi.mocked(uiSpendLogsCall).mockReset(); + testQueryClient.clear(); + }); + + it("opens the row whose request_id is the clicked log id even when a newer row carries that id as its call id", async () => { + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: [ + spendLog({ request_id: "provider-attacker", litellm_call_id: "provider-victim" }), + spendLog({ request_id: "provider-victim", litellm_call_id: "call-victim" }), + ], + total: 2, + }); + + renderWithProviders(); + await userEvent.click(screen.getByText("victim prompt")); + + await waitFor(() => { + expect(screen.getByTestId("log-details-drawer")).toHaveAttribute("data-log-id", "provider-victim"); + }); + expect(vi.mocked(uiSpendLogsCall)).toHaveBeenCalledWith( + expect.objectContaining({ params: { request_id: "provider-victim" } }), + ); + }); + + it("falls back to the first returned row when none carries the clicked id as its request_id", async () => { + vi.mocked(uiSpendLogsCall).mockResolvedValue({ + data: [spendLog({ request_id: "provider-other", litellm_call_id: "provider-victim" })], + total: 1, + }); + + renderWithProviders(); + await userEvent.click(screen.getByText("victim prompt")); + + await waitFor(() => { + expect(screen.getByTestId("log-details-drawer")).toHaveAttribute("data-log-id", "provider-other"); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx index 8d073feae82..0703c94c2ed 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/LogViewer.tsx @@ -92,7 +92,8 @@ export function LogViewer({ enabled: Boolean(accessToken && selectedRequestId && drawerOpen), }); - const selectedLog: ViewLogsLogEntry | null = fullLogResponse?.data?.[0] ?? null; + const selectedLog: ViewLogsLogEntry | null = + fullLogResponse?.data?.find((log) => log.request_id === selectedRequestId) ?? fullLogResponse?.data?.[0] ?? null; const handleLogClick = (log: LogEntry) => { setSelectedRequestId(log.id); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx index e760d0b8dc3..fda3b08b30b 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx @@ -577,6 +577,49 @@ describe("RequestLogsPanel", () => { expect(byIdCall.params?.group_by_session).toBeUndefined(); }); + it("opens the drawer when ?log_id= is the log's litellm_call_id rather than its request_id", async () => { + respondWith([logEntry({ request_id: "chatcmpl-provider", litellm_call_id: "call-1" })]); + renderPanel("?log_id=call-1"); + + await waitFor(() => { + expect(drawer()).toHaveTextContent("open"); + }); + expect(drawer()).toHaveAttribute("data-log-id", "chatcmpl-provider"); + }); + + it("fetches by litellm_call_id and opens the drawer when that log is not in the loaded page", async () => { + vi.mocked(uiSpendLogsCall).mockImplementation(async ({ params }) => + params?.request_id === "call-old" + ? { + data: [logEntry({ request_id: "chatcmpl-old", litellm_call_id: "call-old" })], + total: 1, + page: 1, + page_size: 1, + total_pages: 1, + } + : { data: [], total: 0, page: 1, page_size: 50, total_pages: 0 }, + ); + renderPanel("?log_id=call-old"); + + await waitFor(() => { + expect(drawer()).toHaveTextContent("open"); + }); + expect(drawer()).toHaveAttribute("data-log-id", "chatcmpl-old"); + }); + + it("opens the exact request_id row when another log in the page carries that id as its litellm_call_id", async () => { + respondWith([ + logEntry({ request_id: "chatcmpl-other", litellm_call_id: "victim-req" }), + logEntry({ request_id: "victim-req", litellm_call_id: "victim-call" }), + ]); + renderPanel("?log_id=victim-req"); + + await waitFor(() => { + expect(drawer()).toHaveTextContent("open"); + }); + expect(drawer()).toHaveAttribute("data-log-id", "victim-req"); + }); + it("closing the drawer removes ?log_id= from the URL and closes the drawer", async () => { const user = userEvent.setup(); respondWith([logEntry({ request_id: "req-1" })]); diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx index 4f39bb3b79b..0b5e5ff9616 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx @@ -28,6 +28,9 @@ import { RequestLogsTable } from "./RequestLogsTable"; const PAGE_SIZE = DEFAULT_PAGE_SIZE_OPTIONS[0]; const DEFAULT_INTERVAL = { value: 24, unit: "hours" }; +const matchesLogId = (log: LogEntry, logId: string) => log.request_id === logId || log.litellm_call_id === logId; +const findLogById = (logs: readonly LogEntry[], logId: string): LogEntry | null => + logs.find((log) => log.request_id === logId) ?? logs.find((log) => log.litellm_call_id === logId) ?? null; interface RequestLogsPanelProps { accessToken: string; @@ -141,9 +144,9 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, page_size: 1, params: { request_id: urlLogId }, }); - return response.data.find((log) => log.request_id === urlLogId) ?? null; + return findLogById(response.data, urlLogId); }, - enabled: urlLogId !== null && selectedLog?.request_id !== urlLogId, + enabled: urlLogId !== null && !(selectedLog !== null && matchesLogId(selectedLog, urlLogId)), staleTime: Infinity, }; @@ -151,8 +154,8 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID, const displayLog = useMemo(() => { if (urlLogId === null) return null; - if (selectedLog?.request_id === urlLogId) return selectedLog; - return filteredLogs.data.find((log) => log.request_id === urlLogId) ?? urlLog ?? null; + if (selectedLog !== null && matchesLogId(selectedLog, urlLogId)) return selectedLog; + return findLogById(filteredLogs.data, urlLogId) ?? urlLog ?? null; }, [urlLogId, selectedLog, filteredLogs.data, urlLog]); const displaySessionId = useMemo(() => { diff --git a/ui/litellm-dashboard/src/components/view_logs/columns.tsx b/ui/litellm-dashboard/src/components/view_logs/columns.tsx index b2e29c3a0c1..0a2b22b95e4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/columns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/columns.tsx @@ -12,6 +12,7 @@ export type LogsSortField = keyof typeof LOGS_SORT_FIELD_MAP; export type LogEntry = { request_id: string; + litellm_call_id?: string | null; api_key: string; team_id: string; model: string; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3807286947d..7eadaa6c991 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29639,6 +29639,12 @@ export interface components { aws_web_identity_token?: string | null; /** Azure Ad Token */ azure_ad_token?: string | null; + /** Azure Password */ + azure_password?: string | null; + /** Azure Scope */ + azure_scope?: string | null; + /** Azure Username */ + azure_username?: string | null; /** Bedrock Tags */ bedrock_tags?: unknown[] | null; /** Budget Duration */ @@ -29687,6 +29693,10 @@ export interface components { cache_read_input_token_cost_ultrafast?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; + /** Client Id */ + client_id?: string | null; + /** Client Secret */ + client_secret?: string | null; /** Complexity Router Config */ complexity_router_config?: { [key: string]: unknown; @@ -29900,6 +29910,8 @@ export interface components { tag_regex?: string[] | null; /** Tags */ tags?: string[] | null; + /** Tenant Id */ + tenant_id?: string | null; /** Tiered Pricing */ tiered_pricing?: { [key: string]: unknown; @@ -39841,6 +39853,12 @@ export interface components { aws_web_identity_token?: string | null; /** Azure Ad Token */ azure_ad_token?: string | null; + /** Azure Password */ + azure_password?: string | null; + /** Azure Scope */ + azure_scope?: string | null; + /** Azure Username */ + azure_username?: string | null; /** Bedrock Tags */ bedrock_tags?: unknown[] | null; /** Budget Duration */ @@ -39889,6 +39907,10 @@ export interface components { cache_read_input_token_cost_ultrafast?: number | null; /** Citation Cost Per Token */ citation_cost_per_token?: number | null; + /** Client Id */ + client_id?: string | null; + /** Client Secret */ + client_secret?: string | null; /** Complexity Router Config */ complexity_router_config?: { [key: string]: unknown; @@ -40102,6 +40124,8 @@ export interface components { tag_regex?: string[] | null; /** Tags */ tags?: string[] | null; + /** Tenant Id */ + tenant_id?: string | null; /** Tiered Pricing */ tiered_pricing?: { [key: string]: unknown;