From 0684a1e27521ada35cf8a4afbdff7aeaa87ff41d Mon Sep 17 00:00:00 2001 From: Hendrik Jaks Date: Tue, 31 Mar 2026 22:39:13 +0300 Subject: [PATCH] revert: remove unrelated changes from HttpOnly cookie fix branch Reset files not related to the login cookie fix back to main: - prometheus.py, bedrock converse, guardrail handler - auth_checks.py, reset_budget_job.py, audit_logs.py - test_user_api_key_auth.py Co-Authored-By: Claude Opus 4.6 (1M context) --- litellm/integrations/prometheus.py | 420 +----------------- .../bedrock/chat/converse_transformation.py | 68 +-- .../chat/guardrail_translation/handler.py | 43 +- litellm/proxy/auth/auth_checks.py | 302 +++++-------- .../proxy/common_utils/reset_budget_job.py | 86 +--- .../proxy/management_helpers/audit_logs.py | 128 +----- .../proxy/auth/test_user_api_key_auth.py | 342 ++++---------- 7 files changed, 264 insertions(+), 1125 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 5d1ff832930..357e0229fc6 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -65,17 +65,6 @@ def _get_cached_end_user_id_for_cost_tracking(): class PrometheusLogger(CustomLogger): # Class variables or attributes - - @staticmethod - def get_instance() -> Optional["PrometheusLogger"]: - """Find the PrometheusLogger instance from litellm.callbacks, if registered.""" - import litellm - - for cb in litellm.callbacks: - if isinstance(cb, PrometheusLogger): - return cb - return None - def __init__( # noqa: PLR0915 self, **kwargs, @@ -191,31 +180,6 @@ class PrometheusLogger(CustomLogger): ), ) - # Remaining Budget for Org - self.litellm_remaining_org_budget_metric = self._gauge_factory( - "litellm_remaining_org_budget_metric", - "Remaining budget for org", - labelnames=self.get_labels_for_metric( - "litellm_remaining_org_budget_metric" - ), - ) - - # Max Budget for Org - self.litellm_org_max_budget_metric = self._gauge_factory( - "litellm_org_max_budget_metric", - "Maximum budget set for org", - labelnames=self.get_labels_for_metric("litellm_org_max_budget_metric"), - ) - - # Org Budget Reset At - self.litellm_org_budget_remaining_hours_metric = self._gauge_factory( - "litellm_org_budget_remaining_hours_metric", - "Remaining hours for org budget to be reset", - labelnames=self.get_labels_for_metric( - "litellm_org_budget_remaining_hours_metric" - ), - ) - # Remaining Budget for API Key self.litellm_remaining_api_key_budget_metric = self._gauge_factory( "litellm_remaining_api_key_budget_metric", @@ -476,76 +440,6 @@ class PrometheusLogger(CustomLogger): labelnames=[], ) - ######################################## - # Managed Batch Metrics - ######################################## - self.litellm_managed_batch_created_total = self._counter_factory( - name="litellm_managed_batch_created_total", - documentation="Total number of managed batches created", - labelnames=[ - "model", - "api_provider", - "user", - "user_email", - "api_key_alias", - ], - ) - - self.litellm_managed_file_size_bytes = self._gauge_factory( - "litellm_managed_file_size_bytes", - "Size of the most recent managed batch file in bytes (last-seen value per label combination)", - labelnames=["purpose", "file_type", "model", "api_provider", "user"], - ) - - self.litellm_managed_batch_duration_seconds = self._histogram_factory( - "litellm_managed_batch_duration_seconds", - "Duration of completed managed batches in seconds (completed_at - created_at)", - labelnames=["model", "api_provider"], - buckets=BATCH_DURATION_BUCKETS, - ) - - self.litellm_managed_file_created_total = self._counter_factory( - name="litellm_managed_file_created_total", - documentation="Total number of managed files created", - labelnames=[ - "model", - "api_provider", - "user", - "user_email", - "api_key_alias", - ], - ) - - self.litellm_managed_file_deleted_total = self._counter_factory( - name="litellm_managed_file_deleted_total", - documentation="Total number of managed file deletions (success or blocked)", - labelnames=["result"], - ) - - self.litellm_check_batch_cost_jobs_polled = self._gauge_factory( - "litellm_check_batch_cost_jobs_polled", - "Number of unprocessed batches found by the last CheckBatchCost poll", - labelnames=[], - ) - - self.litellm_check_batch_cost_jobs_processed_total = self._counter_factory( - name="litellm_check_batch_cost_jobs_processed_total", - documentation="Total number of batches successfully cost-tracked by CheckBatchCost", - labelnames=["model", "api_provider"], - ) - - self.litellm_check_batch_cost_errors_total = self._counter_factory( - name="litellm_check_batch_cost_errors_total", - documentation="Total number of errors in CheckBatchCost by error type", - labelnames=["error_type"], - ) - - self.litellm_check_batch_cost_last_run_timestamp = self._gauge_factory( - "litellm_check_batch_cost_last_run_timestamp", - "Unix timestamp of the last CheckBatchCost job run", - labelnames=[], - ) - except Exception as e: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e @@ -1028,9 +922,6 @@ class PrometheusLogger(CustomLogger): user_api_team_alias = standard_logging_payload["metadata"][ "user_api_key_team_alias" ] - user_api_key_org_id = standard_logging_payload["metadata"].get( - "user_api_key_org_id" - ) output_tokens = standard_logging_payload["completion_tokens"] tokens_used = standard_logging_payload["total_tokens"] response_cost = standard_logging_payload["response_cost"] @@ -1040,14 +931,10 @@ class PrometheusLogger(CustomLogger): user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[ "metadata" ].get("user_api_key_auth_metadata") - spend_logs_metadata: Optional[dict] = standard_logging_payload["metadata"].get( - "spend_logs_metadata" - ) combined_metadata: Dict[str, Any] = { **(_requester_metadata if _requester_metadata else {}), **(user_api_key_auth_metadata if user_api_key_auth_metadata else {}), - **(spend_logs_metadata if spend_logs_metadata else {}), } if standard_logging_payload is not None and isinstance( standard_logging_payload, dict @@ -1087,11 +974,9 @@ class PrometheusLogger(CustomLogger): ), client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), user_agent=standard_logging_payload["metadata"].get("user_agent"), - stream=( - str(standard_logging_payload.get("stream")) - if litellm.prometheus_emit_stream_label - else None - ), + stream=str(standard_logging_payload.get("stream")) + if litellm.prometheus_emit_stream_label + else None, ) if ( @@ -1141,7 +1026,6 @@ class PrometheusLogger(CustomLogger): litellm_params=litellm_params, response_cost=response_cost, user_id=user_id, - user_api_key_org_id=user_api_key_org_id, ) # set proxy virtual key rpm/tpm metrics @@ -1297,7 +1181,6 @@ class PrometheusLogger(CustomLogger): litellm_params: dict, response_cost: float, user_id: Optional[str] = None, - user_api_key_org_id: Optional[str] = None, ): _metadata = litellm_params.get("metadata") or {} _team_spend = _metadata.get("user_api_key_team_spend", None) @@ -1330,16 +1213,12 @@ class PrometheusLogger(CustomLogger): user_max_budget=_user_max_budget, response_cost=response_cost, ), - self._set_org_budget_metrics_after_api_request( - org_id=user_api_key_org_id, - response_cost=response_cost, - ), return_exceptions=True, ) for i, r in enumerate(results): if isinstance(r, Exception): verbose_logger.debug( - f"[Non-Blocking] Prometheus: Budget metric lookup {['key', 'team', 'user', 'org'][i]} failed: {r}" + f"[Non-Blocking] Prometheus: Budget metric lookup {['key', 'team', 'user'][i]} failed: {r}" ) def _increment_top_level_request_and_spend_metrics( @@ -1528,9 +1407,6 @@ class PrometheusLogger(CustomLogger): user_api_team_alias = standard_logging_payload["metadata"][ "user_api_key_team_alias" ] - user_api_key_org_id = standard_logging_payload["metadata"].get( - "user_api_key_org_id" - ) try: self.litellm_llm_api_failed_requests_metric.labels( @@ -1546,10 +1422,6 @@ class PrometheusLogger(CustomLogger): ), ).inc() self.set_llm_deployment_failure_metrics(kwargs) - await self._set_org_budget_metrics_after_api_request( - org_id=user_api_key_org_id, - response_cost=0, - ) except Exception as e: verbose_logger.exception( "prometheus Layer Error(): Exception occured - {}".format(str(e)) @@ -1757,11 +1629,9 @@ class PrometheusLogger(CustomLogger): client_ip=_metadata.get("requester_ip_address"), user_agent=_metadata.get("user_agent"), model_id=model_id, - stream=( - str(request_data.get("stream")) - if litellm.prometheus_emit_stream_label - else None - ), + stream=str(request_data.get("stream")) + if litellm.prometheus_emit_stream_label + else None, ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( @@ -2284,127 +2154,6 @@ class PrometheusLogger(CustomLogger): except Exception as e: verbose_logger.debug(f"Error recording guardrail metrics: {str(e)}") - ######################################## - # Managed Batch Metric Recording Methods - ######################################## - - def record_managed_batch_created( - self, - model: Optional[str], - api_provider: Optional[str], - user: Optional[str], - user_email: Optional[str], - api_key_alias: Optional[str], - ): - try: - self.litellm_managed_batch_created_total.labels( - model=model, - api_provider=api_provider, - user=user, - user_email=user_email, - api_key_alias=api_key_alias, - ).inc() - except Exception as e: - verbose_logger.warning(f"Error recording batch created metric: {e}") - - def record_managed_file_size( - self, - size_bytes: int, - purpose: str, - file_type: str, - model: Optional[str] = None, - api_provider: Optional[str] = None, - user: Optional[str] = None, - ): - """Record the size of a managed file. Uses a gauge (last-seen value per label combination).""" - try: - self.litellm_managed_file_size_bytes.labels( - purpose=purpose, - file_type=file_type, - model=model or "", - api_provider=api_provider or "", - user=user or "", - ).set(size_bytes) - except Exception as e: - verbose_logger.warning(f"Error recording file size metric: {e}") - - def record_managed_batch_duration( - self, - duration_seconds: float, - model: Optional[str] = None, - api_provider: Optional[str] = None, - ): - try: - self.litellm_managed_batch_duration_seconds.labels( - model=model or "", - api_provider=api_provider or "", - ).observe(duration_seconds) - except Exception as e: - verbose_logger.warning(f"Error recording batch duration metric: {e}") - - def record_managed_file_created( - self, - model: Optional[str], - api_provider: Optional[str], - user: Optional[str], - user_email: Optional[str], - api_key_alias: Optional[str], - ): - try: - self.litellm_managed_file_created_total.labels( - model=model, - api_provider=api_provider, - user=user, - user_email=user_email, - api_key_alias=api_key_alias, - ).inc() - except Exception as e: - verbose_logger.warning(f"Error recording file created metric: {e}") - - def record_managed_file_deleted(self, result: str): - """Record a managed file deletion attempt. result is 'success' or 'blocked'.""" - try: - self.litellm_managed_file_deleted_total.labels(result=result).inc() - except Exception as e: - verbose_logger.warning(f"Error recording file deleted metric: {e}") - - def record_check_batch_cost_run( - self, - jobs_polled: int, - processed_models: Optional[List[Tuple[Optional[str], Optional[str]]]] = None, - ): - """ - Record CheckBatchCost polling metrics. - - Args: - jobs_polled: Number of unprocessed batches found - processed_models: List of (model, api_provider) tuples for processed jobs - """ - import time - - try: - self.litellm_check_batch_cost_last_run_timestamp.set(time.time()) - self.litellm_check_batch_cost_jobs_polled.set(jobs_polled) - - if processed_models: - for model, api_provider in processed_models: - self.litellm_check_batch_cost_jobs_processed_total.labels( - model=model or "", - api_provider=api_provider or "", - ).inc() - except Exception as e: - verbose_logger.warning(f"Error recording check batch cost metrics: {e}") - - def record_check_batch_cost_error(self, error_type: str): - try: - self.litellm_check_batch_cost_errors_total.labels( - error_type=error_type, - ).inc() - except Exception as e: - verbose_logger.warning( - f"Error recording check batch cost error metric: {e}" - ) - @staticmethod def _get_exception_class_name(exception: Exception) -> str: exception_class_name = "" @@ -2781,35 +2530,6 @@ class PrometheusLogger(CustomLogger): data_type="users", ) - async def _initialize_org_budget_metrics(self): - """ - Initialize org budget metrics by reusing the generic pagination logic. - """ - from litellm.proxy.proxy_server import prisma_client - - if prisma_client is None: - verbose_logger.debug( - "Prometheus: skipping org metrics initialization, DB not initialized" - ) - return - - async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]: - skip = (page - 1) * page_size - orgs = await prisma_client.db.litellm_organizationtable.find_many( - skip=skip, - take=page_size, - order={"created_at": "desc"}, - include={"litellm_budget_table": True}, - ) - total_count = await prisma_client.db.litellm_organizationtable.count() - return orgs, total_count - - await self._initialize_budget_metrics( - data_fetch_function=fetch_orgs, - set_metrics_function=self._set_org_list_budget_metrics, - data_type="orgs", - ) - async def initialize_remaining_budget_metrics(self): """ Handler for initializing remaining budget metrics for all teams to avoid metric discrepancies. @@ -2844,11 +2564,10 @@ class PrometheusLogger(CustomLogger): """ Helper to initialize remaining budget metrics for all teams, API keys, and users. """ - verbose_logger.debug("Emitting key, team, user, org budget metrics....") + verbose_logger.debug("Emitting key, team, user budget metrics....") await self._initialize_team_budget_metrics() await self._initialize_api_key_budget_metrics() await self._initialize_user_budget_metrics() - await self._initialize_org_budget_metrics() await self._initialize_user_and_team_count_metrics() async def _initialize_user_and_team_count_metrics(self): @@ -2904,22 +2623,6 @@ class PrometheusLogger(CustomLogger): for user in users: self._set_user_budget_metrics(user) - async def _set_org_list_budget_metrics(self, orgs: list): - """Helper function to set budget metrics for a list of orgs""" - for org in orgs: - budget_table = getattr(org, "litellm_budget_table", None) - self._set_org_budget_metrics( - org_id=org.organization_id or "", - org_alias=org.organization_alias or "", - spend=org.spend or 0.0, - max_budget=budget_table.max_budget if budget_table else None, - budget_reset_at=( - getattr(budget_table, "budget_reset_at", None) - if budget_table - else None - ), - ) - async def _set_team_budget_metrics_after_api_request( self, user_api_team: Optional[str], @@ -3041,113 +2744,6 @@ class PrometheusLogger(CustomLogger): ) ) - async def _set_org_budget_metrics_after_api_request( - self, - org_id: Optional[str], - response_cost: float, - ): - """ - Set org budget metrics after an LLM API request - - - Fetches org info via cache (get_org_object) - - Sets org budget metrics - """ - if not org_id: - return - - from litellm.proxy.auth.auth_checks import get_org_object - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - if prisma_client is None: - return - - try: - org_info = await get_org_object( - org_id=org_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - include_budget_table=True, - ) - except Exception as e: - verbose_logger.debug( - f"[Non-Blocking] Prometheus: Error getting org info: {str(e)}" - ) - return - - if org_info is None: - return - - org_alias = org_info.organization_alias or "" - _total_org_spend = (org_info.spend or 0.0) + response_cost - budget_table = org_info.litellm_budget_table - max_budget = budget_table.max_budget if budget_table else None - budget_reset_at = ( - getattr(budget_table, "budget_reset_at", None) if budget_table else None - ) - - self._set_org_budget_metrics( - org_id=org_id, - org_alias=org_alias, - spend=_total_org_spend, - max_budget=max_budget, - budget_reset_at=budget_reset_at, - ) - - def _set_org_budget_metrics( - self, - org_id: str, - org_alias: str, - spend: float, - max_budget: Optional[float], - budget_reset_at: Optional[datetime], - ): - """ - Set org budget metrics for a single org - - - Remaining Budget - - Max Budget - - Budget Reset At - """ - enum_values = UserAPIKeyLabelValues( - org_id=org_id, - org_alias=org_alias, - ) - - _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_remaining_org_budget_metric" - ), - enum_values=enum_values, - ) - self.litellm_remaining_org_budget_metric.labels(**_labels).set( - self._safe_get_remaining_budget( - max_budget=max_budget, - spend=spend, - ) - ) - - if max_budget is not None: - _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_org_max_budget_metric" - ), - enum_values=enum_values, - ) - self.litellm_org_max_budget_metric.labels(**_labels).set(max_budget) - - if budget_reset_at is not None: - _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_org_budget_remaining_hours_metric" - ), - enum_values=enum_values, - ) - self.litellm_org_budget_remaining_hours_metric.labels(**_labels).set( - self._get_remaining_hours_for_budget_reset( - budget_reset_at=budget_reset_at - ) - ) - def _set_key_budget_metrics(self, user_api_key_dict: UserAPIKeyAuth): """ Set virtual key budget metrics diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 5cfb00d69b6..229457a73b4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -91,6 +91,34 @@ UNSUPPORTED_BEDROCK_CONVERSE_BETA_PATTERNS = [ "compact-2026-01-12", # The compact beta feature is not currently supported on the Converse and ConverseStream APIs ] +# Models that support Bedrock's native structured outputs API (outputConfig.textFormat) +# Uses substring matching against the Bedrock model ID +# Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html +BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS = { + # Anthropic Claude 4.5+ + "claude-haiku-4-5", + "claude-sonnet-4-5", + "claude-opus-4-5", + "claude-opus-4-6", + # Qwen3 + "qwen3", + # DeepSeek + "deepseek-v3.1", + # Gemma 3 + "gemma-3", + # MiniMax + "minimax-m2", + # Mistral (magistral-small excluded: broken constrained decoding on Bedrock) + "ministral", + "mistral-large-3", + "voxtral", + # Moonshot + "kimi-k2", + # NVIDIA + "nemotron-nano", + # OpenAI (gpt-oss excluded: broken constrained decoding, works via tool-call fallback) +} + class AmazonConverseConfig(BaseConfig): """ @@ -465,7 +493,8 @@ class AmazonConverseConfig(BaseConfig): budget = thinking.get("budget_tokens") if isinstance(budget, int) and budget < BEDROCK_MIN_THINKING_BUDGET_TOKENS: verbose_logger.debug( - "Bedrock requires thinking.budget_tokens >= %d, got %d. Clamping to minimum.", + "Bedrock requires thinking.budget_tokens >= %d, got %d. " + "Clamping to minimum.", BEDROCK_MIN_THINKING_BUDGET_TOKENS, budget, ) @@ -734,20 +763,10 @@ class AmazonConverseConfig(BaseConfig): return _tool @staticmethod - def _supports_native_structured_outputs( - model: str, custom_llm_provider: Optional[str] = None - ) -> bool: - """Check if the Bedrock model supports native structured outputs (outputConfig.textFormat). - - Delegates to the standard ``supports_native_structured_output`` utility - which looks up the flag in ``litellm.model_cost`` via - ``_get_model_info_helper``. - Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/structured-output.html - """ - from litellm.utils import supports_native_structured_output - - return supports_native_structured_output( - model=model, custom_llm_provider=custom_llm_provider + def _supports_native_structured_outputs(model: str) -> bool: + """Check if the Bedrock model supports native structured outputs (outputConfig.textFormat).""" + return any( + substring in model for substring in BEDROCK_NATIVE_STRUCTURED_OUTPUT_MODELS ) @staticmethod @@ -894,9 +913,7 @@ class AmazonConverseConfig(BaseConfig): ) if param == "tool_choice": _tool_choice_value = self.map_tool_choice_values( - model=model, - tool_choice=value, - drop_params=drop_params, # type: ignore + model=model, tool_choice=value, drop_params=drop_params # type: ignore ) if _tool_choice_value is not None: optional_params["tool_choice"] = _tool_choice_value @@ -989,10 +1006,7 @@ class AmazonConverseConfig(BaseConfig): if "type" in value and value["type"] == "text": return optional_params - if ( - self._supports_native_structured_outputs(model, self.custom_llm_provider) - and json_schema is not None - ): + if self._supports_native_structured_outputs(model) and json_schema is not None: # Use Bedrock's native structured outputs API (outputConfig.textFormat) # No synthetic tool injection, no fake_stream needed. # Requires an explicit schema — json_object with no schema falls through @@ -1432,16 +1446,6 @@ class AmazonConverseConfig(BaseConfig): original_tools, model, headers, additional_request_params ) - # Append cachePoint to tools if cache_control_injection_points has tool_config - cache_injection_points = additional_request_params.pop( - "cache_control_injection_points", None - ) - if cache_injection_points and len(bedrock_tools) > 0: - for point in cache_injection_points: - if point.get("location") == "tool_config": - bedrock_tools.append({"cachePoint": {"type": "default"}}) - break - bedrock_tool_config: Optional[ToolConfigBlock] = None if len(bedrock_tools) > 0: tool_choice_values: ToolChoiceValuesBlock = inference_params.pop( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index f854cdb13d0..bab4c3b5eb7 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -260,7 +260,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, - request_data: Optional[dict] = None, ) -> Any: """ Process output response by applying guardrails to text content. @@ -309,21 +308,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 2: Apply guardrail to all texts and tool calls in batch if texts_to_check or tool_calls_to_check: - # Use the real request_data if provided (proxy path), otherwise - # create a standalone dict (SDK / direct-call path). - if request_data is None: - request_data = {"response": response} - else: - if "response" not in request_data: - request_data["response"] = response + # Create a request_data dict with response info and user API key metadata + request_data: dict = {"response": response} # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -371,7 +364,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional[Any] = None, user_api_key_dict: Optional[Any] = None, - request_data: Optional[dict] = None, ) -> List["ModelResponseStream"]: """ Process output streaming responses by applying guardrails to text content. @@ -410,7 +402,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, - request_data=request_data, ) return responses_so_far @@ -445,21 +436,15 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 3: Apply guardrail to all combined texts in batch if texts_to_check: - # Use the real request_data if provided (proxy path), otherwise - # create a standalone dict (SDK / direct-call path). - if request_data is None: - request_data = {"responses": responses_so_far} - else: - if "responses" not in request_data: - request_data["responses"] = responses_so_far + # Create a request_data dict with response info and user API key metadata + request_data: dict = {"responses": responses_so_far} # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata = self.transform_user_api_key_dict_to_metadata( - user_api_key_dict - ) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index ed600b26e83..b4c123b361b 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -8,7 +8,6 @@ Run checks for: 2. If user is in budget 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget """ - import asyncio import re import time @@ -30,7 +29,6 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, ) -from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.proxy._types import ( RBAC_ROLES, @@ -40,7 +38,6 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_JWTAuth, - LiteLLM_ManagedVectorStoresTable, LiteLLM_ObjectPermissionTable, LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, @@ -405,26 +402,23 @@ async def common_checks( # noqa: PLR0915 # 1. If team is blocked if team_object is not None and team_object.blocked is True: raise Exception( - f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if you're an admin." + f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if your admin." ) # 2. If team can call model if _model and team_object: - with tracer.trace("litellm.proxy.auth.common_checks.can_team_access_model"): - if not await can_team_access_model( - model=_model, - team_object=team_object, - llm_router=llm_router, - team_model_aliases=( - valid_token.team_model_aliases if valid_token else None - ), - ): - raise ProxyException( - message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=status.HTTP_401_UNAUTHORIZED, - ) + if not await can_team_access_model( + model=_model, + team_object=team_object, + llm_router=llm_router, + team_model_aliases=valid_token.team_model_aliases if valid_token else None, + ): + raise ProxyException( + message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=status.HTTP_401_UNAUTHORIZED, + ) # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent if valid_token is not None and valid_token.agent_id: @@ -449,62 +443,54 @@ async def common_checks( # noqa: PLR0915 ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: - with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): - await can_user_call_model( - model=_model, - llm_router=llm_router, - user_object=user_object, - ) + await can_user_call_model( + model=_model, + llm_router=llm_router, + user_object=user_object, + ) # 1.1 - 2.2 - 3.0.2 - 3.0.3: Project checks (blocked, model access, budget) - with tracer.trace("litellm.proxy.auth.common_checks.run_project_checks"): - await _run_project_checks( - project_object=project_object, - _model=_model, - llm_router=llm_router, - skip_budget_checks=skip_budget_checks, - valid_token=valid_token, - proxy_logging_obj=proxy_logging_obj, - ) + await _run_project_checks( + project_object=project_object, + _model=_model, + llm_router=llm_router, + skip_budget_checks=skip_budget_checks, + valid_token=valid_token, + proxy_logging_obj=proxy_logging_obj, + ) # If this is a free model, skip all budget checks if not skip_budget_checks: # 3. If team is in budget - with tracer.trace("litellm.proxy.auth.common_checks.team_max_budget_check"): - await _team_max_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ) + await _team_max_budget_check( + team_object=team_object, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) # 3.0.5. If team is over soft budget (alert only, doesn't block) - with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"): - await _team_soft_budget_check( - team_object=team_object, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ) + await _team_soft_budget_check( + team_object=team_object, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) # 3.1. If organization is in budget - with tracer.trace( - "litellm.proxy.auth.common_checks.organization_max_budget_check" - ): - await _organization_max_budget_check( - valid_token=valid_token, - team_object=team_object, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + await _organization_max_budget_check( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) - with tracer.trace("litellm.proxy.auth.common_checks.tag_max_budget_check"): - await _tag_max_budget_check( - request_body=request_body, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - valid_token=valid_token, - ) + await _tag_max_budget_check( + request_body=request_body, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + valid_token=valid_token, + ) # 4. If user is in budget ## 4.1 check personal budget, if personal key @@ -522,15 +508,14 @@ async def common_checks( # noqa: PLR0915 ) ## 4.2 check team member budget, if team key - with tracer.trace("litellm.proxy.auth.common_checks.check_team_member_budget"): - await _check_team_member_budget( - team_object=team_object, - user_object=user_object, - valid_token=valid_token, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget if ( @@ -569,21 +554,19 @@ async def common_checks( # noqa: PLR0915 ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store - with tracer.trace("litellm.proxy.auth.common_checks.vector_store_access_check"): - await vector_store_access_check( - request_body=request_body, - team_object=team_object, - valid_token=valid_token, - ) + await vector_store_access_check( + request_body=request_body, + team_object=team_object, + valid_token=valid_token, + ) # 12. [OPTIONAL] Tool allowlist - key/team allowed_tools (no DB in hot path) - with tracer.trace("litellm.proxy.auth.common_checks.check_tools_allowlist"): - await check_tools_allowlist( - request_body=request_body, - valid_token=valid_token, - team_object=team_object, - route=route, - ) + await check_tools_allowlist( + request_body=request_body, + valid_token=valid_token, + team_object=team_object, + route=route, + ) return True @@ -2296,71 +2279,6 @@ async def get_object_permission( return None -@log_db_metrics -async def get_managed_vector_store_rows_by_uuids( - uuids: List[str], - prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, - parent_otel_span: Optional[Span] = None, - proxy_logging_obj: Optional[ProxyLogging] = None, -) -> List[LiteLLM_ManagedVectorStoresTable]: - """ - Fetch managed vector store rows by their internal UUIDs. - - Follows the get_team_object / get_key_object / get_object_permission pattern: - cache-first lookup (in-memory / Redis), DB fallback only on cache miss. - Critical-path DB access must go through this helper to avoid raw Prisma - calls on the hot request path. - """ - if not uuids or prisma_client is None: - return [] - - result: List[LiteLLM_ManagedVectorStoresTable] = [] - cache_misses: List[str] = [] - - for uuid in uuids: - key = "managed_vector_store_id:{}".format(uuid) - cached = await user_api_key_cache.async_get_cache(key=key) - if cached is not None: - if isinstance(cached, dict): - result.append(LiteLLM_ManagedVectorStoresTable(**cached)) - elif isinstance(cached, LiteLLM_ManagedVectorStoresTable): - result.append(cached) - else: - cache_misses.append(uuid) - else: - cache_misses.append(uuid) - - if not cache_misses: - return result - - rows = await prisma_client.db.litellm_managedvectorstorestable.find_many( - where={"vector_store_id": {"in": cache_misses}}, - take=len(cache_misses), - ) - - for row in rows: - row_dict = ( - row.model_dump() - if hasattr(row, "model_dump") - else (row.dict() if hasattr(row, "dict") else None) - ) - if not isinstance(row_dict, dict) or not row_dict: - row_dict = dict(row) if hasattr(row, "__dict__") else {} - if not row_dict: - continue - cached_obj = LiteLLM_ManagedVectorStoresTable(**row_dict) - key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id) - await user_api_key_cache.async_set_cache( - key=key, - value=row_dict, - ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - ) - result.append(cached_obj) - - return result - - @log_db_metrics async def get_org_object( org_id: str, @@ -2877,15 +2795,7 @@ async def _virtual_key_max_budget_check( Triggers a budget alert if the token is over it's max budget. """ - if valid_token.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - spend = await get_current_spend( - counter_key=f"spend:key:{valid_token.token}", - fallback_spend=valid_token.spend or 0.0, - ) - + if valid_token.spend is not None and valid_token.max_budget is not None: #################################### # collect information for alerting # #################################### @@ -2897,7 +2807,7 @@ async def _virtual_key_max_budget_check( call_info = CallInfo( token=valid_token.token, - spend=spend, + spend=valid_token.spend, max_budget=valid_token.max_budget, soft_budget=valid_token.soft_budget, user_id=valid_token.user_id, @@ -2918,9 +2828,9 @@ async def _virtual_key_max_budget_check( # collect information for alerting # #################################### - if spend >= valid_token.max_budget: + if valid_token.spend >= valid_token.max_budget: raise litellm.BudgetExceededError( - current_cost=spend, + current_cost=valid_token.spend, max_budget=valid_token.max_budget, ) @@ -3051,14 +2961,6 @@ async def _check_team_member_budget( team_member_budget = team_membership.litellm_budget_table.max_budget team_member_spend = team_membership.spend or 0.0 - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = await get_current_spend( - counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", - fallback_spend=team_member_spend, - ) - if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( current_cost=team_member_spend, @@ -3079,39 +2981,35 @@ async def _team_max_budget_check( BudgetExceededError if the team is over it's max budget. Triggers a budget alert if the team is over it's max budget. """ - if team_object is not None and team_object.max_budget is not None: - from litellm.proxy.proxy_server import get_current_spend - - # Read spend from cross-pod counter (Redis-first) or cached object (fallback) - spend = await get_current_spend( - counter_key=f"spend:team:{team_object.team_id}", - fallback_spend=team_object.spend or 0.0, - ) - - if spend > team_object.max_budget: - if valid_token: - call_info = CallInfo( - token=valid_token.token, - spend=spend, - max_budget=team_object.max_budget, - user_id=valid_token.user_id, - team_id=valid_token.team_id, - team_alias=valid_token.team_alias, - organization_id=valid_token.org_id, - event_group=Litellm_EntityType.TEAM, - ) - asyncio.create_task( - proxy_logging_obj.budget_alerts( - type="team_budget", - user_info=call_info, - ) - ) - - raise litellm.BudgetExceededError( - current_cost=spend, + if ( + team_object is not None + and team_object.max_budget is not None + and team_object.spend is not None + and team_object.spend > team_object.max_budget + ): + if valid_token: + call_info = CallInfo( + token=valid_token.token, + spend=team_object.spend, max_budget=team_object.max_budget, - message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}", + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.TEAM, ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="team_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=team_object.spend, + max_budget=team_object.max_budget, + message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}", + ) async def _team_soft_budget_check( diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed24398..674214b19e5 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -54,47 +54,16 @@ class ResetBudgetJob: """ Resets the budget for all LiteLLM Team Members if their budget has expired """ - budget_ids = [ - budget.budget_id - for budget in budgets_to_reset - if budget.budget_id is not None - ] - - # Reset spend counters for affected team members. - # Reset Redis directly so a transient failure doesn't leave stale - # counters that get_current_spend would read as authoritative. - try: - from litellm.proxy.proxy_server import spend_counter_cache - - memberships = await self.prisma_client.db.litellm_teammembership.find_many( - where={"budget_id": {"in": budget_ids}} - ) - for m in memberships: - counter_key = f"spend:team_member:{m.user_id}:{m.team_id}" - # Always reset in-memory - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 - ) - # Explicitly reset Redis with warning on failure - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset team member spend counter in Redis %s: %s. " - "Budget may be over-enforced until counter expires.", - counter_key, - redis_err, - ) - except Exception as e: - verbose_proxy_logger.warning( - "Failed to reset team member spend counters: %s", e - ) - return await self.prisma_client.db.litellm_teammembership.update_many( - where={"budget_id": {"in": budget_ids}}, + where={ + "budget_id": { + "in": [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + } + }, data={ "spend": 0, }, @@ -562,43 +531,6 @@ class ResetBudgetJob: """ try: item.spend = 0.0 - - # Reset the cross-pod spend counter. - # Reset Redis directly (not via DualCache) so a Redis failure - # doesn't silently leave a stale counter that get_current_spend - # would read as authoritative, permanently blocking the user. - from litellm.proxy.proxy_server import spend_counter_cache - - counter_key = None - if item_type == "key" and hasattr(item, "token") and item.token is not None: - counter_key = f"spend:key:{item.token}" - elif ( - item_type == "team" - and hasattr(item, "team_id") - and item.team_id is not None - ): - counter_key = f"spend:team:{item.team_id}" - - if counter_key is not None: - # Always reset in-memory (local fallback) - spend_counter_cache.in_memory_cache.set_cache( - key=counter_key, value=0.0 - ) - # Explicitly reset Redis with warning on failure - if spend_counter_cache.redis_cache is not None: - try: - await spend_counter_cache.redis_cache.async_set_cache( - key=counter_key, value=0.0 - ) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to reset spend counter in Redis for %s key=%s: %s. " - "Budget may be over-enforced until counter expires.", - item_type, - counter_key, - redis_err, - ) - if hasattr(item, "budget_duration") and item.budget_duration is not None: # Get standardized reset time based on budget duration from litellm.proxy.common_utils.timezone_utils import ( diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 7599e11bdef..ea082f468ae 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -2,15 +2,12 @@ Functions to create audit logs for LiteLLM Proxy """ -import asyncio import json +from litellm._uuid import uuid from datetime import datetime, timezone -from typing import Dict import litellm from litellm._logging import verbose_proxy_logger -from litellm._uuid import uuid -from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( AUDIT_ACTIONS, LiteLLM_AuditLogs, @@ -18,99 +15,6 @@ from litellm.proxy._types import ( Optional, UserAPIKeyAuth, ) -from litellm.types.utils import StandardAuditLogPayload - -_audit_log_callback_cache: Dict[str, CustomLogger] = {} - - -def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]: - """Resolve a string callback name to a CustomLogger instance, with caching.""" - if name in _audit_log_callback_cache: - return _audit_log_callback_cache[name] - - from litellm.litellm_core_utils.litellm_logging import ( - _init_custom_logger_compatible_class, - ) - - instance = _init_custom_logger_compatible_class( - logging_integration=name, # type: ignore - internal_usage_cache=None, - llm_router=None, - ) - - if instance is not None: - _audit_log_callback_cache[name] = instance - return instance - - -def _build_audit_log_payload( - request_data: LiteLLM_AuditLogs, -) -> StandardAuditLogPayload: - """Convert LiteLLM_AuditLogs to StandardAuditLogPayload for callback dispatch.""" - updated_at = "" - if request_data.updated_at is not None: - updated_at = request_data.updated_at.isoformat() - - table_name_str: str = ( - request_data.table_name.value - if isinstance(request_data.table_name, LitellmTableNames) - else str(request_data.table_name) - ) - - return StandardAuditLogPayload( - id=request_data.id, - updated_at=updated_at, - changed_by=request_data.changed_by or "", - changed_by_api_key=request_data.changed_by_api_key or "", - action=request_data.action, - table_name=table_name_str, - object_id=request_data.object_id, - before_value=request_data.before_value, - updated_values=request_data.updated_values, - ) - - -def _audit_log_task_done_callback(task: asyncio.Task) -> None: - """Log exceptions from audit log callback tasks so they don't slip through silently.""" - try: - exc = task.exception() - except asyncio.CancelledError: - return - if exc is not None: - verbose_proxy_logger.error( - "Audit log callback task failed: %s", exc, exc_info=exc - ) - - -async def _dispatch_audit_log_to_callbacks( - request_data: LiteLLM_AuditLogs, -) -> None: - """Dispatch audit log to all registered audit_log_callbacks.""" - if not litellm.audit_log_callbacks: - return - - payload = _build_audit_log_payload(request_data) - - for callback in litellm.audit_log_callbacks: - try: - resolved: Optional[CustomLogger] = ( - callback if isinstance(callback, CustomLogger) else None - ) - if isinstance(callback, str): - resolved = _resolve_audit_log_callback(callback) - if resolved is None: - verbose_proxy_logger.warning( - "Could not resolve audit log callback: %s", callback - ) - continue - - if isinstance(resolved, CustomLogger): - task = asyncio.create_task(resolved.async_log_audit_log_event(payload)) - task.add_done_callback(_audit_log_task_done_callback) - except Exception as e: - verbose_proxy_logger.error( - "Failed dispatching audit log to callback: %s", e - ) async def create_object_audit_log( @@ -136,22 +40,20 @@ async def create_object_audit_log( """ from litellm.secret_managers.main import get_secret_bool - _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool( + store_audit_logs = litellm.store_audit_logs or get_secret_bool( "LITELLM_STORE_AUDIT_LOGS" ) - if _store_audit_logs is not True: + if store_audit_logs is not True: return - _changed_by = ( - litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name - ) - await create_audit_log_for_update( request_data=LiteLLM_AuditLogs( id=str(uuid.uuid4()), updated_at=datetime.now(timezone.utc), - changed_by=_changed_by, + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, changed_by_api_key=user_api_key_dict.api_key, table_name=table_name, object_id=object_id, @@ -168,10 +70,10 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): """ from litellm.secret_managers.main import get_secret_bool - _store_audit_logs: Optional[bool] = litellm.store_audit_logs or get_secret_bool( + store_audit_logs = litellm.store_audit_logs or get_secret_bool( "LITELLM_STORE_AUDIT_LOGS" ) - if _store_audit_logs is not True: + if store_audit_logs is not True: return from litellm.proxy.proxy_server import premium_user, prisma_client @@ -179,6 +81,9 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): if premium_user is not True: return + if prisma_client is None: + raise Exception("prisma_client is None, no DB connected") + verbose_proxy_logger.debug("creating audit log for %s", request_data) if isinstance(request_data.updated_values, dict): @@ -187,15 +92,6 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): if isinstance(request_data.before_value, dict): request_data.before_value = json.dumps(request_data.before_value) - # Dispatch to external audit log callbacks regardless of DB availability - await _dispatch_audit_log_to_callbacks(request_data) - - if prisma_client is None: - verbose_proxy_logger.error( - "prisma_client is None, cannot write audit log to DB" - ) - return - _request_data = request_data.model_dump(exclude_none=True) try: @@ -207,3 +103,5 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): except Exception as e: # [Non-Blocking Exception. Do not allow blocking LLM API call] verbose_proxy_logger.error(f"Failed Creating audit log {e}") + + return 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 74c7f9bca5f..f3f0ba56cb9 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 @@ -72,7 +72,7 @@ def test_get_api_key_with_custom_litellm_key_header( def test_team_metadata_with_tags_flows_through_jwt_auth(): """ Test that team_metadata (specifically tags) flows through JWT authentication. - + This is a regression test for the issue where JWT auth was not populating team_metadata, causing team-level tags to be missing in litellm_pre_call_utils.py """ @@ -87,7 +87,7 @@ def test_team_metadata_with_tags_flows_through_jwt_auth(): rpm_limit=100, models=["gpt-4", "gpt-3.5-turbo"], ) - + # Simulate constructing UserAPIKeyAuth like we do in JWT auth # This is the pattern from user_api_key_auth.py lines 552-587 user_api_key_auth = UserAPIKeyAuth( @@ -100,16 +100,14 @@ def test_team_metadata_with_tags_flows_through_jwt_auth(): user_role="internal_user", user_id="test-user", ) - + # Verify team_metadata is set - assert ( - user_api_key_auth.team_metadata is not None - ), "team_metadata should be populated" + assert user_api_key_auth.team_metadata is not None, "team_metadata should be populated" assert user_api_key_auth.team_metadata == team_object.metadata, ( f"team_metadata not correctly mapped. " f"Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}" ) - + # Specifically verify tags are present assert "tags" in user_api_key_auth.team_metadata, "tags should be in team_metadata" assert user_api_key_auth.team_metadata["tags"] == ["production", "high-priority"], ( @@ -120,7 +118,7 @@ def test_team_metadata_with_tags_flows_through_jwt_auth(): def test_route_checks_is_llm_api_route(): """Test RouteChecks.is_llm_api_route() correctly identifies LLM API routes including passthrough endpoints""" - + # Test OpenAI routes openai_routes = [ "/v1/chat/completions", @@ -144,22 +142,18 @@ def test_route_checks_is_llm_api_route(): "/v1/realtime", "/realtime", ] - + for route in openai_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test Anthropic routes anthropic_routes = [ "/v1/messages", "/v1/messages/count_tokens", ] - + for route in anthropic_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test passthrough routes (this is the key improvement over the old route checking) passthrough_routes = [ @@ -177,11 +171,9 @@ def test_route_checks_is_llm_api_route(): "/vllm/v1/chat/completions", "/mistral/v1/chat/completions", ] - + for route in passthrough_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test MCP routes mcp_routes = [ @@ -189,11 +181,9 @@ def test_route_checks_is_llm_api_route(): "/mcp/", "/mcp/test", ] - + for route in mcp_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test LiteLLM native RAG routes rag_routes = [ @@ -203,9 +193,7 @@ def test_route_checks_is_llm_api_route(): "/v1/rag/query", ] for route in rag_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test routes with placeholders placeholder_routes = [ @@ -218,11 +206,9 @@ def test_route_checks_is_llm_api_route(): "/v1/batches/batch_123", "/batches/batch_123", ] - + for route in placeholder_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test Azure OpenAI routes azure_routes = [ @@ -231,11 +217,9 @@ def test_route_checks_is_llm_api_route(): "/engines/gpt-4/chat/completions", "/engines/gpt-3.5-turbo/completions", ] - + for route in azure_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test non-LLM routes (should return False) non_llm_routes = [ @@ -252,11 +236,9 @@ def test_route_checks_is_llm_api_route(): "/debug", "/test", ] - + for route in non_llm_routes: - assert not RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should NOT be identified as LLM API route" + assert not RouteChecks.is_llm_api_route(route=route), f"Route {route} should NOT be identified as LLM API route" # Test invalid inputs invalid_inputs = [ @@ -266,11 +248,9 @@ def test_route_checks_is_llm_api_route(): {}, "", ] - + for invalid_input in invalid_inputs: - assert not RouteChecks.is_llm_api_route( - route=invalid_input - ), f"Invalid input {invalid_input} should return False" + assert not RouteChecks.is_llm_api_route(route=invalid_input), f"Invalid input {invalid_input} should return False" @pytest.mark.asyncio @@ -279,7 +259,7 @@ async def test_proxy_admin_expired_key_from_cache(): Test that PROXY_ADMIN keys retrieved from cache are checked for expiration before being returned. This prevents expired keys from bypassing expiration checks when retrieved from cache (which normally happens at lines 1014-1036). - + Regression test for issue where PROXY_ADMIN keys from cache skipped expiration check. """ from datetime import datetime, timedelta, timezone @@ -300,42 +280,39 @@ async def test_proxy_admin_expired_key_from_cache(): api_key = "sk-test-proxy-admin-key" hashed_key = hash_token(api_key) expired_time = datetime.now(timezone.utc) - timedelta(hours=1) # Expired 1 hour ago - + expired_token = UserAPIKeyAuth( api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN, expires=expired_time, token=hashed_key, ) - + # Mock cache to return the expired token mock_cache = AsyncMock() mock_cache.async_get_cache = AsyncMock(return_value=expired_token) mock_cache.delete_cache = MagicMock() - + # Mock proxy_logging_obj mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() # Mock post_call_failure_hook as async function returning None (no transformation) mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - + # Mock prisma_client mock_prisma_client = MagicMock() - + # Mock get_key_object to return expired token from cache with patch( "litellm.proxy.auth.user_api_key_auth.get_key_object", new_callable=AsyncMock, - ) as mock_get_key_object, patch( - "litellm.proxy.auth.user_api_key_auth._delete_cache_key_object", - new_callable=AsyncMock, - ) as mock_delete_cache: + ) as mock_get_key_object, \ + patch("litellm.proxy.auth.user_api_key_auth._delete_cache_key_object", new_callable=AsyncMock) as mock_delete_cache: + mock_get_key_object.return_value = expired_token - + # Set attributes on proxy_server module (these are imported inside _user_api_key_auth_builder) import litellm.proxy.proxy_server as _proxy_server_mod @@ -354,12 +331,14 @@ async def test_proxy_admin_expired_key_from_cache(): "litellm_proxy_admin_name": "admin", } _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set + attr: getattr(_proxy_server_mod, attr, None) + for attr in _attrs_to_set } try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) + # Create a mock request request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") @@ -379,41 +358,38 @@ async def test_proxy_admin_expired_key_from_cache(): ) # Verify that ProxyException was raised with expired_key type - assert hasattr( - exc_info.value, "type" - ), "Exception should have 'type' attribute" - assert ( - exc_info.value.type == ProxyErrorTypes.expired_key - ), f"Expected expired_key error type, got {exc_info.value.type}" - assert "Expired Key" in str( - exc_info.value.message - ), f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" + assert hasattr(exc_info.value, "type"), "Exception should have 'type' attribute" + assert exc_info.value.type == ProxyErrorTypes.expired_key, ( + f"Expected expired_key error type, got {exc_info.value.type}" + ) + assert "Expired Key" in str(exc_info.value.message), ( + f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" + ) # Verify that the param field does NOT leak the full API key (Issue #18731) # The param should be abbreviated like "sk-...XXXX" not the full plaintext key - assert ( - exc_info.value.param is not None - ), "Exception should have 'param' attribute" + assert exc_info.value.param is not None, "Exception should have 'param' attribute" assert exc_info.value.param != api_key, ( f"SECURITY: Full API key should NOT be in param field! " f"Got: {exc_info.value.param}, Expected abbreviated format like 'sk-...XXXX'" ) - assert exc_info.value.param.startswith( - "sk-..." - ), f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + assert exc_info.value.param.startswith("sk-..."), ( + f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + ) # Verify that cache deletion was called mock_delete_cache.assert_called_once() call_args = mock_delete_cache.call_args - assert ( - call_args[1]["hashed_token"] == hashed_key - ), "Cache deletion should be called with the hashed key" + assert call_args[1]["hashed_token"] == hashed_key, ( + "Cache deletion should be called with the hashed key" + ) finally: # Restore all module-level attributes so subsequent tests are not affected for attr, val in _original_values.items(): setattr(_proxy_server_mod, attr, val) + @pytest.mark.asyncio async def test_return_user_api_key_auth_obj_user_spend_and_budget(): """ @@ -424,7 +400,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj - + user_obj = type( "LiteLLM_UserTable", (), @@ -437,7 +413,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): "user_role": "internal_user", }, ) - + api_key = "sk-test-key" valid_token_dict = { "user_id": "test-user", @@ -445,10 +421,10 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): } route = "/chat/completions" start_time = datetime.now() - + mock_service_logger = MagicMock() mock_service_logger.async_service_success_hook = AsyncMock() - + with patch( "litellm.proxy.auth.user_api_key_auth.user_api_key_service_logger_obj", new=mock_service_logger, @@ -462,7 +438,7 @@ async def test_return_user_api_key_auth_obj_user_spend_and_budget(): start_time=start_time, user_role=None, ) - + assert isinstance(result, UserAPIKeyAuth) assert result.user_spend == 250.0 assert result.user_max_budget == 1000.0 @@ -494,7 +470,9 @@ def test_proxy_admin_jwt_auth_includes_identity_fields(): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="user-abc", team_id="team-123", - team_alias=(team_object.team_alias if team_object is not None else None), + team_alias=( + team_object.team_alias if team_object is not None else None + ), team_metadata=team_object.metadata if team_object is not None else None, org_id="org-456", end_user_id="end-user-789", @@ -525,7 +503,9 @@ def test_proxy_admin_jwt_auth_handles_no_team_object(): user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user", team_id=None, - team_alias=(team_object.team_alias if team_object is not None else None), + team_alias=( + team_object.team_alias if team_object is not None else None + ), team_metadata=team_object.metadata if team_object is not None else None, org_id=None, end_user_id=None, @@ -554,10 +534,7 @@ class TestJWTOAuth2Coexistence: def test_is_jwt_detects_jwt_tokens(self): """JWT tokens have 3 dot-separated parts.""" assert JWTHandler.is_jwt("header.payload.signature") is True - assert ( - JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") - is True - ) + assert JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") is True def test_is_jwt_rejects_opaque_tokens(self): """Opaque OAuth2 tokens do not have 3 dot-separated parts.""" @@ -566,10 +543,6 @@ class TestJWTOAuth2Coexistence: assert JWTHandler.is_jwt("Bearer token") is False assert JWTHandler.is_jwt("two.parts") is False - def test_is_jwt_returns_false_for_none(self): - """None token (missing Authorization header) should not be treated as JWT.""" - assert JWTHandler.is_jwt(None) is False - @pytest.mark.asyncio async def test_both_enabled_opaque_token_uses_oauth2(self): """ @@ -594,20 +567,13 @@ class TestJWTOAuth2Coexistence: mock_request.headers = {"authorization": f"Bearer {opaque_token}"} mock_request.query_params = {} - with patch( - "litellm.proxy.proxy_server.general_settings", general_settings - ), patch("litellm.proxy.proxy_server.premium_user", True), patch( - "litellm.proxy.proxy_server.master_key", "sk-master" - ), patch( - "litellm.proxy.proxy_server.prisma_client", None - ), patch( - "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", - new_callable=AsyncMock, - return_value=mock_oauth2_response, - ) as mock_oauth2, patch( - "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", - new_callable=AsyncMock, - ) as mock_jwt_auth: + with patch("litellm.proxy.proxy_server.general_settings", general_settings), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch("litellm.proxy.proxy_server.master_key", "sk-master"), \ + patch("litellm.proxy.proxy_server.prisma_client", None), \ + patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock, return_value=mock_oauth2_response) as mock_oauth2, \ + patch("litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock) as mock_jwt_auth: + litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -658,20 +624,13 @@ class TestJWTOAuth2Coexistence: mock_request.headers = {"authorization": f"Bearer {jwt_token}"} mock_request.query_params = {} - with patch( - "litellm.proxy.proxy_server.general_settings", general_settings - ), patch("litellm.proxy.proxy_server.premium_user", True), patch( - "litellm.proxy.proxy_server.master_key", "sk-master" - ), patch( - "litellm.proxy.proxy_server.prisma_client", None - ), patch( - "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", - new_callable=AsyncMock, - ) as mock_oauth2, patch( - "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", - new_callable=AsyncMock, - return_value=mock_jwt_result, - ) as mock_jwt_auth: + with patch("litellm.proxy.proxy_server.general_settings", general_settings), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch("litellm.proxy.proxy_server.master_key", "sk-master"), \ + patch("litellm.proxy.proxy_server.prisma_client", None), \ + patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock) as mock_oauth2, \ + patch("litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock, return_value=mock_jwt_result) as mock_jwt_auth: + litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -712,17 +671,12 @@ class TestJWTOAuth2Coexistence: mock_request.headers = {"authorization": f"Bearer {jwt_like_token}"} mock_request.query_params = {} - with patch( - "litellm.proxy.proxy_server.general_settings", general_settings - ), patch("litellm.proxy.proxy_server.premium_user", True), patch( - "litellm.proxy.proxy_server.master_key", "sk-master" - ), patch( - "litellm.proxy.proxy_server.prisma_client", None - ), patch( - "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", - new_callable=AsyncMock, - return_value=mock_oauth2_response, - ) as mock_oauth2: + with patch("litellm.proxy.proxy_server.general_settings", general_settings), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch("litellm.proxy.proxy_server.master_key", "sk-master"), \ + patch("litellm.proxy.proxy_server.prisma_client", None), \ + patch("litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", new_callable=AsyncMock, return_value=mock_oauth2_response) as mock_oauth2: + result = await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_like_token}", @@ -731,131 +685,3 @@ class TestJWTOAuth2Coexistence: # OAuth2 should handle it since JWT auth is disabled mock_oauth2.assert_called_once_with(token=jwt_like_token) assert result.user_id == "oauth2-user" - - -@pytest.mark.asyncio -async def test_user_api_key_auth_builder_no_blocking_calls(): - """ - _user_api_key_auth_builder must never call any synchronous DualCache method - (set_cache, get_cache, batch_get_cache, increment_cache, delete_cache) on - the hot auth path — those methods call Redis synchronously and block the - event loop. Only async_* variants are allowed. - """ - from starlette.datastructures import URL - from starlette.requests import Request - - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder - - _blocking_methods = [ - "set_cache", - "get_cache", - "batch_get_cache", - "increment_cache", - "delete_cache", - ] - - api_key = "sk-test-no-blocking-cache" - valid_token = UserAPIKeyAuth( - api_key=api_key, - token=api_key, - user_role=LitellmUserRoles.INTERNAL_USER, - team_id="team-abc", - ) - - mock_cache = AsyncMock() - mock_cache.async_get_cache = AsyncMock(return_value=valid_token) - mock_cache.async_set_cache = AsyncMock(return_value=None) - # Wire sync methods on the instance as plain MagicMocks (no side_effect) so - # calls are recorded but not raised — the function's broad except Exception - # would swallow a raised error. We assert not_called() after the run instead. - for _m in _blocking_methods: - setattr(mock_cache, _m, MagicMock()) - - mock_proxy_logging_obj = MagicMock() - mock_proxy_logging_obj.internal_usage_cache = MagicMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) - mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - - import litellm.proxy.proxy_server as _proxy_server_mod - - _attrs = { - "prisma_client": MagicMock(), - "user_api_key_cache": mock_cache, - "proxy_logging_obj": mock_proxy_logging_obj, - "master_key": "sk-master-key", - "general_settings": {}, - "llm_model_list": [], - "llm_router": None, - "open_telemetry_logger": None, - "model_max_budget_limiter": MagicMock(), - "user_custom_auth": None, - "jwt_handler": None, - "litellm_proxy_admin_name": "admin", - } - _originals = {k: getattr(_proxy_server_mod, k, None) for k in _attrs} - - try: - for k, v in _attrs.items(): - setattr(_proxy_server_mod, k, v) - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - import contextlib - - from litellm.caching.dual_cache import DualCache - - blocking_patches = [ - patch.object( - DualCache, - m, - MagicMock( - side_effect=AssertionError( - f"Blocking DualCache.{m}() called on async hot path — use async_{m}() instead" - ) - ), - ) - for m in _blocking_methods - ] - - with contextlib.ExitStack() as stack: - for p in blocking_patches: - stack.enter_context(p) - stack.enter_context( - patch( - "litellm.proxy.auth.user_api_key_auth.get_key_object", - new_callable=AsyncMock, - return_value=valid_token, - ) - ) - stack.enter_context( - patch( - "litellm.proxy.auth.user_api_key_auth.get_team_object", - new_callable=AsyncMock, - return_value=None, - ) - ) - await _user_api_key_auth_builder( - request=request, - api_key=f"Bearer {api_key}", - azure_api_key_header="", - anthropic_api_key_header=None, - google_ai_studio_api_key_header=None, - azure_apim_header=None, - request_data={}, - ) - - for _m in _blocking_methods: - mock = getattr(mock_cache, _m) - assert mock.call_count == 0, ( - f"Blocking DualCache.{_m}() was called {mock.call_count} time(s) " - f"on the async hot path — use async_{_m}() instead" - ) - - finally: - for k, v in _originals.items(): - setattr(_proxy_server_mod, k, v)