diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 210f232b17..f9ce9e5dcb 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -2,6 +2,10 @@ +## Linear ticket + + + ## Pre-Submission checklist **Please complete all items before asking a LiteLLM maintainer to review your PR** diff --git a/.github/workflows/create-release-branch.yml b/.github/workflows/create-release-branch.yml index 13b76c94df..ec2651306f 100644 --- a/.github/workflows/create-release-branch.yml +++ b/.github/workflows/create-release-branch.yml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: tag: - description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/" + description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted) — branch will be named release/" required: true type: string commit_hash: @@ -14,7 +14,7 @@ on: workflow_call: inputs: tag: - description: "Release tag" + description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)" required: true type: string commit_hash: @@ -40,8 +40,8 @@ jobs: echo "::error::commit_hash must be a full 40-character commit SHA" exit 1 fi - if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then - echo "::error::tag must start with vX.Y.Z" + if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then + echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable" exit 1 fi diff --git a/.github/workflows/create-release.yml b/.github/workflows/create-release.yml index 68ab397d82..39d078267f 100644 --- a/.github/workflows/create-release.yml +++ b/.github/workflows/create-release.yml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: tag: - description: "Release tag (e.g. v1.83.0-stable)" + description: "Release tag (e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, 1.84.0.post1; legacy v1.83.10-stable still accepted)" required: true type: string commit_hash: @@ -30,8 +30,8 @@ jobs: echo "::error::commit_hash must be a full 40-character commit SHA" exit 1 fi - if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then - echo "::error::tag must start with vX.Y.Z" + if ! echo "${TAG}" | grep -qE '^v?[0-9]+\.[0-9]+\.[0-9]+'; then + echo "::error::tag must start with X.Y.Z (optional leading v), e.g. 1.84.0, 1.84.0rc1, 1.84.0.dev42, or v1.83.10-stable" exit 1 fi @@ -45,6 +45,11 @@ jobs: const tag = process.env.TAG; const commitHash = process.env.COMMIT_HASH; + // Mark RC / dev / nightly / alpha / beta tags as GitHub pre-releases. + // PEP 440 post-releases (e.g. `1.84.0.post1`) and legacy `-stable[.patch.N]` + // are stable maintenance releases, not pre-releases. + const isPrerelease = /(?:rc|nightly|alpha|beta|\.dev)/i.test(tag); + const cosignSection = [ `## Verify Docker Image Signature`, ``, @@ -89,7 +94,7 @@ jobs: target_commitish: commitHash, name: tag, owner: context.repo.owner, - prerelease: false, + prerelease: isPrerelease, repo: context.repo.repo, tag_name: tag, }); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql new file mode 100644 index 0000000000..6454f26765 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429161855_workflow_runs_tables/migration.sql @@ -0,0 +1,75 @@ +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowRun" ( + "run_id" TEXT NOT NULL, + "session_id" TEXT NOT NULL, + "workflow_type" TEXT NOT NULL, + "status" TEXT NOT NULL DEFAULT 'pending', + "created_by" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + "input" JSONB, + "output" JSONB, + "metadata" JSONB, + + CONSTRAINT "LiteLLM_WorkflowRun_pkey" PRIMARY KEY ("run_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowEvent" ( + "event_id" TEXT NOT NULL, + "run_id" TEXT NOT NULL, + "event_type" TEXT NOT NULL, + "step_name" TEXT NOT NULL, + "sequence_number" INTEGER NOT NULL, + "data" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_WorkflowEvent_pkey" PRIMARY KEY ("event_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_WorkflowMessage" ( + "message_id" TEXT NOT NULL, + "run_id" TEXT NOT NULL, + "role" TEXT NOT NULL, + "content" TEXT NOT NULL, + "sequence_number" INTEGER NOT NULL, + "session_id" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_WorkflowMessage_pkey" PRIMARY KEY ("message_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_session_id_key" ON "LiteLLM_WorkflowRun"("session_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_workflow_type_status_idx" ON "LiteLLM_WorkflowRun"("workflow_type", "status"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_session_id_idx" ON "LiteLLM_WorkflowRun"("session_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_created_at_idx" ON "LiteLLM_WorkflowRun"("created_at"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowRun_created_by_idx" ON "LiteLLM_WorkflowRun"("created_by"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowEvent_run_id_idx" ON "LiteLLM_WorkflowEvent"("run_id"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowEvent_run_id_sequence_number_key" ON "LiteLLM_WorkflowEvent"("run_id", "sequence_number"); + +-- CreateIndex +CREATE INDEX "LiteLLM_WorkflowMessage_run_id_idx" ON "LiteLLM_WorkflowMessage"("run_id"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_WorkflowMessage_run_id_sequence_number_key" ON "LiteLLM_WorkflowMessage"("run_id", "sequence_number"); + +-- AddForeignKey +ALTER TABLE "LiteLLM_WorkflowEvent" ADD CONSTRAINT "LiteLLM_WorkflowEvent_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_WorkflowMessage" ADD CONSTRAINT "LiteLLM_WorkflowMessage_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8f07c5afa3..fd54ed5243 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession { @@id([session_id, router_name, model_name]) @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } + +// --------------------------------------------------------------------------- +// Workflow Run Tracking +// +// Generic durable state tracking for any agent or automated workflow. +// Design: three tables — run (header + materialized status), event (append-only +// source of truth for state transitions), message (conversation inbox/outbox). +// +// Usage: +// - Set `workflow_type` to identify the owning system (e.g. "shin-builder"). +// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.). +// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to +// the proxy — all spend logs for this run are automatically tagged. +// --------------------------------------------------------------------------- + +// One instance of work being done. `status` is a materialized cache of the +// latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowRun { + run_id String @id @default(uuid()) + session_id String @unique @default(uuid()) + workflow_type String + status String @default("pending") + created_by String? // user_id of the key that created this run; null = created by master key + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + input Json? + output Json? + metadata Json? + + events LiteLLM_WorkflowEvent[] + messages LiteLLM_WorkflowMessage[] + + @@index([workflow_type, status]) + @@index([session_id]) + @@index([created_at]) + @@index([created_by]) +} + +// Append-only log of state transitions. Never mutate rows here. +// `step_name` and `event_type` are caller-defined strings — no hardcoded enums. +// Status auto-update rules (applied by the append endpoint): +// step.started → run.status = running +// step.failed → run.status = failed +// hook.waiting → run.status = paused +// hook.received → run.status = running +model LiteLLM_WorkflowEvent { + event_id String @id @default(uuid()) + run_id String + event_type String + step_name String + sequence_number Int + data Json? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} + +// Conversation inbox/outbox — full message content, separate from the durable +// event log. Spend logs truncate messages; this table stores them in full. +// `session_id` here is the Claude --resume session ID (or similar). +model LiteLLM_WorkflowMessage { + message_id String @id @default(uuid()) + run_id String + role String + content String + sequence_number Int + session_id String? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 6a68ba8c4d..ce1bc26c5e 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -650,7 +650,10 @@ class Cache: verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") def _convert_to_cached_embedding( - self, embedding_response: Any, model: Optional[str] + self, + embedding_response: Any, + model: Optional[str], + prompt_tokens_details: Optional[dict] = None, ) -> CachedEmbedding: """ Convert any embedding response into the standardized CachedEmbedding TypedDict format. @@ -662,6 +665,7 @@ class Cache: "index": embedding_response.get("index"), "object": embedding_response.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } elif hasattr(embedding_response, "model_dump"): data = embedding_response.model_dump() @@ -670,6 +674,7 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } else: data = vars(embedding_response) @@ -678,10 +683,54 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens_details": prompt_tokens_details, } except KeyError as e: raise ValueError(f"Missing expected key in embedding response: {e}") + def _get_per_item_prompt_tokens_details( + self, + result: EmbeddingResponse, + idx_in_result_data: int, + ) -> Optional[dict]: + """ + Extract per-item prompt_tokens_details from a response for caching. + + For single-item responses (common for multimodal providers like Bedrock Titan, + Nova, Vertex AI), returns the full prompt_tokens_details. + For multi-item responses, distributes integer fields evenly across items + so that summing all per-item details reconstructs the original totals. + """ + if result.usage is None or result.usage.prompt_tokens_details is None: + return None + + details = result.usage.prompt_tokens_details + if hasattr(details, "model_dump"): + details_dict = details.model_dump(exclude_none=True) + elif isinstance(details, dict): + details_dict = {k: v for k, v in details.items() if v is not None} + else: + return None + + if not details_dict: + return None + + num_items = len(result.data) + if num_items <= 1: + return details_dict + + # Distribute integer/float fields evenly across items + per_item: dict = {} + for key, value in details_dict.items(): + if isinstance(value, int): + quotient, remainder = divmod(value, num_items) + per_item[key] = quotient + (1 if idx_in_result_data < remainder else 0) + elif isinstance(value, float): + per_item[key] = value / num_items + else: + per_item[key] = value + return per_item if per_item else None + def add_embedding_response_to_cache( self, result: EmbeddingResponse, @@ -693,10 +742,18 @@ class Cache: kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx_in_result_data] + # Extract per-item prompt_tokens_details from response usage + prompt_tokens_details = self._get_per_item_prompt_tokens_details( + result=result, + idx_in_result_data=idx_in_result_data, + ) + # Always convert to properly typed CachedEmbedding model_name = result.model embedding_dict: CachedEmbedding = self._convert_to_cached_embedding( - embedding_response, model_name + embedding_response, + model_name, + prompt_tokens_details=prompt_tokens_details, ) cache_key, cached_data, kwargs = self._add_cache_logic( diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 2bec705946..7d514e648f 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -59,6 +59,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import PromptTokensDetailsWrapper else: LiteLLMLoggingObj = Any @@ -415,6 +416,7 @@ class LLMCachingHandler: final_embedding_cached_response._hidden_params["cache_hit"] = True prompt_tokens = 0 + aggregated_details: Optional[dict] = None for val in non_null_list: idx, cr = val # (idx, cr) tuple if cr is not None: @@ -431,11 +433,35 @@ class LLMCachingHandler: prompt_tokens += token_counter( text=kwargs_input_as_list[idx], count_response_tokens=True ) + # Aggregate prompt_tokens_details from cached items + item_details = cr.get("prompt_tokens_details") + if item_details: + if aggregated_details is None: + aggregated_details = {} + for key, value in item_details.items(): + if isinstance(value, (int, float)): + aggregated_details[key] = ( + aggregated_details.get(key, 0) + value + ) + else: + aggregated_details[key] = value + ## USAGE + prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None + if aggregated_details: + from litellm.types.utils import PromptTokensDetailsWrapper + + try: + prompt_tokens_details = PromptTokensDetailsWrapper( + **aggregated_details + ) + except Exception: + prompt_tokens_details = None usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens, + prompt_tokens_details=prompt_tokens_details, ) final_embedding_cached_response.usage = usage if len(remaining_list) == 0: @@ -478,8 +504,70 @@ class LLMCachingHandler: prompt_tokens=usage1.prompt_tokens + usage2.prompt_tokens, completion_tokens=usage1.completion_tokens + usage2.completion_tokens, total_tokens=usage1.total_tokens + usage2.total_tokens, + prompt_tokens_details=self._merge_prompt_tokens_details( + usage1.prompt_tokens_details, + usage2.prompt_tokens_details, + ), ) + def _merge_prompt_tokens_details( + self, + details1: Optional["PromptTokensDetailsWrapper"], + details2: Optional["PromptTokensDetailsWrapper"], + ) -> Optional["PromptTokensDetailsWrapper"]: + """Merge two PromptTokensDetailsWrapper objects by summing numeric fields.""" + if details1 is None and details2 is None: + return None + if details1 is None: + return details2 + if details2 is None: + return details1 + + dict1 = ( + details1.model_dump(exclude_none=True) + if hasattr(details1, "model_dump") + else {} + ) + dict2 = ( + details2.model_dump(exclude_none=True) + if hasattr(details2, "model_dump") + else {} + ) + + merged: dict = {} + for key in set(dict1.keys()) | set(dict2.keys()): + v1 = dict1.get(key, 0) + v2 = dict2.get(key, 0) + if isinstance(v1, (int, float)) and isinstance(v2, (int, float)): + merged[key] = v1 + v2 + elif isinstance(v1, dict) and isinstance(v2, dict): + # Recursively merge nested dicts (e.g. cache_creation_token_details) + nested: dict = {} + for nk in set(v1.keys()) | set(v2.keys()): + nv1 = v1.get(nk, 0) + nv2 = v2.get(nk, 0) + if isinstance(nv1, (int, float)) and isinstance(nv2, (int, float)): + nested[nk] = nv1 + nv2 + elif nv1: + nested[nk] = nv1 + else: + nested[nk] = nv2 + merged[key] = nested + elif v1: + merged[key] = v1 + else: + merged[key] = v2 + + if not merged: + return None + + from litellm.types.utils import PromptTokensDetailsWrapper + + try: + return PromptTokensDetailsWrapper(**merged) + except Exception: + return None + def _combine_cached_embedding_response_with_api_result( self, _caching_handler_response: CachingHandlerResponse, diff --git a/litellm/constants.py b/litellm/constants.py index 32e016cd7a..a0e99dd16b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -224,6 +224,16 @@ AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int( ) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) +# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs +# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT +# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing +# during a long provider call and the NAT reaps the flow before the response +# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset +# the NAT idle timer. +AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true" +AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60)) +AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30)) +AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5)) # enable_cleanup_closed is only needed for Python versions with the SSL leak bug # Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960) # Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78 @@ -1396,6 +1406,15 @@ LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) ) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv( + "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false" +) +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS = int( + os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400) +) # 24 hours default +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int( + os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000) +) LITELLM_PROXY_ADMIN_NAME = "default_user_id" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### @@ -1425,6 +1444,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( ) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" +EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 8cb2ee0937..3e97b48077 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -29,12 +29,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: litellm_logging_obj: LiteLLMLoggingObj, request_body: dict, model: str, + hidden_params: Optional[Dict[str, Any]] = None, ): self.litellm_logging_obj = litellm_logging_obj self.request_body = request_body self.start_time = datetime.now() self.collected_chunks: List[bytes] = [] self.model = model + self._hidden_params: Dict[str, Any] = hidden_params or {} async def _handle_async_streaming_logging( self, @@ -76,11 +78,13 @@ class GoogleGenAIGenerateContentStreamingIterator( litellm_metadata: dict, custom_llm_provider: str, request_body: Optional[dict] = None, + hidden_params: Optional[Dict[str, Any]] = None, ): super().__init__( litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + hidden_params=hidden_params, ) self.response = response self.model = model @@ -130,11 +134,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator( litellm_metadata: dict, custom_llm_provider: str, request_body: Optional[dict] = None, + hidden_params: Optional[Dict[str, Any]] = None, ): super().__init__( litellm_logging_obj=logging_obj, request_body=request_body or {}, model=model, + hidden_params=hidden_params, ) self.response = response self.model = model diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 9a8060520d..2982df8fda 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -11,8 +11,9 @@ import json import os import re import traceback -from typing import Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union +import httpx import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -103,6 +104,9 @@ class GenericAPILogger(CustomBatchLogger): event_types: Optional[List[API_EVENT_TYPES]] = None, callback_name: Optional[str] = None, log_format: Optional[LOG_FORMAT_TYPES] = None, + max_retries: int = 0, + retry_delay: float = 1.0, + timeout: Optional[Union[float, httpx.Timeout]] = None, **kwargs, ): """ @@ -114,6 +118,9 @@ class GenericAPILogger(CustomBatchLogger): event_types: Optional[List[API_EVENT_TYPES]] = None, callback_name: Optional[str] = None - If provided, loads config from generic_api_compatible_callbacks.json log_format: Optional[LOG_FORMAT_TYPES] = None - Format for log output: "json_array" (default), "ndjson", or "single" + max_retries: Number of retry attempts after the initial request fails. Defaults to 0. + retry_delay: Initial retry delay in seconds. Retries use exponential backoff. + timeout: Optional timeout to use for Generic API callback requests. """ ######################################################### # Check if callback_name is provided and load config @@ -162,6 +169,10 @@ class GenericAPILogger(CustomBatchLogger): self.endpoint: str = endpoint self.event_types: Optional[List[API_EVENT_TYPES]] = event_types self.callback_name: Optional[str] = callback_name + self.max_retries = max(0, int(max_retries or 0)) + retry_delay_value = 0.0 if retry_delay is None else retry_delay + self.retry_delay = max(0.0, float(retry_delay_value)) + self.timeout = timeout # Validate and store log_format if log_format is not None and log_format not in [ @@ -226,6 +237,53 @@ class GenericAPILogger(CustomBatchLogger): return headers_dict + def _should_retry_exception(self, exception: Exception) -> bool: + if isinstance(exception, (litellm.Timeout, httpx.TransportError)): + return True + + if isinstance(exception, httpx.HTTPStatusError): + return exception.response.status_code >= 500 + + return False + + async def _sleep_before_retry(self, attempt: int) -> None: + if self.retry_delay <= 0: + return + + delay = self.retry_delay * (2**attempt) + await asyncio.sleep(delay) + + async def _post_with_retries(self, data: str) -> httpx.Response: + post_kwargs: Dict[str, Any] = { + "url": self.endpoint, + "headers": self.headers, + "data": data, + } + if self.timeout is not None: + post_kwargs["timeout"] = self.timeout + + total_attempts = self.max_retries + 1 + for attempt in range(total_attempts): + try: + return await self.async_httpx_client.post(**post_kwargs) + except Exception as e: + is_last_attempt = attempt == self.max_retries + should_retry = self._should_retry_exception(e) + if is_last_attempt or not should_retry: + raise + + verbose_logger.warning( + "Generic API Logger - retrying request to %s after error: %s " + "(attempt %s/%s)", + self.endpoint, + str(e), + attempt + 1, + total_attempts, + ) + await self._sleep_before_retry(attempt) + + raise RuntimeError("Generic API Logger retry loop exited unexpectedly") + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Generic API Endpoint @@ -325,11 +383,7 @@ class GenericAPILogger(CustomBatchLogger): # Send each log as individual HTTP request in parallel tasks = [] for log_entry in self.log_queue: - task = self.async_httpx_client.post( - url=self.endpoint, - headers=self.headers, - data=safe_dumps(log_entry), - ) + task = self._post_with_retries(data=safe_dumps(log_entry)) tasks.append(task) # Execute all requests in parallel @@ -356,11 +410,7 @@ class GenericAPILogger(CustomBatchLogger): raise ValueError(f"Unknown log_format: {self.log_format}") # Make POST request - response = await self.async_httpx_client.post( - url=self.endpoint, - headers=self.headers, - data=data, - ) + response = await self._post_with_retries(data=data) verbose_logger.debug( f"Generic API Logger - sent batch to {self.endpoint}, " diff --git a/litellm/integrations/prompt_management_base.py b/litellm/integrations/prompt_management_base.py index 71da650dc4..9c626aea84 100644 --- a/litellm/integrations/prompt_management_base.py +++ b/litellm/integrations/prompt_management_base.py @@ -87,9 +87,7 @@ class PromptManagementBase(ABC): try: messages = compiled_prompt_client["prompt_template"] + client_messages except Exception as e: - raise ValueError( - f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}" - ) + raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}") compiled_prompt_client["completed_messages"] = messages return compiled_prompt_client @@ -116,9 +114,7 @@ class PromptManagementBase(ABC): try: messages = compiled_prompt_client["prompt_template"] + client_messages except Exception as e: - raise ValueError( - f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}" - ) + raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}") compiled_prompt_client["completed_messages"] = messages return compiled_prompt_client diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 95bcd4d718..c0ca6835ee 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -1,4 +1,5 @@ from typing import Optional, Tuple +from urllib.parse import urlparse import litellm from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH @@ -8,6 +9,43 @@ from litellm.secret_managers.main import get_secret, get_secret_str from ..types.router import LiteLLM_Params +def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool: + """ + Match a registered openai-compatible endpoint against a caller-supplied + ``api_base`` using parsed-URL semantics, not unanchored substring search. + + Both inputs may be a bare hostname (``api.perplexity.ai``), host+path + (``api.deepinfra.com/v1/openai``), or a full URL + (``https://api.cerebras.ai/v1``). Hostnames must match exactly + (case-insensitive); if the registered endpoint has a non-trivial path, + the api_base path must start with it on a segment boundary. + + The naive ``endpoint in api_base`` shape lets a caller pass + ``https://attacker.com/api.groq.com/openai/v1`` to coerce the proxy + into reading the server's GROQ_API_KEY from the environment and + forwarding it to the attacker's host as a Bearer credential. + """ + + def _parse(value: str): + # Ensure urlparse sees a scheme so it populates hostname / path. + normalized = value if "://" in value else f"https://{value}" + return urlparse(normalized) + + parsed_endpoint = _parse(endpoint) + parsed_url = _parse(api_base) + + endpoint_host = (parsed_endpoint.hostname or "").lower() + url_host = (parsed_url.hostname or "").lower() + if not endpoint_host or endpoint_host != url_host: + return False + + endpoint_path = parsed_endpoint.path.rstrip("/") + if not endpoint_path: + return True + url_path = parsed_url.path.rstrip("/") + return url_path == endpoint_path or url_path.startswith(endpoint_path + "/") + + def _is_non_openai_azure_model(model: str) -> bool: try: model_name = model.split("/", 1)[1] @@ -210,7 +248,7 @@ def get_llm_provider( # noqa: PLR0915 # check if api base is a known openai compatible endpoint if api_base: for endpoint in litellm.openai_compatible_endpoints: - if endpoint in api_base: + if _endpoint_matches_api_base(endpoint, api_base): if endpoint == "api.perplexity.ai": custom_llm_provider = "perplexity" dynamic_api_key = get_secret_str("PERPLEXITYAI_API_KEY") @@ -348,6 +386,7 @@ def get_llm_provider( # noqa: PLR0915 or "ft:gpt-3.5-turbo" in model or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o or model in litellm.openai_image_generation_models + or model.startswith("gpt-image") or model in litellm.openai_video_generation_models ): custom_llm_provider = "openai" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 57341472b4..829c1c9ca0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1467,6 +1467,8 @@ class Logging(LiteLLMLoggingBaseClass): LiteLLMRealtimeStreamLoggingObject, OpenAIModerationResponse, "SearchResponse", + dict, + list, ], cache_hit: Optional[bool] = None, litellm_model_name: Optional[str] = None, @@ -1725,12 +1727,18 @@ class Logging(LiteLLMLoggingBaseClass): return if self.model_call_details.get("litellm_params") is None: return - self.model_call_details["litellm_params"].setdefault("metadata", {}) - if self.model_call_details["litellm_params"]["metadata"] is None: - self.model_call_details["litellm_params"]["metadata"] = {} - self.model_call_details["litellm_params"]["metadata"]["hidden_params"] = ( - getattr(logging_result, "_hidden_params", {}) - ) + metadata_hidden_params = hidden_params.copy() + response_cost = self.model_call_details.get("response_cost") + if ( + metadata_hidden_params.get("response_cost") is None + and response_cost is not None + ): + metadata_hidden_params["response_cost"] = response_cost + + litellm_params = self.model_call_details["litellm_params"] + metadata = litellm_params.get("metadata") or {} + litellm_params["metadata"] = metadata + metadata["hidden_params"] = metadata_hidden_params def _process_hidden_params_and_response_cost( self, @@ -1738,6 +1746,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, ): + """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params = getattr(logging_result, "_hidden_params", {}) if hidden_params: if self.model_call_details.get("litellm_params") is not None: @@ -1871,24 +1880,12 @@ class Logging(LiteLLMLoggingBaseClass): ): if self._is_recognized_call_type_for_logging( logging_result=logging_result - ): + ) or isinstance(logging_result, (dict, list)): self._process_hidden_params_and_response_cost( logging_result=logging_result, start_time=start_time, end_time=end_time, ) - elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details["standard_logging_object"] = ( - self._build_standard_logging_payload( - result, start_time, end_time - ) - ) - if ( - standard_logging_payload := self.model_call_details.get( - "standard_logging_object" - ) - ) is not None: - emit_standard_logging_payload(standard_logging_payload) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = ( standard_logging_object @@ -5438,11 +5435,6 @@ def get_standard_logging_object_payload( completion_start_time_float=completion_start_time_float, stream=kwargs.get("stream", False), ) - # clean up litellm hidden params - clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( - hidden_params - ) - # clean up litellm metadata clean_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( metadata=metadata, @@ -5476,6 +5468,18 @@ def get_standard_logging_object_payload( ## Get model cost information ## base_model = _get_base_model_from_metadata(model_call_details=kwargs) custom_pricing = use_custom_pricing_for_model(litellm_params=litellm_params) + raw_response_cost = kwargs.get("response_cost") + response_cost: float = raw_response_cost or 0.0 + + # clean up litellm hidden params + clean_hidden_params = StandardLoggingPayloadSetup.get_hidden_params( + hidden_params + ) + if ( + clean_hidden_params["response_cost"] is None + and raw_response_cost is not None + ): + clean_hidden_params["response_cost"] = response_cost model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, @@ -5484,7 +5488,6 @@ def get_standard_logging_object_payload( init_response_obj=init_response_obj, api_base=litellm_params.get("api_base"), ) - response_cost: float = kwargs.get("response_cost", 0) or 0.0 error_information = StandardLoggingPayloadSetup.get_error_information( original_exception=original_exception, diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 888999504f..59d0465e6d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -982,9 +982,9 @@ class CostCalculatorUtils: image_response=completion_response, ) elif custom_llm_provider == litellm.LlmProviders.OPENAI.value: - # Check if this is a gpt-image model (token-based pricing) + # gpt-image models use token-based pricing. model_lower = model.lower() - if "gpt-image-1" in model_lower: + if "gpt-image" in model_lower: from litellm.llms.openai.image_generation.cost_calculator import ( cost_calculator as openai_gpt_image_cost_calculator, ) @@ -1004,9 +1004,9 @@ class CostCalculatorUtils: optional_params=optional_params, ) elif custom_llm_provider == litellm.LlmProviders.AZURE.value: - # Check if this is a gpt-image model (token-based pricing) + # gpt-image models use token-based pricing. model_lower = model.lower() - if "gpt-image-1" in model_lower: + if "gpt-image" in model_lower: from litellm.llms.openai.image_generation.cost_calculator import ( cost_calculator as openai_gpt_image_cost_calculator, ) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 78378faa26..5fd42fe0d3 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -824,8 +824,6 @@ def convert_to_model_response_object( # noqa: PLR0915 stream=stream, start_time=start_time, end_time=end_time, - hidden_params=hidden_params, - _response_headers=_response_headers, convert_tool_call_to_json_mode=convert_tool_call_to_json_mode, ) raise Exception( diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index c5c150274c..6c749118de 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -221,6 +221,13 @@ class LoggingCallbackManager: headers = callback_config.get("headers") event_types = callback_config.get("event_types") log_format = callback_config.get("log_format") + max_retries = max(0, int(callback_config.get("max_retries", 0) or 0)) + retry_delay_value = callback_config.get("retry_delay") + retry_delay = max( + 0.0, + float(0.0 if retry_delay_value is None else retry_delay_value), + ) + timeout = callback_config.get("timeout") if endpoint is None or headers is None: verbose_logger.warning( @@ -236,6 +243,9 @@ class LoggingCallbackManager: and cached_logger.headers == headers and cached_logger.event_types == event_types and cached_logger.log_format == log_format + and cached_logger.max_retries == max_retries + and cached_logger.retry_delay == retry_delay + and cached_logger.timeout == timeout ): return cached_logger @@ -244,6 +254,9 @@ class LoggingCallbackManager: headers=headers, event_types=event_types, log_format=log_format, + max_retries=max_retries, + retry_delay=retry_delay, + timeout=timeout, ) _generic_api_logger_cache[callback] = new_logger return new_logger diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index fcdf49f291..a9cf151464 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -24,6 +24,6 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: return AzureDallE3ImageGenerationConfig() else: verbose_logger.debug( - f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image-1 model format." + f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image model format." ) return AzureGPTImageGenerationConfig() diff --git a/litellm/llms/azure/image_generation/gpt_transformation.py b/litellm/llms/azure/image_generation/gpt_transformation.py index 1f5f65f693..2d46592e3f 100644 --- a/litellm/llms/azure/image_generation/gpt_transformation.py +++ b/litellm/llms/azure/image_generation/gpt_transformation.py @@ -3,7 +3,7 @@ from litellm.llms.openai.image_generation import GPTImageGenerationConfig class AzureGPTImageGenerationConfig(GPTImageGenerationConfig): """ - Azure gpt-image-1 image generation config + Azure gpt-image image generation config """ pass diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index b62acb6516..d1b93c9e7a 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -95,6 +95,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 5fbf0a4b19..85a9c83826 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -59,6 +59,7 @@ class BaseVectorStoreConfig: api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: pass @@ -70,6 +71,7 @@ class BaseVectorStoreConfig: api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Optional async version of transform_search_vector_store_request. @@ -84,6 +86,7 @@ class BaseVectorStoreConfig: api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, + extra_body=extra_body, ) @abstractmethod diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index a27153365d..61a7d4c08d 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1942,8 +1942,8 @@ class AmazonConverseConfig(BaseConfig): completion_response = ConverseResponseBlock(**response.json()) # type: ignore except Exception as e: raise BedrockError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - response.text, str(e) + message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( + str(e) ), status_code=422, ) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 4da0a7c779..f028503c6a 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,14 +1,16 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from copy import deepcopy +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import urlparse import httpx +from litellm._logging import verbose_logger from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, - BedrockKBResponse, BedrockKBRetrievalConfiguration, + BedrockKBResponse, BedrockKBRetrievalQuery, ) from litellm.types.router import GenericLiteLLMParams @@ -202,6 +204,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: if isinstance(query, list): query = " ".join(query) @@ -213,24 +216,46 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): } retrieval_config: Dict[str, Any] = {} + + if isinstance(extra_body, dict): + retrieval_config = deepcopy( + extra_body.get("retrievalConfiguration") + or extra_body.get("retrieval_configuration") + or {} + ) max_results = vector_store_search_optional_params.get("max_num_results") if max_results is not None: + existing_number_of_results = retrieval_config.get( + "vectorSearchConfiguration", {} + ).get("numberOfResults") + if ( + existing_number_of_results is not None + and existing_number_of_results != max_results + ): + verbose_logger.debug( + "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s", + existing_number_of_results, + max_results, + ) retrieval_config.setdefault("vectorSearchConfiguration", {})[ "numberOfResults" ] = max_results filters = vector_store_search_optional_params.get("filters") if filters is not None: + existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get( + "filter" + ) + if existing_filter is not None and existing_filter != filters: + verbose_logger.debug( + "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" + ) retrieval_config.setdefault("vectorSearchConfiguration", {})[ "filter" ] = filters if retrieval_config: - # Create a properly typed retrieval configuration - typed_retrieval_config: BedrockKBRetrievalConfiguration = {} - if "vectorSearchConfiguration" in retrieval_config: - typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config[ - "vectorSearchConfiguration" - ] - request_body["retrievalConfiguration"] = typed_retrieval_config + request_body["retrievalConfiguration"] = cast( + BedrockKBRetrievalConfiguration, retrieval_config + ) litellm_logging_obj.model_call_details["query"] = query return url, request_body diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 03d2af7232..dd955c23d2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1,5 +1,7 @@ import asyncio +import inspect import os +import socket import ssl import sys import time @@ -29,6 +31,10 @@ from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, AIOHTTP_NEEDS_CLEANUP_CLOSED, + AIOHTTP_SO_KEEPALIVE, + AIOHTTP_TCP_KEEPCNT, + AIOHTTP_TCP_KEEPIDLE, + AIOHTTP_TCP_KEEPINTVL, AIOHTTP_TTL_DNS_CACHE, COMPLETION_HTTP_FALLBACK_SECONDS, DEFAULT_SSL_CIPHERS, @@ -54,6 +60,57 @@ except Exception: version = "0.0.0" +# aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older +# versions don't — detect once and skip the keep-alive wiring there. +# https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector +_AIOHTTP_SUPPORTS_SOCKET_FACTORY = ( + "socket_factory" in inspect.signature(TCPConnector.__init__).parameters +) + + +def _build_aiohttp_keepalive_socket_factory() -> ( + Optional[Callable[[Tuple[Any, ...]], socket.socket]] +): + """ + Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets. + + Why: by default, aiohttp creates sockets without SO_KEEPALIVE, so the kernel + sends nothing during a long idle TCP connection. NAT/LB hops (e.g. AWS NAT + Gateway, 350s idle timeout) reap the flow well before slow provider + responses (OpenAI/Azure: up to 600s) arrive. Enabling SO_KEEPALIVE makes + the kernel emit TCP probes that reset the NAT idle timer. + + Returns None when AIOHTTP_SO_KEEPALIVE is disabled or aiohttp is too old. + """ + if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY: + return None + + def factory(addr_info: Tuple[Any, ...]) -> socket.socket: + family, type_, proto = addr_info[0], addr_info[1], addr_info[2] + sock = socket.socket(family=family, type=type_, proto=proto) + sock.setblocking(False) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + # Linux: TCP_KEEPIDLE is idle-before-first-probe. + # macOS/Darwin: TCP_KEEPALIVE is the equivalent. + if hasattr(socket, "TCP_KEEPIDLE"): + sock.setsockopt( + socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE + ) + elif hasattr(socket, "TCP_KEEPALIVE"): + sock.setsockopt( + socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE + ) + if hasattr(socket, "TCP_KEEPINTVL"): + sock.setsockopt( + socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL + ) + if hasattr(socket, "TCP_KEEPCNT"): + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT) + return sock + + return factory + + def get_default_headers() -> dict: """ Get default headers for HTTP requests. @@ -935,6 +992,11 @@ class AsyncHTTPHandler: transport_connector_kwargs["limit_per_host"] = ( AIOHTTP_CONNECTOR_LIMIT_PER_HOST ) + # Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to + # accept socket_factory — version detection lives inside the builder. + socket_factory = _build_aiohttp_keepalive_socket_factory() + if socket_factory is not None: + transport_connector_kwargs["socket_factory"] = socket_factory return LiteLLMAiohttpTransport( client=lambda: ClientSession( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b9ada079f6..a34b73b531 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -155,6 +155,30 @@ else: LiteLLMLoggingObj = Any +def _google_genai_streaming_hidden_params( + *, + api_base: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + response_headers: httpx.Headers, +) -> Dict[str, Any]: + """Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params).""" + from litellm.litellm_core_utils.core_helpers import process_response_headers + + _model_info: Dict[str, Any] = dict( + getattr(litellm_params, "model_info", None) or {} + ) + _raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or "" + _model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id) + return { + "model_id": _model_id, + "api_base": api_base, + "cache_key": "", + "response_cost": "", + "additional_headers": process_response_headers(response_headers), + } + + class BaseLLMHTTPHandler: async def _make_common_async_call( self, @@ -8585,6 +8609,7 @@ class BaseLLMHTTPHandler: api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) else: ( @@ -8597,6 +8622,7 @@ class BaseLLMHTTPHandler: api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) @@ -8697,6 +8723,7 @@ class BaseLLMHTTPHandler: api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) all_optional_params: Dict[str, Any] = dict(litellm_params) @@ -10425,6 +10452,12 @@ class BaseLLMHTTPHandler: litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, request_body=data, + hidden_params=_google_genai_streaming_hidden_params( + api_base=api_base, + litellm_params=litellm_params, + logging_obj=logging_obj, + response_headers=response.headers, + ), ) else: response = sync_httpx_client.post( @@ -10534,6 +10567,12 @@ class BaseLLMHTTPHandler: litellm_metadata=litellm_metadata or {}, custom_llm_provider=custom_llm_provider, request_body=data, + hidden_params=_google_genai_streaming_hidden_params( + api_base=api_base, + litellm_params=litellm_params, + logging_obj=logging_obj, + response_headers=response.headers, + ), ) else: response = await async_httpx_client.post( diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index e6e8369643..35d83bd2ad 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -118,6 +118,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Transform search request to Gemini's generateContent format. diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index fcf5d14db7..af78cd8dbd 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -130,6 +130,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py index 8bca75172f..d009a085fa 100644 --- a/litellm/llms/openai/image_generation/cost_calculator.py +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -1,5 +1,5 @@ """ -Cost calculator for OpenAI image generation models (gpt-image-1, gpt-image-1-mini) +Cost calculator for OpenAI image generation models (gpt-image family) These models use token-based pricing instead of pixel-based pricing like DALL-E. """ @@ -17,13 +17,13 @@ def cost_calculator( custom_llm_provider: Optional[str] = None, ) -> float: """ - Calculate cost for OpenAI gpt-image-1 and gpt-image-1-mini models. + Calculate cost for OpenAI gpt-image models. Uses the same usage format as Responses API, so we reuse the helper to transform to chat completion format and use generic_cost_per_token. Args: - model: The model name (e.g., "gpt-image-1", "gpt-image-1-mini") + model: The model name (e.g., "gpt-image-1", "gpt-image-2") image_response: The ImageResponse containing usage data custom_llm_provider: Optional provider name diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index c106d7f17b..68f799e574 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -15,7 +15,7 @@ if TYPE_CHECKING: class GPTImageGenerationConfig(BaseImageGenerationConfig): """ - OpenAI gpt-image-1 image generation config + OpenAI gpt-image image generation config """ def get_supported_openai_params( diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index c763ed1c8d..2c11d13748 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -106,6 +106,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 275c352b39..5dd1247001 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -101,5 +101,10 @@ "param_mappings": { "max_completion_tokens": "max_tokens" } + }, + "aihubmix": { + "base_url": "https://aihubmix.com/v1", + "api_key_env": "AIHUBMIX_API_KEY", + "api_base_env": "AIHUBMIX_API_BASE" } } diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index ba87a8f2b0..7b22edd867 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -80,6 +80,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( @@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, + extra_body=extra_body, ) return url, request_body diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ed5397eef0..3238d3e9c1 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -102,6 +102,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError( diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 19b5976986..8270e99d45 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -79,6 +79,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Sync version - generates embedding synchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name @@ -140,6 +141,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Async version - generates embedding asynchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 028e02eb0c..7436bfef58 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -4,6 +4,7 @@ from typing import Any, Coroutine, Dict, Optional, Union import httpx import litellm +from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -224,8 +225,14 @@ class VertexAIBatchPrediction(VertexLLM): }, ) - response = sync_handler.get( - url=api_base, + # ``api_base`` here can come from caller-supplied request kwargs + # (clientside override). Wrap the fetch in ``safe_get`` so DNS + # rebind / private / cloud-metadata targets are rejected; the + # proxy auth gate already blocks malicious clientside ``api_base`` + # at the boundary — this is defense-in-depth for SDK callers. + response = safe_get( + sync_handler, + api_base, headers=headers, ) @@ -270,8 +277,13 @@ class VertexAIBatchPrediction(VertexLLM): }, ) - response = await client.get( - url=api_base, + # Mirror the sync path: ``api_base`` may come from caller-supplied + # request kwargs, so wrap the fetch in ``async_safe_get`` to reject + # DNS-rebind / private / cloud-metadata targets. Defense-in-depth + # behind the proxy auth gate's clientside ``api_base`` check. + response = await async_safe_get( + client, + api_base, headers=headers, ) if response.status_code != 200: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index ccd4d4f293..b4bfde5f54 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -597,7 +597,14 @@ def process_items(schema, depth=0): f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting." ) if isinstance(schema, dict): - if "items" in schema and schema["items"] == {}: + # Vertex requires `items` whenever `type == "array"` (even inside anyOf). + # Normalize: empty `items: {}` and missing-items both become {"type": "object"}. + type_val = schema.get("type") + if ( + isinstance(type_val, str) + and type_val.lower() == "array" + and ("items" not in schema or schema.get("items") == {}) + ): schema["items"] = {"type": "object"} for key, value in schema.items(): if isinstance(value, dict): @@ -710,14 +717,10 @@ def convert_anyof_null_to_nullable(schema, depth=0): if contains_null: # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python + # Empty `items: {}` on array branches is left in place; downstream + # process_items() converts it to {"type": "object"}, which Vertex + # requires whenever type == "array" (even inside anyOf). for atype in anyof: - # Remove items field if type is array and items is empty - if ( - atype.get("type") == "array" - and "items" in atype - and not atype["items"] - ): - atype.pop("items") atype["nullable"] = True properties = schema.get("properties", None) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a90e05919f..474ddb402a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -2395,8 +2395,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_response = GenerateContentResponseBody(**raw_response.json()) # type: ignore except Exception as e: raise VertexAIError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - raw_response.text, str(e) + message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( + str(e) ), status_code=422, headers=raw_response.headers, @@ -2530,8 +2530,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): except Exception as e: raise VertexAIError( - message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( - completion_response, str(e) + message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( + str(e) ), status_code=422, headers=raw_response.headers, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 4baa5774c4..d31e1f6c8f 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -100,6 +100,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 179bd7aeff..6cb7a86bea 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -107,6 +107,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5cccd5f00a..e4268fac81 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4735,17 +4735,17 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -4774,17 +4774,17 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -5103,6 +5103,38 @@ "/v1/images/edits" ] }, + "azure/gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/low/1024-x-1024/gpt-image-1-mini": { "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", @@ -19083,6 +19115,38 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "low/1024-x-1024/gpt-image-1.5": { "input_cost_per_image": 0.009, "litellm_provider": "openai", @@ -19898,21 +19962,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -19941,21 +20005,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 251f271903..25c47aacf0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -50,8 +50,11 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mc from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, add_server_prefix_to_name, + compute_short_server_prefix, get_server_prefix, + is_short_mcp_tool_prefix_enabled, is_tool_name_prefixed, + iter_known_server_prefixes, merge_mcp_headers, normalize_server_name, split_server_prefix_from_name, @@ -106,6 +109,12 @@ if not _separator_probe.is_valid: SEP_986_URL, ) +_AZURE_ENTRA_HOSTS = { + "login.microsoftonline.com", # Global + "login.microsoftonline.us", # US Government + "login.chinacloudapi.cn", # China +} + def _warn_on_server_name_fields( *, @@ -364,6 +373,7 @@ class MCPServerManager: aws_session_name=server_config.get("aws_session_name", None), instructions=server_config.get("instructions", None), ) + self._assign_unique_short_prefix(new_server) self.config_mcp_servers[server_id] = new_server # Check if this is an OpenAPI-based server @@ -726,6 +736,7 @@ class MCPServerManager: try: if mcp_server.server_id not in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) + self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) verbose_logger.debug(f"Added MCP Server: {new_server.name}") @@ -738,6 +749,12 @@ class MCPServerManager: try: if mcp_server.server_id in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) + # Carry the previously-resolved short prefix across so the + # tool names stay stable for clients holding cached lists. + existing_prefix = self.registry[mcp_server.server_id].short_prefix + if existing_prefix and not new_server.short_prefix: + new_server.short_prefix = existing_prefix + self._assign_unique_short_prefix(new_server) self.registry[mcp_server.server_id] = new_server await self._maybe_register_openapi_tools(new_server) verbose_logger.debug(f"Updated MCP Server: {new_server.name}") @@ -1236,7 +1253,11 @@ class MCPServerManager: ## HANDLE OPENAPI TOOLS if server.spec_path: - _tools = global_mcp_tool_registry.list_tools(tool_prefix=server.name) + # OpenAPI tools were stored in the registry under the prefix + # active at registration time — fetch by that same prefix. + _tools = global_mcp_tool_registry.list_tools( + tool_prefix=get_server_prefix(server) + ) tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type( _tools ) @@ -1488,11 +1509,28 @@ class MCPServerManager: client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) response = await client.get(server_url) response.raise_for_status() - verbose_logger.warning( - "MCP OAuth discovery unexpectedly succeeded for %s; server did not challenge", - server_url, + ( + authorization_servers, + resource_scopes, + ) = await self._attempt_well_known_discovery(server_url) + metadata = await self._fetch_authorization_server_metadata( + authorization_servers ) - raise RuntimeError("OAuth discovery must not succeed without a challenge") + if ( + metadata is None + and not resource_scopes + and authorization_servers + and response.status_code == 200 + ): + verbose_logger.warning( + "MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.", + server_url, + ) + if metadata is None and resource_scopes: + return MCPOAuthMetadata(scopes=resource_scopes) + if metadata is not None and resource_scopes: + metadata.scopes = resource_scopes + return metadata except HTTPStatusError as exc: verbose_logger.debug( "MCP OAuth discovery for %s received status error: %s", @@ -1510,8 +1548,8 @@ class MCPServerManager: header_value ) - authorization_servers: List[str] = [] - resource_scopes: Optional[List[str]] = None + authorization_servers = [] + resource_scopes = None if resource_metadata_url: ( authorization_servers, @@ -1674,6 +1712,9 @@ class MCPServerManager: f"{base}/.well-known/oauth-authorization-server/{path}" ) candidate_urls.append(f"{base}/.well-known/openid-configuration/{path}") + candidate_urls.append( + f"{issuer_url.rstrip('/')}/.well-known/openid-configuration" + ) candidate_urls.append(f"{base}/.well-known/oauth-authorization-server") candidate_urls.append(f"{base}/.well-known/openid-configuration") candidate_urls.append(issuer_url.rstrip("/")) @@ -1713,7 +1754,28 @@ class MCPServerManager: ): return metadata - return None + return self._build_azure_authorization_server_metadata(parsed) + + @staticmethod + def _build_azure_authorization_server_metadata( + parsed_issuer_url: Any, + ) -> Optional[MCPOAuthMetadata]: + path_parts = [ + part for part in (parsed_issuer_url.path or "").split("/") if part + ] + if ( + parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS + or len(path_parts) != 2 + or path_parts[1] != "v2.0" + ): + return None + + tenant = path_parts[0] + base = f"{parsed_issuer_url.scheme}://{parsed_issuer_url.netloc}/{tenant}" + return MCPOAuthMetadata( + authorization_url=f"{base}/oauth2/v2.0/authorize", + token_url=f"{base}/oauth2/v2.0/token", + ) @staticmethod def _decrypt_credential_field( @@ -1810,6 +1872,63 @@ class MCPServerManager: verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}") return [] + _SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024 + + def _assign_unique_short_prefix(self, server: MCPServer) -> None: + """Resolve and cache a collision-free short tool prefix on ``server``. + + Called at registration time for every MCP server entering the + registry. Mutates ``server.short_prefix`` in place. No-ops when + ``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` is disabled, when the server + has no ``server_id`` (synthetic temp-server objects), or when a + prefix is already cached. + + Collision strategy: take the natural hash; if it's already used by + a *different* server in the combined registry, rehash with an + incrementing attempt counter until we find an unused slot. The + attempt counter is folded into the hash so the resulting prefix is + still deterministic for a given (server_id, set-of-other-server-ids) + pair within one process. + """ + if not is_short_mcp_tool_prefix_enabled(): + return + if server.short_prefix: + return + if not server.server_id: + return + + used: Dict[str, str] = {} + for other in self.get_registry().values(): + if other.server_id == server.server_id: + continue + if other.short_prefix: + used[other.short_prefix] = other.server_id + + for attempt in range(self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS): + candidate = compute_short_server_prefix(server.server_id, attempt=attempt) + if candidate not in used: + server.short_prefix = candidate + if attempt > 0: + verbose_logger.info( + "MCP short-prefix collision resolved for server %s: " + "natural hash collided with %s, using rehashed prefix " + "%s (attempt=%d).", + server.server_id, + used.get( + compute_short_server_prefix(server.server_id, attempt=0), + "", + ), + candidate, + attempt, + ) + return + + raise RuntimeError( + f"Unable to assign a unique short MCP tool prefix for server " + f"{server.server_id} after {self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS} " + "attempts; the 3-character prefix space is too crowded." + ) + def _create_prefixed_tools( self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True ) -> List[MCPTool]: @@ -1838,9 +1957,13 @@ class MCPServerManager: tool_copy.name = name_to_use prefixed_tools.append(tool_copy) - # Update tool to server mapping for resolution (support both forms) + # Register every known prefix form (alias, server_name, server_id, + # short ID) so call_tool can resolve regardless of which form a + # caller / cached client is using. self.tool_name_to_mcp_server_name_mapping[original_name] = prefix - self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix + for known_prefix in iter_known_server_prefixes(server): + qualified = add_server_prefix_to_name(original_name, known_prefix) + self.tool_name_to_mcp_server_name_mapping[qualified] = prefix verbose_logger.info( f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}" @@ -2601,37 +2724,43 @@ class MCPServerManager: Returns: MCPServer if found, None otherwise """ + registry_servers = list(self.get_registry().values()) + + # Build prefix → server lookup covering every known form a tool name + # may take (alias / server_name / server_id / short ID). This is what + # makes the short-prefix mode work without breaking historical names. + prefix_to_server: Dict[str, MCPServer] = {} + for server in registry_servers: + for known_prefix in iter_known_server_prefixes(server): + normalised = normalize_server_name(known_prefix) + prefix_to_server.setdefault(normalised, server) + # First try with the original tool name if tool_name in self.tool_name_to_mcp_server_name_mapping: server_name = self.tool_name_to_mcp_server_name_mapping[tool_name] - for server in self.get_registry().values(): - if normalize_server_name(server.name) == normalize_server_name( - server_name - ): + normalised_lookup = normalize_server_name(server_name) + if normalised_lookup in prefix_to_server: + return prefix_to_server[normalised_lookup] + for server in registry_servers: + if normalize_server_name(server.name) == normalised_lookup: return server - # If not found and tool name is prefixed, try extracting server name from prefix - known_prefixes = { - normalize_server_name(get_server_prefix(s)) - for s in self.get_registry().values() - if get_server_prefix(s) - } - if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes): + # If not found and tool name is prefixed, extract the prefix and + # match against any known form. + if is_tool_name_prefixed( + tool_name, known_server_prefixes=set(prefix_to_server.keys()) + ): ( original_tool_name, server_name_from_prefix, ) = split_server_prefix_from_name(tool_name) - if original_tool_name in self.tool_name_to_mcp_server_name_mapping: - for server in self.get_registry().values(): - if server.server_name is None: - if normalize_server_name(server.name) == normalize_server_name( - server_name_from_prefix - ): - return server - elif normalize_server_name( - server.server_name - ) == normalize_server_name(server_name_from_prefix): - return server + normalised_prefix = normalize_server_name(server_name_from_prefix) + matched_server = prefix_to_server.get(normalised_prefix) + if matched_server is not None and ( + original_tool_name in self.tool_name_to_mcp_server_name_mapping + or tool_name in self.tool_name_to_mcp_server_name_mapping + ): + return matched_server return None @@ -2666,6 +2795,9 @@ class MCPServerManager: previous_registry = self.registry new_registry: Dict[str, MCPServer] = {} + # Stage one: build every server. Stage two assigns short prefixes + # against the *full* set so dedup is deterministic regardless of + # iteration order. for server in db_mcp_servers: existing_server = previous_registry.get(server.server_id) @@ -2689,10 +2821,21 @@ class MCPServerManager: f"Building server from DB: {server.server_id} ({server.server_name})" ) new_server = await self.build_mcp_server_from_table(server) + # Carry the cached short_prefix from the previous registry entry + # (if any) so the prefix is stable across reloads. + if existing_server is not None and existing_server.short_prefix: + new_server.short_prefix = existing_server.short_prefix new_registry[server.server_id] = new_server - await self._maybe_register_openapi_tools(new_server) + # Swap in the new registry first so _assign_unique_short_prefix + # sees the complete set when checking for collisions. self.registry = new_registry + for new_server in new_registry.values(): + self._assign_unique_short_prefix(new_server) + # Register OpenAPI tools *after* the final short prefix is assigned + # so the tools are stored in the global registry under the same + # prefix that lookups will use. + await self._maybe_register_openapi_tools(new_server) verbose_logger.debug( "MCP registry refreshed (%s servers in registry)", len(new_registry) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c2e998f01e..ae6055217b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -49,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_VERSION, add_server_prefix_to_name, get_server_prefix, + iter_known_server_prefixes, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.ip_address_utils import IPAddressUtils @@ -711,13 +712,7 @@ if MCP_AVAILABLE: for server in allowed_mcp_servers: if server: match_list = [ - s.lower() - for s in [ - server.alias, - server.server_name, - server.server_id, - ] - if s is not None + s.lower() for s in iter_known_server_prefixes(server) if s ] if server_or_group.lower() in match_list: @@ -2031,11 +2026,13 @@ if MCP_AVAILABLE: # Remove prefix from tool name for logging and processing original_tool_name, server_name = split_server_prefix_from_name(name) - # If tool name is unprefixed, resolve its server so we can enforce permissions - if not server_name: - mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) - if mcp_server: - server_name = mcp_server.name + # Resolve the actual MCP server up-front so the permission check uses + # the canonical server.name even when the tool name is prefixed with a + # short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the + # server's display name directly. + mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server is not None: + server_name = mcp_server.name # Only enforce server-level permissions when we can resolve a server if server_name: diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 146ba10bb7..df5705c342 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -2,10 +2,11 @@ MCP Server Utilities """ -from typing import Any, Dict, Mapping, Optional, Tuple +from typing import Any, Dict, Iterator, Mapping, Optional, Tuple -import os +import hashlib import importlib +import os # Constants LITELLM_MCP_SERVER_NAME = "litellm-mcp-server" @@ -14,6 +15,89 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM" MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-") MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}" +# --------------------------------------------------------------------------- +# Short-ID tool prefix (opt-in) +# --------------------------------------------------------------------------- +# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP +# tool / prompt / resource / resource-template names switches from the +# (potentially long) human-readable server name to a deterministic three +# character ID derived from the server's ``server_id``. +# +# Why three characters? +# * The first character is restricted to 52 alphabetic characters +# ([A-Za-z]) and the remaining two characters use the full base62 +# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts +# with a digit so it remains a valid identifier for every model API +# (some providers historically required a leading alphabetic char). +# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local +# tool name happening to begin with the exact prefix LiteLLM assigned +# to a given MCP server is negligible in practice. +# * The IDs are short enough that prefixed tool names stay well under +# the 60-character upper bound enforced by some model APIs (Anthropic +# etc.) even for long upstream tool names. +# * The mapping is deterministic (SHA-256 of ``server_id`` → three +# characters drawn from the alphabets above), so the prefix is stable +# across processes, workers and restarts without any persistence +# layer. Two servers with different ``server_id`` values can in +# principle hash to the same three chars; that natural-hash collision +# IS a routing-correctness issue (the second registrant would otherwise +# have its tools misrouted to the first), so registration goes through +# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with +# a deterministic attempt counter until it finds an unused prefix and +# caches the result on ``MCPServer.short_prefix``. A collision is +# logged at INFO when it happens. +# +# This flag is intentionally opt-in for the first release so customers can +# migrate. It will become the default in a future release. +SHORT_MCP_TOOL_PREFIX_LENGTH = 3 +_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" +# Subset of _BASE62_ALPHABET used for the *first* character only, to +# guarantee the prefix never starts with a digit. +_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + + +def is_short_mcp_tool_prefix_enabled() -> bool: + """Return True when the short-ID tool prefix mode is enabled. + + Read at call time (not import time) so tests and runtime config changes + take effect without reimporting the module. + """ + raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "") + return raw.strip().lower() in ("1", "true", "yes", "on") + + +def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str: + """Derive the deterministic three-character prefix for a server. + + Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight + bytes into a fixed-length string whose first character is drawn from + ``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit) + and whose remaining characters are drawn from the full base62 + alphabet. Pass ``attempt > 0`` to rehash to a different prefix when + the natural hash collides with a prefix already assigned to another + server (see ``MCPServerManager._assign_unique_short_prefix``). An + empty ``server_id`` raises ``ValueError`` — short prefixes require a + stable identifier to be deterministic. + """ + if not server_id: + raise ValueError("compute_short_server_prefix requires a non-empty server_id") + + seed = server_id if attempt == 0 else f"{server_id}#{attempt}" + digest = hashlib.sha256(seed.encode("utf-8")).digest() + value = int.from_bytes(digest[:8], "big") + + # Build chars from least-significant to most-significant; we reverse + # at the end so the first emitted char comes from the high-order + # bits of the digest (which is the position we constrain to be + # alphabetic). + chars = [] + for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH): + is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1 + alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET + value, idx = divmod(value, len(alphabet)) + chars.append(alphabet[idx]) + return "".join(reversed(chars)) + def is_mcp_available() -> bool: """ @@ -82,7 +166,25 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str: def get_server_prefix(server: Any) -> str: - """Return the prefix for a server: alias if present, else server_name, else server_id""" + """Return the prefix for a server. + + When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``) + a three-character base62 ID is returned. We prefer the cached + ``server.short_prefix`` value when set — that field is populated at + registration time by ``MCPServerManager._assign_unique_short_prefix`` + and resolves natural-hash collisions deterministically — and only fall + back to the natural hash for ad-hoc / temp-server objects without a + cached value. In default mode the historical behaviour is preserved: + alias if present, else server_name, else server_id. + """ + if is_short_mcp_tool_prefix_enabled(): + cached = getattr(server, "short_prefix", None) + if cached: + return cached + server_id = getattr(server, "server_id", None) + if server_id: + return compute_short_server_prefix(server_id) + if hasattr(server, "alias") and server.alias: return server.alias if hasattr(server, "server_name") and server.server_name: @@ -92,6 +194,36 @@ def get_server_prefix(server: Any) -> str: return "" +def iter_known_server_prefixes(server: Any) -> Iterator[str]: + """Yield every prefix form that may appear in tool names for ``server``. + + Always includes the *current* prefix returned by ``get_server_prefix``. + Additionally yields the historical (alias / server_name / server_id) and + short-ID forms so the routing layer can resolve tool names regardless of + which prefix mode was active when the client first observed them. + """ + seen = set() + + def _emit(value: Optional[str]) -> Iterator[str]: + if value and value not in seen: + seen.add(value) + yield value + + yield from _emit(get_server_prefix(server)) + yield from _emit(getattr(server, "short_prefix", None)) + + server_id = getattr(server, "server_id", None) + if server_id: + try: + yield from _emit(compute_short_server_prefix(server_id)) + except ValueError: + pass + + yield from _emit(getattr(server, "alias", None)) + yield from _emit(getattr(server, "server_name", None)) + yield from _emit(server_id) + + def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: """Return the unprefixed name plus the server name used as prefix.""" if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name: diff --git a/litellm/proxy/_experimental/out/chat.html b/litellm/proxy/_experimental/out/chat.html deleted file mode 100644 index 6ca94ee94f..0000000000 --- a/litellm/proxy/_experimental/out/chat.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_experimental/out/chat.txt b/litellm/proxy/_experimental/out/chat.txt deleted file mode 100644 index 8ee24df074..0000000000 --- a/litellm/proxy/_experimental/out/chat.txt +++ /dev/null @@ -1,22 +0,0 @@ -1:"$Sreact.fragment" -2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"] -7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"] -a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"] -b:"$Sreact.suspense" -d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"] -f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"] -11:I[168027,[],"default"] -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}] -0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true} -8:{} -9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params" -e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]] -12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"] -c:null -10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]] diff --git a/litellm/proxy/_experimental/out/chat/__next._full.txt b/litellm/proxy/_experimental/out/chat/__next._full.txt deleted file mode 100644 index 8ee24df074..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._full.txt +++ /dev/null @@ -1,22 +0,0 @@ -1:"$Sreact.fragment" -2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"] -7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"] -a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"] -b:"$Sreact.suspense" -d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"] -f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"] -11:I[168027,[],"default"] -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}] -0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true} -8:{} -9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params" -e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]] -12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"] -c:null -10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]] diff --git a/litellm/proxy/_experimental/out/chat/__next._head.txt b/litellm/proxy/_experimental/out/chat/__next._head.txt deleted file mode 100644 index 635c8e398b..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._head.txt +++ /dev/null @@ -1,6 +0,0 @@ -1:"$Sreact.fragment" -2:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"] -3:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"] -4:"$Sreact.suspense" -5:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","h",{"children":[null,["$","$L2",null,{"children":[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]}],["$","div",null,{"hidden":true,"children":["$","$L3",null,{"children":["$","$4",null,{"name":"Next.Metadata","children":[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L5","4",{}]]}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],"loading":null,"isPartial":false} diff --git a/litellm/proxy/_experimental/out/chat/__next._index.txt b/litellm/proxy/_experimental/out/chat/__next._index.txt deleted file mode 100644 index d026a48036..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._index.txt +++ /dev/null @@ -1,8 +0,0 @@ -1:"$Sreact.fragment" -2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"] -4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","template":["$","$L5",null,{}],"notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]]}]}]}]}]}]]}],"loading":null,"isPartial":false} diff --git a/litellm/proxy/_experimental/out/chat/__next._tree.txt b/litellm/proxy/_experimental/out/chat/__next._tree.txt deleted file mode 100644 index afb8782562..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next._tree.txt +++ /dev/null @@ -1,4 +0,0 @@ -:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"] -:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"] -:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","tree":{"name":"","paramType":null,"paramKey":"","hasRuntimePrefetch":false,"slots":{"children":{"name":"chat","paramType":null,"paramKey":"chat","hasRuntimePrefetch":false,"slots":{"children":{"name":"__PAGE__","paramType":null,"paramKey":"__PAGE__","hasRuntimePrefetch":false,"slots":null,"isRootLayout":false}},"isRootLayout":false}},"isRootLayout":true},"staleTime":300} diff --git a/litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt b/litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt deleted file mode 100644 index 5d19f68046..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next.chat.__PAGE__.txt +++ /dev/null @@ -1,9 +0,0 @@ -1:"$Sreact.fragment" -2:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"] -3:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"] -6:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"] -7:"$Sreact.suspense" -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[["$","$L2",null,{"Component":"$3","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@4","$@5"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true}]],["$","$L6",null,{"children":["$","$7",null,{"name":"Next.MetadataOutlet","children":"$@8"}]}]]}],"loading":null,"isPartial":false} -4:{} -5:"$0:rsc:props:children:0:props:serverProvidedParams:params" -8:null diff --git a/litellm/proxy/_experimental/out/chat/__next.chat.txt b/litellm/proxy/_experimental/out/chat/__next.chat.txt deleted file mode 100644 index 12f813af35..0000000000 --- a/litellm/proxy/_experimental/out/chat/__next.chat.txt +++ /dev/null @@ -1,4 +0,0 @@ -1:"$Sreact.fragment" -2:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -3:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"] -0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[null,["$","$L2",null,{"parallelRouterKey":"children","template":["$","$L3",null,{}]}]]}],"loading":null,"isPartial":false} diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 448c975d12..91c8f2dd7c 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -6,9 +6,11 @@ from typing import Any, List, Optional, Tuple from fastapi import HTTPException, Request, status +import litellm from litellm import Router, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS +from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import * from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS @@ -53,6 +55,12 @@ def _check_valid_ip( def check_complete_credentials(request_body: dict) -> bool: """ if 'api_base' in request body. Check if complete credentials given. Prevent malicious attacks. + + Supplying an ``api_key`` is necessary but not sufficient: even with + credentials supplied, an ``api_base`` / ``base_url`` that resolves to a + private/internal/cloud-metadata address would still allow the proxy to + be used as an SSRF pivot. Validate any URL fields here so the gate + can't be bypassed with ``api_key=anything`` plus a malicious target. """ given_model: Optional[str] = None @@ -70,10 +78,27 @@ def check_complete_credentials(request_body: dict) -> bool: return False api_key_value = request_body.get("api_key") - if api_key_value and isinstance(api_key_value, str) and api_key_value.strip(): - return True + if not (api_key_value and isinstance(api_key_value, str) and api_key_value.strip()): + return False - return False + # ``validate_url`` itself doesn't consult the toggle; ``safe_get`` / + # ``async_safe_get`` do. Mirror that here so admins who explicitly + # disabled URL validation (e.g. for an internal Ollama endpoint they + # accept the SSRF risk for) aren't blocked at the proxy boundary. + if getattr(litellm, "user_url_validation", False): + for url_field in ("api_base", "base_url"): + url_value = request_body.get(url_field) + if not url_value or not isinstance(url_value, str): + continue + try: + validate_url(url_value) + except SSRFError as e: + raise ValueError( + f"Rejected request: client-side {url_field}={url_value!r} " + f"is rejected by the SSRF guard ({e})." + ) + + return True def check_regex_or_str_match(request_body_value: Any, regex_str: str) -> bool: @@ -159,15 +184,42 @@ def is_request_body_safe( "aws_web_identity_token", "aws_role_name", "vertex_credentials", + # Endpoint-targeting fields that retarget the outbound request or + # an observability callback. An attacker-controlled value either + # exfiltrates the request payload (incl. messages + admin-set + # tokens) to the attacker's host, or coerces the proxy into + # authenticating against the attacker's host with admin secrets. + "aws_bedrock_runtime_endpoint", + "langsmith_base_url", + "langfuse_host", + "posthog_host", + "braintrust_host", + "slack_webhook_url", + # Provider-specific endpoint overrides that flow into the outbound + # request via ``optional_params``. Same threat as ``api_base``: + # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker + # S3; ``sagemaker_base_url`` redirects all SageMaker traffic; + # ``deployment_url`` redirects SAP deployments. + "s3_endpoint_url", + "sagemaker_base_url", + "deployment_url", ] + # The blocklist is enforced unconditionally. Legitimate clientside + # credential / endpoint passthrough goes through one of the two + # explicit admin opt-ins (``general_settings.allow_client_side_credentials`` + # proxy-wide or ``configurable_clientside_auth_params`` per deployment). + # Historically there was a third, *implicit*, *caller-controlled* path: + # ``check_complete_credentials`` returned True when the caller supplied + # any non-empty ``api_key``, which made the entire blocklist a no-op. + # That bypass turned every missing entry on the blocklist into an + # exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp, + # GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l, + # b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now + # has a single, predictable failure mode for missing entries (a 400), + # not a credential leak. for param in banned_params: - if ( - param in request_body - and not check_complete_credentials( # allow client-credentials to be passed to proxy - request_body=request_body - ) - ): + if param in request_body: if general_settings.get("allow_client_side_credentials") is True: return True elif ( @@ -182,7 +234,10 @@ def is_request_body_safe( return True raise ValueError( f"Rejected Request: {param} is not allowed in request body. " - "Enable with `general_settings::allow_client_side_credentials` on proxy config.yaml. " + "Clientside passthrough requires explicit admin opt-in via " + "either `general_settings.allow_client_side_credentials = true` " + "(proxy-wide) or `configurable_clientside_auth_params` on the " + "deployment in your proxy config.yaml. " "Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997", ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c8fab4be4a..76c52f83ee 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -619,6 +619,67 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.error(f"Error setting custom headers: {e}") return {} + @staticmethod + async def build_litellm_proxy_success_headers_from_llm_response( + *, + response: Any, + request_data: dict, + request: Request, + user_api_key_dict: UserAPIKeyAuth, + logging_obj: LiteLLMLoggingObj, + version: Optional[str], + proxy_logging_obj: ProxyLogging, + ) -> Dict[str, str]: + """ + Build LiteLLM proxy response headers for routes that call the LLM directly + (e.g. Google native :generateContent) instead of base_process_llm_request. + """ + if isinstance(response, dict): + hidden_params = response.get("_hidden_params") or {} + else: + hidden_params = getattr(response, "_hidden_params", None) or {} + if not isinstance(hidden_params, dict): + hidden_params = {} + + model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response( + hidden_params, request_data + ) + + cache_key = hidden_params.get("cache_key", None) or "" + api_base = hidden_params.get("api_base", None) or "" + response_cost = hidden_params.get("response_cost", None) or "" + fastest_response_batch_completion = hidden_params.get( + "fastest_response_batch_completion", None + ) + additional_headers = hidden_params.get("additional_headers", {}) or {} + + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=logging_obj.litellm_call_id, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=fastest_response_batch_completion, + request_data=request_data, + hidden_params=hidden_params, + litellm_logging_obj=logging_obj, + **additional_headers, + ) + + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=request_data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) + + return custom_headers + async def common_processing_pre_call_logic( self, request: Request, @@ -875,7 +936,7 @@ class ProxyBaseLLMRequestProcessing: else: verbose_proxy_logger.debug( "Request received by LiteLLM:\n%s", - json.dumps(self.data, indent=4, default=str), + _payload_str, ) async def base_process_llm_request( # noqa: PLR0915 @@ -1511,9 +1572,7 @@ class ProxyBaseLLMRequestProcessing: _response = assembled_response try: from litellm.proxy.proxy_server import llm_router as _global_llm_router - from litellm.proxy.utils import ( - _check_and_merge_model_level_guardrails, - ) + from litellm.proxy.utils import _check_and_merge_model_level_guardrails guardrail_data = _check_and_merge_model_level_guardrails( data=captured_data, llm_router=_global_llm_router @@ -1690,11 +1749,12 @@ class ProxyBaseLLMRequestProcessing: elif isinstance(e, httpx.HTTPStatusError): # Handle httpx.HTTPStatusError - extract actual error from response # This matches the original behavior before the refactor in commit 511d435f6f - error_body = await e.response.aread() + http_status_error: httpx.HTTPStatusError = e + error_body = await http_status_error.response.aread() error_text = error_body.decode("utf-8") raise HTTPException( - status_code=e.response.status_code, + status_code=http_status_error.response.status_code, detail={"error": error_text}, ) error_msg = f"{str(e)}" diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py new file mode 100644 index 0000000000..c25d853312 --- /dev/null +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -0,0 +1,156 @@ +""" +Expired UI session key cleanup manager. + +Deletes expired virtual keys created for LiteLLM dashboard sessions. +""" + +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache +from litellm.constants import ( + EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + UI_SESSION_TOKEN_TEAM_ID, +) +from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth +from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks +from litellm.proxy.management_endpoints.key_management_endpoints import ( + delete_verification_tokens, +) +from litellm.proxy.utils import PrismaClient + + +class ExpiredUISessionKeyCleanupManager: + """ + Cleans up expired UI session keys. + """ + + def __init__( + self, + prisma_client: PrismaClient, + user_api_key_cache: DualCache, + pod_lock_manager=None, + ): + self.prisma_client = prisma_client + self.user_api_key_cache = user_api_key_cache + self.pod_lock_manager = pod_lock_manager + + async def cleanup_expired_keys(self) -> int: + """ + Main entry point for deleting expired UI session keys. + Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments. + """ + lock_acquired = False + try: + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + lock_acquired = ( + await self.pod_lock_manager.acquire_lock( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + or False + ) + if not lock_acquired: + verbose_proxy_logger.debug( + "Expired UI session key cleanup: another pod is already " + "running cleanup or Redis lock acquisition failed - " + "skipping this cycle." + ) + return 0 + + verbose_proxy_logger.info("Starting expired UI session key cleanup...") + + expired_keys = await self._find_expired_ui_session_keys() + if not expired_keys: + verbose_proxy_logger.debug("No expired UI session keys found") + return 0 + + tokens = [key.token for key in expired_keys if key.token is not None] + if not tokens: + return 0 + + system_user = UserAPIKeyAuth.get_litellm_internal_jobs_user_api_key_auth() + response, keys_being_deleted = await delete_verification_tokens( + tokens=tokens, + user_api_key_cache=self.user_api_key_cache, + user_api_key_dict=system_user, + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + ) + await KeyManagementEventHooks.async_key_deleted_hook( + data=KeyRequest(keys=tokens), + keys_being_deleted=keys_being_deleted, + response=response or {}, + user_api_key_dict=system_user, + litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + ) + deleted_count = self._get_deleted_token_count( + tokens=tokens, + response=response, + ) + verbose_proxy_logger.info( + "Deleted %s expired UI session key(s)", deleted_count + ) + return deleted_count + except Exception as e: + if getattr(e, "status_code", None) == 404: + verbose_proxy_logger.debug( + "Expired UI session key cleanup skipped because selected keys " + "were already deleted: %s", + e, + ) + return 0 + verbose_proxy_logger.error(f"Expired UI session key cleanup failed: {e}") + return 0 + finally: + if ( + lock_acquired + and self.pod_lock_manager + and self.pod_lock_manager.redis_cache + ): + await self.pod_lock_manager.release_lock( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + + @staticmethod + def _get_deleted_token_count( + tokens: List[str], + response: Optional[Dict[str, Any]], + ) -> int: + """ + Return the number of tokens actually deleted from the delete helper response. + """ + if response is None: + return len(tokens) + + deleted_keys = response.get("deleted_keys") + if isinstance(deleted_keys, list): + return len(deleted_keys) + if isinstance(deleted_keys, int): + return deleted_keys + if isinstance(deleted_keys, dict): + nested_deleted_keys = deleted_keys.get("deleted_keys") + if isinstance(nested_deleted_keys, list): + return len(nested_deleted_keys) + if isinstance(nested_deleted_keys, int): + return nested_deleted_keys + + failed_tokens = response.get("failed_tokens") or [] + if failed_tokens: + return max(len(tokens) - len(set(failed_tokens)), 0) + + return len(tokens) + + async def _find_expired_ui_session_keys(self) -> List[LiteLLM_VerificationToken]: + """ + Find expired LiteLLM dashboard session keys. + """ + now = datetime.now(timezone.utc) + return await self.prisma_client.db.litellm_verificationtoken.find_many( + where={ + "team_id": UI_SESSION_TOKEN_TEAM_ID, + "expires": {"lt": now}, + }, + take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + ) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index cfa90a48ee..ab9d341aa5 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,5 +1,6 @@ -from typing import Union +from typing import Any, Awaitable, Callable, Optional, Union +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, ProxyErrorTypes, @@ -123,3 +124,138 @@ class PrismaDBExceptionHandler: ): return None raise e + + +# Default fallback timeouts when neither the caller nor the prisma_client +# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`. +# Match the auth path's existing defaults so behavior is uniform across read paths. +_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0 +_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1 + + +def _coerce_timeout(value: Any, fallback: float) -> float: + """Return `value` if it is a real int/float, else `fallback`. Guards + against tests that mock `prisma_client` and leave the timeout slots as + MagicMock instances.""" + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return fallback + + +async def call_with_db_reconnect_retry( + prisma_client: Any, + coro_factory: Callable[[], Awaitable[Any]], + *, + reason: str, + timeout_seconds: Optional[float] = None, + lock_timeout_seconds: Optional[float] = None, +) -> Any: + """Run a Prisma read coroutine with one transport-reconnect-and-retry. + + The canonical "self-heal a transient DB transport blip" wrapper used by + `PrismaClient.get_generic_data` and other read paths. Mirrors the inline + pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we + have a single implementation rather than three drifting copies. + + Behavior: + 1. Await `coro_factory()`. On success, return its value. + 2. On exception, if it is NOT a transport error (per + `is_database_transport_error`), re-raise — data-layer errors like + `UniqueViolationError` mean the DB is reachable, reconnect would be + pointless. + 3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise. + This guards against partial stand-ins / older clients in tests. + 4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns + False (cooldown / lock contention / reconnect failure), re-raise. + 5. Otherwise await `coro_factory()` a second time and return / propagate + its result. At-most-one retry by construction — no infinite loop. + + `coro_factory` MUST be a zero-arg callable that returns a fresh awaitable + on each call. Passing an already-awaited coroutine would fail on retry + with `RuntimeError: cannot reuse already awaited coroutine`. + + `reason` should follow `___failure` so + telemetry distinguishes between fan-out callers (e.g. + `_update_config_from_db` issues four concurrent reads). + + Args: + prisma_client: The `PrismaClient` (or stand-in) that owns + `attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults. + coro_factory: Zero-arg callable returning the read awaitable. + reason: Telemetry tag forwarded to `attempt_db_reconnect`. + timeout_seconds: Optional override for the reconnect cycle timeout. + Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`, + then to 2.0s. + lock_timeout_seconds: Optional override for how long the helper will + wait to acquire the reconnect lock. Defaults to + `prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to + 0.1s. + + Returns: + Whatever `coro_factory()` returns (on first or second attempt). + + Raises: + Whatever `coro_factory()` raises if the failure is not a transport + error, or if the reconnect attempt does not succeed, or if the retry + also fails. + """ + try: + return await coro_factory() + except Exception as first_exc: + if not PrismaDBExceptionHandler.is_database_transport_error(first_exc): + raise + if not hasattr(prisma_client, "attempt_db_reconnect"): + raise + + resolved_timeout = _coerce_timeout( + ( + timeout_seconds + if timeout_seconds is not None + else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None) + ), + _DEFAULT_RECONNECT_TIMEOUT_SECONDS, + ) + resolved_lock_timeout = _coerce_timeout( + ( + lock_timeout_seconds + if lock_timeout_seconds is not None + else getattr( + prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None + ) + ), + _DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS, + ) + + verbose_proxy_logger.warning( + "DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s", + reason, + first_exc, + ) + + # Preserve the original transport error in telemetry. If + # `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer + # error, unexpected internal failure), surfacing that exception + # instead of `first_exc` would mask the actual DB transport problem + # in `failure_handler` / `db_exceptions` alerts. Chain the reconnect + # error as the cause for debuggability without losing the original. + try: + did_reconnect = await prisma_client.attempt_db_reconnect( + reason=reason, + timeout_seconds=resolved_timeout, + lock_timeout_seconds=resolved_lock_timeout, + ) + except Exception as reconnect_exc: + verbose_proxy_logger.warning( + "DB reconnect attempt raised; preserving original transport error. " + "reason=%s reconnect_error=%s", + reason, + reconnect_exc, + ) + raise first_exc from reconnect_exc + if not did_reconnect: + raise + + # At most one retry. If the retry also raises a transport error, we + # propagate — repeated reconnect-loops are the watchdog's job, not + # this helper's. + return await coro_factory() diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 73735796eb..d112e22230 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -52,18 +52,25 @@ class PrismaWrapper: engine = self._original_prisma._engine process = getattr(engine, "process", None) if engine is not None else None if process is not None: - return process.pid + pid = process.pid + if isinstance(pid, int): + return pid except (AttributeError, TypeError): pass return 0 @staticmethod async def _kill_engine_process(pid: int) -> None: - """Force-kill an orphaned engine subprocess to prevent DB connection pool leaks. + """Force-kill the engine subprocess to prevent DB connection pool leaks. - Called when disconnect() fails and the old engine process may still be - holding open connections. Sends SIGTERM for graceful shutdown, waits - briefly, then SIGKILL as a backstop. + Called on every reconnect (in `recreate_prisma_client`) to retire the + old query-engine subprocess without invoking prisma-client-py's + synchronous `disconnect()` — which blocks the asyncio event loop on + `subprocess.Popen.wait()` for 30-120+ seconds when the engine is + stuck on TCP close. + + Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as + a backstop. """ if pid <= 0: return @@ -72,7 +79,7 @@ class PrismaWrapper: except (ProcessLookupError, PermissionError, OSError): return # Already dead or inaccessible verbose_proxy_logger.warning( - "Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.", + "Sent SIGTERM to prisma-query-engine PID %s during reconnect.", pid, ) # Brief wait for graceful shutdown, then force-kill @@ -217,15 +224,18 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, http_client: Optional[Any] = None ): - """Disconnect and reconnect the Prisma client with a new database URL.""" + """Disconnect and reconnect the Prisma client with a new database URL. + + Kills the old engine subprocess directly (SIGTERM → SIGKILL) rather than + calling `disconnect()`. prisma-client-py's `disconnect()` calls a + synchronous `subprocess.Popen.wait()` that can freeze the asyncio event + loop for 30-120+ seconds when the engine is stuck on TCP close, + breaking `/health/liveliness` and causing Kubernetes pod restarts. + """ from prisma import Prisma # type: ignore old_engine_pid = self._get_engine_pid() - - try: - await self._original_prisma.disconnect() - except Exception as e: - verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}") + if old_engine_pid > 0: await self._kill_engine_process(old_engine_pid) if http_client is not None: diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 9768d93e92..6ada8f5878 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -35,6 +35,7 @@ async def google_generate_content( general_settings, llm_router, proxy_config, + proxy_logging_obj, version, ) @@ -73,6 +74,16 @@ async def google_generate_content( if llm_router is None: raise HTTPException(status_code=500, detail="Router not initialized") response = await llm_router.agenerate_content(**data) + success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data=data, + request=request, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + proxy_logging_obj=proxy_logging_obj, + ) + fastapi_response.headers.update(success_headers) return response @@ -95,6 +106,7 @@ async def google_stream_generate_content( general_settings, llm_router, proxy_config, + proxy_logging_obj, version, ) @@ -137,9 +149,24 @@ async def google_stream_generate_content( raise HTTPException(status_code=500, detail="Router not initialized") response = await llm_router.agenerate_content_stream(**data) + success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=response, + request_data=data, + request=request, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, + version=version, + proxy_logging_obj=proxy_logging_obj, + ) + # Check if response is an async iterator (streaming response) if response is not None and hasattr(response, "__aiter__"): - return StreamingResponse(content=response, media_type="text/event-stream") + return StreamingResponse( + content=response, + media_type="text/event-stream", + headers=success_headers, + ) + fastapi_response.headers.update(success_headers) return response diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py new file mode 100644 index 0000000000..a19af4dd48 --- /dev/null +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -0,0 +1,492 @@ +""" +WORKFLOW RUN MANAGEMENT + +Generic durable state tracking for agents and automated workflows. + +POST /v1/workflows/runs - Create a workflow run +GET /v1/workflows/runs - List runs (filter by type, status) +GET /v1/workflows/runs/{run_id} - Get run with latest event +PATCH /v1/workflows/runs/{run_id} - Update status, metadata, output +POST /v1/workflows/runs/{run_id}/events - Append event (updates run status) +GET /v1/workflows/runs/{run_id}/events - Full event log +POST /v1/workflows/runs/{run_id}/messages - Append conversation message +GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history +""" + +import json +from typing import Any, Dict, Literal, Optional + +from fastapi import APIRouter, Depends, HTTPException, Query + +try: + from prisma.errors import UniqueViolationError +except ImportError: + UniqueViolationError = None # type: ignore +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + +_MAX_SEQUENCE_RETRIES = 5 + + +def _json(value: Any) -> str: + """Serialize a Python value for prisma-client-py Json fields (must be a string).""" + return json.dumps(value) + + +def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + + +def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + """Return the hashed key token that identifies this caller, or None for master key.""" + return user_api_key_dict.token + + +# Status transitions driven by event_type +_EVENT_STATUS_MAP: Dict[str, str] = { + "step.started": "running", + "step.failed": "failed", + "hook.waiting": "paused", + "hook.received": "running", +} + + +# --------------------------------------------------------------------------- +# Request / Response models +# --------------------------------------------------------------------------- + + +class WorkflowRunCreateRequest(BaseModel): + workflow_type: str + input: Optional[Dict[str, Any]] = None + metadata: Optional[Dict[str, Any]] = None + + +WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed"] + + +class WorkflowRunUpdateRequest(BaseModel): + status: Optional[WorkflowRunStatus] = None + output: Optional[Dict[str, Any]] = None + metadata: Optional[Dict[str, Any]] = None + + +class WorkflowEventCreateRequest(BaseModel): + event_type: str + step_name: str + data: Optional[Dict[str, Any]] = None + + +class WorkflowMessageCreateRequest(BaseModel): + role: str + content: str + session_id: Optional[str] = None + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int: + """Return MAX(sequence_number) + 1 for the given run, for either events or messages.""" + if table == "events": + rows = await prisma_client.db.litellm_workflowevent.find_many( + where={"run_id": run_id}, + order={"sequence_number": "desc"}, + take=1, + ) + else: + rows = await prisma_client.db.litellm_workflowmessage.find_many( + where={"run_id": run_id}, + order={"sequence_number": "desc"}, + take=1, + ) + return (rows[0].sequence_number + 1) if rows else 0 + + +async def _require_run( + prisma_client: Any, + run_id: str, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, +) -> Any: + """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" + run = await prisma_client.db.litellm_workflowrun.find_unique( + where={"run_id": run_id} + ) + if run is None: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + if user_api_key_dict is not None and not _is_admin(user_api_key_dict): + caller = _caller_key(user_api_key_dict) + if not caller or run.created_by != caller: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + return run + + +# --------------------------------------------------------------------------- +# Endpoints +# --------------------------------------------------------------------------- + + +@router.post( + "/v1/workflows/runs", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def create_workflow_run( + data: WorkflowRunCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Create a new workflow run. Returns run_id and session_id. + + The caller's API key token is stored as created_by so that non-admin keys + can only see and modify their own runs. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + create_data: Dict[str, Any] = { + "workflow_type": data.workflow_type, + "created_by": _caller_key(user_api_key_dict), + } + if data.input is not None: + create_data["input"] = _json(data.input) + if data.metadata is not None: + create_data["metadata"] = _json(data.metadata) + run = await prisma_client.db.litellm_workflowrun.create(data=create_data) + return run + except Exception as e: + verbose_proxy_logger.exception("Error creating workflow run: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/workflows/runs", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_workflow_runs( + workflow_type: Optional[str] = Query(None), + status: Optional[str] = Query(None), + limit: int = Query(50, ge=1, le=250), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """List workflow runs. Filter by workflow_type and/or status. + + Non-admin callers only see runs created by their own API key. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + where: Dict[str, Any] = {} + if workflow_type: + where["workflow_type"] = workflow_type + if status: + statuses = [s.strip() for s in status.split(",")] + where["status"] = {"in": statuses} if len(statuses) > 1 else statuses[0] + + # Non-admin callers are scoped to their own key. + if not _is_admin(user_api_key_dict): + caller = _caller_key(user_api_key_dict) + if caller: + where["created_by"] = caller + + try: + runs = await prisma_client.db.litellm_workflowrun.find_many( + where=where, + order={"created_at": "desc"}, + take=limit, + ) + return {"runs": runs, "count": len(runs)} + except Exception as e: + verbose_proxy_logger.exception("Error listing workflow runs: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/workflows/runs/{run_id}", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_workflow_run( + run_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Get a workflow run with its most recent event.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + run = await prisma_client.db.litellm_workflowrun.find_unique( + where={"run_id": run_id}, + include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}}, + ) + if run is None: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + if not _is_admin(user_api_key_dict): + caller = _caller_key(user_api_key_dict) + if not caller or run.created_by != caller: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + return run + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting workflow run: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.patch( + "/v1/workflows/runs/{run_id}", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_workflow_run( + run_id: str, + data: WorkflowRunUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Update status, metadata, or output on a workflow run.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + update: Dict[str, Any] = {} + if data.status is not None: + update["status"] = data.status + if data.output is not None: + update["output"] = _json(data.output) + if data.metadata is not None: + update["metadata"] = _json(data.metadata) + + if not update: + raise HTTPException(status_code=400, detail="No fields to update") + + # Enforce ownership before writing. + await _require_run(prisma_client, run_id, user_api_key_dict) + + try: + run = await prisma_client.db.litellm_workflowrun.update( + where={"run_id": run_id}, + data=update, + ) + if run is None: + raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found") + return run + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error updating workflow run: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/workflows/runs/{run_id}/events", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def append_workflow_event( + run_id: str, + data: WorkflowEventCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Append an event to the run's event log. Also updates run.status if event_type maps to a status. + + Sequence numbers use optimistic concurrency: on a unique-constraint collision + (concurrent append), retries up to _MAX_SEQUENCE_RETRIES times with a fresh MAX+1. + The event+status update is atomic in a single DB transaction. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + new_status = _EVENT_STATUS_MAP.get(data.event_type) + + for attempt in range(_MAX_SEQUENCE_RETRIES): + try: + seq = await _get_next_sequence_number(prisma_client, run_id, "events") + event_data: Dict[str, Any] = { + "run_id": run_id, + "event_type": data.event_type, + "step_name": data.step_name, + "sequence_number": seq, + } + if data.data is not None: + event_data["data"] = _json(data.data) + + async with prisma_client.db.tx() as tx: + event = await tx.litellm_workflowevent.create(data=event_data) + if new_status: + await tx.litellm_workflowrun.update( + where={"run_id": run_id}, + data={"status": new_status}, + ) + + return event + + except Exception as e: + if UniqueViolationError is not None and isinstance(e, UniqueViolationError): + if attempt == _MAX_SEQUENCE_RETRIES - 1: + verbose_proxy_logger.exception( + "Sequence number collision after %d retries for run %s", + _MAX_SEQUENCE_RETRIES, + run_id, + ) + raise HTTPException( + status_code=409, + detail="Concurrent write conflict — please retry", + ) + continue + verbose_proxy_logger.exception("Error appending workflow event: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + raise HTTPException( + status_code=500, detail="Failed to append event" + ) # pragma: no cover + + +@router.get( + "/v1/workflows/runs/{run_id}/events", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_workflow_events( + run_id: str, + limit: int = Query(100, ge=1, le=500), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Fetch event log for a run, ordered by sequence_number. Default limit 100, max 500.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + try: + events = await prisma_client.db.litellm_workflowevent.find_many( + where={"run_id": run_id}, + order={"sequence_number": "asc"}, + take=limit, + ) + return {"events": events, "count": len(events)} + except Exception as e: + verbose_proxy_logger.exception("Error listing workflow events: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/workflows/runs/{run_id}/messages", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def append_workflow_message( + run_id: str, + data: WorkflowMessageCreateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Append a conversation message. Stores full content (not truncated). + + Uses optimistic concurrency for sequence numbers. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + for attempt in range(_MAX_SEQUENCE_RETRIES): + try: + seq = await _get_next_sequence_number(prisma_client, run_id, "messages") + msg_data: Dict[str, Any] = { + "run_id": run_id, + "role": data.role, + "content": data.content, + "sequence_number": seq, + } + if data.session_id is not None: + msg_data["session_id"] = data.session_id + msg = await prisma_client.db.litellm_workflowmessage.create(data=msg_data) + return msg + + except Exception as e: + if UniqueViolationError is not None and isinstance(e, UniqueViolationError): + if attempt == _MAX_SEQUENCE_RETRIES - 1: + verbose_proxy_logger.exception( + "Sequence number collision after %d retries for run %s", + _MAX_SEQUENCE_RETRIES, + run_id, + ) + raise HTTPException( + status_code=409, + detail="Concurrent write conflict — please retry", + ) + continue + verbose_proxy_logger.exception("Error appending workflow message: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + raise HTTPException( + status_code=500, detail="Failed to append message" + ) # pragma: no cover + + +@router.get( + "/v1/workflows/runs/{run_id}/messages", + tags=["workflow management"], + dependencies=[Depends(user_api_key_auth)], +) +async def list_workflow_messages( + run_id: str, + limit: int = Query(100, ge=1, le=500), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Fetch conversation history for a run, ordered by sequence_number. Default limit 100, max 500.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + await _require_run(prisma_client, run_id, user_api_key_dict) + + try: + messages = await prisma_client.db.litellm_workflowmessage.find_many( + where={"run_id": run_id}, + order={"sequence_number": "asc"}, + take=limit, + ) + return {"messages": messages, "count": len(messages)} + except Exception as e: + verbose_proxy_logger.exception("Error listing workflow messages: %s", e) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 214f2df8a1..2088ddf45f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -426,6 +426,9 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.tool_management_endpoints import ( router as tool_management_router, ) +from litellm.proxy.management_endpoints.workflow_management_endpoints import ( + router as workflow_management_router, +) from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, @@ -497,14 +500,18 @@ from litellm.proxy.utils import ( _get_redoc_url, _is_projected_spend_over_limit, _is_valid_team_configs, + get_config_param, get_custom_url, get_error_message_str, get_server_root_path, handle_exception_on_proxy, hash_password, hash_token, + invalidate_config_param, + litellm_config_cache, migrate_passwords_to_scrypt_async, model_dump_with_preserved_fields, + prefetch_config_params, update_spend, ) from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router @@ -738,6 +745,10 @@ async def _initialize_shared_aiohttp_session(): try: from aiohttp import ClientSession, TCPConnector + from litellm.llms.custom_httpx.http_handler import ( + _build_aiohttp_keepalive_socket_factory, + ) + connector_kwargs: Dict[str, Any] = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, @@ -748,6 +759,9 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST + socket_factory = _build_aiohttp_keepalive_socket_factory() + if socket_factory is not None: + connector_kwargs["socket_factory"] = socket_factory connector = TCPConnector(**connector_kwargs) session = ClientSession(connector=connector) @@ -2929,8 +2943,13 @@ class ProxyConfig: ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache spend_counter_cache.redis_cache = redis_usage_cache + litellm_config_cache.redis_cache = redis_usage_cache # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. + elif litellm_config_cache.redis_cache is None: + verbose_proxy_logger.info( + "litellm_config_cache: no Redis configured; cluster-wide cache sharing disabled." + ) def switch_on_llm_response_caching(self): """ @@ -4846,10 +4865,7 @@ class ProxyConfig: "environment_variables", ] for k in keys: - response = prisma_client.get_generic_data( - key="param_name", value=k, table_name="config" - ) - _tasks.append(response) + _tasks.append(get_config_param(prisma_client, k)) responses = await asyncio.gather(*_tasks) for response in responses: @@ -4931,6 +4947,19 @@ class ProxyConfig: global llm_router, llm_model_list, master_key, general_settings try: + # warm the config cache so the per-param reads below all hit + await prefetch_config_params( + prisma_client, + [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + "model_cost_map_reload_config", + "anthropic_beta_headers_reload_config", + ], + ) + # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) if self._should_load_db_object(object_type="models"): new_models = await self._get_models_from_db(prisma_client=prisma_client) @@ -4940,8 +4969,8 @@ class ProxyConfig: new_models=new_models, proxy_logging_obj=proxy_logging_obj ) - db_general_settings = await prisma_client.db.litellm_config.find_first( - where={"param_name": "general_settings"} + db_general_settings = await get_config_param( + prisma_client, "general_settings" ) # update general settings @@ -5034,10 +5063,7 @@ class ProxyConfig: from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook try: - # Load litellm_settings from DB - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "litellm_settings"} - ) + config_record = await get_config_param(prisma_client, "litellm_settings") if config_record is None or config_record.param_value is None: return @@ -5192,8 +5218,8 @@ class ProxyConfig: """ try: # Get model cost map reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "model_cost_map_reload_config"} + config_record = await get_config_param( + prisma_client, "model_cost_map_reload_config" ) if config_record is None or config_record.param_value is None: @@ -5288,6 +5314,7 @@ class ProxyConfig: }, }, ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" @@ -5307,8 +5334,8 @@ class ProxyConfig: """ try: # Get anthropic beta headers reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( - where={"param_name": "anthropic_beta_headers_reload_config"} + config_record = await get_config_param( + prisma_client, "anthropic_beta_headers_reload_config" ) if config_record is None or config_record.param_value is None: @@ -5396,6 +5423,7 @@ class ProxyConfig: }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") # Count providers in config provider_count = sum( @@ -6688,6 +6716,10 @@ class ProxyStartupEvent: Args: scheduler: The scheduler to add the background jobs to """ + global prisma_client + global proxy_logging_obj + global user_api_key_cache + ######################################################## # CloudZero Background Job ######################################################## @@ -6761,8 +6793,6 @@ class ProxyStartupEvent: ) # Get prisma_client and proxy_logging_obj from global scope - global prisma_client - global proxy_logging_obj if prisma_client is not None: # Reuse the PodLockManager from db_spend_update_writer pod_lock_manager = ( @@ -6792,6 +6822,83 @@ class ProxyStartupEvent: "Key rotation disabled (set LITELLM_KEY_ROTATION_ENABLED=true to enable)" ) + await cls._initialize_expired_ui_session_key_cleanup_background_job( + scheduler=scheduler + ) + + @classmethod + async def _initialize_expired_ui_session_key_cleanup_background_job( + cls, scheduler: AsyncIOScheduler + ): + """ + Initialize the expired UI session key cleanup background job. + """ + global prisma_client + global proxy_logging_obj + global user_api_key_cache + + ######################################################## + # Expired UI Session Key Cleanup Background Job + ######################################################## + from litellm.constants import ( + EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS, + ) + + expired_ui_session_key_cleanup_enabled: Optional[bool] = str_to_bool( + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED + ) + verbose_proxy_logger.debug( + "expired_ui_session_key_cleanup_enabled: " + f"{expired_ui_session_key_cleanup_enabled}" + ) + + if expired_ui_session_key_cleanup_enabled is True: + try: + from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import ( + ExpiredUISessionKeyCleanupManager, + ) + + if prisma_client is not None: + pod_lock_manager = ( + proxy_logging_obj.db_spend_update_writer.pod_lock_manager + ) + expired_ui_session_key_cleanup_manager = ( + ExpiredUISessionKeyCleanupManager( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + pod_lock_manager=pod_lock_manager, + ) + ) + verbose_proxy_logger.debug( + "Expired UI session key cleanup background job scheduled " + "every " + f"{LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS} " + "seconds " + "(LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true)" + ) + scheduler.add_job( + expired_ui_session_key_cleanup_manager.cleanup_expired_keys, + "interval", + seconds=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS, + id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + else: + verbose_proxy_logger.warning( + "Expired UI session key cleanup enabled but prisma_client " + "not available" + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to setup expired UI session key cleanup job: {e}" + ) + else: + verbose_proxy_logger.debug( + "Expired UI session key cleanup disabled (set " + "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" + ) + @classmethod async def _initialize_slack_alerting_jobs( cls, @@ -12633,6 +12740,7 @@ async def update_config( # noqa: PLR0915 updated_litellm_settings["success_callback"] = normalize_callback_names( incoming_cb ) + await invalidate_config_param(k) merged = {**existing, **updated_litellm_settings} @@ -12770,6 +12878,7 @@ async def update_config_general_settings( "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) + await invalidate_config_param("general_settings") return response @@ -13053,6 +13162,7 @@ async def delete_config_general_settings( "update": {"param_value": json.dumps(general_settings)}, # type: ignore }, ) + await invalidate_config_param("general_settings") return response @@ -13418,6 +13528,7 @@ async def reload_model_cost_map( }, }, ) + await invalidate_config_param("model_cost_map_reload_config") models_count = len(new_model_cost_map) if new_model_cost_map else 0 verbose_proxy_logger.info( @@ -13487,6 +13598,7 @@ async def schedule_model_cost_map_reload( }, }, ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reload scheduled for every {hours} hours" @@ -13540,6 +13652,7 @@ async def cancel_model_cost_map_reload( await prisma_client.db.litellm_config.delete( where={"param_name": "model_cost_map_reload_config"} ) + await invalidate_config_param("model_cost_map_reload_config") verbose_proxy_logger.info("Model cost map reload schedule cancelled") @@ -13770,6 +13883,7 @@ async def reload_anthropic_beta_headers( }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") provider_count = sum( 1 for k in new_config.keys() if k not in ["provider_aliases", "description"] @@ -13843,6 +13957,7 @@ async def schedule_anthropic_beta_headers_reload( }, }, ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info( f"Anthropic beta headers reload scheduled for every {hours} hours" @@ -13896,6 +14011,7 @@ async def cancel_anthropic_beta_headers_reload( await prisma_client.db.litellm_config.delete( where={"param_name": "anthropic_beta_headers_reload_config"} ) + await invalidate_config_param("anthropic_beta_headers_reload_config") verbose_proxy_logger.info("Anthropic beta headers reload schedule cancelled") @@ -14158,6 +14274,7 @@ app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) app.include_router(tool_management_router) +app.include_router(workflow_management_router) app.include_router(memory_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8f07c5afa3..fd54ed5243 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession { @@id([session_id, router_name, model_name]) @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } + +// --------------------------------------------------------------------------- +// Workflow Run Tracking +// +// Generic durable state tracking for any agent or automated workflow. +// Design: three tables — run (header + materialized status), event (append-only +// source of truth for state transitions), message (conversation inbox/outbox). +// +// Usage: +// - Set `workflow_type` to identify the owning system (e.g. "shin-builder"). +// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.). +// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to +// the proxy — all spend logs for this run are automatically tagged. +// --------------------------------------------------------------------------- + +// One instance of work being done. `status` is a materialized cache of the +// latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowRun { + run_id String @id @default(uuid()) + session_id String @unique @default(uuid()) + workflow_type String + status String @default("pending") + created_by String? // user_id of the key that created this run; null = created by master key + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + input Json? + output Json? + metadata Json? + + events LiteLLM_WorkflowEvent[] + messages LiteLLM_WorkflowMessage[] + + @@index([workflow_type, status]) + @@index([session_id]) + @@index([created_at]) + @@index([created_by]) +} + +// Append-only log of state transitions. Never mutate rows here. +// `step_name` and `event_type` are caller-defined strings — no hardcoded enums. +// Status auto-update rules (applied by the append endpoint): +// step.started → run.status = running +// step.failed → run.status = failed +// hook.waiting → run.status = paused +// hook.received → run.status = running +model LiteLLM_WorkflowEvent { + event_id String @id @default(uuid()) + run_id String + event_type String + step_name String + sequence_number Int + data Json? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} + +// Conversation inbox/outbox — full message content, separate from the durable +// event log. Spend logs truncate messages; this table stores them in full. +// `session_id` here is the Claude --resume session ID (or similar). +model LiteLLM_WorkflowMessage { + message_id String @id @default(uuid()) + run_id String + role String + content String + sequence_number Int + session_id String? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 712853a33c..d2dfa17751 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import ( should_create_missing_views, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + call_with_db_reconnect_retry, +) from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( @@ -2442,6 +2445,92 @@ async def _lookup_deprecated_key( return None +# DualCache for LiteLLM_Config param_name reads. +# Redis layer is attached in proxy_server._init_cache. +LITELLM_CONFIG_CACHE_TTL_SECONDS: int = int( + os.environ.get("LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS", "60") +) +_CONFIG_CACHE_MISS: str = "__litellm_config_param_miss__" + +litellm_config_cache: DualCache = DualCache( + default_in_memory_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS, + default_redis_ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS, +) + + +class _ConfigRow: + """Mimics the Prisma litellm_config row shape for cached entries.""" + + __slots__ = ("param_name", "param_value") + + def __init__(self, param_name: str, param_value: Any) -> None: + self.param_name = param_name + self.param_value = param_value + + +def _config_cache_key(param_name: str) -> str: + return f"litellm_config:param:{param_name}" + + +def _pack_config_row(row: Any) -> Dict[str, Any]: + return {"param_name": row.param_name, "param_value": row.param_value} + + +def _unpack_config_row(cached: Any) -> Optional[_ConfigRow]: + if cached is None or cached == _CONFIG_CACHE_MISS: + return None + if isinstance(cached, dict): + return _ConfigRow(cached["param_name"], cached["param_value"]) + return None + + +async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any]: + """Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None.""" + cache_key = _config_cache_key(param_name) + cached = await litellm_config_cache.async_get_cache(cache_key) + if cached is not None: + return _unpack_config_row(cached) + + row = await prisma_client.get_generic_data( + key="param_name", value=param_name, table_name="config" + ) + cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + await litellm_config_cache.async_set_cache( + cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS + ) + return row + + +async def invalidate_config_param(param_name: str) -> None: + """Evict from both cache layers; call after every LiteLLM_Config write.""" + await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + + +async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: + """Batch-load LiteLLM_Config rows into the cache with one find_many.""" + if not param_names: + return + try: + rows = await prisma_client.db.litellm_config.find_many( + where={"param_name": {"in": param_names}} # type: ignore + ) + except Exception as e: + verbose_proxy_logger.debug( + "prefetch_config_params failed, falling through to per-param queries: %s", + e, + ) + return + by_name = {row.param_name: row for row in rows} + for name in param_names: + row = by_name.get(name) + cache_value: Any = ( + _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + ) + await litellm_config_cache.async_set_cache( + _config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS + ) + + class PrismaClient: spend_log_transactions: List = [] _spend_log_transactions_lock = asyncio.Lock() @@ -2693,30 +2782,42 @@ class PrismaClient: table_name: Literal["users", "keys", "config", "spend"], ): """ - Generic implementation of get data + Generic implementation of get data. + + Self-heals across a single transient transport blip via + `call_with_db_reconnect_retry`: on `httpx.ReadError` / + `ClientNotConnectedError` / similar, attempt one DB reconnect and + retry once before surfacing the failure. Restores the 1.82.6 behavior + that was lost in 1.83.x — see issue #25143. """ start_time = time.time() - try: + + async def _do_query(): if table_name == "users": - response = await self.db.litellm_usertable.find_first( + return await self.db.litellm_usertable.find_first( where={key: value} # type: ignore ) elif table_name == "keys": - response = await self.db.litellm_verificationtoken.find_first( # type: ignore + return await self.db.litellm_verificationtoken.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "config": - response = await self.db.litellm_config.find_first( # type: ignore + return await self.db.litellm_config.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "spend": - response = await self.db.l.find_first( # type: ignore + return await self.db.l.find_first( # type: ignore where={key: value} # type: ignore ) - return response - except Exception as e: - import traceback + return None + try: + return await call_with_db_reconnect_retry( + self, + _do_query, + reason=f"prisma_get_generic_data_{table_name}_lookup_failure", + ) + except Exception as e: error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}" verbose_proxy_logger.error(error_msg) error_msg = error_msg + "\nException Type: {}".format(type(e)) @@ -3310,6 +3411,9 @@ class PrismaClient: tasks.append(updated_table_row) await asyncio.gather(*tasks) + # invalidate cache so other pods see writes from save_config + for k in data.keys(): + await invalidate_config_param(k) verbose_proxy_logger.info("Data Inserted into Config Table") elif table_name == "spend": db_data = self.jsonify_object(data=data) @@ -4094,8 +4198,11 @@ class PrismaClient: Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll handlers) to choose between heavy reconnect (engine dead -- recreate - Prisma client, re-arm watcher) and lightweight reconnect (network - blip -- disconnect, connect, SELECT 1). + Prisma client, re-arm watcher) and direct reconnect (network blip -- + recreate Prisma client, re-arm watcher, SELECT 1). Both paths recreate + the client via the non-blocking kill-then-construct flow rather than + calling disconnect(), which blocks the event loop on the synchronous + subprocess.Popen.wait() inside prisma-client-py (see issue #26191). """ effective_timeout = ( timeout_seconds @@ -4115,7 +4222,6 @@ class PrismaClient: ) self._reap_all_zombies() self._cleanup_engine_watcher() - self._engine_confirmed_dead = False async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") @@ -4128,23 +4234,32 @@ class PrismaClient: await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + # Only clear the "dead engine" flag after the heavy reconnect + # actually completed. If `_do_heavy_reconnect()` raises (timeout, + # missing DATABASE_URL, recreate failure), the flag stays True so + # the next attempt re-enters the heavy branch instead of silently + # demoting to the lightweight path. + self._engine_confirmed_dead = False else: verbose_proxy_logger.debug( "Performing Prisma DB reconnect (engine alive or unknown)." ) async def _do_direct_reconnect() -> None: - old_pid = self._get_engine_pid() - try: - await self.db.disconnect() - except Exception as disconnect_err: - verbose_proxy_logger.warning( - "Prisma DB disconnect before reconnect failed: %s", - disconnect_err, + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error( + "DATABASE_URL not set; cannot reconnect Prisma client." ) - await PrismaWrapper._kill_engine_process(old_pid) - - await self.db.connect() + raise RuntimeError("DATABASE_URL not set") + # Fresh Prisma client + new engine subprocess. The previous + # "lightweight" path called `disconnect()` which blocks the + # event loop on `subprocess.Popen.wait()`; since that call + # ends up killing the engine anyway, we do it non-blockingly + # via `_kill_engine_process` inside `recreate_prisma_client`. + self._cleanup_engine_watcher() + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() await self.db.query_raw("SELECT 1") await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) diff --git a/litellm/proxy/workflows/README.md b/litellm/proxy/workflows/README.md new file mode 100644 index 0000000000..f452066afb --- /dev/null +++ b/litellm/proxy/workflows/README.md @@ -0,0 +1,150 @@ +# Workflow Run Tracking + +Generic durable state tracking for agents and automated workflows built on the LiteLLM proxy. + +## The Problem + +Agents like [shin-builder](https://github.com/BerriAI/shin-builder) run multi-stage pipelines (triage → plan → implement → PR). Their task state and conversation history lived in memory — a process restart lost everything. + +## Three-Table Design + +``` +WorkflowRun one instance of work (header + materialized status) +WorkflowEvent append-only state transitions (source of truth for replay) +WorkflowMessage conversation inbox/outbox (full content, not truncated) +``` + +**WorkflowEvent is the source of truth.** `WorkflowRun.status` is a materialized cache updated automatically when events are appended. If you need to debug a run, replay its events. + +## API + +All endpoints require a valid LiteLLM API key (`Authorization: Bearer sk-...`). + +### Runs + +``` +POST /v1/workflows/runs Create a run +GET /v1/workflows/runs List runs (?workflow_type=&status=) +GET /v1/workflows/runs/{run_id} Get run + latest event +PATCH /v1/workflows/runs/{run_id} Update status / metadata / output +``` + +### Events + +``` +POST /v1/workflows/runs/{run_id}/events Append event (auto-updates run status) +GET /v1/workflows/runs/{run_id}/events Full event log (ordered by sequence) +``` + +### Messages + +``` +POST /v1/workflows/runs/{run_id}/messages Append message +GET /v1/workflows/runs/{run_id}/messages Conversation history (ordered by sequence) +``` + +## Quick Start + +```bash +# Create a run +curl -X POST http://localhost:4000/v1/workflows/runs \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{"workflow_type": "shin-builder", "metadata": {"title": "Fix login bug"}}' + +# {"run_id": "abc-123", "session_id": "xyz-456", "status": "pending", ...} + +# Mark step started (sets status → running) +curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/events \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{"event_type": "step.started", "step_name": "grill", "data": {"claude_session_id": "sess-789"}}' + +# Store a conversation message +curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/messages \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{"role": "user", "content": "What is the expected behavior?", "session_id": "sess-789"}' + +# Restart recovery: fetch active runs and resume from last event's data.claude_session_id +curl "http://localhost:4000/v1/workflows/runs?status=running,paused&workflow_type=shin-builder" \ + -H "Authorization: Bearer sk-1234" +``` + +## Status Auto-Update Rules + +When you append an event, the run's status is updated automatically: + +| event_type | run.status | +|-----------------|------------| +| `step.started` | `running` | +| `step.failed` | `failed` | +| `hook.waiting` | `paused` | +| `hook.received` | `running` | + +Set `status = completed` explicitly via PATCH when the workflow finishes. + +## Linking to Spend Logs + +`WorkflowRun.session_id` is generated automatically (UUID). Pass it as the `x-litellm-session-id` header when making completions through the proxy: + +```python +headers = {"x-litellm-session-id": run.session_id} +``` + +All spend log entries for this run are then tagged automatically. Query cost per run: + +``` +POST /ui/spend_logs/view_session_spend_logs?session_id={run.session_id} +``` + +## Sequence Numbers + +Sequence numbers on events and messages are assigned server-side (`MAX + 1` per run). Callers never supply them. This guarantees ordering even under concurrent writes. + +## Using from shin-builder + +Replace the in-memory `tasks.py` dict with calls to these endpoints: + +```python +import httpx + +class WorkflowRunClient: + def __init__(self, base_url: str, api_key: str): + self._client = httpx.AsyncClient( + base_url=base_url, + headers={"Authorization": f"Bearer {api_key}"}, + ) + + async def create_task(self, title: str, **metadata) -> dict: + r = await self._client.post("/v1/workflows/runs", json={ + "workflow_type": "shin-builder", + "metadata": {"title": title, **metadata}, + }) + r.raise_for_status() + return r.json() + + async def list_active_tasks(self) -> list: + r = await self._client.get( + "/v1/workflows/runs", + params={"workflow_type": "shin-builder", "status": "running,paused"}, + ) + r.raise_for_status() + return r.json()["runs"] + + async def transition(self, run_id: str, step_name: str, event_type: str, data: dict = None): + r = await self._client.post(f"/v1/workflows/runs/{run_id}/events", json={ + "event_type": event_type, + "step_name": step_name, + "data": data or {}, + }) + r.raise_for_status() + + async def append_message(self, run_id: str, role: str, content: str, session_id: str = None): + r = await self._client.post(f"/v1/workflows/runs/{run_id}/messages", json={ + "role": role, "content": content, "session_id": session_id, + }) + r.raise_for_status() +``` + +On startup, call `list_active_tasks()` to restore in-flight runs. The last `step.started` event's `data.claude_session_id` gives you the `--resume` ID. diff --git a/litellm/router_utils/clientside_credential_handler.py b/litellm/router_utils/clientside_credential_handler.py index c98f614335..45ade81b2d 100644 --- a/litellm/router_utils/clientside_credential_handler.py +++ b/litellm/router_utils/clientside_credential_handler.py @@ -11,9 +11,60 @@ If given, generate a unique model_id for the deployment. Ensures cooldowns are applied correctly. """ +from typing import List + clientside_credential_keys = ["api_key", "api_base", "base_url"] +def _admin_config_fields_to_clear_on_base_override() -> List[str]: + """ + Provider-specific credential / endpoint-targeting fields that must NOT + flow through to a client-redirected upstream. + + Built dynamically from ``CredentialLiteLLMParams.model_fields`` so any + new provider field added there (Bedrock endpoint, Watsonx region, etc.) + is gated automatically — plus a fixed list of kwargs-only fields that + aren't declared on the typed model. + """ + from litellm.types.router import CredentialLiteLLMParams + + typed_fields = [ + f + for f in CredentialLiteLLMParams.model_fields + if f not in clientside_credential_keys + ] + kwargs_only_fields = [ + # Caller-supplied via **kwargs, not declared on CredentialLiteLLMParams. + "organization", + "extra_body", + "extra_headers", + "default_headers", + "api_type", + "azure_ad_token", + "azure_ad_token_provider", + "aws_session_token", + "aws_sts_endpoint", + "aws_web_identity_token", + "aws_role_name", + # OCI provider — consumed by litellm/llms/oci/* via optional_params + # and not declared on CredentialLiteLLMParams. Without these here, + # an admin's OCI signing key / tenancy / fingerprint would flow + # through to an attacker-redirected upstream. + "oci_signer", + "oci_user", + "oci_fingerprint", + "oci_tenancy", + "oci_key", + "oci_key_file", + ] + return typed_fields + kwargs_only_fields + + +_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE = ( + _admin_config_fields_to_clear_on_base_override() +) + + def is_clientside_credential(request_kwargs: dict) -> bool: """ Check if the credential is a clientside credential. @@ -34,4 +85,20 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di for key in clientside_credential_keys: if key in request_kwargs: litellm_params[key] = request_kwargs[key] + + # If the caller redirected api_base/base_url to a client-controlled value, + # don't forward the admin's organization / extra_body / region / token / + # vertex / aws fields — those were meant for the original upstream. + # Always drop the admin's value first, then write the caller's value back + # if they resupplied the field. The naive + # ``if field not in request_kwargs: pop`` shape lets a caller *echo* a + # field name (with any value, including an empty string) to keep the + # admin's value in ``litellm_params`` and have it forwarded to the + # redirected upstream. + if "api_base" in request_kwargs or "base_url" in request_kwargs: + for field in _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE: + litellm_params.pop(field, None) + if field in request_kwargs: + litellm_params[field] = request_kwargs[field] + return litellm_params diff --git a/litellm/types/caching.py b/litellm/types/caching.py index c8194ce2e7..f8050b292c 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -118,3 +118,4 @@ class CachedEmbedding(TypedDict): index: Optional[int] object: Optional[str] model: Optional[str] + prompt_tokens_details: Optional[dict] diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index ace8c8a418..8f8673b0a7 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -81,6 +81,12 @@ class MCPServer(BaseModel): # Defaults to the token's expires_in minus the expiry buffer, or # MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. token_storage_ttl_seconds: Optional[int] = None + # Resolved short-ID tool prefix when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is + # enabled. Set by ``MCPServerManager._assign_unique_short_prefix`` at + # registration time so that natural-hash collisions between two + # different ``server_id`` values are bumped deterministically. Left + # ``None`` in default-prefix mode. + short_prefix: Optional[str] = None model_config = ConfigDict(arbitrary_types_allowed=True) @property diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a212d56c1a..ed29d49fc2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2659,7 +2659,7 @@ class StandardLoggingHiddenParams(TypedDict): ] # id of the model in the router, separates multiple models with the same name but different credentials cache_key: Optional[str] api_base: Optional[str] - response_cost: Optional[str] + response_cost: Optional[Union[str, float]] litellm_overhead_time_ms: Optional[float] additional_headers: Optional[StandardLoggingAdditionalHeaders] batch_models: Optional[List[str]] diff --git a/litellm/utils.py b/litellm/utils.py index e1ad1db63e..027c9fedce 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6526,6 +6526,7 @@ def validate_environment( # noqa: PLR0915 or model in litellm.open_ai_text_completion_models or model in litellm.open_ai_embedding_models or model in litellm.openai_image_generation_models + or model.startswith("gpt-image") ): if "OPENAI_API_KEY" in os.environ: keys_in_environment = True @@ -8410,6 +8411,17 @@ class ProviderConfigManager: model: str, provider: LlmProviders, ) -> Optional[BaseAnthropicMessagesConfig]: + return ProviderConfigManager._get_provider_anthropic_messages_config_cached( + model=model, provider=provider + ) + + @staticmethod + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) + def _get_provider_anthropic_messages_config_cached( + model: str, + provider: LlmProviders, + ) -> Optional[BaseAnthropicMessagesConfig]: + model_lower = model.lower() if litellm.LlmProviders.ANTHROPIC == provider: return litellm.AnthropicMessagesConfig() # The 'BEDROCK' provider corresponds to Amazon's implementation of Anthropic Claude v3. @@ -8419,14 +8431,14 @@ class ProviderConfigManager: return BedrockModelInfo.get_bedrock_provider_config_for_messages_api(model) elif litellm.LlmProviders.VERTEX_AI == provider: - if "claude" in model.lower(): + if "claude" in model_lower: from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) return VertexAIPartnerModelsAnthropicMessagesConfig() elif litellm.LlmProviders.AZURE_AI == provider: - if "claude" in model.lower(): + if "claude" in model_lower: from litellm.llms.azure_ai.anthropic.messages_transformation import ( AzureAnthropicMessagesConfig, ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4d8f3a984f..ca7d323ad6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4749,17 +4749,17 @@ "supports_web_search": true }, "azure/gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -4788,17 +4788,17 @@ "supports_low_reasoning_effort": false }, "azure/gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, "supported_endpoints": [ "/v1/batch", "/v1/responses" @@ -5117,6 +5117,38 @@ "/v1/images/edits" ] }, + "azure/gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/low/1024-x-1024/gpt-image-1-mini": { "input_cost_per_pixel": 2.0751953125e-09, "litellm_provider": "azure", @@ -19097,6 +19129,38 @@ "supports_vision": true, "supports_pdf_input": true }, + "gpt-image-2": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "gpt-image-2-2026-04-21": { + "cache_read_input_image_token_cost": 2e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "openai", + "mode": "image_generation", + "output_cost_per_token": 1e-05, + "input_cost_per_image_token": 8e-06, + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "low/1024-x-1024/gpt-image-1.5": { "input_cost_per_image": 0.009, "litellm_provider": "openai", @@ -19912,21 +19976,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" @@ -19955,21 +20019,21 @@ "supports_minimal_reasoning_effort": true }, "gpt-5.5-pro-2026-04-23": { - "cache_read_input_token_cost": 6e-06, - "cache_read_input_token_cost_above_272k_tokens": 1.2e-05, - "input_cost_per_token": 6e-05, - "input_cost_per_token_above_272k_tokens": 0.00012, - "input_cost_per_token_flex": 3e-05, - "input_cost_per_token_batches": 3e-05, + "cache_read_input_token_cost": 3e-06, + "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "input_cost_per_token": 3e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_cost_per_token_flex": 1.5e-05, + "input_cost_per_token_batches": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", - "output_cost_per_token": 0.00036, - "output_cost_per_token_above_272k_tokens": 0.00054, - "output_cost_per_token_flex": 0.00018, - "output_cost_per_token_batches": 0.00018, + "output_cost_per_token": 0.00018, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_cost_per_token_flex": 9e-05, + "output_cost_per_token_batches": 9e-05, "supported_endpoints": [ "/v1/responses", "/v1/batch" diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 6f23c87f91..ed49c14621 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -193,6 +193,23 @@ "a2a": false } }, + "aihubmix": { + "display_name": "AIHubMix (`aihubmix`)", + "url": "https://docs.litellm.ai/docs/providers/aihubmix", + "endpoints": { + "chat_completions": true, + "messages": true, + "responses": true, + "embeddings": true, + "image_generations": true, + "audio_transcriptions": true, + "audio_speech": true, + "moderations": true, + "batches": false, + "rerank": true, + "a2a": false + } + }, "assemblyai": { "display_name": "AssemblyAI (`assemblyai`)", "url": "https://docs.litellm.ai/docs/pass_through/assembly_ai", diff --git a/pyproject.toml b/pyproject.toml index a15fa5a06a..657632d69e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.83.14" +version = "1.84.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -236,7 +236,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.83.14" +version = "1.84.0" version_files = [ "pyproject.toml:^version", ] diff --git a/schema.prisma b/schema.prisma index 8f07c5afa3..fd54ed5243 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession { @@id([session_id, router_name, model_name]) @@index([last_activity_at], map: "idx_adaptive_router_session_activity") } + +// --------------------------------------------------------------------------- +// Workflow Run Tracking +// +// Generic durable state tracking for any agent or automated workflow. +// Design: three tables — run (header + materialized status), event (append-only +// source of truth for state transitions), message (conversation inbox/outbox). +// +// Usage: +// - Set `workflow_type` to identify the owning system (e.g. "shin-builder"). +// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.). +// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to +// the proxy — all spend logs for this run are automatically tagged. +// --------------------------------------------------------------------------- + +// One instance of work being done. `status` is a materialized cache of the +// latest event; the event log is the authoritative source of truth. +model LiteLLM_WorkflowRun { + run_id String @id @default(uuid()) + session_id String @unique @default(uuid()) + workflow_type String + status String @default("pending") + created_by String? // user_id of the key that created this run; null = created by master key + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + input Json? + output Json? + metadata Json? + + events LiteLLM_WorkflowEvent[] + messages LiteLLM_WorkflowMessage[] + + @@index([workflow_type, status]) + @@index([session_id]) + @@index([created_at]) + @@index([created_by]) +} + +// Append-only log of state transitions. Never mutate rows here. +// `step_name` and `event_type` are caller-defined strings — no hardcoded enums. +// Status auto-update rules (applied by the append endpoint): +// step.started → run.status = running +// step.failed → run.status = failed +// hook.waiting → run.status = paused +// hook.received → run.status = running +model LiteLLM_WorkflowEvent { + event_id String @id @default(uuid()) + run_id String + event_type String + step_name String + sequence_number Int + data Json? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} + +// Conversation inbox/outbox — full message content, separate from the durable +// event log. Spend logs truncate messages; this table stores them in full. +// `session_id` here is the Claude --resume session ID (or similar). +model LiteLLM_WorkflowMessage { + message_id String @id @default(uuid()) + run_id String + role String + content String + sequence_number Int + session_id String? + created_at DateTime @default(now()) + + run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id]) + + @@unique([run_id, sequence_number]) + @@index([run_id]) +} diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py index 786167b948..0d241f7574 100644 --- a/tests/litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( @pytest.mark.asyncio -async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( +async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive( engine_client, ) -> None: - """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" - engine_client._engine_pid = 1234 + """Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1. - with patch.object(engine_client, "_is_engine_alive", return_value=True): + The old "lightweight" path called `disconnect()` + `connect()`, which + blocks the event loop on the sync `process.wait()` inside aclose(). + The fix routes both engine-alive and engine-dead paths through + `recreate_prisma_client`, which non-blockingly kills the old engine. + """ + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=True), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await engine_client._run_reconnect_cycle(timeout_seconds=5.0) - engine_client.db.connect.assert_awaited_once() + engine_client.db.recreate_prisma_client.assert_awaited_once_with( + "postgresql://test" + ) engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") - engine_client.db.recreate_prisma_client.assert_not_awaited() + engine_client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio -async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( +async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown( engine_client, ) -> None: - """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + """When the engine PID is not tracked, direct reconnect still runs.""" engine_client._engine_pid = 0 + engine_client._start_engine_watcher = AsyncMock() - await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) - engine_client.db.connect.assert_awaited_once() + engine_client.db.recreate_prisma_client.assert_awaited_once_with( + "postgresql://test" + ) engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") - engine_client.db.recreate_prisma_client.assert_not_awaited() + engine_client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio @@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): @pytest.mark.asyncio -async def test_escalation_after_consecutive_lightweight_failures(engine_client): - """After N consecutive lightweight reconnect failures, _engine_confirmed_dead +async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client): + """After N consecutive direct reconnect failures, _engine_confirmed_dead is set to True so _run_reconnect_cycle takes the heavy reconnect path.""" engine_client._reconnect_escalation_threshold = 3 engine_client._consecutive_reconnect_failures = 0 engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test + engine_client._start_engine_watcher = AsyncMock(return_value=None) - # Make lightweight reconnect fail every time - engine_client.db.disconnect = AsyncMock(return_value=None) - engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed")) + # Make direct reconnect fail every time + engine_client.db.recreate_prisma_client = AsyncMock( + side_effect=Exception("recreate failed") + ) # Run 3 failed reconnect attempts - for i in range(3): - result = await engine_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=5.0 - ) - assert result is False + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + for _ in range(3): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) + assert result is False assert engine_client._consecutive_reconnect_failures == 3 - # Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle + # Next attempt should escalate to the heavy path (recreate_prisma_client still + # the call, but via the _engine_confirmed_dead branch that also re-arms the watcher). engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) - engine_client._start_engine_watcher = AsyncMock(return_value=None) with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): result = await engine_client._attempt_reconnect_inside_lock( force=True, reason="test_escalation", timeout_seconds=5.0 ) - # Heavy reconnect should have been attempted (recreate_prisma_client called) engine_client.db.recreate_prisma_client.assert_awaited_once() @@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client): """A successful reconnect resets _consecutive_reconnect_failures to 0.""" engine_client._consecutive_reconnect_failures = 2 engine_client._db_reconnect_cooldown_seconds = 0 + engine_client._start_engine_watcher = AsyncMock() # Make reconnect succeed - engine_client.db.disconnect = AsyncMock(return_value=None) - engine_client.db.connect = AsyncMock(return_value=None) + engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) - result = await engine_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=5.0 - ) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) assert result is True assert engine_client._consecutive_reconnect_failures == 0 diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index a41907722e..45c6a04ad5 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -314,11 +314,11 @@ def test_update_litellm_params_for_health_check(): # Issue #15807: Fixes health checks sending "region/model" as model ID to AWS model_info = {} litellm_params = { - "model": "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0", "api_key": "fake_key", } updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) - assert updated_params["model"] == "anthropic.claude-3-7-sonnet-20250219-v1:0" + assert updated_params["model"] == "anthropic.claude-sonnet-4-5-20250929-v1:0" # Test with Bedrock cross-region inference profile - should preserve the inference profile prefix # AWS requires inference profile IDs like "us.anthropic.claude..." for cross-region routing diff --git a/tests/litellm_utils_tests/test_logging_callback_manager.py b/tests/litellm_utils_tests/test_logging_callback_manager.py index 88ae07fd81..d9540f8f85 100644 --- a/tests/litellm_utils_tests/test_logging_callback_manager.py +++ b/tests/litellm_utils_tests/test_logging_callback_manager.py @@ -366,3 +366,40 @@ def test_generic_api_compatible_callbacks_json_unknown_callback(): # Should return the string unchanged assert result == "unknown_callback", "Unknown callback should be returned as-is" assert isinstance(result, str), "Unknown callback should remain a string" + + +@pytest.mark.asyncio +async def test_generic_api_callback_settings_retry_config(): + """ + Test that generic_api callback_settings are passed to GenericAPILogger. + """ + from litellm.integrations.generic_api.generic_api_callback import GenericAPILogger + from litellm.litellm_core_utils.logging_callback_manager import ( + _generic_api_logger_cache, + ) + + callback_name = "test_generic_api_retry_config" + _generic_api_logger_cache.pop(callback_name, None) + litellm.callback_settings[callback_name] = { + "callback_type": "generic_api", + "endpoint": "https://example.com/api/logs", + "headers": {"Content-Type": "application/json"}, + "max_retries": 2, + "retry_delay": 0.5, + "timeout": 3, + } + + try: + result = LoggingCallbackManager._add_custom_callback_generic_api_str( + callback_name + ) + + assert isinstance(result, GenericAPILogger) + assert result.endpoint == "https://example.com/api/logs" + assert result.headers == {"Content-Type": "application/json"} + assert result.max_retries == 2 + assert result.retry_delay == 0.5 + assert result.timeout == 3 + finally: + litellm.callback_settings.pop(callback_name, None) + _generic_api_logger_cache.pop(callback_name, None) diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index d5df4ef75a..20af6e1023 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -2309,11 +2309,11 @@ def test_get_provider_audio_transcription_config(): @pytest.mark.parametrize( "model, expected_bool", [ - ("anthropic.claude-3-7-sonnet-20250219-v1:0", True), - ("us.anthropic.claude-3-7-sonnet-20250219-v1:0", True), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", True), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", True), ], ) -def test_claude_3_7_sonnet_supports_pdf_input(model, expected_bool): +def test_claude_sonnet_4_5_supports_pdf_input(model, expected_bool): from litellm.utils import supports_pdf_input assert supports_pdf_input(model) == expected_bool diff --git a/tests/llm_translation/test_bedrock_anthropic_regression.py b/tests/llm_translation/test_bedrock_anthropic_regression.py index 5928ca0223..8b8ce0a6cc 100644 --- a/tests/llm_translation/test_bedrock_anthropic_regression.py +++ b/tests/llm_translation/test_bedrock_anthropic_regression.py @@ -134,7 +134,7 @@ class TestBedrockAnthropicPromptCachingRegression: if "converse" in model_prefix: config = AmazonConverseConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, @@ -162,7 +162,7 @@ class TestBedrockAnthropicPromptCachingRegression: else: config = AmazonAnthropicClaudeConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, @@ -227,7 +227,7 @@ class TestBedrockAnthropicPromptCachingRegression: if "converse" in model_prefix: config = AmazonConverseConfig() result = config._transform_request_helper( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", system_content_blocks=[], optional_params={}, messages=messages, @@ -236,7 +236,7 @@ class TestBedrockAnthropicPromptCachingRegression: else: config = AmazonAnthropicClaudeConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, @@ -498,7 +498,7 @@ class TestBedrockAnthropicCombinedRegressions: if "converse" in model_prefix: config = AmazonConverseConfig() result = config._transform_request_helper( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", system_content_blocks=[], optional_params={}, messages=messages, @@ -518,7 +518,7 @@ class TestBedrockAnthropicCombinedRegressions: else: config = AmazonAnthropicClaudeConfig() result = config.transform_request( - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, optional_params={}, litellm_params={}, diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index ddfe383f2a..15f950224d 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -1323,7 +1323,7 @@ def test_base_aws_llm_get_credentials(): def test_bedrock_completion_test_2(): litellm.set_verbose = True data = { - "model": "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [ { "role": "system", @@ -1630,7 +1630,7 @@ def test_bedrock_completion_test_4(modify_params): litellm.modify_params = modify_params data = { - "model": "anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [ { "role": "user", @@ -2115,7 +2115,7 @@ class TestBedrockConverseAnthropicUnitTests(BaseAnthropicChatTest): def get_base_completion_call_args_with_thinking(self) -> dict: return { - "model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "thinking": {"type": "enabled", "budget_tokens": 16000}, } @@ -2828,7 +2828,7 @@ async def test_bedrock_thinking_in_assistant_message(sync_mode): client = AsyncHTTPHandler() params = { - "model": "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "model": "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [ { "role": "assistant", @@ -2887,7 +2887,7 @@ async def test_bedrock_stream_thinking_content_openwebui(): ``` """ response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[{"role": "user", "content": "Hello who is this?"}], stream=True, max_tokens=1080, diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index 8fc961d12d..8b6f37bfbc 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -580,7 +580,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): def test_litellm_gateway_from_sdk_with_thinking_param(): try: response = litellm.completion( - model="litellm_proxy/anthropic.claude-3-7-sonnet-20250219-v1:0", + model="litellm_proxy/anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[{"role": "user", "content": "Hello world"}], api_base="http://0.0.0.0:4000", api_key="sk-PIp1h0RekR", diff --git a/tests/llm_translation/test_optional_params.py b/tests/llm_translation/test_optional_params.py index 82a3d96b02..b40ce11bb9 100644 --- a/tests/llm_translation/test_optional_params.py +++ b/tests/llm_translation/test_optional_params.py @@ -1828,7 +1828,7 @@ def test_azure_response_format_param(): "model, provider", [ ("claude-3-7-sonnet-20240620-v1:0", "anthropic"), - ("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), ("invoke/anthropic.claude-3-7-sonnet-20240620-v1:0", "bedrock"), ("claude-3-7-sonnet@20250219", "vertex_ai"), ], diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 9070a9feab..9782bf3c2a 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3493,8 +3493,14 @@ def test_litellm_api_base(monkeypatch, provider, route): def test_gemini_tool_calling_working_demo(): - load_vertex_ai_credentials() - litellm._turn_on_debug() + """ + Regression test: tool params with anyOf containing a `{"type": "array"}` + branch (no items field at all) must synthesize items before the request + is sent to Vertex (Vertex rejects array types missing items). + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + args = { "messages": [ { @@ -3564,13 +3570,75 @@ def test_gemini_tool_calling_working_demo(): ], "vertex_location": "global", } - response = completion(model="vertex_ai/gemini-3-flash-preview", **args) - print(response) + + client = HTTPHandler() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello!"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + with ( + patch.object(client, "post", return_value=mock_response) as mock_post, + patch.object( + VertexBase, + "_ensure_access_token", + return_value=("fake-token", "fake-project"), + ), + ): + completion( + model="vertex_ai/gemini-3-flash-preview", + client=client, + **args, + ) + + sent_body = mock_post.call_args.kwargs.get( + "json" + ) or mock_post.call_args.kwargs.get("data") + assert sent_body is not None, "expected request body to be sent" + if isinstance(sent_body, str): + sent_body = json.loads(sent_body) + + function_decl = sent_body["tools"][0]["function_declarations"][0] + callbacks_schema = function_decl["parameters"]["properties"]["config"][ + "properties" + ]["callbacks"] + array_branches = [ + branch + for branch in callbacks_schema["anyOf"] + if branch.get("type", "").lower() == "array" + ] + assert array_branches, "expected an array branch in callbacks anyOf" + for branch in array_branches: + assert "items" in branch and branch["items"], ( + f"array branch in callbacks.anyOf must include non-empty items " + f"(Vertex rejects array types missing items). Got: {branch}" + ) def test_gemini_tool_calling_not_working(): - load_vertex_ai_credentials() - litellm._turn_on_debug() + """ + Regression test: tool params with anyOf containing both an empty-items + array branch and a null branch must serialize with items present on the + array branch (Vertex rejects array types missing `items`). + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + args = { "messages": [ { @@ -3637,8 +3705,64 @@ def test_gemini_tool_calling_not_working(): ], "vertex_location": "global", } - response = completion(model="vertex_ai/gemini-3-flash-preview", **args) - print(response) + + client = HTTPHandler() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Type": "application/json"} + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Hello!"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + with ( + patch.object(client, "post", return_value=mock_response) as mock_post, + patch.object( + VertexBase, + "_ensure_access_token", + return_value=("fake-token", "fake-project"), + ), + ): + completion( + model="vertex_ai/gemini-3-flash-preview", + client=client, + **args, + ) + + sent_body = mock_post.call_args.kwargs.get( + "json" + ) or mock_post.call_args.kwargs.get("data") + assert sent_body is not None, "expected request body to be sent" + if isinstance(sent_body, str): + sent_body = json.loads(sent_body) + + function_decl = sent_body["tools"][0]["function_declarations"][0] + callbacks_schema = function_decl["parameters"]["properties"]["config"][ + "properties" + ]["callbacks"] + array_branches = [ + branch + for branch in callbacks_schema["anyOf"] + if branch.get("type", "").lower() == "array" + ] + assert array_branches, "expected an array branch in callbacks anyOf" + for branch in array_branches: + assert "items" in branch and branch["items"], ( + f"array branch in callbacks.anyOf must include non-empty items " + f"(Vertex rejects array types missing items). Got: {branch}" + ) def test_vertex_ai_llama_tool_calling(): diff --git a/tests/local_testing/test_function_calling.py b/tests/local_testing/test_function_calling.py index b52805c066..02affa1d57 100644 --- a/tests/local_testing/test_function_calling.py +++ b/tests/local_testing/test_function_calling.py @@ -159,7 +159,7 @@ def test_aaparallel_function_call(model): "model", [ "anthropic/claude-4-sonnet-20250514", - "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", ], ) @pytest.mark.flaky(retries=3, delay=1) diff --git a/tests/local_testing/test_model_alias_map.py b/tests/local_testing/test_model_alias_map.py index cf731d6628..14c1de2f6a 100644 --- a/tests/local_testing/test_model_alias_map.py +++ b/tests/local_testing/test_model_alias_map.py @@ -30,10 +30,9 @@ def test_model_alias_map(caplog): ) print(response.model) - captured_logs = [rec.levelname for rec in caplog.records] - - for log in captured_logs: - assert "ERROR" not in log + for rec in caplog.records: + if rec.levelname == "ERROR" and rec.name.startswith("LiteLLM"): + pytest.fail(f"Unexpected litellm ERROR log: {rec.getMessage()}") assert "llama-3.1-8b-instant" in response.model except litellm.ServiceUnavailableError: diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index abe96b2ea2..3e8d59b299 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=litellm_params_dict, + extra_body=None, ) ) captured_request_body["url"] = url diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py index 528a5101df..6984b6fa00 100644 --- a/tests/logging_callback_tests/test_generic_api_callback.py +++ b/tests/logging_callback_tests/test_generic_api_callback.py @@ -8,6 +8,7 @@ sys.path.insert(0, os.path.abspath("../..")) import asyncio import litellm import gzip +import httpx import json import logging import time @@ -470,3 +471,96 @@ async def test_generic_api_callback_invalid_log_format(): endpoint=test_endpoint, log_format="invalid_format", # type: ignore # Intentionally invalid for testing ) + + +@pytest.mark.asyncio +async def test_generic_api_callback_retries_timeout_then_succeeds(): + """ + Test that GenericAPILogger retries LiteLLM timeout errors when configured. + """ + test_endpoint = "https://example.com/api/logs" + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + max_retries=1, + retry_delay=0, + timeout=0.2, + ) + + mock_post = AsyncMock() + mock_post.side_effect = [ + litellm.Timeout( + message="Connection timed out", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + type("Response", (), {"status_code": 200})(), + ] + generic_logger.async_httpx_client.post = mock_post + generic_logger.log_queue = [{"event": "timeout-retry"}] + + await generic_logger.async_send_batch() + + assert mock_post.call_count == 2 + first_call = mock_post.call_args_list[0][1] + assert first_call["url"] == test_endpoint + assert first_call["timeout"] == 0.2 + assert json.loads(first_call["data"]) == [{"event": "timeout-retry"}] + + +@pytest.mark.asyncio +async def test_generic_api_callback_retries_5xx_then_succeeds(): + """ + Test that GenericAPILogger retries transient HTTP 5xx errors when configured. + """ + test_endpoint = "https://example.com/api/logs" + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + max_retries=1, + retry_delay=0, + ) + + request = httpx.Request("POST", test_endpoint) + response = httpx.Response(status_code=503, request=request) + mock_post = AsyncMock() + mock_post.side_effect = [ + httpx.HTTPStatusError( + "Server error", + request=request, + response=response, + ), + type("Response", (), {"status_code": 200})(), + ] + generic_logger.async_httpx_client.post = mock_post + generic_logger.log_queue = [{"event": "5xx-retry"}] + + await generic_logger.async_send_batch() + + assert mock_post.call_count == 2 + + +@pytest.mark.asyncio +async def test_generic_api_callback_does_not_retry_4xx(): + """ + Test that GenericAPILogger does not retry non-transient HTTP 4xx errors. + """ + test_endpoint = "https://example.com/api/logs" + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + max_retries=2, + retry_delay=0, + ) + + request = httpx.Request("POST", test_endpoint) + response = httpx.Response(status_code=401, request=request) + mock_post = AsyncMock() + mock_post.side_effect = httpx.HTTPStatusError( + "Unauthorized", + request=request, + response=response, + ) + generic_logger.async_httpx_client.post = mock_post + generic_logger.log_queue = [{"event": "4xx-no-retry"}] + + await generic_logger.async_send_batch() + + mock_post.assert_called_once() diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py index 3c71af97c9..5fc4ecefb3 100644 --- a/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py +++ b/tests/pass_through_unit_tests/base_anthropic_messages_prompt_caching_test.py @@ -96,8 +96,8 @@ class BaseAnthropicMessagesPromptCachingTest(ABC): Returns the model string to use for tests. Examples: - - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0" - - "bedrock/invoke/anthropic.claude-3-7-sonnet-20250219-v1:0" + - "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" + - "bedrock/invoke/anthropic.claude-sonnet-4-5-20250929-v1:0" """ pass diff --git a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py index 83a47a0149..a194ded12f 100644 --- a/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py +++ b/tests/pass_through_unit_tests/test_anthropic_messages_prompt_caching.py @@ -31,7 +31,7 @@ class TestBedrockConversePromptCaching(BaseAnthropicMessagesPromptCachingTest): """ def get_model(self) -> str: - return "bedrock/converse/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + return "bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0" class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest): @@ -43,4 +43,4 @@ class TestBedrockInvokePromptCaching(BaseAnthropicMessagesPromptCachingTest): """ def get_model(self) -> str: - return "bedrock/invoke/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + return "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0" diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 837ce7d405..742a4f410d 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -52,3 +52,183 @@ async def test_process_async_embedding_cached_response(): print(f"response: {response}") assert len(response.data) == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_preserves_prompt_tokens_details(): + """Test that prompt_tokens_details (including image_count) survives a full cache hit.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": "base64imagedata"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_backward_compat_no_prompt_tokens_details(): + """Test that old cached items without prompt_tokens_details still work.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Old-format cached item — no prompt_tokens_details field + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "test"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens_details is None + + +@pytest.mark.asyncio +async def test_embedding_cache_aggregates_multiple_image_counts(): + """Test that image_count is summed correctly across multiple cached items.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + { + "embedding": [0.031, 0.042], + "index": 1, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens_details": {"image_count": 1}, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={ + "model": "amazon.titan-embed-image-v1", + "input": ["img1", "img2"], + }, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage.prompt_tokens_details is not None + assert response.usage.prompt_tokens_details.image_count == 2 + + +def test_combine_usage_merges_prompt_tokens_details(): + """Test that combine_usage merges prompt_tokens_details from both Usage objects.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + usage1 = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + usage2 = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=2), + ) + + combined = llm_caching_handler.combine_usage(usage1, usage2) + + assert combined.prompt_tokens == 30 + assert combined.total_tokens == 30 + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 3 + + +def test_combine_usage_handles_none_details(): + """Test that combine_usage works when one or both sides have null prompt_tokens_details.""" + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # Both null + usage_a = Usage(prompt_tokens=10, completion_tokens=0, total_tokens=10) + usage_b = Usage(prompt_tokens=20, completion_tokens=0, total_tokens=20) + combined = llm_caching_handler.combine_usage(usage_a, usage_b) + assert combined.prompt_tokens_details is None + + # Only first has details + usage_c = Usage( + prompt_tokens=10, + completion_tokens=0, + total_tokens=10, + prompt_tokens_details=PromptTokensDetailsWrapper(image_count=1), + ) + combined = llm_caching_handler.combine_usage(usage_c, usage_b) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 + + # Only second has details + combined = llm_caching_handler.combine_usage(usage_a, usage_c) + assert combined.prompt_tokens_details is not None + assert combined.prompt_tokens_details.image_count == 1 diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json b/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json index 913e3bfeda..818e4fa3ea 100644 --- a/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json +++ b/tests/test_litellm/integrations/open_telemetry/data/captured_kwargs.json @@ -1 +1 @@ -{"litellm_trace_id": null, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "input": [{"role": "user", "content": "What is the capital of France?"}], "litellm_params": {"acompletion": true, "api_key": null, "force_timeout": 600, "logger_fn": null, "verbose": false, "custom_llm_provider": "bedrock", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse", "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "model_alias_map": {}, "completion_call_id": null, "aembedding": null, "metadata": {"requester_metadata": {}, "user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_user_id": null, "user_api_key_org_id": null, "user_api_key_team_alias": null, "user_api_key_end_user_id": null, "user_api_key_user_email": null, "user_api_key": "unused-for-aws-bedrock", "user_api_end_user_max_budget": null, "litellm_api_version": "1.72.3", "global_max_parallel_requests": null, "user_api_key_team_max_budget": null, "user_api_key_team_spend": null, "user_api_key_spend": 0.0, "user_api_key_max_budget": null, "user_api_key_model_max_budget": {}, "user_api_key_metadata": {}, "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "endpoint": "http://0.0.0.0:44444/chat/completions", "litellm_parent_otel_span": null, "requester_ip_address": "", "model_group": "claude-3-7-sonnet", "model_group_size": 1, "deployment": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "api_base": null, "caching_groups": null, "hidden_params": {"custom_llm_provider": "bedrock", "region_name": null, "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "api_base": null, "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "response_cost": 0.001047, "additional_headers": {"x-litellm-model-group": "claude-3-7-sonnet", "x-litellm-attempted-retries": 0, "x-litellm-attempted-fallbacks": 0}, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "litellm_overhead_time_ms": 231.156, "_response_ms": 236.798}}, "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "proxy_server_request": {"url": "http://0.0.0.0:44444/chat/completions", "method": "POST", "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "body": {"messages": [{"role": "user", "content": "What is the capital of France?"}], "model": "claude-3-7-sonnet", "stream": false}}, "preset_cache_key": null, "no-log": null, "stream_response": {}, "input_cost_per_token": null, "input_cost_per_second": null, "output_cost_per_token": null, "output_cost_per_second": null, "cooldown_time": null, "text_completion": null, "azure_ad_token_provider": null, "user_continue_message": null, "base_model": null, "litellm_trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "litellm_session_id": null, "hf_model_name": null, "custom_prompt_dict": {}, "litellm_metadata": null, "disable_add_transform_inline_image_block": null, "drop_params": null, "prompt_id": null, "prompt_variables": null, "async_call": null, "ssl_verify": null, "merge_reasoning_content_in_choices": false, "api_version": null, "azure_ad_token": null, "tenant_id": null, "client_id": null, "client_secret": null, "azure_username": null, "azure_password": null, "max_retries": 0, "timeout": 6000.0, "bucket_name": null, "vertex_credentials": null, "vertex_project": null, "use_litellm_proxy": false}, "applied_guardrails": [], "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "messages": [{"role": "user", "content": "What is the capital of France?"}], "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "start_time": "2025-06-22 10:59:08.159939", "stream": false, "user": null, "call_type": "acompletion", "completion_start_time": "2025-06-22 10:59:08.399523", "standard_callback_dynamic_params": {}, "stream_options": null, "max_retries": 0, "provider": "aws", "region": "us-west-2", "custom_llm_provider": "bedrock", "api_key": "", "additional_args": {"complete_input_dict": "{\"messages\": [{\"role\": \"user\", \"content\": [{\"text\": \"What is the capital of France?\"}]}], \"additionalModelRequestFields\": {\"provider\": \"aws\", \"region\": \"us-west-2\"}, \"system\": [], \"inferenceConfig\": {}}"}, "log_event_type": "post_api_call", "api_call_start_time": "2025-06-22 10:59:08.387641", "llm_api_duration_ms": 5.642, "original_response": "{\"metrics\":{\"latencyMs\":1513},\"output\":{\"message\":{\"content\":[{\"text\":\"The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.\"}],\"role\":\"assistant\"}},\"stopReason\":\"end_turn\",\"usage\":{\"cacheReadInputTokenCount\":0,\"cacheReadInputTokens\":0,\"cacheWriteInputTokenCount\":0,\"cacheWriteInputTokens\":0,\"inputTokens\":14,\"outputTokens\":67,\"totalTokens\":81}}", "end_time": "2025-06-22 10:59:08.399523", "cache_hit": null, "response_cost": 0.001047, "standard_logging_object": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "call_type": "acompletion", "cache_hit": null, "stream": true, "status": "success", "custom_llm_provider": "bedrock", "saved_cache_cost": 0.0, "startTime": 1750615148.162725, "endTime": 1750615148.399523, "completionStartTime": 1750615148.399523, "response_time": 0.23679804801940918, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "metadata": {"user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_org_id": null, "user_api_key_user_id": null, "user_api_key_team_alias": null, "user_api_key_user_email": null, "spend_logs_metadata": null, "requester_ip_address": "", "requester_metadata": {}, "user_api_key_end_user_id": null, "prompt_management_metadata": null, "applied_guardrails": [], "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "requester_custom_headers": {"x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600"}}, "cache_key": null, "response_cost": 0.001047, "total_tokens": 81, "prompt_tokens": 14, "completion_tokens": 67, "request_tags": [], "end_user": "", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse", "model_group": "claude-3-7-sonnet", "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "requester_ip_address": "", "messages": [{"role": "user", "content": "What is the capital of France?"}], "response": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}, "model_parameters": {"stream": false}, "hidden_params": {"model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "cache_key": null, "api_base": null, "response_cost": 0.001047, "additional_headers": {}, "litellm_overhead_time_ms": 231.156, "batch_models": null, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "usage_object": null}, "model_map_information": {"model_map_key": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "model_map_value": {"key": "anthropic.claude-3-7-sonnet-20250219-v1:0", "max_tokens": 8192, "max_input_tokens": 200000, "max_output_tokens": 8192, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_reasoning_token": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "bedrock_converse", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": null, "supports_audio_output": null, "supports_pdf_input": true, "supports_embedding_image_input": null, "supports_native_streaming": null, "supports_web_search": null, "supports_url_context": null, "supports_reasoning": true, "supports_computer_use": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["max_tokens", "max_completion_tokens", "stream", "stream_options", "stop", "temperature", "top_p", "extra_headers", "response_format", "tools", "tool_choice", "thinking", "reasoning_effort"]}}, "error_str": null, "error_information": {"error_code": "", "error_class": "", "llm_provider": "", "traceback": "", "error_message": ""}, "response_cost_failure_debug_info": null, "guardrail_information": null, "standard_built_in_tools_params": {"web_search_options": null, "file_search": null}}, "async_complete_streaming_response": "ModelResponse(id='chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1', created=1750615148, model='arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None))], usage=Usage(completion_tokens=67, prompt_tokens=14, total_tokens=81, completion_tokens_details=None, prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None), cache_creation_input_tokens=0, cache_read_input_tokens=0))"} \ No newline at end of file +{"litellm_trace_id": null, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "input": [{"role": "user", "content": "What is the capital of France?"}], "litellm_params": {"acompletion": true, "api_key": null, "force_timeout": 600, "logger_fn": null, "verbose": false, "custom_llm_provider": "bedrock", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-sonnet-4-5-20250929-v1%3A0/converse", "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "model_alias_map": {}, "completion_call_id": null, "aembedding": null, "metadata": {"requester_metadata": {}, "user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_user_id": null, "user_api_key_org_id": null, "user_api_key_team_alias": null, "user_api_key_end_user_id": null, "user_api_key_user_email": null, "user_api_key": "unused-for-aws-bedrock", "user_api_end_user_max_budget": null, "litellm_api_version": "1.72.3", "global_max_parallel_requests": null, "user_api_key_team_max_budget": null, "user_api_key_team_spend": null, "user_api_key_spend": 0.0, "user_api_key_max_budget": null, "user_api_key_model_max_budget": {}, "user_api_key_metadata": {}, "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "endpoint": "http://0.0.0.0:44444/chat/completions", "litellm_parent_otel_span": null, "requester_ip_address": "", "model_group": "claude-3-7-sonnet", "model_group_size": 1, "deployment": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "api_base": null, "caching_groups": null, "hidden_params": {"custom_llm_provider": "bedrock", "region_name": null, "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "litellm_call_id": "dbecd23a-e71a-49cf-90d4-712a8a8e29c5", "api_base": null, "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "response_cost": 0.001047, "additional_headers": {"x-litellm-model-group": "claude-3-7-sonnet", "x-litellm-attempted-retries": 0, "x-litellm-attempted-fallbacks": 0}, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "litellm_overhead_time_ms": 231.156, "_response_ms": 236.798}}, "model_info": {"id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "db_model": false}, "proxy_server_request": {"url": "http://0.0.0.0:44444/chat/completions", "method": "POST", "headers": {"host": "0.0.0.0:44444", "accept-encoding": "gzip, deflate, zstd", "connection": "keep-alive", "accept": "application/json", "content-type": "application/json", "user-agent": "AsyncOpenAI/Python 1.84.0", "x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600", "content-length": "116"}, "body": {"messages": [{"role": "user", "content": "What is the capital of France?"}], "model": "claude-3-7-sonnet", "stream": false}}, "preset_cache_key": null, "no-log": null, "stream_response": {}, "input_cost_per_token": null, "input_cost_per_second": null, "output_cost_per_token": null, "output_cost_per_second": null, "cooldown_time": null, "text_completion": null, "azure_ad_token_provider": null, "user_continue_message": null, "base_model": null, "litellm_trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "litellm_session_id": null, "hf_model_name": null, "custom_prompt_dict": {}, "litellm_metadata": null, "disable_add_transform_inline_image_block": null, "drop_params": null, "prompt_id": null, "prompt_variables": null, "async_call": null, "ssl_verify": null, "merge_reasoning_content_in_choices": false, "api_version": null, "azure_ad_token": null, "tenant_id": null, "client_id": null, "client_secret": null, "azure_username": null, "azure_password": null, "max_retries": 0, "timeout": 6000.0, "bucket_name": null, "vertex_credentials": null, "vertex_project": null, "use_litellm_proxy": false}, "applied_guardrails": [], "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "messages": [{"role": "user", "content": "What is the capital of France?"}], "optional_params": {"stream": false, "max_retries": 0, "provider": "aws", "region": "us-west-2"}, "start_time": "2025-06-22 10:59:08.159939", "stream": false, "user": null, "call_type": "acompletion", "completion_start_time": "2025-06-22 10:59:08.399523", "standard_callback_dynamic_params": {}, "stream_options": null, "max_retries": 0, "provider": "aws", "region": "us-west-2", "custom_llm_provider": "bedrock", "api_key": "", "additional_args": {"complete_input_dict": "{\"messages\": [{\"role\": \"user\", \"content\": [{\"text\": \"What is the capital of France?\"}]}], \"additionalModelRequestFields\": {\"provider\": \"aws\", \"region\": \"us-west-2\"}, \"system\": [], \"inferenceConfig\": {}}"}, "log_event_type": "post_api_call", "api_call_start_time": "2025-06-22 10:59:08.387641", "llm_api_duration_ms": 5.642, "original_response": "{\"metrics\":{\"latencyMs\":1513},\"output\":{\"message\":{\"content\":[{\"text\":\"The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.\"}],\"role\":\"assistant\"}},\"stopReason\":\"end_turn\",\"usage\":{\"cacheReadInputTokenCount\":0,\"cacheReadInputTokens\":0,\"cacheWriteInputTokenCount\":0,\"cacheWriteInputTokens\":0,\"inputTokens\":14,\"outputTokens\":67,\"totalTokens\":81}}", "end_time": "2025-06-22 10:59:08.399523", "cache_hit": null, "response_cost": 0.001047, "standard_logging_object": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "trace_id": "4c97150b-b1a3-4dec-bd7a-734786b1b3bc", "call_type": "acompletion", "cache_hit": null, "stream": true, "status": "success", "custom_llm_provider": "bedrock", "saved_cache_cost": 0.0, "startTime": 1750615148.162725, "endTime": 1750615148.399523, "completionStartTime": 1750615148.399523, "response_time": 0.23679804801940918, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "metadata": {"user_api_key_hash": "unused-for-aws-bedrock", "user_api_key_alias": null, "user_api_key_team_id": null, "user_api_key_org_id": null, "user_api_key_user_id": null, "user_api_key_team_alias": null, "user_api_key_user_email": null, "spend_logs_metadata": null, "requester_ip_address": "", "requester_metadata": {}, "user_api_key_end_user_id": null, "prompt_management_metadata": null, "applied_guardrails": [], "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "usage_object": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "requester_custom_headers": {"x-stainless-lang": "python", "x-stainless-package-version": "1.84.0", "x-stainless-os": "MacOS", "x-stainless-arch": "arm64", "x-stainless-runtime": "CPython", "x-stainless-runtime-version": "3.12.10", "x-stainless-async": "async:asyncio", "x-stainless-retry-count": "0", "x-stainless-read-timeout": "600"}}, "cache_key": null, "response_cost": 0.001047, "total_tokens": 81, "prompt_tokens": 14, "completion_tokens": 67, "request_tags": [], "end_user": "", "api_base": "https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A1234567890123%3Ainference-profile%2Fus.anthropic.claude-sonnet-4-5-20250929-v1%3A0/converse", "model_group": "claude-3-7-sonnet", "model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "requester_ip_address": "", "messages": [{"role": "user", "content": "What is the capital of France?"}], "response": {"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}, "model_parameters": {"stream": false}, "hidden_params": {"model_id": "6bace4d6db0105943b3b0bfe7eb1a62c06e6f16f008cc4673fdf918eb3e9e62a", "cache_key": null, "api_base": null, "response_cost": 0.001047, "additional_headers": {}, "litellm_overhead_time_ms": 231.156, "batch_models": null, "litellm_model_name": "bedrock/arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "usage_object": null}, "model_map_information": {"model_map_key": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "model_map_value": {"key": "anthropic.claude-sonnet-4-5-20250929-v1:0", "max_tokens": 8192, "max_input_tokens": 200000, "max_output_tokens": 8192, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_reasoning_token": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "bedrock_converse", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": null, "supports_audio_output": null, "supports_pdf_input": true, "supports_embedding_image_input": null, "supports_native_streaming": null, "supports_web_search": null, "supports_url_context": null, "supports_reasoning": true, "supports_computer_use": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["max_tokens", "max_completion_tokens", "stream", "stream_options", "stop", "temperature", "top_p", "extra_headers", "response_format", "tools", "tool_choice", "thinking", "reasoning_effort"]}}, "error_str": null, "error_information": {"error_code": "", "error_class": "", "llm_provider": "", "traceback": "", "error_message": ""}, "response_cost_failure_debug_info": null, "guardrail_information": null, "standard_built_in_tools_params": {"web_search_options": null, "file_search": null}}, "async_complete_streaming_response": "ModelResponse(id='chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1', created=1750615148, model='arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.', role='assistant', tool_calls=None, function_call=None, provider_specific_fields=None))], usage=Usage(completion_tokens=67, prompt_tokens=14, total_tokens=81, completion_tokens_details=None, prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None), cache_creation_input_tokens=0, cache_read_input_tokens=0))"} \ No newline at end of file diff --git a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json b/tests/test_litellm/integrations/open_telemetry/data/captured_response.json index 3cf77781cc..1fa1889909 100644 --- a/tests/test_litellm/integrations/open_telemetry/data/captured_response.json +++ b/tests/test_litellm/integrations/open_telemetry/data/captured_response.json @@ -1 +1 @@ -{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}} \ No newline at end of file +{"id": "chatcmpl-fa9be5b7-9487-46ab-86de-6462d578fea1", "created": 1750615148, "model": "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0", "object": "chat.completion", "system_fingerprint": null, "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "The capital of France is Paris. Paris has been the capital city of France since 987 CE when Hugh Capet, the first king of the Capetian dynasty, made the city his seat of government. Today, Paris is not only the political capital but also the cultural and economic center of France.", "role": "assistant", "tool_calls": null, "function_call": null}}], "usage": {"completion_tokens": 67, "prompt_tokens": 14, "total_tokens": 81, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}} \ No newline at end of file diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 271d58061a..1a4d03528e 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -220,7 +220,7 @@ async def test_anthropic_cache_control_hook_negative_indices(): with patch.object(client, "post", return_value=mock_response) as mock_post: # Test with multiple messages and negative indices response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "system", @@ -352,7 +352,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds_logging(): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, cache_control_injection_points=[ {"location": "message", "index": 10} @@ -420,7 +420,7 @@ async def test_anthropic_cache_control_hook_negative_out_of_bounds_logging(): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, cache_control_injection_points=[ { @@ -486,7 +486,7 @@ async def test_anthropic_cache_control_hook_multiple_user_messages(): with patch.object(client, "post", return_value=mock_response) as mock_post: # Test with multiple user messages and negative indices response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", @@ -586,7 +586,7 @@ async def test_anthropic_cache_control_hook_out_of_bounds(bad_index): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, cache_control_injection_points=[ {"location": "message", "index": bad_index} @@ -651,7 +651,7 @@ async def test_anthropic_cache_control_hook_single_message(message_list): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=message_list, cache_control_injection_points=[{"location": "message", "index": -1}], client=client, @@ -691,7 +691,7 @@ async def test_anthropic_cache_control_hook_empty_message_list(): match="bedrock requires at least one non-system message", ): await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[], cache_control_injection_points=[ {"location": "message", "index": -1} @@ -742,7 +742,7 @@ async def test_anthropic_cache_control_hook_no_op(): ] await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=messages, # No cache_control_injection_points parameter client=client, @@ -799,7 +799,7 @@ async def test_anthropic_cache_control_hook_multiple_content_items_last_only(): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", @@ -874,7 +874,7 @@ async def test_anthropic_cache_control_hook_document_analysis_multiple_pages(): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: response = await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ { "role": "user", @@ -1057,7 +1057,7 @@ async def test_anthropic_cache_control_hook_string_negative_index(): client = AsyncHTTPHandler() with patch.object(client, "post", return_value=mock_response) as mock_post: await litellm.acompletion( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", messages=[ {"role": "user", "content": "First message"}, {"role": "assistant", "content": "First response"}, diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index f710647189..b31bbca889 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -262,7 +262,7 @@ class TestOpenTelemetryProviderInitialization(unittest.TestCase): class TestOpenTelemetry(unittest.TestCase): POLL_INTERVAL = 0.05 POLL_TIMEOUT = 2.0 - MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0" + MODEL = "arn:aws:bedrock:us-west-2:1234567890123:inference-profile/us.anthropic.claude-sonnet-4-5-20250929-v1:0" HERE = os.path.dirname(__file__) @patch.dict(os.environ, {}, clear=True) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 3016762043..77284a64cf 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -369,7 +369,7 @@ def test_generic_cost_per_token_gpt55(): def test_generic_cost_per_token_gpt55_pro(): - """gpt-5.5-pro: responses-only model — $60/1M input, $360/1M output, $6/1M cached input.""" + """gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input.""" model = "gpt-5.5-pro" custom_llm_provider = "openai" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" @@ -378,18 +378,18 @@ def test_generic_cost_per_token_gpt55_pro(): model_cost_map = litellm.model_cost[model] # Sanity-check the map values match OpenAI's published pricing. - assert model_cost_map["input_cost_per_token"] == 6e-5 - assert model_cost_map["output_cost_per_token"] == 3.6e-4 - assert model_cost_map["cache_read_input_token_cost"] == 6e-6 + assert model_cost_map["input_cost_per_token"] == 3e-5 + assert model_cost_map["output_cost_per_token"] == 1.8e-4 + assert model_cost_map["cache_read_input_token_cost"] == 3e-6 assert model_cost_map["litellm_provider"] == "openai" # gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint). assert model_cost_map["mode"] == "responses" assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"] assert "/v1/responses" in model_cost_map["supported_endpoints"] - # Inherits GPT-5.4-pro's long-context window + tiered pricing (scaled 2x). + # Inherits GPT-5.4-pro's long-context window + tiered pricing. assert model_cost_map["max_input_tokens"] == 1050000 - assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1.2e-4 - assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 5.4e-4 + assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5 + assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4 prompt_tokens = 1000 completion_tokens = 500 @@ -454,8 +454,8 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities( [ ("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7), ("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7), - ("azure/gpt-5.5-pro", "responses", 6e-5, 3.6e-4, 6e-6), - ("azure/gpt-5.5-pro-2026-04-23", "responses", 6e-5, 3.6e-4, 6e-6), + ("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6), + ("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6), ], ) def test_azure_gpt55_entries_present_with_correct_pricing( @@ -464,7 +464,7 @@ def test_azure_gpt55_entries_present_with_correct_pricing( """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure. Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page - on 2026-04-24): $5/$30 input/output per 1M for chat, $60/$360 for pro. + on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro. Cache discount is 10% of input. """ os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 72cfd89408..f8708dd2f7 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -77,7 +77,7 @@ async def test_anthropic_bedrock_thinking_blocks_with_none_content(): # test _bedrock_converse_messages_pt_async result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( messages=messages, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", llm_provider="bedrock", ) diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py new file mode 100644 index 0000000000..fc5b39a2fd --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py @@ -0,0 +1,138 @@ +""" +Regression tests for the parsed-URL hostname match used to identify a +caller-supplied ``api_base`` as a known openai-compatible provider. + +The previous shape (``if endpoint in api_base:``) used unanchored +substring search, which let a caller pass +``https://attacker.com/api.groq.com/openai/v1`` and have the proxy +return ``GROQ_API_KEY`` as the dynamic credential — exfiltrating the +server's real provider key to an attacker-controlled host on the +outbound request. +""" + +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.litellm_core_utils.get_llm_provider_logic import ( + _endpoint_matches_api_base, + get_llm_provider, +) + + +class TestEndpointMatchesApiBase: + """Direct unit tests on the parsed-URL matcher.""" + + @pytest.mark.parametrize( + "endpoint, api_base", + [ + # Bare hostname endpoint, exact host match. + ("api.perplexity.ai", "https://api.perplexity.ai/v1"), + # Endpoint includes a path; api_base path starts with it. + ("api.groq.com/openai/v1", "https://api.groq.com/openai/v1"), + # Endpoint with full URL scheme. + ("https://api.cerebras.ai/v1", "https://api.cerebras.ai/v1/chat"), + # Trailing-slash on registered endpoint must not break match. + ("https://llm.chutes.ai/v1/", "https://llm.chutes.ai/v1/chat"), + # Case-insensitive on hostname. + ("api.groq.com/openai/v1", "https://API.GROQ.COM/openai/v1"), + ], + ) + def test_legitimate_provider_urls_match(self, endpoint, api_base): + assert _endpoint_matches_api_base(endpoint, api_base) is True + + @pytest.mark.parametrize( + "endpoint, api_base", + [ + # Attacker host, registered endpoint smuggled into path. + ( + "api.groq.com/openai/v1", + "https://attacker.com/api.groq.com/openai/v1", + ), + # Attacker host, registered endpoint smuggled into a path segment. + ( + "api.groq.com/openai/v1", + "https://attacker.com/foo/api.groq.com/openai/v1", + ), + # Lookalike host that contains the registered host as a suffix label. + ( + "api.groq.com/openai/v1", + "https://api.groq.com.attacker.com/openai/v1", + ), + # Lookalike host with the registered host as a prefix. + ( + "api.groq.com/openai/v1", + "https://api.groq.com.evil.example/openai/v1", + ), + # Right host, wrong path — endpoint requires ``/openai/v1`` prefix. + ("api.groq.com/openai/v1", "https://api.groq.com/v1"), + # Path-segment lookalike: ``/openai/v10`` must not match ``/openai/v1``. + ("api.groq.com/openai/v1", "https://api.groq.com/openai/v10"), + # Userinfo / @-injection trick — the ``hostname`` after ``@`` is + # what httpx connects to. + ( + "api.groq.com/openai/v1", + "https://api.groq.com@attacker.com/openai/v1", + ), + ], + ) + def test_attacker_smuggling_does_not_match(self, endpoint, api_base): + assert _endpoint_matches_api_base(endpoint, api_base) is False + + +class TestGetLlmProviderRejectsAttackerSmuggledApiBase: + """ + End-to-end: ``get_llm_provider`` must NOT return the server's stored + secret (e.g. ``GROQ_API_KEY``) for an api_base whose hostname is + attacker-controlled, even when the registered endpoint string appears + elsewhere in the URL. + """ + + def test_attacker_host_does_not_yield_groq_secret(self): + # The function may either fall through (different provider) or + # raise BadRequestError because the model can't be identified. + # The invariant under test is that ``GROQ_API_KEY`` is never + # looked up against an attacker-controlled hostname. + import litellm + + with patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str", + return_value="server-real-groq-key", + ) as mocked_secret: + try: + _, _, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://attacker.com/api.groq.com/openai/v1", + ) + # If it returned, the dynamic key must not be the secret. + assert dynamic_api_key != "server-real-groq-key" + except litellm.exceptions.BadRequestError: + # Acceptable outcome: provider unidentifiable, no secret + # was returned. + pass + + # Regardless of return / raise, the secret must never have been + # read against this attacker-controlled api_base. + groq_lookups = [ + call + for call in mocked_secret.call_args_list + if call.args and call.args[0] == "GROQ_API_KEY" + ] + assert groq_lookups == [] + + def test_legitimate_groq_api_base_still_resolves(self): + with patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str", + return_value="server-real-groq-key", + ): + _, provider, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://api.groq.com/openai/v1", + ) + + assert provider == "groq" + assert dynamic_api_key == "server-real-groq-key" 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 cf7be6bf1c..1764d9c609 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2337,6 +2337,104 @@ def test_merge_hidden_params_from_response_into_metadata_populates_metadata(): assert meta["hidden_params"]["model_id"] == "mid-test" +def test_merge_hidden_params_from_response_into_metadata_backfills_response_cost(): + """Streaming metadata should include the already-calculated response cost.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="merge-hp-cost-test", + function_id="merge-hp-cost-fn", + ) + logging_obj.model_call_details = { + "litellm_params": {"metadata": {}}, + "response_cost": 0.002, + } + + class _Resp: + _hidden_params = {"response_cost": None, "model_id": "mid-test"} + + response = _Resp() + logging_obj._merge_hidden_params_from_response_into_metadata(response) + meta = logging_obj.model_call_details["litellm_params"]["metadata"] + assert meta["hidden_params"]["response_cost"] == 0.002 + assert meta["hidden_params"]["model_id"] == "mid-test" + assert response._hidden_params["response_cost"] is None + + +def test_standard_logging_hidden_params_backfills_response_cost_without_mutating_response(): + """Streaming standard logging payload should expose the calculated response cost.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import Usage + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="standard-hp-cost-test", + function_id="standard-hp-cost-fn", + ) + logging_obj.model_call_details = { + "litellm_params": {"metadata": {}, "proxy_server_request": {}}, + "litellm_call_id": "standard-hp-cost-test", + "call_type": "acompletion", + "stream": True, + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "optional_params": {"stream": True}, + "response_cost": 0.002, + } + response = ModelResponse( + id="standard-hp-cost-response", + model="gpt-4o-mini", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + response._hidden_params = {"response_cost": None, "model_id": "mid-test"} + + payload = logging_obj._build_standard_logging_payload( + response, datetime.now(), datetime.now() + ) + + assert payload is not None + assert payload["hidden_params"]["response_cost"] == 0.002 + assert response._hidden_params["response_cost"] is None + + +def test_merge_hidden_params_from_response_into_metadata_preserves_response_cost(): + """Do not overwrite provider-supplied response cost when it already exists.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="merge-hp-preserve-cost-test", + function_id="merge-hp-preserve-cost-fn", + ) + logging_obj.model_call_details = { + "litellm_params": {"metadata": {}}, + "response_cost": 0.002, + } + + class _Resp: + _hidden_params = {"response_cost": 0.001, "model_id": "mid-test"} + + logging_obj._merge_hidden_params_from_response_into_metadata(_Resp()) + meta = logging_obj.model_call_details["litellm_params"]["metadata"] + assert meta["hidden_params"]["response_cost"] == 0.001 + assert meta["hidden_params"]["model_id"] == "mid-test" + + def test_merge_hidden_params_from_response_into_metadata_no_op_when_empty(): from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -2436,3 +2534,141 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload is not None assert payload["litellm_call_id"] == call_id + + +def _make_dict_logging_obj(): + """Build a Logging instance configured for a non-streaming dict result.""" + obj = LitellmLogging( + model="claude-haiku-4-5@20251001", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + litellm_call_id="test-call-id", + start_time=time.time(), + function_id="test-fn", + ) + obj.model_call_details = { + "model": "claude-haiku-4-5@20251001", + "custom_llm_provider": "vertex_ai", + "litellm_params": {"metadata": {}}, + "response_cost": None, + } + return obj + + +def test_success_handler_computes_cost_for_dict_response(): + """Non-streaming dict responses run through the cost calculator.""" + logging_obj = _make_dict_logging_obj() + expected_cost = 0.42 + with ( + patch.object( + logging_obj, + "_response_cost_calculator", + return_value=expected_cost, + ) as mock_calc, + patch.object( + logging_obj, + "_build_standard_logging_payload", + return_value={"response_cost": expected_cost}, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), + patch.object( + logging_obj, + "_is_recognized_call_type_for_logging", + return_value=False, + ), + patch.object( + logging_obj, + "_transform_usage_objects", + side_effect=lambda result: result, + ), + ): + logging_obj.success_handler( + result={"id": "msg_1"}, + start_time=time.time(), + end_time=time.time(), + ) + mock_calc.assert_called_once() + assert logging_obj.model_call_details["response_cost"] == expected_cost + + +def test_success_handler_preserves_precomputed_cost_for_dict_response(): + """Precomputed response_cost on model_call_details must not be overwritten.""" + logging_obj = _make_dict_logging_obj() + precomputed_cost = 1.23 + logging_obj.model_call_details["response_cost"] = precomputed_cost + with ( + patch.object( + logging_obj, + "_response_cost_calculator", + return_value=9.99, + ) as mock_calc, + patch.object( + logging_obj, + "_build_standard_logging_payload", + return_value={"response_cost": precomputed_cost}, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), + patch.object( + logging_obj, + "_is_recognized_call_type_for_logging", + return_value=False, + ), + patch.object( + logging_obj, + "_transform_usage_objects", + side_effect=lambda result: result, + ), + ): + logging_obj.success_handler( + result={"id": "msg_2"}, + start_time=time.time(), + end_time=time.time(), + ) + mock_calc.assert_not_called() + assert logging_obj.model_call_details["response_cost"] == precomputed_cost + + +def test_success_handler_unified_helper_runs_for_typed_results(): + """Recognized typed responses still flow through the unified helper.""" + logging_obj = _make_dict_logging_obj() + expected_cost = 0.10 + typed_result = MagicMock() + typed_result._hidden_params = {} + + with ( + patch.object( + logging_obj, + "_response_cost_calculator", + return_value=expected_cost, + ) as mock_calc, + patch.object( + logging_obj, + "_build_standard_logging_payload", + return_value={"response_cost": expected_cost}, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.emit_standard_logging_payload" + ), + patch.object( + logging_obj, + "_is_recognized_call_type_for_logging", + return_value=True, + ), + patch.object( + logging_obj, + "_transform_usage_objects", + side_effect=lambda result: result, + ), + ): + logging_obj.success_handler( + result=typed_result, + start_time=time.time(), + end_time=time.time(), + ) + mock_calc.assert_called_once() + assert logging_obj.model_call_details["response_cost"] == expected_cost diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 38a59c694e..633aed38a0 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -279,7 +279,7 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): } optional_params = config.map_openai_params( - model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", non_default_params=non_default_params, optional_params={}, drop_params=False, @@ -2797,7 +2797,7 @@ def test_thinking_with_max_completion_tokens(): result = config.map_openai_params( non_default_params=non_default_params_with_max_completion, optional_params=optional_params, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) @@ -2819,7 +2819,7 @@ def test_thinking_with_max_completion_tokens(): result = config.map_openai_params( non_default_params=non_default_params_with_max_tokens, optional_params=optional_params, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) @@ -2842,7 +2842,7 @@ def test_thinking_with_max_completion_tokens(): result = config.map_openai_params( non_default_params=non_default_params_without_max, optional_params=optional_params, - model="us.anthropic.claude-3-7-sonnet-20250219-v1:0", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) @@ -3617,7 +3617,7 @@ class TestBedrockMinThinkingBudgetTokens: """Test that thinking.budget_tokens is clamped to the Bedrock minimum (1024).""" def _map_params( - self, thinking_value, model="anthropic.claude-3-7-sonnet-20250219-v1:0" + self, thinking_value, model="anthropic.claude-sonnet-4-5-20250929-v1:0" ): """Helper to call map_openai_params with the given thinking value.""" config = AmazonConverseConfig() @@ -3651,7 +3651,7 @@ class TestBedrockMinThinkingBudgetTokens: result = config.map_openai_params( non_default_params={}, optional_params={}, - model="anthropic.claude-3-7-sonnet-20250219-v1:0", + model="anthropic.claude-sonnet-4-5-20250929-v1:0", drop_params=False, ) assert "thinking" not in result or result.get("thinking") is None @@ -4146,3 +4146,39 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools( # finish_reason must be "stop", not "tool_calls" assert result.choices[0].finish_reason == "stop" + + +def test_transform_response_does_not_leak_body_on_parse_failure(): + from litellm.llms.bedrock.common_utils import BedrockError + + leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}} + + class MockResponse: + def json(self): + return leaky_body + + @property + def text(self): + return json.dumps(leaky_body) + + with patch( + "litellm.llms.bedrock.chat.converse_transformation.ConverseResponseBlock", + side_effect=KeyError("missing required field"), + ): + with pytest.raises(BedrockError) as exc_info: + AmazonConverseConfig()._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + msg = str(exc_info.value) + assert "secret content" not in msg + assert "Error converting to valid response block" in msg diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index bcd4c62005..d60d0487d0 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -21,7 +21,112 @@ def test_transform_search_request(): api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=None, ) assert url.endswith("/kb123/retrieve") assert body["retrievalQuery"].get("text") == "hello" + + +def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + + url, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + extra_body={ + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "overrideSearchType": "HYBRID", + "numberOfResults": 8, + } + }, + "unrelatedField": {"should_not": "be_forwarded"}, + }, + ) + + assert url.endswith("/kb123/retrieve") + assert body["retrievalQuery"].get("text") == "hello" + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"][ + "overrideSearchType" + ] + == "HYBRID" + ) + assert "unrelatedField" not in body + + +def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + extra_body = { + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "overrideSearchType": "HYBRID", + "numberOfResults": 8, + } + } + } + + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"max_num_results": 10}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + extra_body=extra_body, + ) + + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"] + == 10 + ) + assert ( + extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][ + "numberOfResults" + ] + == 8 + ) + + +def test_transform_search_request_overrides_filter_without_mutating_extra_body(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + extra_body = { + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "filter": {"equals": {"key": "tenant", "value": "a"}} + } + } + } + new_filter = {"equals": {"key": "tenant", "value": "b"}} + + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"filters": new_filter}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + extra_body=extra_body, + ) + + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"] + == new_filter + ) + assert ( + extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][ + "equals" + ]["value"] + == "a" + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py new file mode 100644 index 0000000000..5515e6ce81 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py @@ -0,0 +1,161 @@ +import socket +from unittest.mock import MagicMock, patch + + +def _invoke_connector_factory(http_handler_module): + """ + Drive the lambda factory installed on the transport so TCPConnector is + actually constructed. _create_aiohttp_transport returns a transport whose + _client_factory is the lambda that builds (TCPConnector → ClientSession); + invoking it directly avoids relying on _get_valid_client_session's internal + branching to trigger connector construction. + """ + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( + shared_session=None + ) + transport._client_factory() + return transport + + +def test_socket_factory_omitted_when_disabled(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + + with patch.object( + http_handler_module, "TCPConnector", return_value=connector_mock + ) as mock_tcp_connector: + with patch.object( + http_handler_module, "ClientSession", return_value=session_mock + ): + _invoke_connector_factory(http_handler_module) + + assert mock_tcp_connector.call_count >= 1 + assert "socket_factory" not in mock_tcp_connector.call_args.kwargs + + +def test_socket_factory_attached_when_enabled(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + + with patch.object( + http_handler_module, "TCPConnector", return_value=connector_mock + ) as mock_tcp_connector: + with patch.object( + http_handler_module, "ClientSession", return_value=session_mock + ): + _invoke_connector_factory(http_handler_module) + + assert mock_tcp_connector.call_count >= 1 + factory = mock_tcp_connector.call_args.kwargs.get("socket_factory") + assert callable(factory) + + +def test_socket_factory_skipped_on_old_aiohttp(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False) + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + + with patch.object( + http_handler_module, "TCPConnector", return_value=connector_mock + ) as mock_tcp_connector: + with patch.object( + http_handler_module, "ClientSession", return_value=session_mock + ): + _invoke_connector_factory(http_handler_module) + + assert mock_tcp_connector.call_count >= 1 + assert "socket_factory" not in mock_tcp_connector.call_args.kwargs + + +def test_socket_factory_sets_keepalive_options(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4) + + factory = http_handler_module._build_aiohttp_keepalive_socket_factory() + assert factory is not None + + addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0)) + + fake_sock = MagicMock(spec=socket.socket) + with patch("socket.socket", return_value=fake_sock) as sock_ctor: + returned = factory(addr_info) + + sock_ctor.assert_called_once_with( + family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP + ) + assert returned is fake_sock + fake_sock.setblocking.assert_called_once_with(False) + + setsockopt_calls = { + (call.args[0], call.args[1]): call.args[2] + for call in fake_sock.setsockopt.call_args_list + } + assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1 + + if hasattr(socket, "TCP_KEEPIDLE"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45 + elif hasattr(socket, "TCP_KEEPALIVE"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45 + if hasattr(socket, "TCP_KEEPINTVL"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15 + if hasattr(socket, "TCP_KEEPCNT"): + assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4 + + +def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch): + """ + Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE + is present, the factory should fall back to TCP_KEEPALIVE for the idle timer. + Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to + simulate the BSD-derived environment. + """ + from litellm.llms.custom_httpx import http_handler as http_handler_module + + monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True) + monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True) + monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60) + + factory = http_handler_module._build_aiohttp_keepalive_socket_factory() + assert factory is not None + + fake_socket_module = MagicMock(spec=[]) + fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET + fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE + fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP + fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10) + fake_sock = MagicMock(spec=socket.socket) + fake_socket_module.socket = MagicMock(return_value=fake_sock) + + addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0)) + + with patch.object(http_handler_module, "socket", fake_socket_module): + factory(addr_info) + + setsockopt_calls = { + (call.args[0], call.args[1]): call.args[2] + for call in fake_sock.setsockopt.call_args_list + } + assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1 + assert ( + setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60 + ) + assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 6924eb8d3d..752b5ff090 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -2,12 +2,16 @@ import os import sys from unittest.mock import AsyncMock, Mock, patch +import httpx import pytest sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path -from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.custom_httpx.llm_http_handler import ( + BaseLLMHTTPHandler, + _google_genai_streaming_hidden_params, +) from litellm.types.router import GenericLiteLLMParams @@ -320,3 +324,29 @@ async def test_async_anthropic_messages_handler_header_priority(): assert captured_headers["X-Forwarded-Only"] == "keep" assert captured_headers["X-Extra-Only"] == "also-keep" assert captured_headers["X-Provider-Only"] == "keep-this-too" + + +def test_google_genai_streaming_hidden_params_model_info_and_router_fallback(): + logging_obj = Mock() + logging_obj.get_router_model_id = Mock(return_value="router-model-id") + + from_model_info = _google_genai_streaming_hidden_params( + api_base="https://generativelanguage.googleapis.com/v1beta", + litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}), + logging_obj=logging_obj, + response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}), + ) + assert from_model_info["model_id"] == "info-id" + assert ( + from_model_info["api_base"] + == "https://generativelanguage.googleapis.com/v1beta" + ) + assert isinstance(from_model_info["additional_headers"], dict) + + from_router = _google_genai_streaming_hidden_params( + api_base="https://x", + litellm_params=GenericLiteLLMParams(), + logging_obj=logging_obj, + response_headers=httpx.Headers({}), + ) + assert from_router["model_id"] == "router-model-id" diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 7c06823d16..7085e45cdc 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -59,6 +59,7 @@ class TestS3VectorsVectorStoreConfig: api_base="https://s3vectors.us-west-2.api.aws", litellm_logging_obj=mock_logging_obj, litellm_params={}, + extra_body=None, ) def test_transform_search_response(self): diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0118b39ddd..353d19b019 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -4261,3 +4261,32 @@ def test_sync_streaming_uses_custom_client(): # Verify that gemini_client is in the partial's keywords assert "gemini_client" in partial_make_sync_call.keywords assert partial_make_sync_call.keywords["gemini_client"] is mock_client + + +def test_transform_response_does_not_leak_body_on_parse_failure(): + leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]} + raw_response = MagicMock() + raw_response.json.return_value = leaky_body + raw_response.text = json.dumps(leaky_body) + raw_response.headers = {} + + with patch( + "litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody", + side_effect=KeyError("missing required field"), + ): + with pytest.raises(VertexAIError) as exc_info: + VertexGeminiConfig().transform_response( + model="gemini-pro", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + msg = str(exc_info.value) + assert "secret content" not in msg + assert "Error converting to valid response block" in msg diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index ef93375c3c..95507390df 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -225,7 +225,11 @@ def test_build_vertex_schema(): "metadata": {"type": "object"}, "callbacks": { "anyOf": [ - {"type": "array", "nullable": True}, + { + "type": "array", + "items": {"type": "object"}, + "nullable": True, + }, {"type": "object", "nullable": True}, ] }, @@ -288,6 +292,43 @@ def test_process_items_basic(): process_items(schema) assert schema["properties"]["nested"]["items"] == {"type": "object"} + # Vertex rejects array types missing `items` entirely (not just empty). + # Synthesize {"type": "object"} so the request validates. + schema = {"type": "array"} + process_items(schema) + assert schema["items"] == {"type": "object"} + + +def test_build_vertex_schema_array_branch_missing_items_in_anyof(): + """ + Regression: an `anyOf` branch with `{"type": "array"}` (no items) must + end up with synthesized `items: {"type": "object"}` after the schema + transform — Vertex returns INVALID_ARGUMENT otherwise. + """ + from litellm.llms.vertex_ai.common_utils import _build_vertex_schema + + parameters = { + "properties": { + "callbacks": { + "anyOf": [ + {"type": "array"}, + {"type": "object"}, + {"type": "null"}, + ] + } + }, + "type": "object", + } + + result = _build_vertex_schema(parameters) + callbacks_anyof = result["properties"]["callbacks"]["anyOf"] + array_branches = [b for b in callbacks_anyof if b.get("type") == "array"] + assert array_branches, "expected an array branch to remain after transform" + for branch in array_branches: + assert branch.get("items") == { + "type": "object" + }, f"array branch must have items synthesized; got {branch}" + def test_vertex_ai_complex_response_schema(): import json diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index 214e5f0797..b8cd65d3c9 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -311,3 +311,29 @@ def test_transform_anthropic_messages_request_removes_scope_from_cache_control() # scope removed from message content assert "scope" not in result["messages"][0]["content"][0]["cache_control"] assert result["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" + + +def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance(): + """ + Regression test: repeated provider config lookups for the same Vertex Claude model + should return the same config instance (which preserves auth cache state). + """ + import litellm + from litellm.utils import ProviderConfigManager + + ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() + try: + first_config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="claude-opus-4-6", + provider=litellm.LlmProviders.VERTEX_AI, + ) + second_config = ProviderConfigManager.get_provider_anthropic_messages_config( + model="claude-opus-4-6", + provider=litellm.LlmProviders.VERTEX_AI, + ) + + assert isinstance(first_config, VertexAIPartnerModelsAnthropicMessagesConfig) + assert isinstance(second_config, VertexAIPartnerModelsAnthropicMessagesConfig) + assert first_config is second_config + finally: + ProviderConfigManager._get_provider_anthropic_messages_config_cached.cache_clear() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 447c28078e..8614cb3ba7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -728,6 +728,136 @@ class TestMCPServerManager: ] assert scopes == ["read", "write"] + @pytest.mark.asyncio + async def test_descovery_metadata_probes_well_known_when_server_does_not_challenge( + self, + ): + manager = MCPServerManager() + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + mock_metadata = MCPOAuthMetadata( + scopes=None, + authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize", + token_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/token", + registration_url=None, + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ), + patch.object( + manager, + "_attempt_well_known_discovery", + AsyncMock( + return_value=( + ["https://login.microsoftonline.com/test-tenant-id/v2.0"], + ["api://some-scope/.default"], + ) + ), + ) as mock_well_known, + patch.object( + manager, + "_fetch_authorization_server_metadata", + AsyncMock(return_value=mock_metadata), + ) as mock_fetch_auth, + ): + result = await manager._descovery_metadata("http://localhost:8001/mcp") + + mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp") + mock_fetch_auth.assert_awaited_once_with( + ["https://login.microsoftonline.com/test-tenant-id/v2.0"] + ) + assert result is mock_metadata + assert result.scopes == ["api://some-scope/.default"] + + @pytest.mark.asyncio + async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path( + self, + ): + manager = MCPServerManager() + issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0" + + def build_response(url: str): + mock_response = MagicMock() + if url == f"{issuer}/.well-known/openid-configuration": + mock_response.json.return_value = { + "authorization_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize", + "token_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token", + "scopes_supported": ["api://some-scope/.default"], + } + mock_response.raise_for_status = MagicMock() + else: + request = httpx.Request("GET", url) + response_obj = httpx.Response(status_code=404, request=request) + mock_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError( + "not found", request=request, response=response_obj + ) + ) + return mock_response + + mock_client = MagicMock() + mock_client.get = AsyncMock(side_effect=build_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata(issuer) + + assert result is not None + assert ( + result.authorization_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" + ) + assert ( + result.token_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" + ) + assert result.scopes == ["api://some-scope/.default"] + + @pytest.mark.asyncio + async def test_fetch_single_authorization_server_metadata_derives_azure_metadata( + self, + ): + manager = MCPServerManager() + issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0" + + request = httpx.Request("GET", issuer) + response_obj = httpx.Response(status_code=404, request=request) + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock( + side_effect=httpx.HTTPStatusError( + "not found", request=request, response=response_obj + ) + ) + + mock_client = MagicMock() + mock_client.get = AsyncMock(return_value=mock_response) + + with patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client", + return_value=mock_client, + ): + result = await manager._fetch_single_authorization_server_metadata(issuer) + + assert result is not None + assert ( + result.authorization_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize" + ) + assert ( + result.token_url + == "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token" + ) + @pytest.mark.asyncio async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self): manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py new file mode 100644 index 0000000000..bf90e9ebef --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_short_mcp_tool_prefix.py @@ -0,0 +1,302 @@ +""" +Tests for the short-ID MCP tool prefix (LITELLM_USE_SHORT_MCP_TOOL_PREFIX). + +The short-prefix mode swaps the historical alias/server_name prefix on +tool names for a deterministic three-character base62 ID derived from the +server's ``server_id``. This keeps tool names well below the 60-char +upper bound enforced by some model APIs while remaining stable across +processes/restarts and tolerant of mixed-version clients. +""" + +from typing import List + +import pytest +from mcp.types import Tool as MCPTool + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager +from litellm.proxy._experimental.mcp_server.utils import ( + SHORT_MCP_TOOL_PREFIX_LENGTH, + add_server_prefix_to_name, + compute_short_server_prefix, + get_server_prefix, + is_short_mcp_tool_prefix_enabled, + iter_known_server_prefixes, +) +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +def _make_server( + *, + server_id: str = "abcdef-1234", + server_name: str = "github_onprem", + alias: str = "github_onprem", +) -> MCPServer: + return MCPServer( + server_id=server_id, + name=alias or server_name, + alias=alias, + server_name=server_name, + transport="http", + ) + + +@pytest.fixture(autouse=True) +def _reset_env(monkeypatch): + monkeypatch.delenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", raising=False) + yield + + +# --------------------------------------------------------------------------- +# Pure helpers +# --------------------------------------------------------------------------- + + +class TestShortPrefixHelpers: + def test_short_prefix_is_three_base62_chars(self): + prefix = compute_short_server_prefix("any-server-id") + assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH + assert prefix.isalnum() and prefix.isascii() + + def test_short_prefix_first_char_is_alphabetic(self): + """The first char must be [A-Za-z] so the prefix is a valid identifier + on every model API (some providers historically required the first + character of a function name to be alphabetic).""" + # Sweep many server_ids and rehash attempts to give us coverage of + # every position the high-order bits can land on. + for i in range(200): + for attempt in range(4): + prefix = compute_short_server_prefix(f"server-{i}", attempt=attempt) + assert prefix[0].isalpha(), ( + f"prefix {prefix!r} for server-{i} (attempt={attempt}) " + f"starts with a non-alphabetic character" + ) + + def test_short_prefix_is_deterministic(self): + assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc") + assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd") + + def test_short_prefix_requires_server_id(self): + with pytest.raises(ValueError): + compute_short_server_prefix("") + + def test_flag_defaults_to_false(self): + assert is_short_mcp_tool_prefix_enabled() is False + + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "On"]) + def test_flag_truthy_values(self, monkeypatch, value): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value) + assert is_short_mcp_tool_prefix_enabled() is True + + @pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) + def test_flag_falsey_values(self, monkeypatch, value): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value) + assert is_short_mcp_tool_prefix_enabled() is False + + +# --------------------------------------------------------------------------- +# get_server_prefix behaviour +# --------------------------------------------------------------------------- + + +class TestGetServerPrefix: + def test_default_mode_uses_alias(self): + server = _make_server(alias="github_onprem", server_name="github_onprem") + assert get_server_prefix(server) == "github_onprem" + + def test_short_mode_uses_short_id(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = _make_server(server_id="abcdef-1234") + prefix = get_server_prefix(server) + assert prefix == compute_short_server_prefix("abcdef-1234") + assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH + + def test_short_mode_falls_back_when_no_server_id(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + + class _Bare: + alias = "fallback_alias" + server_name = None + server_id = None + + assert get_server_prefix(_Bare()) == "fallback_alias" + + +# --------------------------------------------------------------------------- +# iter_known_server_prefixes — covers reverse-lookup tolerance +# --------------------------------------------------------------------------- + + +class TestIterKnownServerPrefixes: + def test_default_mode_includes_short_id_too(self): + server = _make_server() + prefixes = list(iter_known_server_prefixes(server)) + # Contains the live prefix and every known form so that mixed-mode + # clients can be resolved. + assert "github_onprem" in prefixes + assert compute_short_server_prefix(server.server_id) in prefixes + + def test_short_mode_still_yields_long_forms(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = _make_server() + prefixes = list(iter_known_server_prefixes(server)) + assert "github_onprem" in prefixes + assert compute_short_server_prefix(server.server_id) in prefixes + + +# --------------------------------------------------------------------------- +# Manager-level behaviour: list + reverse-lookup +# --------------------------------------------------------------------------- + + +def _stub_tools() -> List[MCPTool]: + return [ + MCPTool(name="get_repo", description="", inputSchema={"type": "object"}), + MCPTool(name="list_issues", description="", inputSchema={"type": "object"}), + ] + + +class TestManagerShortPrefix: + def test_list_tools_uses_short_prefix_when_flag_on(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server() + + out = manager._create_prefixed_tools(_stub_tools(), server) + + short = compute_short_server_prefix(server.server_id) + assert {t.name for t in out} == {f"{short}-get_repo", f"{short}-list_issues"} + + def test_call_tool_lookup_resolves_short_prefix(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server() + manager.registry[server.server_id] = server + manager._create_prefixed_tools(_stub_tools(), server) + + short = compute_short_server_prefix(server.server_id) + resolved = manager._get_mcp_server_from_tool_name(f"{short}-get_repo") + assert resolved is server + + def test_call_tool_lookup_resolves_long_prefix_in_short_mode(self, monkeypatch): + """Old clients that cached the long-prefix name must still route.""" + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server() + manager.registry[server.server_id] = server + manager._create_prefixed_tools(_stub_tools(), server) + + resolved = manager._get_mcp_server_from_tool_name("github_onprem-get_repo") + assert resolved is server + + def test_default_mode_unchanged(self): + manager = MCPServerManager() + server = _make_server() + + out = manager._create_prefixed_tools(_stub_tools(), server) + + assert {t.name for t in out} == { + "github_onprem-get_repo", + "github_onprem-list_issues", + } + assert ( + manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is None + ) # registry empty + manager.registry[server.server_id] = server + assert ( + manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is server + ) + + def test_total_tool_name_length_short_enough(self, monkeypatch): + """The short prefix keeps tool names under the 60-char limit even + when the upstream tool name is itself reasonably long.""" + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + long_server_name = "a" * 50 + server = _make_server( + server_id="server-id-1", + server_name=long_server_name, + alias=long_server_name, + ) + prefix = get_server_prefix(server) + full = add_server_prefix_to_name("get_repo", prefix) + assert len(full) < 60 + + +# --------------------------------------------------------------------------- +# Collision-resolution at registration time +# --------------------------------------------------------------------------- + + +class TestShortPrefixCollisionResolution: + """``_assign_unique_short_prefix`` must rehash on collision. + + The dedup path is exercised by forcing two distinct ``server_id`` + values to both hash to the same natural prefix via a monkeypatched + ``compute_short_server_prefix``. + """ + + def test_no_op_when_flag_off(self): + manager = MCPServerManager() + server = _make_server(server_id="abc") + manager._assign_unique_short_prefix(server) + assert server.short_prefix is None + + def test_assigns_natural_hash_when_no_collision(self, monkeypatch): + from litellm.proxy._experimental.mcp_server import utils as mcp_utils + + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server(server_id="abc") + manager._assign_unique_short_prefix(server) + + assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc") + + def test_rehashes_when_natural_hash_collides(self, monkeypatch): + """Two server_ids that natural-hash to the same prefix get + deterministic, distinct short prefixes.""" + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + + # Force every attempt=0 hash to "AAA" and attempt=1 to "AAB". + # That way the second server registered must rehash to "AAB". + from litellm.proxy._experimental.mcp_server import utils as mcp_utils + + def _fake_hash(server_id: str, attempt: int = 0) -> str: + return "AAA" if attempt == 0 else f"AA{chr(ord('A') + attempt)}" + + monkeypatch.setattr(mcp_utils, "compute_short_server_prefix", _fake_hash) + # Also patch the symbol that the manager imported at module load. + from litellm.proxy._experimental.mcp_server import ( + mcp_server_manager as mgr_module, + ) + + monkeypatch.setattr(mgr_module, "compute_short_server_prefix", _fake_hash) + + manager = MCPServerManager() + first = _make_server(server_id="server-1", alias="srv1") + second = _make_server(server_id="server-2", alias="srv2") + + # Pretend both are already in the registry so dedup sees both. + manager.registry[first.server_id] = first + manager._assign_unique_short_prefix(first) + manager.registry[second.server_id] = second + manager._assign_unique_short_prefix(second) + + assert first.short_prefix == "AAA" + assert second.short_prefix == "AAB" + assert first.short_prefix != second.short_prefix + + def test_cached_prefix_is_reused(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + manager = MCPServerManager() + server = _make_server(server_id="abc") + server.short_prefix = "ZZZ" # pretend a previous registration set this + + manager._assign_unique_short_prefix(server) + + assert server.short_prefix == "ZZZ" + + def test_get_server_prefix_prefers_cached(self, monkeypatch): + monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true") + server = _make_server(server_id="abc") + server.short_prefix = "Q9q" + + assert get_server_prefix(server) == "Q9q" diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 15cfc84c23..91f300b88c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -5,6 +5,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext from typing import Optional from unittest.mock import MagicMock, patch +import pytest + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( _get_customer_id_from_standard_headers, @@ -15,6 +17,7 @@ from litellm.proxy.auth.auth_utils import ( get_key_model_tpm_limit, get_project_model_rpm_limit, get_project_model_tpm_limit, + is_request_body_safe, ) @@ -660,3 +663,304 @@ class TestCheckCompleteCredentials: def test_returns_true_when_api_key_is_valid(self): result = check_complete_credentials({"model": "gpt-4", "api_key": "sk-valid"}) assert result is True + + +class TestCheckCompleteCredentialsBlocksSSRF: + """ + Even with credentials supplied, ``api_base`` / ``base_url`` must not + point at private / internal / cloud-metadata addresses. Without this + the gate accepts ``api_key=anything`` plus a malicious target and the + proxy is used as an SSRF pivot. + + The check only runs when ``litellm.user_url_validation`` is True, so + every test in this class flips the toggle. Tests stay mock-only — no + real DNS is performed. + """ + + @pytest.fixture(autouse=True) + def _enable_url_validation(self, monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", True, raising=False) + + @pytest.mark.parametrize( + "url_field", + ["api_base", "base_url"], + ) + @pytest.mark.parametrize( + "blocked_url", + [ + "http://169.254.169.254/latest/meta-data/iam/security-credentials/", + "http://metadata.google.internal/computeMetadata/v1/", + "http://127.0.0.1:8080/admin", + "http://10.0.0.1/", + "http://192.168.1.1/", + ], + ) + def test_rejects_private_or_metadata_targets(self, url_field, blocked_url): + from litellm.litellm_core_utils.url_utils import SSRFError + + with patch( + "litellm.proxy.auth.auth_utils.validate_url", + side_effect=SSRFError(f"blocked: {blocked_url}"), + ): + with pytest.raises(ValueError) as exc_info: + check_complete_credentials( + { + "model": "gpt-4", + "api_key": "sk-some-clientside-key", + url_field: blocked_url, + } + ) + assert url_field in str(exc_info.value) + assert "SSRF" in str(exc_info.value) + + def test_allows_public_target_when_validate_url_passes(self): + # ``validate_url`` is mocked so no real DNS is performed. + with patch( + "litellm.proxy.auth.auth_utils.validate_url", + return_value=("https://api.openai.com/v1", "api.openai.com"), + ): + result = check_complete_credentials( + { + "model": "gpt-4", + "api_key": "sk-some-clientside-key", + "api_base": "https://api.openai.com/v1", + } + ) + assert result is True + + def test_skips_url_validation_when_toggle_is_off(self, monkeypatch): + # Admins who disable ``user_url_validation`` (default) should not + # have requests rejected at the proxy boundary even if the URL + # would fail the SSRF guard. + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + with patch( + "litellm.proxy.auth.auth_utils.validate_url", + ) as mocked: + result = check_complete_credentials( + { + "model": "gpt-4", + "api_key": "sk-some-clientside-key", + "api_base": "http://127.0.0.1:8080/admin", + } + ) + assert result is True + mocked.assert_not_called() + + +class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride: + """ + When the caller redirects ``api_base`` / ``base_url`` to their own + server, admin-set fields like ``OpenAI-Organization``, ``extra_body``, + AWS / Vertex / Azure tokens, and per-deployment ``api_version`` must + NOT flow through to that destination. + """ + + def test_clears_admin_organization_and_extra_body_on_base_override(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + admin_params = { + "model": "gpt-4", + "api_key": "sk-admin-key", + "api_base": "https://admin.upstream/v1", + "organization": "org-admin-corp", + "extra_body": {"x-admin-secret": "super-secret"}, + "api_version": "2026-04-01", + } + out = get_dynamic_litellm_params( + litellm_params=dict(admin_params), + request_kwargs={ + "api_key": "sk-attacker", + "api_base": "https://attacker.example", + }, + ) + assert out["api_base"] == "https://attacker.example" + assert out["api_key"] == "sk-attacker" + assert "organization" not in out + assert "extra_body" not in out + assert "api_version" not in out + + def test_clears_aws_and_vertex_secrets_on_base_override(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + admin_params = { + "model": "bedrock/claude-3", + "aws_access_key_id": "AKIA-EXAMPLE", + "aws_secret_access_key": "secret-example", + "aws_session_token": "session-example", + "vertex_credentials": '{"private_key":"-----BEGIN..."}', + "vertex_project": "admin-gcp-project", + } + out = get_dynamic_litellm_params( + litellm_params=dict(admin_params), + request_kwargs={"base_url": "https://attacker.example"}, + ) + assert "aws_access_key_id" not in out + assert "aws_secret_access_key" not in out + assert "aws_session_token" not in out + assert "vertex_credentials" not in out + assert "vertex_project" not in out + + def test_caller_resupplied_value_overrides_admin_value_on_base_override(self): + # When the caller redirects ``api_base`` and *also* supplies their + # own value for one of the admin fields (e.g. ``organization``), + # the caller's value must win — never the admin's. The naive + # ``if field not in request_kwargs: pop`` shape lets a caller echo + # the field name with any value (or empty string) to keep the + # admin's value forwarded, which is the exfiltration vector this + # test guards against. + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + out = get_dynamic_litellm_params( + litellm_params={ + "api_base": "https://admin.upstream/v1", + "organization": "org-admin", + "extra_body": {"admin": "value"}, + }, + request_kwargs={ + "api_base": "https://attacker.example", + "organization": "org-attacker", + "extra_body": {"attacker": "value"}, + }, + ) + assert out["organization"] == "org-attacker" + assert out["extra_body"] == {"attacker": "value"} + + def test_field_echo_does_not_preserve_admin_value(self): + # Regression: a caller that echoes an admin-config field name with + # an *empty* value (or any value) must not be able to keep the + # admin's value in ``litellm_params``. + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + out = get_dynamic_litellm_params( + litellm_params={ + "api_base": "https://admin.upstream/v1", + "organization": "org-admin-secret", + "extra_body": {"x-admin-only": "secret"}, + }, + request_kwargs={ + "api_base": "https://attacker.example", + "organization": "", + "extra_body": "", + }, + ) + assert out["organization"] == "" + assert out["extra_body"] == "" + assert "org-admin-secret" not in str(out) + + def test_no_clearing_when_only_api_key_overridden(self): + from litellm.router_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, + ) + + # Caller only overrides api_key (BYOK pattern); admin's organization / + # extra_body / region still apply because the destination is unchanged. + out = get_dynamic_litellm_params( + litellm_params={ + "api_base": "https://admin.upstream/v1", + "organization": "org-admin", + "api_version": "2026-04-01", + }, + request_kwargs={"api_key": "sk-byok"}, + ) + assert out["organization"] == "org-admin" + assert out["api_version"] == "2026-04-01" + assert out["api_base"] == "https://admin.upstream/v1" + + +class TestIsRequestBodySafeBlocksEndpointTargetingFields: + """ + ``is_request_body_safe`` rejects request-body fields that retarget the + outbound request to a caller-controlled host. Beyond the original + ``api_base`` / ``base_url``, the same protection must apply to: + + * ``aws_bedrock_runtime_endpoint`` — Bedrock endpoint redirect; an + attacker-controlled value coerces the proxy to authenticate against + their host with the admin's AWS creds. + * ``langsmith_base_url`` — Langsmith callback host; attacker-controlled + values exfiltrate the entire request payload (incl. message content) + via the observability hook. + * ``langfuse_host`` — same exfil vector via the Langfuse hook. + """ + + @pytest.fixture(autouse=True) + def _disable_url_validation(self, monkeypatch): + # The new banned-params entries should be rejected even when + # ``user_url_validation`` is off — the gate isn't the URL guard, + # it's the banned-params list. + import litellm + + monkeypatch.setattr(litellm, "user_url_validation", False, raising=False) + + @pytest.mark.parametrize( + "field", + [ + "aws_bedrock_runtime_endpoint", + "langsmith_base_url", + "langfuse_host", + "posthog_host", + "braintrust_host", + "slack_webhook_url", + "s3_endpoint_url", + "sagemaker_base_url", + "deployment_url", + ], + ) + def test_endpoint_targeting_field_in_request_body_is_rejected(self, field): + with pytest.raises(ValueError) as exc: + is_request_body_safe( + request_body={"model": "gpt-4", field: "https://attacker.example"}, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + # The function lists the offending param name in the error. + assert field in str(exc.value) + + @pytest.mark.parametrize( + "field", + ["api_base", "base_url", "user_config", "langfuse_host", "slack_webhook_url"], + ) + def test_api_key_does_not_bypass_blocklist(self, field): + # Regression: the historical ``check_complete_credentials`` clause + # made the entire blocklist a no-op for any caller that supplied + # a non-empty ``api_key``. That bypass turned every missing entry + # on the blocklist into an SSRF / credential-exfil hole. Verify + # that supplying an api_key (alongside the banned param) does NOT + # bypass the gate — it can only be opened by an admin opt-in. + with pytest.raises(ValueError) as exc: + is_request_body_safe( + request_body={ + "model": "gpt-4", + "api_key": "sk-anything", + field: "https://attacker.example", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + assert field in str(exc.value) + + def test_admin_opt_in_proxy_wide_still_allows(self): + # ``general_settings.allow_client_side_credentials = True`` remains + # the documented proxy-wide BYOK opt-in. + assert ( + is_request_body_safe( + request_body={"model": "gpt-4", "api_base": "https://my-byok.example"}, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) diff --git a/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py new file mode 100644 index 0000000000..3efeeee9a2 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_expired_ui_session_key_cleanup_manager.py @@ -0,0 +1,348 @@ +""" +Test expired UI session key cleanup manager functionality. +""" + +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException, status + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.constants import ( + EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, + UI_SESSION_TOKEN_TEAM_ID, +) +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager import ( + ExpiredUISessionKeyCleanupManager, +) + + +class TestExpiredUISessionKeyCleanupManager: + """Test the ExpiredUISessionKeyCleanupManager class functionality.""" + + @pytest.mark.asyncio + async def test_find_expired_ui_session_keys_filters_dashboard_team_and_expiry(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + now = datetime(2026, 4, 25, 12, 0, 0, tzinfo=timezone.utc) + mock_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=now - timedelta(seconds=1), + ) + ] + mock_prisma_client.db.litellm_verificationtoken.find_many.return_value = ( + mock_keys + ) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.datetime" + ) as mock_datetime: + mock_datetime.now.return_value = now + mock_datetime.side_effect = lambda *args, **kwargs: datetime( + *args, **kwargs + ) + + keys = await manager._find_expired_ui_session_keys() + + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( + where={ + "team_id": UI_SESSION_TOKEN_TEAM_ID, + "expires": {"lt": now}, + }, + take=LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, + ) + assert keys == mock_keys + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_uses_existing_delete_path(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_key = LiteLLM_VerificationToken( + token="expired-dashboard-token", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key]) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + {"deleted_keys": ["expired-dashboard-token"], "failed_tokens": []}, + [expired_key], + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ) as mock_key_deleted_hook: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 1 + mock_delete_verification_tokens.assert_called_once() + call_kwargs = mock_delete_verification_tokens.call_args.kwargs + assert call_kwargs["tokens"] == ["expired-dashboard-token"] + assert call_kwargs["user_api_key_cache"] == mock_cache + assert ( + call_kwargs["litellm_changed_by"] + == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + ) + assert call_kwargs["user_api_key_dict"].user_id == "system" + mock_key_deleted_hook.assert_called_once() + hook_kwargs = mock_key_deleted_hook.call_args.kwargs + assert hook_kwargs["data"].keys == ["expired-dashboard-token"] + assert hook_kwargs["keys_being_deleted"] == [expired_key] + assert hook_kwargs["response"] == { + "deleted_keys": ["expired-dashboard-token"], + "failed_tokens": [], + } + assert ( + hook_kwargs["litellm_changed_by"] + == LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + ) + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_deletes_multiple_keys(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + LiteLLM_VerificationToken( + token="expired-dashboard-token-2", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + ] + tokens = [key.token for key in expired_keys] + manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + {"deleted_keys": tokens, "failed_tokens": []}, + expired_keys, + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ) as mock_key_deleted_hook: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 2 + assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens + hook_kwargs = mock_key_deleted_hook.call_args.kwargs + assert hook_kwargs["data"].keys == tokens + assert hook_kwargs["keys_being_deleted"] == expired_keys + assert hook_kwargs["response"] == {"deleted_keys": tokens, "failed_tokens": []} + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_returns_successful_delete_count(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + LiteLLM_VerificationToken( + token="expired-dashboard-token-2", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + ] + tokens = [key.token for key in expired_keys] + manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + { + "deleted_keys": ["expired-dashboard-token-1"], + "failed_tokens": ["expired-dashboard-token-2"], + }, + [expired_keys[0]], + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ): + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 1 + assert mock_delete_verification_tokens.call_args.kwargs["tokens"] == tokens + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_counts_nested_delete_response(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_keys = [ + LiteLLM_VerificationToken( + token="expired-dashboard-token-1", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + LiteLLM_VerificationToken( + token="expired-dashboard-token-2", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ), + ] + tokens = [key.token for key in expired_keys] + manager._find_expired_ui_session_keys = AsyncMock(return_value=expired_keys) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.return_value = ( + { + "deleted_keys": {"deleted_keys": 2}, + "failed_tokens": tokens, + }, + expired_keys, + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ): + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 2 + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_treats_missing_keys_as_noop(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + expired_key = LiteLLM_VerificationToken( + token="expired-dashboard-token", + team_id=UI_SESSION_TOKEN_TEAM_ID, + expires=datetime.now(timezone.utc) - timedelta(seconds=1), + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[expired_key]) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + mock_delete_verification_tokens.side_effect = HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": "No keys found"}, + ) + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.KeyManagementEventHooks.async_key_deleted_hook", + new_callable=AsyncMock, + ) as mock_key_deleted_hook: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_key_deleted_hook.assert_not_called() + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_noops_when_no_keys_found(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[]) + + with patch( + "litellm.proxy.common_utils.expired_ui_session_key_cleanup_manager.delete_verification_tokens", + new_callable=AsyncMock, + ) as mock_delete_verification_tokens: + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_delete_verification_tokens.assert_not_called() + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_skips_when_lock_held(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + pod_lock_manager=mock_pod_lock_manager, + ) + manager._find_expired_ui_session_keys = AsyncMock() + + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_pod_lock_manager.acquire_lock.assert_called_once_with( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) + manager._find_expired_ui_session_keys.assert_not_called() + mock_pod_lock_manager.release_lock.assert_not_called() + + @pytest.mark.asyncio + async def test_cleanup_expired_keys_releases_acquired_lock(self): + mock_prisma_client = AsyncMock() + mock_cache = MagicMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = ExpiredUISessionKeyCleanupManager( + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + pod_lock_manager=mock_pod_lock_manager, + ) + manager._find_expired_ui_session_keys = AsyncMock(return_value=[]) + + deleted_count = await manager.cleanup_expired_keys() + + assert deleted_count == 0 + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id=EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, + ) diff --git a/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py new file mode 100644 index 0000000000..ae0e1f845b --- /dev/null +++ b/tests/test_litellm/proxy/db/test_exception_handler_reconnect_retry.py @@ -0,0 +1,255 @@ +""" +Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read, +on transport error reconnect once and retry once" helper. + +Covers the regression in issue #25143 where read paths (e.g. +`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in +LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient +`httpx.ReadError` flaps that used to self-heal in 1.82.6. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from prisma.errors import UniqueViolationError + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry + + +def _make_client( + *, + attempt_db_reconnect_return: bool = True, + has_attempt_db_reconnect: bool = True, +): + """Build a minimal stand-in for PrismaClient that exposes only the surface + `call_with_db_reconnect_retry` actually pokes at.""" + client = MagicMock() + if has_attempt_db_reconnect: + client.attempt_db_reconnect = AsyncMock( + return_value=attempt_db_reconnect_return + ) + else: + # `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock + # auto-creates attributes, so we wipe it out via `spec`. + client = MagicMock(spec=[]) + client._db_auth_reconnect_timeout_seconds = 2.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.1 + return client + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_returns_value_on_first_success(): + """Happy path: factory succeeds first call, no reconnect attempted.""" + client = _make_client() + + async def _factory(): + return {"id": 1} + + result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path") + + assert result == {"id": 1} + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_retries_after_transport_error(): + """Transport error on first call → reconnect → second call succeeds.""" + client = _make_client(attempt_db_reconnect_return=True) + + invocations = [] + + async def _factory(): + invocations.append(None) + if len(invocations) == 1: + raise httpx.ReadError("transport blip") + return {"id": 1} + + result = await call_with_db_reconnect_retry( + client, _factory, reason="prisma_get_generic_data_config_lookup_failure" + ) + + assert result == {"id": 1} + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure" + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error(): + """Data-layer errors (e.g. UniqueViolationError) are NOT transport errors — + propagate immediately, do not reconnect.""" + client = _make_client() + + async def _factory(): + raise UniqueViolationError( + data={"user_facing_error": {"meta": {}}}, + message="Unique constraint failed", + ) + + with pytest.raises(UniqueViolationError): + await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test") + + client.attempt_db_reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails(): + """Transport error, but reconnect returns False → propagate the original + exception. Do not call factory a second time.""" + client = _make_client(attempt_db_reconnect_return=False) + + invocations = [] + + async def _factory(): + invocations.append(None) + raise httpx.ReadError("transport blip") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails") + + assert len(invocations) == 1 + client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error(): + """Transport error, reconnect succeeds, retry also raises transport error → + propagate. At most one retry by construction (no infinite loop).""" + client = _make_client(attempt_db_reconnect_return=True) + + invocations = [] + + async def _factory(): + invocations.append(None) + raise httpx.ReadError("still failing") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry( + client, _factory, reason="second_transport_error" + ) + + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr(): + """Older PrismaClient stand-ins / partial mocks may not expose + `attempt_db_reconnect`. The helper must not crash — just propagate the + original exception. Mirrors the `hasattr` guard from + `auth_checks._fetch_key_object_from_db_with_reconnect`.""" + client = _make_client(has_attempt_db_reconnect=False) + + async def _factory(): + raise httpx.ReadError("transport blip") + + with pytest.raises(httpx.ReadError): + await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr") + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro(): + """Guard against the obvious bug of awaiting the same coroutine twice + (`RuntimeError: cannot reuse already awaited coroutine`). The helper must + call the factory a fresh time on retry, not cache an awaitable.""" + client = _make_client(attempt_db_reconnect_return=True) + + factory_call_count = 0 + + async def _factory(): + nonlocal factory_call_count + factory_call_count += 1 + if factory_call_count == 1: + raise httpx.ReadError("transport blip") + return "ok" + + result = await call_with_db_reconnect_retry( + client, _factory, reason="fresh_coro_on_retry" + ) + + assert result == "ok" + assert factory_call_count == 2 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_passes_explicit_timeouts(): + """Explicit timeout_seconds / lock_timeout_seconds override the auth + defaults read off the prisma_client object.""" + client = _make_client(attempt_db_reconnect_return=True) + + async def _factory(): + if not hasattr(_factory, "_called"): + _factory._called = True # type: ignore[attr-defined] + raise httpx.ReadError("transport blip") + return "ok" + + result = await call_with_db_reconnect_retry( + client, + _factory, + reason="explicit_timeouts", + timeout_seconds=5.5, + lock_timeout_seconds=0.25, + ) + + assert result == "ok" + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["timeout_seconds"] == 5.5 + assert call_kwargs["lock_timeout_seconds"] == 0.25 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset(): + """When timeouts are not provided, helper reads + `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds` + off the prisma_client (matching the auth path's existing convention).""" + client = _make_client(attempt_db_reconnect_return=True) + client._db_auth_reconnect_timeout_seconds = 3.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.5 + + async def _factory(): + if not hasattr(_factory, "_called"): + _factory._called = True # type: ignore[attr-defined] + raise httpx.ReadError("transport blip") + return "ok" + + await call_with_db_reconnect_retry(client, _factory, reason="defaults") + + call_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert call_kwargs["timeout_seconds"] == 3.0 + assert call_kwargs["lock_timeout_seconds"] == 0.5 + + +@pytest.mark.asyncio +async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises(): + """If `attempt_db_reconnect` itself raises (lock cancellation, timer + error, unexpected internal failure), the helper must surface the + *original* transport error to telemetry — not the reconnect exception. + Otherwise `failure_handler` / `db_exceptions` alerts log the wrong + error string and the actual DB transport problem becomes invisible. + + The reconnect error is chained as the `__cause__` for debuggability.""" + client = MagicMock() + reconnect_exc = RuntimeError("simulated reconnect lock cancellation") + client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc) + client._db_auth_reconnect_timeout_seconds = 2.0 + client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + original_exc = httpx.ReadError("transport blip") + + async def _factory(): + raise original_exc + + with pytest.raises(httpx.ReadError) as exc_info: + await call_with_db_reconnect_retry( + client, _factory, reason="reconnect_itself_raises" + ) + + assert exc_info.value is original_exc + assert exc_info.value.__cause__ is reconnect_exc + client.attempt_db_reconnect.assert_awaited_once() diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index fb215e5477..3f9ba6af3a 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -5,6 +5,7 @@ import sys import time from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest sys.path.insert( @@ -34,18 +35,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging): client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() - result = await client.attempt_db_reconnect( - reason="unit_test_reconnect_success", - force=True, - ) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await client.attempt_db_reconnect( + reason="unit_test_reconnect_success", + force=True, + ) assert result is True - client.db.disconnect.assert_awaited_once() - client.db.connect.assert_awaited_once() + client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") client.db.query_raw.assert_awaited_once_with("SELECT 1") @@ -140,15 +141,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt( ) client._db_last_reconnect_attempt_ts = 0.0 client._db_reconnect_cooldown_seconds = 10 - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() # Use a counter-based mock to avoid StopIteration when time.time() is called # more times than expected (varies by Python version / internal code paths). fake_clock = iter(range(100, 10000)) - with patch( - "litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock)) + with ( + patch( + "litellm.proxy.utils.time.time", + side_effect=lambda: float(next(fake_clock)), + ), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), ): result = await client.attempt_db_reconnect( reason="unit_test_cooldown_timestamp_after_attempt", @@ -162,23 +167,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt( @pytest.mark.asyncio -async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops( +async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client( mock_proxy_logging, ): + """Direct reconnect goes through recreate_prisma_client (which non-blockingly + kills the old engine) instead of calling disconnect() — see issue #26191. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used")) - client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used")) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.disconnect = AsyncMock( + side_effect=AssertionError("disconnect must not be called") + ) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() - await client._run_reconnect_cycle(timeout_seconds=None) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await client._run_reconnect_cycle(timeout_seconds=None) - client.db.disconnect.assert_awaited_once() - client.db.connect.assert_awaited_once() + client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") client.db.query_raw.assert_awaited_once_with("SELECT 1") + client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio @@ -189,19 +199,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) client._db_watchdog_reconnect_timeout_seconds = 0.1 - client.db.disconnect = AsyncMock(return_value=None) + client._start_engine_watcher = AsyncMock() - async def _slow_connect(): + async def _slow_recreate(_db_url): await asyncio.sleep(0.08) async def _slow_query(_query: str): await asyncio.sleep(0.08) return [{"result": 1}] - client.db.connect = AsyncMock(side_effect=_slow_connect) + client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate) client.db.query_raw = AsyncMock(side_effect=_slow_query) - with pytest.raises(asyncio.TimeoutError): + with ( + pytest.raises(asyncio.TimeoutError), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await client._run_reconnect_cycle(timeout_seconds=None) @@ -212,19 +225,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget( client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) + client._start_engine_watcher = AsyncMock() - async def _slow_connect(): + async def _slow_recreate(_db_url): await asyncio.sleep(0.08) async def _slow_query(_query: str): await asyncio.sleep(0.08) return [{"result": 1}] - client.db.connect = AsyncMock(side_effect=_slow_connect) + client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate) client.db.query_raw = AsyncMock(side_effect=_slow_query) - with pytest.raises(asyncio.TimeoutError): + with ( + pytest.raises(asyncio.TimeoutError), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await client._run_reconnect_cycle(timeout_seconds=0.1) @@ -319,42 +335,154 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging): @pytest.mark.asyncio -async def test_lightweight_reconnect_kills_engine_on_disconnect_failure( +async def test_recreate_prisma_client_kills_old_engine_without_disconnect( mock_proxy_logging, ): - """Lightweight reconnect must kill the old engine PID when disconnect() fails.""" + """recreate_prisma_client SIGTERMs the old engine PID directly rather than + calling `disconnect()`, which blocks the asyncio event loop on the sync + `subprocess.Popen.wait()` inside prisma-client-py — see issue #26191. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed")) - client.db.connect = AsyncMock(return_value=None) - client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + disconnect_mock = AsyncMock( + side_effect=AssertionError("disconnect must not be called on reconnect path") + ) + client.db._original_prisma.disconnect = disconnect_mock with ( - patch.object(client, "_get_engine_pid", return_value=9999), - patch("os.kill") as mock_kill, - patch("asyncio.sleep", new_callable=AsyncMock), + patch.object(client.db, "_get_engine_pid", return_value=9999), + patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill, + patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock), ): - await client._run_reconnect_cycle(timeout_seconds=5.0) + # Return a Prisma instance whose connect() is awaitable. + fake_new_prisma = MagicMock() + fake_new_prisma.connect = AsyncMock(return_value=None) + with patch("prisma.Prisma", return_value=fake_new_prisma): + await client.db.recreate_prisma_client("postgresql://test") mock_kill.assert_any_call(9999, signal.SIGTERM) - client.db.connect.assert_awaited_once() - client.db.query_raw.assert_awaited_once_with("SELECT 1") + disconnect_mock.assert_not_awaited() + fake_new_prisma.connect.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# get_generic_data: transport-reconnect-and-retry coverage (issue #25143) +# --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_lightweight_reconnect_skips_kill_on_successful_disconnect( +async def test_get_generic_data_retries_on_transport_error_for_config_table( mock_proxy_logging, ): - """Lightweight reconnect must NOT kill when disconnect() succeeds.""" + """`get_generic_data(table_name="config")` self-heals on a transient + `httpx.ReadError`: reconnect once, retry once, return the row. + + Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry + branch that 1.82.6 had on this method. `_update_config_from_db` fans out + four concurrent `get_generic_data` calls, so a single transport flap used + to surface as four `db_exceptions` alerts and a stale config window. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) - client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) - with patch("os.kill") as mock_kill: - await client._run_reconnect_cycle(timeout_seconds=5.0) + expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}} + invocations: list[None] = [] - mock_kill.assert_not_called() + async def _flaky_find_first(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise httpx.ReadError("simulated transport blip") + return expected_row + + client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first) + client.attempt_db_reconnect = AsyncMock(return_value=True) + + result = await client.get_generic_data( + key="param_name", + value="general_settings", + table_name="config", + ) + + assert result == expected_row + assert len(invocations) == 2 + client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs + assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure" + + # The failure_handler telemetry side-effect must NOT fire on the first + # transport blip — only if the post-retry call also fails. Drain the + # event loop so any spuriously-spawned task would have run by now. + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging): + """If reconnect itself does not succeed, propagate the original transport + error and let the existing failure_handler / db_exceptions telemetry fire.""" + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + + client.db.litellm_config.find_first = AsyncMock( + side_effect=httpx.ReadError("simulated transport blip") + ) + client.attempt_db_reconnect = AsyncMock(return_value=False) + + with pytest.raises(httpx.ReadError): + await client.get_generic_data( + key="param_name", + value="general_settings", + table_name="config", + ) + + client.attempt_db_reconnect.assert_awaited_once() + # Failure telemetry IS expected here — the read genuinely failed. + await asyncio.sleep(0) + mock_proxy_logging.failure_handler.assert_called_once() + + +# --------------------------------------------------------------------------- +# _engine_confirmed_dead flag-reset bug (B2) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect( + mock_proxy_logging, +): + """Regression test for the flag-reset bug. + + Before the fix, `_run_reconnect_cycle` cleared + `self._engine_confirmed_dead = False` *before* awaiting + `_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout, + missing DATABASE_URL, recreate failure), the flag was left cleared and the + next attempt could demote to the lightweight path even though the engine + was genuinely dead. + + The fix moves the reset into the success branch — the flag must stay True + when heavy reconnect raises. + """ + client = PrismaClient( + database_url="mock://test", proxy_logging_obj=mock_proxy_logging + ) + client._engine_confirmed_dead = True + client._engine_pid = 0 # so `_is_engine_alive` is not consulted + + # Make the heavy reconnect path raise. + client.db.recreate_prisma_client = AsyncMock( + side_effect=RuntimeError("simulated heavy reconnect failure") + ) + client._start_engine_watcher = AsyncMock() + client._cleanup_engine_watcher = MagicMock() + client._reap_all_zombies = MagicMock() + + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + with pytest.raises(Exception): + await client._run_reconnect_cycle(timeout_seconds=5.0) + + # The flag must STILL be True so the next attempt re-enters the heavy + # branch instead of silently demoting to the lightweight path. + assert client._engine_confirmed_dead is True diff --git a/tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py new file mode 100644 index 0000000000..a337ff6d88 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_workflow_management_endpoints.py @@ -0,0 +1,611 @@ +""" +Unit tests for workflow management endpoints (/v1/workflows/runs/*). +Uses FastAPI TestClient with a mocked prisma_client. +""" + +import os +import sys +from datetime import datetime, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient +from prisma.errors import UniqueViolationError + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.management_endpoints.workflow_management_endpoints import router + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_run( + run_id: str = "run-1", + session_id: str = "sess-1", + workflow_type: str = "shin-builder", + status: str = "pending", + created_by: Any = "tok-test", +) -> MagicMock: + obj = MagicMock() + obj.run_id = run_id + obj.session_id = session_id + obj.workflow_type = workflow_type + obj.status = status + obj.created_by = created_by + obj.created_at = datetime.now(timezone.utc) + obj.updated_at = datetime.now(timezone.utc) + obj.input = None + obj.output = None + obj.metadata = None + return obj + + +def _make_event( + event_id: str = "evt-1", + run_id: str = "run-1", + event_type: str = "step.started", + step_name: str = "grill", + sequence_number: int = 0, +) -> MagicMock: + obj = MagicMock() + obj.event_id = event_id + obj.run_id = run_id + obj.event_type = event_type + obj.step_name = step_name + obj.sequence_number = sequence_number + obj.data = None + obj.created_at = datetime.now(timezone.utc) + return obj + + +def _make_message( + message_id: str = "msg-1", + run_id: str = "run-1", + role: str = "user", + content: str = "hello", + sequence_number: int = 0, +) -> MagicMock: + obj = MagicMock() + obj.message_id = message_id + obj.run_id = run_id + obj.role = role + obj.content = content + obj.sequence_number = sequence_number + obj.session_id = None + obj.created_at = datetime.now(timezone.utc) + return obj + + +def _make_tx(event_return=None, run_return=None, msg_return=None) -> MagicMock: + """Build an async context-manager mock for prisma_client.db.tx().""" + tx = MagicMock() + tx.litellm_workflowevent = MagicMock() + tx.litellm_workflowevent.create = AsyncMock( + return_value=event_return or _make_event() + ) + tx.litellm_workflowrun = MagicMock() + tx.litellm_workflowrun.update = AsyncMock(return_value=run_return or _make_run()) + tx.litellm_workflowmessage = MagicMock() + tx.litellm_workflowmessage.create = AsyncMock( + return_value=msg_return or _make_message() + ) + tx.__aenter__ = AsyncMock(return_value=tx) + tx.__aexit__ = AsyncMock(return_value=False) + return tx + + +def _make_prisma_client() -> MagicMock: + client = MagicMock() + client.db = MagicMock() + client.db.litellm_workflowrun = MagicMock() + client.db.litellm_workflowevent = MagicMock() + client.db.litellm_workflowmessage = MagicMock() + # default tx() returns a no-op transaction + client.db.tx = MagicMock(return_value=_make_tx()) + return client + + +def _make_app() -> FastAPI: + app = FastAPI() + app.include_router(router) + return app + + +def _override_auth() -> Any: + from litellm.proxy._types import UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-test", user_id="admin") + auth.token = "tok-test" + return auth + + +def _override_auth_admin() -> Any: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-master") + auth.user_role = LitellmUserRoles.PROXY_ADMIN # type: ignore[assignment] + return auth + + +def _override_auth_user_with_token(token: str = "tok-abc") -> Any: + """Return a non-admin caller whose hashed token equals `token`.""" + from litellm.proxy._types import UserAPIKeyAuth + + auth = UserAPIKeyAuth(api_key="sk-user", user_id="user-1") + auth.token = token # override the computed hash with a predictable value + return auth + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestCreateWorkflowRun: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_create_returns_run(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run()) + + resp = self.client.post( + "/v1/workflows/runs", + json={"workflow_type": "shin-builder"}, + ) + assert resp.status_code == 200 + self._prisma.db.litellm_workflowrun.create.assert_awaited_once() + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_create_500_when_no_db(self): + resp = self.client.post( + "/v1/workflows/runs", + json={"workflow_type": "shin-builder"}, + ) + assert resp.status_code == 500 + + +class TestListWorkflowRuns: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_returns_runs(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock( + return_value=[_make_run()] + ) + + resp = self.client.get("/v1/workflows/runs") + assert resp.status_code == 200 + data = resp.json() + assert data["count"] == 1 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_filters_by_status(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs?status=running") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert call_kwargs["where"]["status"] == "running" + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_filters_by_multiple_statuses(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs?status=running,paused") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert call_kwargs["where"]["status"] == {"in": ["running", "paused"]} + + +class TestGetWorkflowRun: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_get_existing_run(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + + resp = self.client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 200 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_get_missing_run_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.get("/v1/workflows/runs/nonexistent") + assert resp.status_code == 404 + + +class TestUpdateWorkflowRun: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_update_status(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + updated = _make_run(status="completed") + self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=updated) + + resp = self.client.patch( + "/v1/workflows/runs/run-1", json={"status": "completed"} + ) + assert resp.status_code == 200 + self._prisma.db.litellm_workflowrun.update.assert_awaited_once() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_update_no_fields_returns_400(self, mock_pc): + mock_pc.db = self._prisma.db + resp = self.client.patch("/v1/workflows/runs/run-1", json={}) + assert resp.status_code == 400 + + +class TestAppendWorkflowEvent: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_event_updates_run_status(self, mock_pc): + mock_pc.db = self._prisma.db + # _require_run check + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + tx = _make_tx( + event_return=_make_event(), run_return=_make_run(status="running") + ) + self._prisma.db.tx = MagicMock(return_value=tx) + + resp = self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "step.started", "step_name": "grill"}, + ) + assert resp.status_code == 200 + # run status updated inside tx + tx.litellm_workflowrun.update.assert_awaited_once() + update_call = tx.litellm_workflowrun.update.call_args[1] + assert update_call["data"]["status"] == "running" + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_event_no_status_update_for_unknown_type(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + tx = _make_tx(event_return=_make_event(event_type="custom.event")) + self._prisma.db.tx = MagicMock(return_value=tx) + + resp = self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "custom.event", "step_name": "grill"}, + ) + assert resp.status_code == 200 + # no status update inside tx for unknown event_type + tx.litellm_workflowrun.update.assert_not_awaited() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_sequence_number_increments(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + existing = _make_event(sequence_number=4) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock( + return_value=[existing] + ) + tx = _make_tx(event_return=_make_event(sequence_number=5)) + self._prisma.db.tx = MagicMock(return_value=tx) + + self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "step.started", "step_name": "plan"}, + ) + create_call = tx.litellm_workflowevent.create.call_args[1] + assert create_call["data"]["sequence_number"] == 5 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_unknown_run_id_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.post( + "/v1/workflows/runs/nonexistent/events", + json={"event_type": "step.started", "step_name": "grill"}, + ) + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_sequence_collision_retries_and_succeeds(self, mock_pc): + """UniqueViolationError on first attempt triggers retry; second attempt succeeds.""" + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + + # First tx raises UniqueViolationError; second succeeds. + tx_fail = _make_tx() + tx_fail.__aenter__ = AsyncMock(return_value=tx_fail) + tx_fail.litellm_workflowevent.create = AsyncMock( + side_effect=UniqueViolationError( + {"user_facing_error": {"message": "unique"}} + ) + ) + tx_fail.__aexit__ = AsyncMock(return_value=False) + + tx_ok = _make_tx(event_return=_make_event(sequence_number=1)) + + self._prisma.db.tx = MagicMock(side_effect=[tx_fail, tx_ok]) + + resp = self.client.post( + "/v1/workflows/runs/run-1/events", + json={"event_type": "step.started", "step_name": "grill"}, + ) + assert resp.status_code == 200 + + +class TestWorkflowMessages: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_message(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) + self._prisma.db.litellm_workflowmessage.create = AsyncMock( + return_value=_make_message() + ) + + resp = self.client.post( + "/v1/workflows/runs/run-1/messages", + json={"role": "user", "content": "fix the bug"}, + ) + assert resp.status_code == 200 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_append_message_unknown_run_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.post( + "/v1/workflows/runs/nonexistent/messages", + json={"role": "user", "content": "hello"}, + ) + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_messages_ordered(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock( + return_value=[ + _make_message(sequence_number=0), + _make_message(sequence_number=1, role="assistant"), + ] + ) + + resp = self.client.get("/v1/workflows/runs/run-1/messages") + assert resp.status_code == 200 + data = resp.json() + assert data["count"] == 2 + call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1] + assert call_kwargs["order"] == {"sequence_number": "asc"} + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_messages_respects_limit(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs/run-1/messages?limit=25") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1] + assert call_kwargs["take"] == 25 + + +class TestListWorkflowEvents: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_events_ordered(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock( + return_value=[ + _make_event(sequence_number=0), + _make_event(sequence_number=1), + ] + ) + + resp = self.client.get("/v1/workflows/runs/run-1/events") + assert resp.status_code == 200 + data = resp.json() + assert data["count"] == 2 + call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1] + assert call_kwargs["order"] == {"sequence_number": "asc"} + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_events_respects_limit(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run() + ) + self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[]) + + resp = self.client.get("/v1/workflows/runs/run-1/events?limit=10") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1] + assert call_kwargs["take"] == 10 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_list_events_unknown_run_returns_404(self, mock_pc): + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None) + + resp = self.client.get("/v1/workflows/runs/nonexistent/events") + assert resp.status_code == 404 + + +class TestTenantIsolation: + """Ownership enforcement: non-admin callers only see their own runs.""" + + def _make_app_with_auth(self, auth_fn): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + self._prisma = _make_prisma_client() + app = _make_app() + app.dependency_overrides[user_api_key_auth] = auth_fn + return TestClient(app, raise_server_exceptions=True) + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_create_stores_caller_token(self, mock_pc): + token = "tok-owner" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.create = AsyncMock( + return_value=_make_run(created_by=token) + ) + + resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"}) + assert resp.status_code == 200 + create_call = self._prisma.db.litellm_workflowrun.create.call_args[1] + assert create_call["data"]["created_by"] == token + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_list_scoped_to_caller_token(self, mock_pc): + token = "tok-owner" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/workflows/runs") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert call_kwargs["where"].get("created_by") == token + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_admin_list_not_scoped(self, mock_pc): + client = self._make_app_with_auth(_override_auth_admin) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[]) + + resp = client.get("/v1/workflows/runs") + assert resp.status_code == 200 + call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1] + assert "created_by" not in call_kwargs["where"] + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_get_other_users_run_returns_404(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + # Run owned by a different key + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by="tok-other-owner") + ) + + resp = client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_get_null_owner_run_returns_404(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by=None) + ) + + resp = client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 404 + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_update_null_owner_run_returns_404(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by=None) + ) + self._prisma.db.litellm_workflowrun.update = AsyncMock( + return_value=_make_run(status="completed") + ) + + resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"}) + assert resp.status_code == 404 + self._prisma.db.litellm_workflowrun.update.assert_not_awaited() + + @patch("litellm.proxy.proxy_server.prisma_client") + def test_non_admin_get_own_run_succeeds(self, mock_pc): + token = "tok-caller" + client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token)) + mock_pc.db = self._prisma.db + self._prisma.db.litellm_workflowrun.find_unique = AsyncMock( + return_value=_make_run(created_by=token) + ) + + resp = client.get("/v1/workflows/runs/run-1") + assert resp.status_code == 200 diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index a635c16e98..d4b4617730 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -218,6 +218,141 @@ class TestProxyBaseLLMRequestProcessing: headers_with_invalid ) + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_from_llm_response(self): + """ + Google native :generateContent uses this helper instead of base_process_llm_request; + ensure x-litellm-* headers and callback hooks merge like the main proxy path. + """ + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + + class _FakeGenaiResponse: + _hidden_params = { + "model_id": "deployment-model-id", + "cache_key": "ck-test", + "api_base": "https://generativelanguage.googleapis.com/v1beta", + "response_cost": 0.001, + "additional_headers": {"llm_provider-ratelimit-requests": "1000"}, + } + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "call-id-test" + + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-ratelimit-remaining-requests": "999"} + ) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_FakeGenaiResponse(), + request_data={"model": "gemini/gemini-1.5-flash"}, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="9.9.9", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-call-id"] == "call-id-test" + assert headers["x-litellm-model-id"] == "deployment-model-id" + assert headers["x-litellm-version"] == "9.9.9" + assert headers["llm_provider-ratelimit-requests"] == "1000" + assert headers["x-ratelimit-remaining-requests"] == "999" + proxy_logging_obj.post_call_response_headers_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_streaming_style_iterator(self): + """AsyncGoogleGenAIGenerateContentStreamingIterator sets _hidden_params at init; headers must propagate.""" + + class _FakeStreamLike: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + _hidden_params = { + "model_id": "stream-model-id", + "api_base": "https://generativelanguage.googleapis.com/v1beta", + "cache_key": "", + "response_cost": "", + "additional_headers": {"llm_provider-x": "y"}, + } + + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + logging_obj = MagicMock() + logging_obj.litellm_call_id = "cid-stream" + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_FakeStreamLike(), + request_data={"model": "gemini/gemini-2.0-flash"}, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="1.0.0", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-model-id"] == "stream-model-id" + assert headers["x-litellm-model-api-base"] == ( + "https://generativelanguage.googleapis.com/v1beta" + ) + assert headers["llm_provider-x"] == "y" + + @pytest.mark.asyncio + async def test_build_litellm_proxy_success_headers_no_hidden_params_metadata_fallback( + self, + ): + """When response has no _hidden_params, model_id can still come from litellm_metadata.""" + + class _BareResponse: + pass + + mock_request = MagicMock(spec=Request) + mock_request.headers = {} + logging_obj = MagicMock() + logging_obj.litellm_call_id = "cid-meta" + mock_user = MagicMock() + mock_user.tpm_limit = None + mock_user.rpm_limit = None + mock_user.max_budget = None + mock_user.spend = 0.0 + mock_user.allowed_model_region = None + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response( + response=_BareResponse(), + request_data={ + "model": "gemini/gemini-1.5-flash", + "litellm_metadata": {"model_info": {"id": "meta-model-id"}}, + }, + request=mock_request, + user_api_key_dict=mock_user, + logging_obj=logging_obj, + version="1.0.0", + proxy_logging_obj=proxy_logging_obj, + ) + + assert headers["x-litellm-model-id"] == "meta-model-id" + @pytest.mark.asyncio async def test_add_litellm_data_to_request_with_stream_timeout_header(self): """ @@ -1158,13 +1293,6 @@ class TestCommonRequestProcessingHelpers: assert mock_tracer.trace.call_count == 4 # Verify that each call was made with the correct operation name - expected_calls = [ - (("streaming.chunk.yield",), {}), - (("streaming.chunk.yield",), {}), - (("streaming.chunk.yield",), {}), - (("streaming.chunk.yield",), {}), - ] - actual_calls = mock_tracer.trace.call_args_list assert len(actual_calls) == 4 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 3349a138ee..1f4f82a64e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2544,6 +2544,14 @@ class TestPriceDataReloadAPI: class TestPriceDataReloadIntegration: """Integration tests for the complete price data reload feature""" + @pytest.fixture(autouse=True) + def _flush_litellm_config_cache(self): + from litellm.proxy.utils import litellm_config_cache + + litellm_config_cache.flush_cache() + yield + litellm_config_cache.flush_cache() + @pytest.fixture def client_with_auth(self): """Create a test client with authentication""" @@ -2601,6 +2609,7 @@ class TestPriceDataReloadIntegration: def test_distributed_reload_check_function(self): """Test the _check_and_reload_model_cost_map function""" from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import litellm_config_cache proxy_config = ProxyConfig() @@ -2609,14 +2618,19 @@ class TestPriceDataReloadIntegration: # Test case 1: No config in database mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + # _check_and_reload_model_cost_map routes through get_config_param, + # which calls prisma.get_generic_data on a cache miss. + mock_prisma.get_generic_data = AsyncMock(return_value=None) # Should return early without reloading asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) # Test case 2: Config with interval but not time to reload + litellm_config_cache.flush_cache() mock_config = MagicMock() mock_config.param_value = {"interval_hours": 6, "force_reload": False} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) # Mock current time and last reload time with patch( @@ -2632,8 +2646,10 @@ class TestPriceDataReloadIntegration: asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) # Test case 3: Config with force reload + litellm_config_cache.flush_cache() mock_config.param_value = {"interval_hours": 6, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost.copy() @@ -2675,6 +2691,8 @@ class TestPriceDataReloadIntegration: mock_config = MagicMock() mock_config.param_value = {"interval_hours": 24, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_model_cost_map now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) original_model_cost = litellm.model_cost.copy() @@ -2770,6 +2788,8 @@ class TestPriceDataReloadIntegration: mock_config = MagicMock() mock_config.param_value = {"interval_hours": 12, "force_reload": True} mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) + # _check_and_reload_anthropic_beta_headers now reads through get_generic_data. + mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) with patch( diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py index 620c073498..6644b1389c 100644 --- a/tests/test_litellm/test_gpt_image_cost_calculator.py +++ b/tests/test_litellm/test_gpt_image_cost_calculator.py @@ -29,8 +29,21 @@ from litellm.types.utils import ( ) +@pytest.fixture(autouse=True) +def _use_local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + class TestGPTImageCostCalculator: - """Test the OpenAI gpt-image-1 cost calculator""" + """Test the OpenAI gpt-image cost calculator""" def test_gpt_image_1_cost_with_text_only(self): """Test cost calculation with only text input tokens""" @@ -149,6 +162,44 @@ class TestGPTImageCostCalculator: assert cost == 0.0 + def test_gpt_image_2_cost_with_text_and_image_tokens(self): + """Test cost calculation for gpt-image-2 token pricing""" + from litellm.llms.openai.image_generation.cost_calculator import cost_calculator + + usage = Usage( + prompt_tokens=600, + completion_tokens=5000, + total_tokens=5600, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=100, + image_tokens=500, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=1000, + image_tokens=4000, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = cost_calculator( + model="gpt-image-2", + image_response=image_response, + custom_llm_provider="openai", + ) + + # GPT Image 2 pricing: + # Text input: 100 * $5/1M = 0.0005 + # Image input: 500 * $8/1M = 0.004 + # Text output: 1000 * $10/1M = 0.01 + # Image output: 4000 * $30/1M = 0.12 + expected_cost = 0.0005 + 0.004 + 0.01 + 0.12 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + class TestGPTImageCostRouting: """Test that gpt-image models are properly routed to the token-based calculator""" @@ -182,6 +233,33 @@ class TestGPTImageCostRouting: expected_cost = 0.0005 + 0.2 assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + def test_openai_gpt_image_2_routes_to_token_calculator(self): + """Test that OpenAI gpt-image-2 routes to token-based calculator""" + from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils + + usage = Usage( + prompt_tokens=100, + completion_tokens=5000, + total_tokens=5100, + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=5000), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="gpt-image-2", + completion_response=image_response, + custom_llm_provider="openai", + ) + + expected_cost = 0.0005 + 0.15 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + def test_openai_dalle_routes_to_pixel_calculator(self): """Test that OpenAI DALL-E still routes to pixel-based calculator""" from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 93c61e003d..f28fe3ed25 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -32,6 +32,19 @@ from litellm.utils import ( # Adds the parent directory to the system path +@pytest.fixture +def local_model_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_check_provider_match_azure_ai_allows_openai_and_azure(): """ Test that azure_ai provider can match openai and azure models. @@ -198,6 +211,72 @@ def test_get_optional_params_image_gen_filters_empty_values(): assert optional_params == {} +def test_gpt_image_provider_detection_covers_existing_family(): + for image_model in ("gpt-image-1", "gpt-image-1-mini", "gpt-image-1.5"): + model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=image_model) + + assert model == image_model + assert custom_llm_provider == "openai" + + +def test_gpt_image_2_provider_and_model_info(local_model_cost_map): + + model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="gpt-image-2") + + assert model == "gpt-image-2" + assert custom_llm_provider == "openai" + + model_info = litellm.get_model_info(model="gpt-image-2") + assert model_info["litellm_provider"] == "openai" + assert model_info["mode"] == "image_generation" + assert model_info["input_cost_per_token"] == 5e-06 + assert model_info["input_cost_per_image_token"] == 8e-06 + assert model_info["output_cost_per_token"] == 1e-05 + assert model_info["output_cost_per_image_token"] == 3e-05 + assert ( + "/v1/images/generations" + in litellm.model_cost["gpt-image-2"]["supported_endpoints"] + ) + assert ( + "/v1/images/edits" in litellm.model_cost["gpt-image-2"]["supported_endpoints"] + ) + assert model_info["supports_vision"] is True + assert model_info["supports_pdf_input"] is True + + +def test_gpt_image_2_snapshot_model_info(local_model_cost_map): + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model="gpt-image-2-2026-04-21" + ) + + assert model == "gpt-image-2-2026-04-21" + assert custom_llm_provider == "openai" + + model_info = litellm.get_model_info(model="gpt-image-2-2026-04-21") + assert model_info["litellm_provider"] == "openai" + assert model_info["mode"] == "image_generation" + assert model_info["output_cost_per_image_token"] == 3e-05 + + +def test_azure_gpt_image_2_model_info(local_model_cost_map): + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model="azure/gpt-image-2" + ) + + assert model == "gpt-image-2" + assert custom_llm_provider == "azure" + + model_info = litellm.get_model_info( + model="gpt-image-2", custom_llm_provider="azure" + ) + assert model_info["litellm_provider"] == "azure" + assert model_info["mode"] == "image_generation" + assert model_info["input_cost_per_token"] == 5e-06 + assert model_info["input_cost_per_image_token"] == 8e-06 + assert model_info["output_cost_per_token"] == 1e-05 + assert model_info["output_cost_per_image_token"] == 3e-05 + + def test_all_model_configs(): from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( VertexAIAi21Config, @@ -1179,7 +1258,7 @@ def test_get_model_info_shows_supports_computer_use(): "model, custom_llm_provider", [ ("gpt-3.5-turbo", "openai"), - ("anthropic.claude-3-7-sonnet-20250219-v1:0", "bedrock"), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), ("gemini-2.5-pro", "vertex_ai"), ], ) @@ -1325,7 +1404,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", False, ), ( @@ -1623,7 +1702,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Bedrock Claude 3 Opus via Converse API", ), @@ -1710,7 +1789,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/staging-claude-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Staging Claude Opus", ), @@ -1722,7 +1801,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/high-performance-claude", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "High-performance Claude deployment", ), @@ -1860,7 +1939,7 @@ class TestProxyFunctionCalling: bedrock_models = [ "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0", "bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", ] for model in bedrock_models: @@ -1892,7 +1971,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Bedrock Claude 3 Opus via Converse API", ), @@ -1979,7 +2058,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/staging-claude-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Staging Claude Opus", ), @@ -1991,7 +2070,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/high-performance-claude", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "High-performance Claude deployment", ), @@ -2129,7 +2208,7 @@ class TestProxyFunctionCalling: bedrock_models = [ "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0", "bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", ] for model in bedrock_models: @@ -2161,7 +2240,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/bedrock-claude-3-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Bedrock Claude 3 Opus via Converse API", ), @@ -2248,7 +2327,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/staging-claude-opus", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "Staging Claude Opus", ), @@ -2260,7 +2339,7 @@ class TestProxyFunctionCalling: ), ( "litellm_proxy/high-performance-claude", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", False, "High-performance Claude deployment", ), @@ -2398,7 +2477,7 @@ class TestProxyFunctionCalling: bedrock_models = [ "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0", "bedrock/converse/anthropic.claude-3-sonnet-20240229-v1:0", - "bedrock/converse/anthropic.claude-3-7-sonnet-20250219-v1:0", + "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0", ] for model in bedrock_models: diff --git a/tests/vector_store_tests/test_ragflow_vector_store.py b/tests/vector_store_tests/test_ragflow_vector_store.py index 4ca3823312..cb4cfd75c1 100644 --- a/tests/vector_store_tests/test_ragflow_vector_store.py +++ b/tests/vector_store_tests/test_ragflow_vector_store.py @@ -267,6 +267,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest): api_base="http://localhost:9380", litellm_logging_obj=logging_obj, litellm_params={}, + extra_body=None, ) def test_transform_search_vector_store_response_not_implemented(self): diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts index a8ce55d274..4ba80df5c6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.test.ts @@ -7,6 +7,7 @@ import { useDeleteProxyConfigField, getProxyConfigCall, deleteProxyConfigFieldCall, + proxyConfigKeys, ConfigType, GeneralSettingsFieldName, type ProxyConfigResponse, @@ -426,6 +427,28 @@ describe("useDeleteProxyConfigField", () => { expect(result.current.error).toBeDefined(); }); + + it("should invalidate proxyConfig queries after a successful delete", async () => { + (fetchSpy as any).mockResolvedValue({ + ok: true, + json: async () => mockDeleteResponse, + }); + + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const { result } = renderHook(() => useDeleteProxyConfigField(), { wrapper }); + + result.current.mutate({ + config_type: ConfigType.GENERAL_SETTINGS, + field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(invalidateSpy).toHaveBeenCalledWith({ queryKey: proxyConfigKeys.all }); + }); }); describe("getProxyConfigCall", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts index b823ce4ffd..485c7cc1f9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts @@ -1,4 +1,4 @@ -import { useQuery, useMutation, UseMutationResult } from "@tanstack/react-query"; +import { useQuery, useMutation, UseMutationResult, useQueryClient } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import useAuthorized from "../useAuthorized"; import { proxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; @@ -101,7 +101,7 @@ export const getProxyConfigCall = async (accessToken: string, configType: Config } }; -const proxyConfigKeys = createQueryKeys("proxyConfig"); +export const proxyConfigKeys = createQueryKeys("proxyConfig"); /** * Network call function to delete a proxy config field @@ -168,6 +168,7 @@ export const useDeleteProxyConfigField = (): UseMutationResult< DeleteProxyConfigFieldRequest > => { const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); return useMutation({ mutationFn: async (request: DeleteProxyConfigFieldRequest) => { @@ -176,5 +177,8 @@ export const useDeleteProxyConfigField = (): UseMutationResult< } return await deleteProxyConfigFieldCall(accessToken, request); }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: proxyConfigKeys.all }); + }, }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts index 9c6211c308..67b52997a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts @@ -1,6 +1,7 @@ -import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { useMutation, UseMutationResult, useQueryClient } from "@tanstack/react-query"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; import useAuthorized from "../useAuthorized"; +import { proxyConfigKeys } from "../proxyConfig/useProxyConfig"; export interface StoreRequestInSpendLogsParams { store_prompts_in_spend_logs: boolean; @@ -51,6 +52,7 @@ export const useStoreRequestInSpendLogs = (): UseMutationResult< StoreRequestInSpendLogsParams > => { const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); return useMutation({ mutationFn: async (params: StoreRequestInSpendLogsParams) => { @@ -59,5 +61,8 @@ export const useStoreRequestInSpendLogs = (): UseMutationResult< } return await performStoreRequestInSpendLogs(accessToken, params); }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: proxyConfigKeys.all }); + }, }); }; diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index d32cf74d6a..a8553d5405 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -40,6 +40,7 @@ import { ProjectsPage } from "@/components/Projects/ProjectsPage"; import VectorStoreManagement from "@/components/vector_store_management"; import ToolPoliciesView from "@/components/ToolPoliciesView"; import { MemoryView } from "@/components/MemoryView"; +import WorkflowRuns from "@/components/workflow_runs"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -631,6 +632,8 @@ function CreateKeyPageContent() { ) : page == "tool-policies" ? ( + ) : page == "workflows" ? ( + ) : page == "memory" ? ( = ({ proxySettings }) => { ), children: , }, + { + key: "logging-settings", + label: "Logging Settings", + children: , + }, { key: "hashicorp-vault", label: "Hashicorp Vault", diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx new file mode 100644 index 0000000000..1adf8039e6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx @@ -0,0 +1,335 @@ +import { useDeleteProxyConfigField, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { useStoreRequestInSpendLogs } from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../../tests/test-utils"; +import LoggingSettings from "./LoggingSettings"; + +vi.mock("@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs"); +vi.mock("@/app/(dashboard)/hooks/proxyConfig/useProxyConfig", async () => { + const actual = await vi.importActual( + "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig", + ); + return { + ...actual, + useProxyConfig: vi.fn(), + useDeleteProxyConfigField: vi.fn(), + }; +}); +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn(), +})); + +const mockUseStoreRequestInSpendLogs = vi.mocked(useStoreRequestInSpendLogs); +const mockUseProxyConfig = vi.mocked(useProxyConfig); +const mockUseDeleteProxyConfigField = vi.mocked(useDeleteProxyConfigField); +const mockNotificationsManager = vi.mocked(NotificationsManager); +const mockParseErrorMessage = vi.mocked(parseErrorMessage); + +describe("LoggingSettings", () => { + const mockMutate = vi.fn(); + const mockDeleteField = vi.fn(); + const mockRefetch = vi.fn(); + + beforeEach(() => { + vi.resetAllMocks(); + mockUseStoreRequestInSpendLogs.mockReturnValue({ + mutate: mockMutate, + isPending: false, + } as any); + mockUseDeleteProxyConfigField.mockReturnValue({ + mutate: mockDeleteField, + isPending: false, + } as any); + mockUseProxyConfig.mockReturnValue({ + data: [], + isLoading: false, + refetch: mockRefetch, + } as any); + mockParseErrorMessage.mockImplementation((error: any) => error?.message || String(error)); + }); + + it("should render the card with title and form fields", () => { + renderWithProviders(); + + expect(screen.getByText("Logging Settings")).toBeInTheDocument(); + expect(screen.getByText("Store Prompts in Spend Logs")).toBeInTheDocument(); + expect(screen.getByLabelText("Maximum Spend Logs Retention Period (Optional)")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("e.g., 7d, 30d")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Save Settings" })).toBeInTheDocument(); + }); + + it("should toggle store prompts switch", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + expect(switchElement).not.toBeChecked(); + + await user.click(switchElement); + + await waitFor(() => { + expect(switchElement).toBeChecked(); + }); + }); + + it("should update retention period input", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); + await user.type(retentionInput, "30d"); + + expect(retentionInput).toHaveValue("30d"); + }); + + it("should submit form with store prompts enabled and retention period", async () => { + const user = userEvent.setup(); + mockMutate.mockImplementation((_params, options) => { + options?.onSuccess?.(); + }); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + await user.click(switchElement); + + const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); + await user.type(retentionInput, "30d"); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockDeleteField).not.toHaveBeenCalled(); + expect(mockMutate).toHaveBeenCalledWith( + { + store_prompts_in_spend_logs: true, + maximum_spend_logs_retention_period: "30d", + }, + expect.any(Object), + ); + }); + }); + + it("should delete retention period field when left empty on submit", async () => { + const user = userEvent.setup(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { + options?.onSuccess?.(); + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockDeleteField).toHaveBeenCalled(); + expect(mockMutate).toHaveBeenCalledWith( + { + store_prompts_in_spend_logs: false, + }, + expect.any(Object), + ); + }); + }); + + it("should show success notification on successful submission", async () => { + const user = userEvent.setup(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { + options?.onSuccess?.(); + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("Spend logs settings updated successfully"); + }); + }); + + it("should show a single error notification via onError callback", async () => { + const user = userEvent.setup(); + const error = new Error("Backend error"); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { + options?.onError?.(error); + }); + mockParseErrorMessage.mockReturnValue("Backend error"); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith( + "Failed to save spend logs settings: Backend error", + ); + }); + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledTimes(1); + }); + + it("should show loading state on save button when update pending", () => { + mockUseStoreRequestInSpendLogs.mockReturnValue({ + mutate: mockMutate, + isPending: true, + } as any); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: /Saving/i }); + expect(saveButton).toBeInTheDocument(); + expect(saveButton.className).toContain("ant-btn-loading"); + }); + + it("should show loading state on save button when delete pending", () => { + mockUseDeleteProxyConfigField.mockReturnValue({ + mutate: mockDeleteField, + isPending: true, + } as any); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: /Saving/i }); + expect(saveButton).toBeInTheDocument(); + expect(saveButton.className).toContain("ant-btn-loading"); + }); + + it("should disable save button while config is loading", () => { + mockUseProxyConfig.mockReturnValue({ + data: undefined, + isLoading: true, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + expect(saveButton).toBeDisabled(); + }); + + it("should render form with initial values from config data", () => { + mockUseProxyConfig.mockReturnValue({ + data: [ + { + field_name: "store_prompts_in_spend_logs", + field_type: "bool", + field_description: "Store prompts in spend logs", + field_value: true, + stored_in_db: true, + field_default_value: false, + }, + { + field_name: "maximum_spend_logs_retention_period", + field_type: "string", + field_description: "Maximum retention period", + field_value: "30d", + stored_in_db: true, + field_default_value: undefined, + }, + ], + isLoading: false, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d"); + + expect(switchElement).toBeChecked(); + expect(retentionInput).toHaveValue("30d"); + }); + + it("should show skeleton loaders when config is loading", () => { + mockUseProxyConfig.mockReturnValue({ + data: undefined, + isLoading: true, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + expect(screen.queryByPlaceholderText("e.g., 7d, 30d")).not.toBeInTheDocument(); + + const skeletons = document.querySelectorAll(".ant-skeleton"); + expect(skeletons.length).toBeGreaterThan(0); + }); + + it("should continue with update even if deleteField fails", async () => { + const user = userEvent.setup(); + const deleteError = new Error("Field does not exist"); + mockDeleteField.mockImplementation((_params, options) => { + options?.onError?.(deleteError); + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { + options?.onSuccess?.(); + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockDeleteField).toHaveBeenCalled(); + expect(mockMutate).toHaveBeenCalledWith( + { + store_prompts_in_spend_logs: false, + }, + expect.any(Object), + ); + expect(mockNotificationsManager.success).toHaveBeenCalled(); + }); + }); + + it("should submit with only store prompts enabled when retention is empty", async () => { + const user = userEvent.setup(); + mockDeleteField.mockImplementation((_params, options) => { + options?.onSettled?.(); + }); + mockMutate.mockImplementation((_params, options) => { + options?.onSuccess?.(); + }); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + await user.click(switchElement); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockDeleteField).toHaveBeenCalled(); + expect(mockMutate).toHaveBeenCalledWith( + { + store_prompts_in_spend_logs: true, + }, + expect.any(Object), + ); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx new file mode 100644 index 0000000000..456240ec5d --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx @@ -0,0 +1,145 @@ +"use client"; + +import { + ConfigType, + GeneralSettingsFieldName, + useDeleteProxyConfigField, + useProxyConfig, +} from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { + StoreRequestInSpendLogsParams, + useStoreRequestInSpendLogs, +} from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { ClockCircleOutlined } from "@ant-design/icons"; +import { Button, Card, Form, Input, Skeleton, Space, Switch, Typography } from "antd"; +import React, { useMemo } from "react"; + +const LoggingSettings: React.FC = () => { + const [form] = Form.useForm(); + const { mutate, isPending } = useStoreRequestInSpendLogs(); + const { mutate: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField(); + const { data: proxyConfigData, isLoading: isLoadingConfig } = useProxyConfig(ConfigType.GENERAL_SETTINGS); + const storePromptsValue = Form.useWatch("store_prompts_in_spend_logs", form); + + const initialValues = useMemo(() => { + if (!proxyConfigData) { + return { + store_prompts_in_spend_logs: false, + maximum_spend_logs_retention_period: undefined, + }; + } + + const storePromptsField = proxyConfigData.find((field) => field.field_name === "store_prompts_in_spend_logs"); + const retentionPeriodField = proxyConfigData.find( + (field) => field.field_name === "maximum_spend_logs_retention_period", + ); + + return { + store_prompts_in_spend_logs: storePromptsField?.field_value ?? false, + maximum_spend_logs_retention_period: retentionPeriodField?.field_value ?? undefined, + }; + }, [proxyConfigData]); + + const handleFormSubmit = (formValues: StoreRequestInSpendLogsParams) => { + const retentionPeriodValue = formValues.maximum_spend_logs_retention_period; + const hasRetentionPeriod = + typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() !== ""; + + const updateParams: StoreRequestInSpendLogsParams = { + store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs, + ...(hasRetentionPeriod && { maximum_spend_logs_retention_period: retentionPeriodValue }), + }; + + const submitUpdate = () => + mutate(updateParams, { + onSuccess: () => NotificationsManager.success("Spend logs settings updated successfully"), + onError: (error) => + NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)), + }); + + if (hasRetentionPeriod) { + submitUpdate(); + return; + } + + deleteField( + { + config_type: ConfigType.GENERAL_SETTINGS, + field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD, + }, + { + onError: (deleteError) => + console.warn("Failed to delete retention period field (may not exist):", deleteError), + onSettled: submitUpdate, + }, + ); + }; + + return ( + + + + Proxy-wide settings that control how request and response data are written to spend logs. + + +
+ f.field_name === "store_prompts_in_spend_logs")?.field_description || + "When enabled, prompts will be stored in spend logs for tracking and analysis purposes." + } + valuePropName="checked" + > + {isLoadingConfig ? ( + + ) : ( + form.setFieldValue("store_prompts_in_spend_logs", checked)} + /> + )} + + + f.field_name === "maximum_spend_logs_retention_period") + ?.field_description || + "Set the maximum retention period for spend logs (e.g., '7d' for 7 days, '30d' for 30 days). Leave empty for no limit." + } + > + {isLoadingConfig ? ( + + ) : ( + } /> + )} + + + + + + + + + ); +}; + +export default LoggingSettings; diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index c168ffa21e..08b66c382b 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -69,7 +69,7 @@ describe("Sidebar (leftnav)", () => { "Virtual Keys", "Playground", "Models + Endpoints", - "Agents", + "Agentic", "MCP Servers", "Guardrails", "Policies", diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index c340b65496..cecc99739b 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -3,6 +3,7 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { ApiOutlined, + ApartmentOutlined, AppstoreOutlined, AuditOutlined, BankOutlined, @@ -120,11 +121,31 @@ const menuGroups: MenuGroup[] = [ roles: rolesWithWriteAccess, }, { - key: "agents", - page: "agents", - label: "Agents", + key: "agentic", + page: "agentic", + label: "Agentic", icon: , - roles: rolesWithWriteAccess, + children: [ + { + key: "agents", + page: "agents", + label: "Agents", + icon: , + roles: rolesWithWriteAccess, + }, + { + key: "workflows", + page: "workflows", + label: "Workflow Runs", + icon: , + }, + { + key: "memory", + page: "memory", + label: "Memory", + icon: , + }, + ], }, { key: "mcp-servers", @@ -139,12 +160,6 @@ const menuGroups: MenuGroup[] = [ icon: , roles: all_admin_roles, }, - { - key: "memory", - page: "memory", - label: "Memory", - icon: , - }, { key: "guardrails", page: "guardrails", diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 63f03f4586..277bb0bd02 100644 --- a/ui/litellm-dashboard/src/components/page_metadata.ts +++ b/ui/litellm-dashboard/src/components/page_metadata.ts @@ -9,6 +9,8 @@ export const pageDescriptions: Record = { "llm-playground": "Interactive playground for testing LLM requests", models: "Configure and manage LLM models and endpoints", agents: "Create and manage AI agents", + agentic: "Manage agentic resources: agents, workflow runs, and memory", + workflows: "Track and inspect durable workflow run history", "mcp-servers": "Configure Model Context Protocol servers", memory: "Inspect and manage agent memory entries stored under /v1/memory", guardrails: "Set up content moderation and safety guardrails", diff --git a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx index 9e28b27cec..ad9b724ba8 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.test.tsx @@ -1,6 +1,5 @@ import { render, screen } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; -import userEvent from "@testing-library/user-event"; +import { describe, expect, it } from "vitest"; import { ConfigInfoMessage } from "./ConfigInfoMessage"; describe("ConfigInfoMessage", () => { @@ -19,23 +18,8 @@ describe("ConfigInfoMessage", () => { expect(screen.getByText(/store_prompts_in_spend_logs: true/)).toBeInTheDocument(); }); - it("should render the settings button when onOpenSettings is provided", () => { - render( {}} />); - expect(screen.getByText("open the settings")).toBeInTheDocument(); - }); - - it("should not render the settings button when onOpenSettings is omitted", () => { + it("should reference Admin Settings \u2192 Logging Settings", () => { render(); - expect(screen.queryByText("open the settings")).not.toBeInTheDocument(); - }); - - it("should call onOpenSettings when the settings button is clicked", async () => { - const user = userEvent.setup(); - const onOpenSettings = vi.fn(); - - render(); - await user.click(screen.getByText("open the settings")); - - expect(onOpenSettings).toHaveBeenCalledOnce(); + expect(screen.getByText(/Admin Settings → Logging Settings/)).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx index 509b1c73de..a231a09971 100644 --- a/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/ConfigInfoMessage.tsx @@ -2,10 +2,9 @@ import React from "react"; interface ConfigInfoMessageProps { show: boolean; - onOpenSettings?: () => void; } -export const ConfigInfoMessage: React.FC = ({ show, onOpenSettings }) => { +export const ConfigInfoMessage: React.FC = ({ show }) => { if (!show) return null; return ( @@ -31,18 +30,8 @@ export const ConfigInfoMessage: React.FC = ({ show, onOp

Request/Response Data Not Available

To view request and response details, enable prompt storage in your LiteLLM configuration by adding the - following to your proxy_config.yaml file - {onOpenSettings && ( - <> or{" "} - - {" "}to configure this directly. - - )} + following to your proxy_config.yaml file, or toggle + the setting in Admin Settings → Logging Settings. 

           {`general_settings:
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx
index 992063fb69..a2da513675 100644
--- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx
@@ -171,27 +171,6 @@ describe("LogDetailContent", () => {
     expect(screen.queryByText("Request/Response Data Not Available")).not.toBeInTheDocument();
   });
 
-  it("should call onOpenSettings when user clicks open settings in ConfigInfoMessage", async () => {
-    const onOpenSettings = vi.fn();
-    const user = userEvent.setup();
-
-    render(
-      ,
-    );
-
-    const settingsButton = screen.getByRole("button", { name: /open the settings/i });
-    await user.click(settingsButton);
-
-    expect(onOpenSettings).toHaveBeenCalledTimes(1);
-  });
-
   it("should display loading state when isLoadingDetails is true", () => {
     render(
        void;
   /** When true, log details (messages/response) are still being lazy-loaded. */
   isLoadingDetails?: boolean;
   accessToken?: string | null;
@@ -52,7 +51,7 @@ export interface LogDetailContentProps {
  * Designed to be placed inside LogDetailsDrawer's right panel so it can
  * be reused for both single-log and session-mode views.
  */
-export function LogDetailContent({ logEntry, onOpenSettings, isLoadingDetails = false, accessToken }: LogDetailContentProps) {
+export function LogDetailContent({ logEntry, isLoadingDetails = false, accessToken }: LogDetailContentProps) {
   const metadata = logEntry.metadata || {};
   const hasError = metadata.status === "failure";
   const errorInfo = hasError ? metadata.error_information : null;
@@ -158,7 +157,7 @@ export function LogDetailContent({ logEntry, onOpenSettings, isLoadingDetails =
       {/* Configuration Info Message */}
       {missingData && (
         

-          
+          
         

       )}
 
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx
index 036a24c045..39315b0b58 100644
--- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx
@@ -26,7 +26,6 @@ export interface LogDetailsDrawerProps {
   logEntry: LogEntry | null;
   sessionId?: string | null;
   accessToken?: string | null;
-  onOpenSettings?: () => void;
   allLogs?: LogEntry[];
   onSelectLog?: (log: LogEntry) => void;
   startTime?: string;
@@ -109,7 +108,6 @@ export function LogDetailsDrawer({
   logEntry,
   sessionId,
   accessToken,
-  onOpenSettings,
   allLogs = [],
   onSelectLog,
   startTime,
@@ -399,7 +397,6 @@ export function LogDetailsDrawer({
             

               
diff --git a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx
deleted file mode 100644
index 40d06d9046..0000000000
--- a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.test.tsx
+++ /dev/null
@@ -1,484 +0,0 @@
-import { useDeleteProxyConfigField, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig";
-import { useStoreRequestInSpendLogs } from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs";
-import NotificationsManager from "@/components/molecules/notifications_manager";
-import { parseErrorMessage } from "@/components/shared/errorUtils";
-import { screen, waitFor } from "@testing-library/react";
-import userEvent from "@testing-library/user-event";
-import { beforeEach, describe, expect, it, vi } from "vitest";
-import { renderWithProviders } from "../../../../tests/test-utils";
-import SpendLogsSettingsModal from "./SpendLogsSettingsModal";
-
-vi.mock("@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs");
-vi.mock("@/app/(dashboard)/hooks/proxyConfig/useProxyConfig");
-vi.mock("@/components/molecules/notifications_manager", () => ({
-  default: {
-    success: vi.fn(),
-    fromBackend: vi.fn(),
-  },
-}));
-vi.mock("@/components/shared/errorUtils", () => ({
-  parseErrorMessage: vi.fn(),
-}));
-
-const mockUseStoreRequestInSpendLogs = vi.mocked(useStoreRequestInSpendLogs);
-const mockUseProxyConfig = vi.mocked(useProxyConfig);
-const mockUseDeleteProxyConfigField = vi.mocked(useDeleteProxyConfigField);
-const mockNotificationsManager = vi.mocked(NotificationsManager);
-const mockParseErrorMessage = vi.mocked(parseErrorMessage);
-
-describe("SpendLogsSettingsModal", () => {
-  const mockOnCancel = vi.fn();
-  const mockOnSuccess = vi.fn();
-  const mockMutateAsync = vi.fn();
-  const mockDeleteField = vi.fn();
-  const mockRefetch = vi.fn();
-
-  const defaultProps = {
-    isVisible: true,
-    onCancel: mockOnCancel,
-    onSuccess: mockOnSuccess,
-  };
-
-  beforeEach(() => {
-    vi.clearAllMocks();
-    mockUseStoreRequestInSpendLogs.mockReturnValue({
-      mutateAsync: mockMutateAsync,
-      isPending: false,
-    } as any);
-    mockUseDeleteProxyConfigField.mockReturnValue({
-      mutateAsync: mockDeleteField,
-      isPending: false,
-    } as any);
-    mockUseProxyConfig.mockReturnValue({
-      data: [],
-      isLoading: false,
-      refetch: mockRefetch,
-    } as any);
-    mockParseErrorMessage.mockImplementation((error: any) => error?.message || String(error));
-  });
-
-  it("should render the modal", () => {
-    renderWithProviders();
-    expect(screen.getByRole("dialog")).toBeInTheDocument();
-    expect(screen.getByText("Spend Logs Settings")).toBeInTheDocument();
-  });
-
-  it("should render form fields with initial values", () => {
-    renderWithProviders();
-
-    expect(screen.getByText("Store Prompts in Spend Logs")).toBeInTheDocument();
-    expect(screen.getByLabelText("Maximum Spend Logs Retention Period (Optional)")).toBeInTheDocument();
-    expect(screen.getByPlaceholderText("e.g., 7d, 30d")).toBeInTheDocument();
-  });
-
-  it("should render cancel and save buttons", () => {
-    renderWithProviders();
-
-    expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument();
-    expect(screen.getByRole("button", { name: "Save Settings" })).toBeInTheDocument();
-  });
-
-  it("should call onCancel when cancel button is clicked", async () => {
-    const user = userEvent.setup();
-    renderWithProviders();
-
-    const cancelButton = screen.getByRole("button", { name: "Cancel" });
-    await user.click(cancelButton);
-
-    expect(mockOnCancel).toHaveBeenCalledTimes(1);
-  });
-
-  it("should call onCancel when modal close button is clicked", async () => {
-    const user = userEvent.setup();
-    renderWithProviders();
-
-    const closeButton = screen.getByRole("button", { name: /close/i });
-    await user.click(closeButton);
-
-    expect(mockOnCancel).toHaveBeenCalledTimes(1);
-  });
-
-  it("should toggle store prompts switch", async () => {
-    const user = userEvent.setup();
-    renderWithProviders();
-
-    const switchElement = screen.getByRole("switch");
-    expect(switchElement).not.toBeChecked();
-
-    await user.click(switchElement);
-
-    await waitFor(() => {
-      expect(switchElement).toBeChecked();
-    });
-  });
-
-  it("should update retention period input", async () => {
-    const user = userEvent.setup();
-    renderWithProviders();
-
-    const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d");
-    await user.type(retentionInput, "30d");
-
-    expect(retentionInput).toHaveValue("30d");
-  });
-
-  it("should submit form with store prompts enabled and retention period", async () => {
-    const user = userEvent.setup();
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const switchElement = screen.getByRole("switch");
-    await user.click(switchElement);
-
-    const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d");
-    await user.type(retentionInput, "30d");
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockDeleteField).not.toHaveBeenCalled();
-      expect(mockMutateAsync).toHaveBeenCalledWith(
-        {
-          store_prompts_in_spend_logs: true,
-          maximum_spend_logs_retention_period: "30d",
-        },
-        expect.any(Object)
-      );
-    });
-  });
-
-  it("should submit form with store prompts disabled and no retention period", async () => {
-    const user = userEvent.setup();
-    mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" });
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockDeleteField).toHaveBeenCalled();
-      expect(mockMutateAsync).toHaveBeenCalledWith(
-        {
-          store_prompts_in_spend_logs: false,
-        },
-        expect.any(Object)
-      );
-    });
-  });
-
-  it("should show success notification and call onSuccess on successful submission", async () => {
-    const user = userEvent.setup();
-    mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" });
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockNotificationsManager.success).toHaveBeenCalledWith("Spend logs settings updated successfully");
-      expect(mockRefetch).toHaveBeenCalled();
-      expect(mockOnSuccess).toHaveBeenCalledTimes(1);
-    });
-  });
-
-  it("should show error notification when submission fails", async () => {
-    const user = userEvent.setup();
-    const error = new Error("Network error");
-    mockMutateAsync.mockRejectedValue(error);
-    mockParseErrorMessage.mockReturnValue("Network error");
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save spend logs settings: Network error");
-    });
-  });
-
-  it("should show error notification from onError callback", async () => {
-    const user = userEvent.setup();
-    const error = new Error("Backend error");
-    mockMutateAsync.mockImplementation((params, options) => {
-      options?.onError?.(error);
-      return Promise.reject(error);
-    });
-    mockParseErrorMessage.mockReturnValue("Backend error");
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save spend logs settings: Backend error");
-    });
-  });
-
-  it("should disable cancel button when pending", () => {
-    mockUseStoreRequestInSpendLogs.mockReturnValue({
-      mutateAsync: mockMutateAsync,
-      isPending: true,
-    } as any);
-
-    renderWithProviders();
-
-    const cancelButton = screen.getByRole("button", { name: "Cancel" });
-    expect(cancelButton).toBeDisabled();
-  });
-
-  it("should disable cancel button when deleting field", () => {
-    mockUseDeleteProxyConfigField.mockReturnValue({
-      mutateAsync: mockDeleteField,
-      isPending: true,
-    } as any);
-
-    renderWithProviders();
-
-    const cancelButton = screen.getByRole("button", { name: "Cancel" });
-    expect(cancelButton).toBeDisabled();
-  });
-
-  it("should disable cancel button when loading config", () => {
-    mockUseProxyConfig.mockReturnValue({
-      data: undefined,
-      isLoading: true,
-      refetch: mockRefetch,
-    } as any);
-
-    renderWithProviders();
-
-    const cancelButton = screen.getByRole("button", { name: "Cancel" });
-    expect(cancelButton).toBeDisabled();
-  });
-
-  it("should show loading state on save button when pending", () => {
-    mockUseStoreRequestInSpendLogs.mockReturnValue({
-      mutateAsync: mockMutateAsync,
-      isPending: true,
-    } as any);
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: /Saving/i });
-    expect(saveButton).toBeInTheDocument();
-    expect(saveButton.className).toContain("ant-btn-loading");
-  });
-
-  it("should show loading state on save button when deleting field", () => {
-    mockUseDeleteProxyConfigField.mockReturnValue({
-      mutateAsync: mockDeleteField,
-      isPending: true,
-    } as any);
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: /Saving/i });
-    expect(saveButton).toBeInTheDocument();
-    expect(saveButton.className).toContain("ant-btn-loading");
-  });
-
-  it("should call onCancel when cancel button is clicked after modifying form", async () => {
-    const user = userEvent.setup();
-    renderWithProviders();
-
-    const switchElement = screen.getByRole("switch");
-    await user.click(switchElement);
-
-    const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d");
-    await user.type(retentionInput, "30d");
-
-    expect(switchElement).toBeChecked();
-    expect(retentionInput).toHaveValue("30d");
-
-    const cancelButton = screen.getByRole("button", { name: "Cancel" });
-    await user.click(cancelButton);
-
-    expect(mockOnCancel).toHaveBeenCalledTimes(1);
-  });
-
-  it("should call refetch after successful submission", async () => {
-    const user = userEvent.setup();
-    mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" });
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const switchElement = screen.getByRole("switch");
-    await user.click(switchElement);
-
-    const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d");
-    await user.type(retentionInput, "30d");
-
-    expect(switchElement).toBeChecked();
-    expect(retentionInput).toHaveValue("30d");
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockNotificationsManager.success).toHaveBeenCalled();
-      expect(mockRefetch).toHaveBeenCalled();
-    });
-  });
-
-  it("should not call onSuccess when it is not provided", async () => {
-    const user = userEvent.setup();
-    mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" });
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockNotificationsManager.success).toHaveBeenCalled();
-    });
-  });
-
-  it("should not render modal when isVisible is false", () => {
-    renderWithProviders();
-
-    expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
-  });
-
-  it("should call refetch when modal opens", () => {
-    renderWithProviders();
-
-    expect(mockRefetch).toHaveBeenCalledTimes(1);
-  });
-
-  it("should render form with initial values from config data", () => {
-    mockUseProxyConfig.mockReturnValue({
-      data: [
-        {
-          field_name: "store_prompts_in_spend_logs",
-          field_type: "bool",
-          field_description: "Store prompts in spend logs",
-          field_value: true,
-          stored_in_db: true,
-          field_default_value: false,
-        },
-        {
-          field_name: "maximum_spend_logs_retention_period",
-          field_type: "string",
-          field_description: "Maximum retention period",
-          field_value: "30d",
-          stored_in_db: true,
-          field_default_value: undefined,
-        },
-      ],
-      isLoading: false,
-      refetch: mockRefetch,
-    } as any);
-
-    renderWithProviders();
-
-    const switchElement = screen.getByRole("switch");
-    const retentionInput = screen.getByPlaceholderText("e.g., 7d, 30d");
-
-    expect(switchElement).toBeChecked();
-    expect(retentionInput).toHaveValue("30d");
-  });
-
-  it("should show skeleton loaders when config is loading", () => {
-    mockUseProxyConfig.mockReturnValue({
-      data: undefined,
-      isLoading: true,
-      refetch: mockRefetch,
-    } as any);
-
-    renderWithProviders();
-
-    // Check that switch and input are not present when loading (skeletons are shown instead)
-    expect(screen.queryByRole("switch")).not.toBeInTheDocument();
-    expect(screen.queryByPlaceholderText("e.g., 7d, 30d")).not.toBeInTheDocument();
-
-    // Check for skeleton elements (Ant Design Skeleton.Input renders with ant-skeleton class)
-    const skeletons = document.querySelectorAll(".ant-skeleton");
-    expect(skeletons.length).toBeGreaterThan(0);
-  });
-
-  it("should continue with update even if deleteField fails", async () => {
-    const user = userEvent.setup();
-    const deleteError = new Error("Field does not exist");
-    mockDeleteField.mockRejectedValue(deleteError);
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockDeleteField).toHaveBeenCalled();
-      expect(mockMutateAsync).toHaveBeenCalledWith(
-        {
-          store_prompts_in_spend_logs: false,
-        },
-        expect.any(Object)
-      );
-      expect(mockNotificationsManager.success).toHaveBeenCalled();
-    });
-  });
-
-  it("should submit form with only store prompts enabled and no retention period", async () => {
-    const user = userEvent.setup();
-    mockDeleteField.mockResolvedValue({ message: "Field deleted successfully" });
-    mockMutateAsync.mockImplementation(async (params, options) => {
-      await Promise.resolve();
-      options?.onSuccess?.();
-      return { message: "Success" };
-    });
-
-    renderWithProviders();
-
-    const switchElement = screen.getByRole("switch");
-    await user.click(switchElement);
-
-    const saveButton = screen.getByRole("button", { name: "Save Settings" });
-    await user.click(saveButton);
-
-    await waitFor(() => {
-      expect(mockDeleteField).toHaveBeenCalled();
-      expect(mockMutateAsync).toHaveBeenCalledWith(
-        {
-          store_prompts_in_spend_logs: true,
-        },
-        expect.any(Object)
-      );
-    });
-  });
-});
diff --git a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx b/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx
deleted file mode 100644
index b3f73af060..0000000000
--- a/ui/litellm-dashboard/src/components/view_logs/SpendLogsSettingsModal/SpendLogsSettingsModal.tsx
+++ /dev/null
@@ -1,156 +0,0 @@
-"use client";
-
-import { ConfigType, GeneralSettingsFieldName, useDeleteProxyConfigField, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig";
-import { StoreRequestInSpendLogsParams, useStoreRequestInSpendLogs } from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs";
-import NotificationsManager from "@/components/molecules/notifications_manager";
-import { parseErrorMessage } from "@/components/shared/errorUtils";
-import { ClockCircleOutlined } from "@ant-design/icons";
-import { Button, Form, Input, Modal, Skeleton, Space, Switch, Typography } from "antd";
-import React, { useEffect, useMemo } from "react";
-
-interface SpendLogsSettingsModalProps {
-  isVisible: boolean;
-  onCancel: () => void;
-  onSuccess?: () => void;
-}
-
-const SpendLogsSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => {
-  const [form] = Form.useForm();
-  const { mutateAsync, isPending } = useStoreRequestInSpendLogs();
-  const { mutateAsync: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField();
-  const { data: proxyConfigData, isLoading: isLoadingConfig, refetch } = useProxyConfig(ConfigType.GENERAL_SETTINGS);
-  const storePromptsValue = Form.useWatch('store_prompts_in_spend_logs', form);
-
-  // Refetch config when modal opens to ensure we have the latest values
-  useEffect(() => {
-    if (isVisible) {
-      refetch();
-    }
-  }, [isVisible, refetch]);
-
-  // Compute initial values from fetched config data
-  const initialValues = useMemo(() => {
-    if (!proxyConfigData) {
-      return {
-        store_prompts_in_spend_logs: false,
-        maximum_spend_logs_retention_period: undefined,
-      };
-    }
-
-    const storePromptsField = proxyConfigData.find(field => field.field_name === 'store_prompts_in_spend_logs');
-    const retentionPeriodField = proxyConfigData.find(field => field.field_name === 'maximum_spend_logs_retention_period');
-
-    return {
-      store_prompts_in_spend_logs: storePromptsField?.field_value ?? false,
-      maximum_spend_logs_retention_period: retentionPeriodField?.field_value ?? undefined,
-    };
-  }, [proxyConfigData]);
-
-  const handleFormSubmit = async (formValues: StoreRequestInSpendLogsParams) => {
-    try {
-      // If maximum_spend_logs_retention_period is empty/null, delete the field first
-      const retentionPeriodValue = formValues.maximum_spend_logs_retention_period;
-      const shouldDeleteRetentionPeriod =
-        !retentionPeriodValue ||
-        (typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() === "");
-
-      if (shouldDeleteRetentionPeriod) {
-        try {
-          await deleteField({
-            config_type: ConfigType.GENERAL_SETTINGS,
-            field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD,
-          });
-        } catch (deleteError) {
-          // If field doesn't exist, that's okay - continue with update
-          console.warn("Failed to delete retention period field (may not exist):", deleteError);
-        }
-      }
-
-      // Update the settings (excluding maximum_spend_logs_retention_period if it's empty)
-      const updateParams: StoreRequestInSpendLogsParams = {
-        store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs,
-        ...(retentionPeriodValue &&
-          typeof retentionPeriodValue === "string" &&
-          retentionPeriodValue.trim() !== "" && {
-          maximum_spend_logs_retention_period: retentionPeriodValue,
-        }),
-      };
-
-      await mutateAsync(updateParams, {
-        onSuccess: () => {
-          NotificationsManager.success("Spend logs settings updated successfully");
-          refetch(); // Refetch config to get updated values
-          onSuccess?.();
-        },
-        onError: (error) => {
-          NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error));
-        },
-      });
-    } catch (error) {
-      NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error));
-    }
-  };
-
-  const handleCancel = () => {
-    form.resetFields();
-    onCancel();
-  };
-
-  return (
-    Spend Logs Settings}
-      open={isVisible}
-      footer={
-        
-          
-          
-        
-      }
-      onCancel={handleCancel}
-    >
-
-      

-         f.field_name === 'store_prompts_in_spend_logs')?.field_description ||
-            "When enabled, prompts will be stored in spend logs for tracking and analysis purposes."
-          }
-          valuePropName="checked"
-        >
-          

-
-            {isLoadingConfig ?  :  form.setFieldValue('store_prompts_in_spend_logs', checked)} />}
-          

-        
-
-         f.field_name === 'maximum_spend_logs_retention_period')?.field_description ||
-            "Set the maximum retention period for spend logs (e.g., '7d' for 7 days, '30d' for 30 days). Leave empty for no limit."
-          }
-        >
-          {isLoadingConfig ?  : }
-          />}
-        
-      
-    
-  );
-};
-
-export default SpendLogsSettingsModal;
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx
index 8205c0b4b8..8a015d8305 100644
--- a/ui/litellm-dashboard/src/components/view_logs/index.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx
@@ -4,7 +4,7 @@ import { useCallback, useDeferredValue, useEffect, useRef, useState } from "reac
 import GuardrailViewer from "@/components/view_logs/GuardrailViewer/GuardrailViewer";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { truncateString } from "@/utils/textUtils";
-import { SettingOutlined, SyncOutlined } from "@ant-design/icons";
+import { SyncOutlined } from "@ant-design/icons";
 import { Row } from "@tanstack/react-table";
 import { Switch, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
 import { Button, Tag, Tooltip } from "antd";
@@ -28,7 +28,6 @@ import { useLogFilterLogic } from "./log_filter_logic";
 import { LogDetailsDrawer } from "./LogDetailsDrawer";
 import { getTimeRangeDisplay } from "./logs_utils";
 import { RequestResponsePanel } from "./RequestResponsePanel";
-import SpendLogsSettingsModal from "./SpendLogsSettingsModal/SpendLogsSettingsModal";
 import { DataTable } from "./table";
 import { VectorStoreViewer } from "./VectorStoreViewer";
 
@@ -85,7 +84,6 @@ export default function SpendLogsTable({
   const [selectedLog, setSelectedLog] = useState(null);
   const [isDrawerOpen, setIsDrawerOpen] = useState(false);
   const [selectedSessionId, setSelectedSessionId] = useState(null);
-  const [isSpendLogsSettingsModalVisible, setIsSpendLogsSettingsModalVisible] = useState(false);
 
   const [sortBy, setSortBy] = useState("startTime");
   const [sortOrder, setSortOrder] = useState<"asc" | "desc">("desc");
@@ -490,11 +488,6 @@ export default function SpendLogsTable({
           
             

               
Request Logs

-              

             {selectedKeyInfo && selectedKeyIdInfoView && selectedKeyInfo.api_key === selectedKeyIdInfoView ? (
               
-                 setIsSpendLogsSettingsModalVisible(false)}
-                  onSuccess={() => setIsSpendLogsSettingsModalVisible(false)}
-                />
                 

                   

                     

@@ -725,7 +713,6 @@ export default function SpendLogsTable({
         logEntry={selectedLog}
         sessionId={selectedSessionId}
         accessToken={accessToken}
-        onOpenSettings={() => setIsSpendLogsSettingsModalVisible(true)}
         allLogs={filteredData}
         onSelectLog={handleSelectLog}
         startTime={moment(startTime).utc().format("YYYY-MM-DD HH:mm:ss")}
@@ -734,7 +721,7 @@ export default function SpendLogsTable({
   );
 }
 
-export function RequestViewer({ row, onOpenSettings }: { row: Row; onOpenSettings?: () => void }) {
+export function RequestViewer({ row }: { row: Row }) {
   // Helper function to clean metadata by removing specific fields
   const formatData = (input: any) => {
     if (typeof input === "string") {
@@ -961,7 +948,7 @@ export function RequestViewer({ row, onOpenSettings }: { row: Row; onO
       />
 
       {/* Configuration Info Message - Show when data is missing */}
-      
+      
 
       {/* Request/Response Panel */}
       

diff --git a/ui/litellm-dashboard/src/components/workflow_runs/index.tsx b/ui/litellm-dashboard/src/components/workflow_runs/index.tsx
new file mode 100644
index 0000000000..a69eac282b
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/workflow_runs/index.tsx
@@ -0,0 +1,751 @@
+import React, { useState, useEffect, useCallback } from "react";
+import { Button, Collapse, Drawer, Empty, Spin, Table, Tooltip, Typography } from "antd";
+import { ReloadOutlined } from "@ant-design/icons";
+import { proxyBaseUrl } from "@/components/networking";
+
+const { Text } = Typography;
+
+interface WorkflowRunsProps {
+  accessToken: string | null;
+}
+
+type RunStatus = "pending" | "running" | "paused" | "completed" | "failed";
+
+interface RunMetadata {
+  title?: string;
+  state?: string;
+  pr_url?: string;
+  worktree_path?: string;
+  plan_text?: string;
+  grill_session_id?: string;
+  session_id?: string;
+  [key: string]: unknown;
+}
+
+interface WorkflowRun {
+  run_id: string;
+  status: RunStatus;
+  workflow_type: string;
+  created_at: string;
+  metadata?: RunMetadata | null;
+}
+
+interface WorkflowRunEvent {
+  event_id: string;
+  event_type: string;
+  step_name: string;
+  sequence_number: number;
+  created_at: string;
+  data?: Record | null;
+}
+
+interface WorkflowRunMessage {
+  message_id: string;
+  role: string;
+  content: string;
+  sequence_number: number;
+  created_at: string;
+}
+
+// ── design tokens ─────────────────────────────────────────────────────────────
+
+const STATUS_DOT: Record = {
+  pending:   "#a1a1aa",
+  running:   "#3b82f6",
+  paused:    "#f59e0b",
+  completed: "#22c55e",
+  failed:    "#ef4444",
+};
+
+const EVENT_COLOR: Record = {
+  "step.started":  { bar: "#f0fdf4", border: "#86efac", text: "#16a34a" },
+  "step.failed":   { bar: "#fef2f2", border: "#fca5a5", text: "#dc2626" },
+  "hook.waiting":  { bar: "#fffbeb", border: "#fcd34d", text: "#d97706" },
+  "hook.received": { bar: "#eff6ff", border: "#93c5fd", text: "#2563eb" },
+};
+
+function eventStyle(type: string) {
+  return EVENT_COLOR[type] ?? { bar: "#f4f4f5", border: "#d4d4d8", text: "#52525b" };
+}
+
+// ── helpers ───────────────────────────────────────────────────────────────────
+
+function timeAgo(iso: string): string {
+  const diff = Date.now() - new Date(iso).getTime();
+  if (isNaN(diff)) return iso;
+  const s = Math.floor(diff / 1000);
+  if (s < 60) return `${s}s ago`;
+  const m = Math.floor(s / 60);
+  if (m < 60) return `${m}m ago`;
+  const h = Math.floor(m / 60);
+  if (h < 24) return `${h}h ago`;
+  return `${Math.floor(h / 24)}d ago`;
+}
+
+function fmtDuration(ms: number): string {
+  if (ms < 0) return "";
+  if (ms < 1000) return `${ms}ms`;
+  return `${(ms / 1000).toFixed(1)}s`;
+}
+
+function runTitle(run: WorkflowRun): string {
+  const t = run.metadata?.title;
+  if (t) return String(t);
+  return run.workflow_type ?? run.run_id.slice(0, 8);
+}
+
+function shortId(id: string): string {
+  return id.slice(0, 8);
+}
+
+// ── status dot ────────────────────────────────────────────────────────────────
+
+const StatusDot: React.FC<{ status: RunStatus; size?: number }> = ({ status, size = 8 }) => (
+  
+);
+
+// ── truncated text value ──────────────────────────────────────────────────────
+
+const TRUNCATE_AT = 120;
+
+const TruncatedValue: React.FC<{ value: string }> = ({ value }) => {
+  const [expanded, setExpanded] = useState(false);
+  if (value.length <= TRUNCATE_AT) {
+    return {value};
+  }
+  return (
+    
+      {expanded ? value : value.slice(0, TRUNCATE_AT) + "…"}
+      
+    
+  );
+};
+
+// ── metadata card ─────────────────────────────────────────────────────────────
+
+const MetadataCard: React.FC<{ run: WorkflowRun }> = ({ run }) => {
+  const meta = run.metadata ?? {};
+
+  const primaryFields: { key: string; label: string }[] = [
+    { key: "state",            label: "state" },
+    { key: "worktree_path",    label: "worktree" },
+    { key: "grill_session_id", label: "grill session" },
+    { key: "session_id",       label: "session" },
+  ];
+
+  const primaryKeys = new Set(["title", ...primaryFields.map((f) => f.key)]);
+  const extraEntries = Object.entries(meta).filter(
+    ([k, v]) => !primaryKeys.has(k) && v !== null && v !== undefined && v !== ""
+  );
+
+  return (
+    

+      {/* title bar */}
+      

+        
+        
+          {runTitle(run)}
+        
+        
+          {shortId(run.run_id)}
+        
+        
+          {run.workflow_type}
+        
+      

+
+      {/* key fields grid */}
+      

+        
+          {run.status}
+        
+        
+          {timeAgo(run.created_at)}
+        
+
+        {meta.pr_url && (
+          
+            
+              {String(meta.pr_url)}
+            
+          
+        )}
+
+        {primaryFields.map(({ key, label }) => {
+          const v = meta[key];
+          if (v === null || v === undefined || v === "") return null;
+          const str = typeof v === "object" ? JSON.stringify(v) : String(v);
+          return (
+            
+              
+            
+          );
+        })}
+
+        {extraEntries.map(([k, v]) => {
+          const str = typeof v === "object" ? JSON.stringify(v) : String(v);
+          return (
+            
+              
+            
+          );
+        })}
+      

+    

+  );
+};
+
+const FieldPair: React.FC<{ label: string; children: React.ReactNode }> = ({
+  label,
+  children,
+}) => (
+  

+    
+      {label}
+    
+    {children}
+  

+);
+
+// ── gantt timeline ────────────────────────────────────────────────────────────
+
+const GanttTimeline: React.FC<{
+  run: WorkflowRun;
+  events: WorkflowRunEvent[];
+}> = ({ run, events }) => {
+  if (events.length === 0) {
+    return (
+      

+        No events recorded
+      

+    );
+  }
+
+  const runStart = new Date(run.created_at).getTime();
+  const eventTimes = events.map((e) => new Date(e.created_at).getTime());
+  const lastTime = Math.max(...eventTimes);
+  const totalSpan = Math.max(lastTime - runStart, 1);
+  const totalDur = fmtDuration(lastTime - runStart);
+
+  return (
+    

+      {/* ruler */}
+      

+        

+        

+          {[0, 100].map((pct) => (
+            
+              {pct === 0 ? "0" : totalDur}
+            
+          ))}
+        

+      

+
+      {/* outer run bar */}
+      

+        

+          {runTitle(run)}
+        

+        

+          {totalDur}
+        

+      

+
+      {/* event rows */}
+      

+        {events.map((ev) => {
+          const evTime = new Date(ev.created_at).getTime();
+          const leftPct = ((evTime - runStart) / totalSpan) * 100;
+
+          const nextIdx = events.findIndex((e) => e.sequence_number > ev.sequence_number);
+          const nextTime =
+            nextIdx >= 0
+              ? new Date(events[nextIdx].created_at).getTime()
+              : lastTime + Math.max(totalSpan * 0.12, 500);
+          const widthPct = Math.max(8, ((nextTime - evTime) / totalSpan) * 100);
+          const style = eventStyle(ev.event_type);
+          const dur = fmtDuration(nextTime - evTime);
+
+          return (
+            
+              

+                {ev.step_name || ev.event_type}
+              

+              

+                
+                      
type: {ev.event_type}

+                      
step: {ev.step_name}

+                      
seq:  {ev.sequence_number}

+                      
time: {timeAgo(ev.created_at)}

+                      {ev.data && Object.keys(ev.data).length > 0 && (
+                        
data: {JSON.stringify(ev.data)}

+                      )}
+                    

+                  }
+                >
+                  

+                    {ev.event_type}
+                    {dur && {dur}}
+                  

+                
+              

+            
+          );
+        })}
+      

+    

+  );
+};
+
+// ── message row ───────────────────────────────────────────────────────────────
+
+const MessageRow: React.FC<{ msg: WorkflowRunMessage }> = ({ msg }) => {
+  const roleColor: Record = {
+    user:        "#2563eb",
+    assistant:   "#16a34a",
+    system:      "#7c3aed",
+    tool_result: "#d97706",
+  };
+  const color = roleColor[msg.role] ?? "#52525b";
+
+  return (
+    

+      [{msg.role}]
+      

+        
+          {msg.content}
+        
+        
+          {timeAgo(msg.created_at)}
+        
+      

+    

+  );
+};
+
+// ── main component ────────────────────────────────────────────────────────────
+
+const WorkflowRuns: React.FC = ({ accessToken }) => {
+  const [runs, setRuns] = useState([]);
+  const [loadingRuns, setLoadingRuns] = useState(false);
+  const [selectedRun, setSelectedRun] = useState(null);
+  const [events, setEvents] = useState([]);
+  const [messages, setMessages] = useState([]);
+  const [loadingDetail, setLoadingDetail] = useState(false);
+  const [drawerOpen, setDrawerOpen] = useState(false);
+
+  const fetchRuns = useCallback(async () => {
+    if (!accessToken) return;
+    setLoadingRuns(true);
+    try {
+      const res = await fetch(`${proxyBaseUrl ?? ""}/v1/workflows/runs?limit=100`, {
+        headers: { Authorization: `Bearer ${accessToken}` },
+      });
+      if (!res.ok) throw new Error(`HTTP ${res.status}`);
+      const data = await res.json();
+      setRuns(data.runs ?? []);
+    } catch (err) {
+      console.error("workflow runs fetch failed:", err);
+    } finally {
+      setLoadingRuns(false);
+    }
+  }, [accessToken]);
+
+  const fetchRunDetail = useCallback(
+    async (run: WorkflowRun) => {
+      if (!accessToken) return;
+      setSelectedRun(run);
+      setDrawerOpen(true);
+      setLoadingDetail(true);
+      setEvents([]);
+      setMessages([]);
+      try {
+        const base = proxyBaseUrl ?? "";
+        const [evRes, msgRes] = await Promise.all([
+          fetch(`${base}/v1/workflows/runs/${run.run_id}/events`, {
+            headers: { Authorization: `Bearer ${accessToken}` },
+          }),
+          fetch(`${base}/v1/workflows/runs/${run.run_id}/messages`, {
+            headers: { Authorization: `Bearer ${accessToken}` },
+          }),
+        ]);
+        const evData = evRes.ok ? await evRes.json() : { events: [] };
+        const msgData = msgRes.ok ? await msgRes.json() : { messages: [] };
+        setEvents(
+          [...(evData.events ?? [])].sort(
+            (a: WorkflowRunEvent, b: WorkflowRunEvent) => a.sequence_number - b.sequence_number
+          )
+        );
+        setMessages(
+          [...(msgData.messages ?? [])].sort(
+            (a: WorkflowRunMessage, b: WorkflowRunMessage) => a.sequence_number - b.sequence_number
+          )
+        );
+      } catch (err) {
+        console.error("workflow run detail fetch failed:", err);
+      } finally {
+        setLoadingDetail(false);
+      }
+    },
+    [accessToken]
+  );
+
+  useEffect(() => {
+    fetchRuns();
+  }, [fetchRuns]);
+
+  const columns = [
+    {
+      title: "Run",
+      dataIndex: "run_id",
+      key: "run",
+      render: (_: string, run: WorkflowRun) => (
+        

+          
+          

+            

+              {runTitle(run)}
+            

+            

+              {shortId(run.run_id)}
+            

+          

+        

+      ),
+    },
+    {
+      title: "Type",
+      dataIndex: "workflow_type",
+      key: "workflow_type",
+      render: (v: string) => (
+        {v}
+      ),
+    },
+    {
+      title: "Status",
+      dataIndex: "status",
+      key: "status",
+      render: (status: RunStatus, run: WorkflowRun) => {
+        const state = run.metadata?.state;
+        return (
+          

+            
+            
+              {state ?? status}
+            
+          

+        );
+      },
+    },
+    {
+      title: "Created",
+      dataIndex: "created_at",
+      key: "created_at",
+      render: (v: string) => (
+        {timeAgo(v)}
+      ),
+    },
+  ];
+
+  return (
+    

+      {/* page header */}
+      

+        

+          
Workflow Runs

+          

+            Durable state tracking for agents and automated workflows
+          

+        

+        
+      

+
+      {/* runs table — matches logs page density */}
+      

+
({ + onClick: () => fetchRunDetail(run), + style: { cursor: "pointer" }, + })} + locale={{ + emptyText: ( + No workflow runs yet} + image={Empty.PRESENTED_IMAGE_SIMPLE} + /> + ), + }} + className="[&_.ant-table-cell]:py-0.5 [&_.ant-table-thead_.ant-table-cell]:py-1" + style={{ border: "none" }} + /> + + + {/* detail drawer */} + setDrawerOpen(false)} + width={680} + title={null} + closable={false} + bodyStyle={{ padding: 0 }} + styles={{ body: { padding: 0 } }} + > + {!selectedRun ? null : loadingDetail ? ( +
+ +
+ ) : ( +
+ {/* drawer close + refresh */} +
+ + +
+ + {/* metadata card — top */} + + + {/* collapsible sections */} + + Timeline + + {events.length} {events.length === 1 ? "event" : "events"} + + + ), + children: ( +
+ +
+ ), + }, + { + key: "messages", + label: ( + + Messages + + {messages.length} + + + ), + children: messages.length === 0 ? ( +
+ No messages +
+ ) : ( +
+ {messages.map((msg) => ( + + ))} +
+ ), + }, + ]} + /> +
+ )} + + + ); +}; + +export default WorkflowRuns; diff --git a/uv.lock b/uv.lock index 53f032cfba..f837e2b5ef 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-23T02:32:27.506663Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" [manifest] @@ -3085,7 +3085,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.83.14" +version = "1.84.0" source = { editable = "." } dependencies = [ { name = "aiohttp" },