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/.npmrc b/.npmrc index 168e81a1c4..7999681cc3 100644 --- a/.npmrc +++ b/.npmrc @@ -2,4 +2,4 @@ # Packages needing lifecycle scripts: npm rebuild ignore-scripts=true # Protects local npm install only — npm ci (used in CI) ignores this -min-release-age=3d +min-release-age=3 diff --git a/litellm-js/proxy/.npmrc b/litellm-js/proxy/.npmrc index 168e81a1c4..7999681cc3 100644 --- a/litellm-js/proxy/.npmrc +++ b/litellm-js/proxy/.npmrc @@ -2,4 +2,4 @@ # Packages needing lifecycle scripts: npm rebuild ignore-scripts=true # Protects local npm install only — npm ci (used in CI) ignores this -min-release-age=3d +min-release-age=3 diff --git a/litellm-js/spend-logs/.npmrc b/litellm-js/spend-logs/.npmrc index 168e81a1c4..7999681cc3 100644 --- a/litellm-js/spend-logs/.npmrc +++ b/litellm-js/spend-logs/.npmrc @@ -2,4 +2,4 @@ # Packages needing lifecycle scripts: npm rebuild ignore-scripts=true # Protects local npm install only — npm ci (used in CI) ignores this -min-release-age=3d +min-release-age=3 diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429120000_search_tools_on_object_permission/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429120000_search_tools_on_object_permission/migration.sql new file mode 100644 index 0000000000..bffdaaebc5 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260429120000_search_tools_on_object_permission/migration.sql @@ -0,0 +1,2 @@ +-- Search tool allowlists live on LiteLLM_ObjectPermissionTable (with agents, MCP, vector stores). +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN IF NOT EXISTS "search_tools" TEXT[] DEFAULT ARRAY[]::TEXT[]; 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..a9d3911c07 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -277,6 +277,7 @@ model LiteLLM_ObjectPermissionTable { models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user + search_tools String[] @default([]) // search_tool_name values this key/team/user may call teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -1290,3 +1291,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/batches/main.py b/litellm/batches/main.py index 23f444b1ce..259439d4d0 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -543,15 +543,17 @@ def _handle_retrieve_batch_providers_without_provider_config( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=( + "LiteLLM doesn't support custom_llm_provider={} for 'retrieve_batch' without a `model` kwarg. " + "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. " + "'bedrock' is supported but requires `model` to be passed so the provider config can be loaded." + ).format(custom_llm_provider), model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( status_code=400, content="Unsupported provider", - request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"), # type: ignore ), ) return response diff --git a/litellm/constants.py b/litellm/constants.py index fc9f5730cd..d78c124d71 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 @@ -1383,6 +1393,10 @@ except (ValueError, TypeError): LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check" LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli" LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs" +# Stable identifier substituted in place of the master key on UserAPIKeyAuth +# objects so the master key (or its hash) never propagates to spend logs, +# Prometheus metrics, audit trails, or any other downstream consumer. +LITELLM_PROXY_MASTER_KEY_ALIAS = "litellm_proxy_master_key" # Key Rotation Constants LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false") 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/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/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/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index fe8387476e..3a83162fb2 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1661,6 +1661,20 @@ def _sanitize_anthropic_tool_use_id(tool_use_id: str) -> str: return sanitized +_ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES = {"application/pdf", "text/plain"} + + +def _is_anthropic_document_data_uri(url: str) -> bool: + # Anthropic's base64 document source accepts only application/pdf and + # text/plain (see select_anthropic_content_block_type_for_file). Routing + # other mimes here would produce a document block the API rejects, so we + # leave them on the image code path. + match = re.match(r"data:([^;,]+)", url) + if not match: + return False + return match.group(1) in _ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES + + def convert_to_anthropic_tool_result( message: Union[ChatCompletionToolMessage, ChatCompletionFunctionMessage], force_base64: bool = False, @@ -1698,14 +1712,24 @@ def convert_to_anthropic_tool_result( """ anthropic_content: Union[ str, - List[Union[AnthropicMessagesToolResultContent, AnthropicMessagesImageParam]], + List[ + Union[ + AnthropicMessagesToolResultContent, + AnthropicMessagesImageParam, + AnthropicMessagesDocumentParam, + ] + ], ] = "" if isinstance(message["content"], str): anthropic_content = message["content"] elif isinstance(message["content"], List): content_list = message["content"] anthropic_content_list: List[ - Union[AnthropicMessagesToolResultContent, AnthropicMessagesImageParam] + Union[ + AnthropicMessagesToolResultContent, + AnthropicMessagesImageParam, + AnthropicMessagesDocumentParam, + ] ] = [] for content in content_list: if content["type"] == "text": @@ -1720,21 +1744,62 @@ def convert_to_anthropic_tool_result( text_content["cache_control"] = cache_control_value anthropic_content_list.append(text_content) elif content["type"] == "image_url": + image_url_value = content["image_url"] format = ( - content["image_url"].get("format") - if isinstance(content["image_url"], dict) + image_url_value.get("format") + if isinstance(image_url_value, dict) else None ) - _anthropic_image_param = create_anthropic_image_param( - content["image_url"], format=format, is_bedrock_invoke=force_base64 + url_str = ( + image_url_value.get("url") + if isinstance(image_url_value, dict) + else image_url_value ) - _anthropic_image_param = add_cache_control_to_content( - anthropic_content_element=_anthropic_image_param, + # Data URIs with non-image mime types (e.g. application/pdf) must + # translate to Anthropic document blocks, not image blocks — + # wrapping a PDF in `type: "image"` is rejected by the API. + if isinstance(url_str, str) and _is_anthropic_document_data_uri( + url_str + ): + synth_file_message: ChatCompletionFileObject = { + "type": "file", + "file": {"file_data": url_str}, + } + _document_block = anthropic_process_openai_file_message( + synth_file_message + ) + _document_block = add_cache_control_to_content( + anthropic_content_element=cast( + AnthropicMessagesDocumentParam, _document_block + ), + original_content_element=content, + ) + anthropic_content_list.append( + cast(AnthropicMessagesDocumentParam, _document_block) + ) + else: + _anthropic_image_param = create_anthropic_image_param( + image_url_value, + format=format, + is_bedrock_invoke=force_base64, + ) + _anthropic_image_param = add_cache_control_to_content( + anthropic_content_element=_anthropic_image_param, + original_content_element=content, + ) + anthropic_content_list.append( + cast(AnthropicMessagesImageParam, _anthropic_image_param) + ) + elif content["type"] == "file": + file_content = cast(ChatCompletionFileObject, content) + _file_block = anthropic_process_openai_file_message(file_content) + _file_block = add_cache_control_to_content( + anthropic_content_element=cast( + AnthropicMessagesDocumentParam, _file_block + ), original_content_element=content, ) - anthropic_content_list.append( - cast(AnthropicMessagesImageParam, _anthropic_image_param) - ) + anthropic_content_list.append(_file_block) anthropic_content = anthropic_content_list anthropic_tool_result: Optional[AnthropicMessagesToolResultParam] = None @@ -3977,6 +4042,55 @@ def _convert_to_bedrock_tool_call_result( tool_result_content_blocks.append( BedrockToolResultContentBlock(image=_block["image"]) ) + elif "document" in _block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(document=_block["document"]) + ) + else: + verbose_logger.warning( + "Bedrock Converse: unrecognized BedrockContentBlock keys " + "%s for image_url tool-result block %s; dropping.", + list(_block.keys()), + content, + ) + elif content["type"] == "file": + # Match the user-message path (_process_file_message): accept + # either file_data (base64 data URI) or file_id (server-side + # reference / URL) and hand off to BedrockImageProcessor. Raise + # BadRequestError on both-None rather than silently dropping. + file_obj = content.get("file") or {} + file_data = file_obj.get("file_data") + file_id = file_obj.get("file_id") + if file_data is None and file_id is None: + raise litellm.BadRequestError( + message="file_data and file_id cannot both be None. Got={}".format( + content + ), + model="", + llm_provider="bedrock", + ) + file_format = file_obj.get("format") + _file_block: BedrockContentBlock = ( + BedrockImageProcessor.process_image_sync( + image_url=cast(str, file_id or file_data), + format=file_format, + ) + ) + if "document" in _file_block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(document=_file_block["document"]) + ) + elif "image" in _file_block: + tool_result_content_blocks.append( + BedrockToolResultContentBlock(image=_file_block["image"]) + ) + else: + verbose_logger.warning( + "Bedrock Converse: unrecognized BedrockContentBlock keys " + "%s for file tool-result block %s; dropping.", + list(_file_block.keys()), + content, + ) message.get("name", "") id = str(message.get("tool_call_id", str(uuid.uuid4()))) diff --git a/litellm/litellm_core_utils/sensitive_data_masker.py b/litellm/litellm_core_utils/sensitive_data_masker.py index 8b88ef9482..d7803455b4 100644 --- a/litellm/litellm_core_utils/sensitive_data_masker.py +++ b/litellm/litellm_core_utils/sensitive_data_masker.py @@ -21,6 +21,8 @@ class SensitiveDataMasker: "auth", "authorization", "credential", + # Plural form: Vertex uses ``vertex_credentials``; segment-exact + # matching otherwise misses it because "credential" != "credentials". "credentials", "access", "private", 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/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 8511d785fb..13a45fd165 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -712,6 +712,7 @@ }, "anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", @@ -735,6 +736,7 @@ }, "anthropic.claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", @@ -955,6 +957,7 @@ }, "anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -982,6 +985,7 @@ }, "anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1011,6 +1015,7 @@ }, "global.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1040,6 +1045,7 @@ }, "us.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", @@ -1127,6 +1133,7 @@ }, "anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1157,6 +1164,7 @@ }, "global.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1187,6 +1195,7 @@ }, "us.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", @@ -1277,6 +1286,7 @@ }, "anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", @@ -1305,6 +1315,7 @@ }, "global.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", @@ -1333,6 +1344,7 @@ }, "us.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, + "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", @@ -1447,11 +1459,13 @@ }, "anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, @@ -5103,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", @@ -17889,11 +17935,13 @@ }, "global.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, @@ -17950,6 +17998,7 @@ }, "global.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", @@ -19083,6 +19132,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", @@ -30052,6 +30133,7 @@ }, "us.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_1hr": 2.2e-06, "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", @@ -30203,11 +30285,13 @@ }, "us.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, + "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "input_cost_per_token_above_200k_tokens": 6.6e-06, "output_cost_per_token_above_200k_tokens": 2.475e-05, "cache_creation_input_token_cost_above_200k_tokens": 8.25e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, @@ -30308,6 +30392,7 @@ }, "us.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", @@ -30335,6 +30420,7 @@ }, "global.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index edee50bdfc..c4c9aea6f6 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -21,6 +21,7 @@ import httpx from httpx._types import CookieTypes, QueryParamTypes, RequestFiles import litellm +from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -390,19 +391,28 @@ def _sync_streaming( ): from litellm.utils import executor + raw_bytes: List[bytes] = [] + flush_scheduled = False try: - raw_bytes: List[bytes] = [] for chunk in response.iter_bytes(): # type: ignore raw_bytes.append(chunk) yield chunk - - executor.submit( - litellm_logging_obj.flush_passthrough_collected_chunks, - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - except Exception as e: - raise e + finally: + if not flush_scheduled and raw_bytes: + flush_scheduled = True + try: + executor.submit( + litellm_logging_obj.flush_passthrough_collected_chunks, + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + except Exception as e: + verbose_logger.exception( + "Failed to schedule passthrough spend-tracking flush " + "in _sync_streaming; %d buffered chunks dropped: %s", + len(raw_bytes), + e, + ) async def _async_streaming( @@ -411,23 +421,45 @@ async def _async_streaming( provider_config: "BasePassthroughConfig", ): iter_response = await response + try: iter_response.raise_for_status() - raw_bytes: List[bytes] = [] - - async for chunk in iter_response.aiter_bytes(): # type: ignore - raw_bytes.append(chunk) - yield chunk - - asyncio.create_task( - litellm_logging_obj.async_flush_passthrough_collected_chunks( - raw_bytes=raw_bytes, - provider_config=provider_config, - ) - ) except Exception: try: await iter_response.aclose() except Exception: pass raise + + raw_bytes: List[bytes] = [] + flush_scheduled = False + try: + async for chunk in iter_response.aiter_bytes(): # type: ignore + raw_bytes.append(chunk) + yield chunk + except Exception: + try: + await iter_response.aclose() + except Exception: + pass + raise + finally: + # GeneratorExit (raised on client disconnect) is not caught by + # `except Exception`; the finally block ensures partial usage + # still gets flushed for spend tracking. See LIT-2642. + if not flush_scheduled and raw_bytes: + flush_scheduled = True + try: + asyncio.create_task( + litellm_logging_obj.async_flush_passthrough_collected_chunks( + raw_bytes=raw_bytes, + provider_config=provider_config, + ) + ) + except Exception as e: + verbose_logger.exception( + "Failed to schedule passthrough spend-tracking flush " + "in _async_streaming; %d buffered chunks dropped: %s", + len(raw_bytes), + e, + ) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index d4799e9f20..756b2ed91d 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -117,7 +117,10 @@ class MCPRequestHandler: return b"{}" request.body = mock_body # type: ignore - if ".well-known" in str(request.url): # public routes + # Only OAuth metadata routes registered under /.well-known/ are public. + # Match on request.url.path (path-only, exact prefix) so the substring + # cannot be smuggled via query string, hostname, or a deeper URL segment. + if request.url.path.startswith("/.well-known/"): validated_user_api_key_auth = UserAPIKeyAuth() elif has_explicit_litellm_key: # Explicit x-litellm-api-key provided - always validate normally @@ -126,27 +129,37 @@ class MCPRequestHandler: ) elif oauth2_headers: # No x-litellm-api-key, but Authorization header present. - # Could be a LiteLLM key (backward compat) OR an OAuth2 token - # from an upstream MCP provider (e.g. Atlassian). - # Try LiteLLM auth first; on auth failure, treat as OAuth2 passthrough. + # Could be a LiteLLM key (backward compat) OR an opaque OAuth2 token + # the operator wants forwarded to an upstream OAuth2-mode MCP server. + # Try LiteLLM auth first; on auth failure, only fall back to anonymous + # passthrough when the request actually targets a server whose operator + # configured ``auth_type=oauth2``. For any other server (api_key, + # bearer_token, basic, etc.), a failed LiteLLM auth is a real failure + # and must propagate — otherwise an attacker can exchange any garbage + # bearer for an anonymous session. try: validated_user_api_key_auth = await user_api_key_auth( api_key=litellm_api_key, request=request ) - except HTTPException as e: - if e.status_code in (401, 403): + except (HTTPException, ProxyException) as e: + # HTTPException.status_code is int; ProxyException.code is + # normalized to str in its __init__ but can be ``"None"`` or any + # non-numeric string when the caller didn't supply a numeric + # code, so we compare against both int and str forms rather + # than coercing (``int("None")`` would raise ValueError and + # rewrite the auth error as a 500). + status = e.status_code if isinstance(e, HTTPException) else e.code + if status in ( + 401, + 403, + "401", + "403", + ) and MCPRequestHandler._target_servers_use_oauth2( + path=request.url.path, mcp_servers=mcp_servers + ): verbose_logger.debug( - "MCP OAuth2: Authorization header is not a valid LiteLLM key, " - "treating as OAuth2 token passthrough" - ) - validated_user_api_key_auth = UserAPIKeyAuth() - else: - raise - except ProxyException as e: - if str(e.code) in ("401", "403"): - verbose_logger.debug( - "MCP OAuth2: Authorization header is not a valid LiteLLM key, " - "treating as OAuth2 token passthrough" + "MCP OAuth2: target server is OAuth2-mode, treating " + "Authorization as upstream OAuth2 token passthrough" ) validated_user_api_key_auth = UserAPIKeyAuth() else: @@ -165,6 +178,62 @@ class MCPRequestHandler: dict(headers), ) + @staticmethod + def _extract_target_server_names_from_path(path: str) -> List[str]: + """ + Extract the target MCP server name from the standard MCP transport + URL patterns: ``/mcp/{server_name}[/...]`` and + ``/{server_name}/mcp[/...]``. Returns ``[]`` for any other path so + callers fail closed when the target cannot be resolved. + + REST/admin endpoints, OAuth2 server endpoints + (``/{server_name}/authorize``, ``/token`` etc.), and ``.well-known`` + discovery routes intentionally fall through — those flows do not need + OAuth2 token passthrough. Clients aggregating multiple servers should + use ``x-mcp-servers``, which takes precedence over path parsing. + """ + segments = [s for s in path.split("/") if s] + if len(segments) >= 2 and segments[0] == "mcp": + return [segments[1]] + if len(segments) >= 2 and segments[1] == "mcp": + return [segments[0]] + return [] + + @staticmethod + def _target_servers_use_oauth2(path: str, mcp_servers: Optional[List[str]]) -> bool: + """ + True only when EVERY MCP server the request targets is configured for + ``auth_type == oauth2``. If any target is non-OAuth2 — or if the target + cannot be resolved at all — return False so the caller fails closed. + + Used to gate the "treat Authorization as opaque OAuth2 token" fallback + in :meth:`process_mcp_request` so a failed LiteLLM-auth cannot be + exchanged for an anonymous session against a non-OAuth2 server. + """ + # Inline imports avoid a circular dependency: mcp_server_manager imports + # from this module. + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.types.mcp import MCPAuth + + # Use the x-mcp-servers header verbatim when present (including the + # explicitly-empty list, which means "no targets" → fail closed). + # Only fall back to path parsing when the header was absent entirely. + target_names = ( + mcp_servers + if mcp_servers is not None + else MCPRequestHandler._extract_target_server_names_from_path(path) + ) + if not target_names: + return False + + for name in target_names: + server = global_mcp_server_manager.get_mcp_server_by_name(name) + if server is None or server.auth_type != MCPAuth.oauth2: + return False + return True + @staticmethod def _get_mcp_auth_header_from_headers(headers: Headers) -> Optional[str]: """ diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 13b19aa2d6..cebd224a1a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -131,6 +131,22 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data +def _get_validated_client_redirect_uri(state_data: Dict[str, Any]) -> str: + """Return a loopback client redirect URI from OAuth state.""" + redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url") + if not redirect_uri or not isinstance(redirect_uri, str): + raise HTTPException(status_code=400, detail="Invalid redirect URI") + validate_loopback_redirect_uri(redirect_uri) + return redirect_uri + + +def _append_query_params(url: str, params: Dict[str, str]) -> str: + parsed = urlparse(url) + query_params = parse_qsl(parsed.query, keep_blank_values=True) + query_params.extend(params.items()) + return urlunparse(parsed._replace(query=urlencode(query_params))) + + def _resolve_oauth2_server_for_root_endpoints( client_ip: Optional[str] = None, ) -> Optional[MCPServer]: @@ -568,7 +584,7 @@ async def authorize( else None ) if mcp_server is None and mcp_server_name is None: - mcp_server = _resolve_oauth2_server_for_root_endpoints() + mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") # Use server's stored client_id when caller doesn't supply one. @@ -630,7 +646,7 @@ async def token_endpoint( lookup_name, client_ip=client_ip ) if mcp_server is None and mcp_server_name is None: - mcp_server = _resolve_oauth2_server_for_root_endpoints() + mcp_server = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if mcp_server is None: raise HTTPException(status_code=404, detail="MCP server not found") return await exchange_token_with_server( @@ -651,7 +667,6 @@ async def token_endpoint( async def callback(code: str, state: str): try: state_data = decode_state_hash(state) - base_url = state_data["base_url"] original_state = state_data["original_state"] # Re-validate loopback at the sink. /authorize rejects non-loopback @@ -659,10 +674,10 @@ async def callback(code: str, state: str): # minted before that check was added have no expiry and remain # valid indefinitely. Validating here blocks the open-redirect + # code-theft primitive even for pre-fix states. - validate_loopback_redirect_uri(base_url) + redirect_uri = _get_validated_client_redirect_uri(state_data) params = {"code": code, "state": original_state} - complete_returned_url = f"{base_url}?{urlencode(params)}" + complete_returned_url = _append_query_params(redirect_uri, params) return RedirectResponse(url=complete_returned_url, status_code=302) except HTTPException: @@ -719,16 +734,16 @@ def _build_oauth_protected_resource_response( ) request_base_url = get_request_base_url(request) + client_ip = IPAddressUtils.get_mcp_client_ip(request) # When no server name provided, try to resolve the single OAuth2 server if mcp_server_name is None: - resolved = _resolve_oauth2_server_for_root_endpoints() + resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: mcp_server_name = resolved.server_name or resolved.name mcp_server: Optional[MCPServer] = None if mcp_server_name: - client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = global_mcp_server_manager.get_mcp_server_by_name( mcp_server_name, client_ip=client_ip ) @@ -835,10 +850,11 @@ def _build_oauth_authorization_server_response( ) request_base_url = get_request_base_url(request) + client_ip = IPAddressUtils.get_mcp_client_ip(request) # When no server name provided, try to resolve the single OAuth2 server if mcp_server_name is None: - resolved = _resolve_oauth2_server_for_root_endpoints() + resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: mcp_server_name = resolved.server_name or resolved.name @@ -855,7 +871,6 @@ def _build_oauth_authorization_server_response( mcp_server: Optional[MCPServer] = None if mcp_server_name: - client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = global_mcp_server_manager.get_mcp_server_by_name( mcp_server_name, client_ip=client_ip ) @@ -1007,8 +1022,9 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non "client_secret": "dummy", "redirect_uris": [f"{request_base_url}/callback"], } + client_ip = IPAddressUtils.get_mcp_client_ip(request) if not mcp_server_name: - resolved = _resolve_oauth2_server_for_root_endpoints() + resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) if resolved: return await register_client_with_server( request=request, @@ -1021,7 +1037,6 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non ) return dummy_return - client_ip = IPAddressUtils.get_mcp_client_ip(request) mcp_server = global_mcp_server_manager.get_mcp_server_by_name( mcp_server_name, client_ip=client_ip ) 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/_types.py b/litellm/proxy/_types.py index 92c920ca59..fd4d4df241 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -904,6 +904,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): agents: Optional[List[str]] = None agent_access_groups: Optional[List[str]] = None models: Optional[List[str]] = None + search_tools: Optional[List[str]] = None class BudgetLimitEntry(LiteLLMPydanticObjectBase): @@ -1934,6 +1935,7 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): agent_access_groups: Optional[List[str]] = [] mcp_toolsets: Optional[List[str]] = None blocked_tools: Optional[List[str]] = [] + search_tools: Optional[List[str]] = [] class LiteLLM_TeamTable(TeamBase): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 333684eae5..9266a917c5 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -376,7 +376,7 @@ def _guardrail_modification_check( coerced = _coerce_to_dict(container) if coerced is None: return False - return any(coerced.get(key) for key in _GUARDRAIL_MODIFICATION_KEYS) + return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) # Check both metadata keys — callers can populate either depending on the # endpoint. Cover the top-level too so root-level injection is rejected. @@ -922,7 +922,8 @@ async def get_team_member_default_budget( Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"]. This budget is applied to team members whose TeamMembership row has no - linked budget. Results are cached for performance. + linked budget, or whose linked budget has max_budget=NULL. Results are + cached for performance. Args: budget_id: The budget_id pulled from team.metadata["team_member_budget_id"] @@ -2995,6 +2996,116 @@ async def can_user_call_model( ) +def _search_tool_names_from_object_permission( + object_permission: Optional[LiteLLM_ObjectPermissionTable], +) -> List[str]: + """Return allowlisted search tool names from object_permission (empty = unrestricted).""" + if object_permission is None: + return [] + raw = object_permission.search_tools + if not raw: + return [] + return list(raw) + + +def _can_object_call_search_tool( + search_tool_name: str, + allowed_search_tools: List[str], + object_type: Literal["key", "team", "project"], +) -> Literal[True]: + """ + Check if an object (key/team/project) can access a specific search tool. + + Similar to _can_object_call_model but for search tools. + + Args: + search_tool_name: The search tool being requested + allowed_search_tools: List of allowed search tool names for this object + object_type: Type of object for error messaging + + Returns: + True if access is allowed + + Raises: + ProxyException if access is denied + """ + # Empty list means all search tools are allowed + if not allowed_search_tools: + return True + + # Check if the search tool is in the allowlist + if search_tool_name in allowed_search_tools: + return True + + # Access denied + raise ProxyException( + message=f"{object_type.capitalize()} not allowed to access search tool: {search_tool_name}. " + f"Allowed search tools: {allowed_search_tools}", + type=ProxyErrorTypes.key_model_access_denied, + param="search_tool_name", + code=status.HTTP_403_FORBIDDEN, + ) + + +async def can_key_call_search_tool( + search_tool_name: str, + valid_token: UserAPIKeyAuth, +) -> Literal[True]: + """ + Check if a key can access a specific search tool. + + Similar to can_key_call_model but for search tools. + + Args: + search_tool_name: The search tool being requested + valid_token: The authenticated key + + Returns: + True if access is allowed + + Raises: + ProxyException if access is denied + """ + return _can_object_call_search_tool( + search_tool_name=search_tool_name, + allowed_search_tools=_search_tool_names_from_object_permission( + valid_token.object_permission + ), + object_type="key", + ) + + +async def can_team_call_search_tool( + search_tool_name: str, + team_object: Optional[LiteLLM_TeamTable], +) -> Literal[True]: + """ + Check if a team can access a specific search tool. + + Similar to can_team_access_model but for search tools. + + Args: + search_tool_name: The search tool being requested + team_object: The team object + + Returns: + True if access is allowed + + Raises: + ProxyException if access is denied + """ + if team_object is None: + return True + + return _can_object_call_search_tool( + search_tool_name=search_tool_name, + allowed_search_tools=_search_tool_names_from_object_permission( + team_object.object_permission + ), + object_type="team", + ) + + async def is_valid_fallback_model( model: str, llm_router: Optional[Router], @@ -3326,6 +3437,7 @@ async def _check_team_member_budget( if ( team_membership is not None and team_membership.litellm_budget_table is not None + and team_membership.litellm_budget_table.max_budget is not None ): team_member_budget = team_membership.litellm_budget_table.max_budget else: 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/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 529d91b204..a94abbc9d4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -20,6 +20,8 @@ from fastapi.security.api_key import APIKeyHeader import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.caching import DualCache +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * @@ -1117,10 +1119,14 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) if is_master_key_valid: + # Substitute a stable alias for the raw master key so neither the + # master key nor its hash propagates into spend logs, Prometheus + # /metrics labels, audit trails, rate-limit buckets, or any other + # downstream consumer of UserAPIKeyAuth.api_key. _user_api_key_obj = await _return_user_api_key_auth_obj( user_obj=None, user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=master_key, + api_key=LITELLM_PROXY_MASTER_KEY_ALIAS, parent_otel_span=parent_otel_span, valid_token_dict={ **end_user_params, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 160e9c23f0..935b96a0e3 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -474,6 +474,10 @@ async def retrieve_batch( # noqa: PLR0915 ) # Fix: The helper sets "file_id" but we need "batch_id" data["batch_id"] = data.pop("file_id", original_batch_id) + # Provider-config providers (e.g. bedrock) require `model` in kwargs + # so litellm.aretrieve_batch can load BedrockBatchesConfig. Without + # it the call falls into the legacy provider switch and 400s. + data["model"] = model_from_id # Retrieve batch using model credentials response = await litellm.aretrieve_batch( 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/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/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a424f1558f..8129fb0de5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -65,6 +65,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( attach_object_permission_to_dict, handle_update_object_permission_common, validate_key_mcp_servers_against_team, + validate_key_search_tools_against_team, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, @@ -768,6 +769,10 @@ async def _common_key_generation_helper( # noqa: PLR0915 object_permission=data_json.get("object_permission"), team_obj=team_table, ) + await validate_key_search_tools_against_team( + object_permission=data_json.get("object_permission"), + team_obj=team_table, + ) data_json = await _set_object_permission( data_json=data_json, @@ -2010,6 +2015,10 @@ async def _validate_mcp_servers_for_key_update( object_permission=object_permission_dict, team_obj=effective_team_obj, ) + await validate_key_search_tools_against_team( + object_permission=object_permission_dict, + team_obj=effective_team_obj, + ) async def _validate_update_key_data( diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 4eec7c6b7c..a6c0c7dcc0 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -4,15 +4,22 @@ Endpoints to control callbacks per team Use this when each team should control its own callbacks """ +import asyncio +import copy import json import traceback -from typing import List, Optional +from datetime import datetime, timezone +from typing import Any, List, Optional from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +import litellm from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.proxy._types import ( AddTeamCallback, + LiteLLM_AuditLogs, + LitellmTableNames, ProxyErrorTypes, ProxyException, TeamCallbackMetadata, @@ -24,6 +31,106 @@ from litellm.proxy.management_helpers.utils import management_endpoint_wrapper router = APIRouter() +_CALLBACK_VARS_REDACTED = "***REDACTED***" + + +def _redact_callback_secrets(metadata: Any) -> Any: + """Strip secret values out of a team-metadata snapshot before audit logging. + + Both ``team_metadata["logging"]`` (list of ``AddTeamCallback`` dicts) and + ``team_metadata["callback_settings"]["callback_vars"]`` carry provider + credentials such as ``langfuse_secret_key``, ``langsmith_api_key``, and + ``gcs_path_service_account``. Persisting them verbatim into + ``LiteLLM_AuditLogs`` would let anyone with read access to the audit + table harvest team callback credentials, so we replace each value with + a fixed marker. The keys themselves are kept so the audit reader can + still see *which* fields changed. + """ + if not isinstance(metadata, dict): + return metadata + redacted = copy.deepcopy(metadata) + logging_entries = redacted.get("logging") + if isinstance(logging_entries, list): + for entry in logging_entries: + if isinstance(entry, dict) and isinstance(entry.get("callback_vars"), dict): + entry["callback_vars"] = { + k: _CALLBACK_VARS_REDACTED for k in entry["callback_vars"] + } + callback_settings = redacted.get("callback_settings") + if isinstance(callback_settings, dict) and isinstance( + callback_settings.get("callback_vars"), dict + ): + callback_settings["callback_vars"] = { + k: _CALLBACK_VARS_REDACTED for k in callback_settings["callback_vars"] + } + return redacted + + +def _log_audit_task_exception(task: "asyncio.Task[None]") -> None: + """Surface a fire-and-forget audit-log task failure. + + ``asyncio.create_task`` swallows exceptions silently — if the audit + write fails (transient DB error etc.) we'd otherwise lose the row + without any signal. Log at warning level so the operator sees there's + a gap in the audit trail. + """ + if task.cancelled(): + return + exc = task.exception() + if exc is not None: + verbose_proxy_logger.warning("Failed to write team-callback audit log: %s", exc) + + +async def _emit_team_callback_audit_log( + *, + team_id: str, + before_metadata: Any, + after_metadata: Any, + user_api_key_dict: UserAPIKeyAuth, + litellm_changed_by: Optional[str], +) -> None: + """Emit an audit-log row for a team-callback mutation. + + Mirrors the ``store_audit_logs``-gated pattern used in + ``team_endpoints.py``: the call is async-fire-and-forget and is a no-op + when audit logging is not enabled on the proxy. Captured under + ``LitellmTableNames.TEAM_TABLE_NAME`` so the row co-locates with other + team mutations in the audit table. + + Callback secrets are redacted before serialization so the audit table + cannot itself become a credential-harvest sink. + """ + if litellm.store_audit_logs is not True: + return + + from litellm.proxy.management_helpers.audit_logs import ( + create_audit_log_for_update, + ) + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + redacted_before = _redact_callback_secrets(before_metadata) + redacted_after = _redact_callback_secrets(after_metadata) + + task = asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + object_id=team_id, + action="updated", + updated_values=json.dumps({"metadata": redacted_after}, default=str), + before_value=json.dumps({"metadata": redacted_before}, default=str), + ) + ) + ) + task.add_done_callback(_log_audit_task_exception) + + @router.post( "/team/{team_id:path}/callback", tags=["team management"], @@ -123,6 +230,7 @@ async def add_team_callbacks( param="callback_name", ) + before_metadata = copy.deepcopy(team_metadata) team_callback_settings.append(data.model_dump()) team_metadata["logging"] = team_callback_settings @@ -132,6 +240,14 @@ async def add_team_callbacks( where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore ) + await _emit_team_callback_audit_log( + team_id=team_id, + before_metadata=before_metadata, + after_metadata=team_metadata, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + return { "status": "success", "data": new_team_row, @@ -165,6 +281,10 @@ async def disable_team_logging( http_request: Request, team_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), ): """ Disable all logging callbacks for a team @@ -198,6 +318,7 @@ async def disable_team_logging( # Update team metadata to disable logging team_metadata = _existing_team.metadata + before_metadata = copy.deepcopy(team_metadata) team_callback_settings = team_metadata.get("callback_settings", {}) team_callback_settings_obj = TeamCallbackMetadata(**team_callback_settings) @@ -222,6 +343,17 @@ async def disable_team_logging( }, ) + # Disabling a team's logging callbacks is itself a logging-control + # action — emit an audit-log row so the action remains traceable + # even though the team's own observability is now off. + await _emit_team_callback_audit_log( + team_id=team_id, + before_metadata=before_metadata, + after_metadata=team_metadata, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + return { "status": "success", "message": f"Logging disabled for team {team_id}", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f254fea3e7..bcafd93a4c 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2003,21 +2003,34 @@ def team_member_add_duplication_check( async def _validate_team_member_add_permissions( user_api_key_dict: UserAPIKeyAuth, complete_team_data: LiteLLM_TeamTable, + data: TeamMemberAddRequest, ) -> None: - """Validate if user has permission to add members to the team.""" + """Validate if user has permission to add members to the team. + + Standard users can self-join an *available team*, but the bypass + must not be allowed to escalate them to ``role=admin`` or to add + other users into the team. When access is granted via the + available-team bypass we therefore enforce that every member in + the request matches the caller's own ``user_id`` and is being + added with ``role="user"``. + """ if ( - hasattr(user_api_key_dict, "user_role") - and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value - and not _is_user_team_admin( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ) - and not await _is_user_org_admin_for_team( - user_api_key_dict=user_api_key_dict, team_obj=complete_team_data - ) - and not _is_available_team( - team_id=complete_team_data.team_id, - user_api_key_dict=user_api_key_dict, - ) + getattr(user_api_key_dict, "user_role", None) + == LitellmUserRoles.PROXY_ADMIN.value + ): + return + if _is_user_team_admin( + user_api_key_dict=user_api_key_dict, team_obj=complete_team_data + ): + return + if await _is_user_org_admin_for_team( + user_api_key_dict=user_api_key_dict, team_obj=complete_team_data + ): + return + + if not _is_available_team( + team_id=complete_team_data.team_id, + user_api_key_dict=user_api_key_dict, ): raise HTTPException( status_code=403, @@ -2029,6 +2042,34 @@ async def _validate_team_member_add_permissions( }, ) + # Available-team self-join: caller may add only themselves, only as a + # standard user. Enforce that here so the bypass cannot be used as a + # privilege-escalation or cross-user-injection primitive. + members = data.member if isinstance(data.member, list) else [data.member] + caller_user_id = getattr(user_api_key_dict, "user_id", None) + for member in members: + if getattr(member, "role", "user") != "user": + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Available-team self-join cannot assign 'admin' role. " + "Only proxy/team/org admins can add admins to a team." + ) + }, + ) + member_user_id = getattr(member, "user_id", None) + if not caller_user_id or not member_user_id or member_user_id != caller_user_id: + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Available-team self-join can only add the caller " + "(user_id must match the authenticated user's user_id)." + ) + }, + ) + async def _process_team_members( data: TeamMemberAddRequest, @@ -2049,8 +2090,11 @@ async def _process_team_members( # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models member_allowed_models = data.allowed_models - if member_allowed_models is None and complete_team_data.default_team_member_models: - member_allowed_models = complete_team_data.default_team_member_models + team_default_member_models = getattr( + complete_team_data, "default_team_member_models", None + ) + if member_allowed_models is None and team_default_member_models: + member_allowed_models = team_default_member_models if isinstance(data.member, Member): try: @@ -2381,6 +2425,7 @@ async def team_member_add( await _validate_team_member_add_permissions( user_api_key_dict=user_api_key_dict, complete_team_data=complete_team_data, + data=data, ) # Validate and populate user_email/user_id for members before processing @@ -4694,6 +4739,8 @@ async def update_team_member_permissions( complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump()) + # Available-team self-join must NOT grant write access to team-wide + # permission policies; only proxy/team/org admins can update them. if ( hasattr(user_api_key_dict, "user_role") and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value @@ -4703,16 +4750,12 @@ async def update_team_member_permissions( and not await _is_user_org_admin_for_team( user_api_key_dict=user_api_key_dict, team_obj=complete_team_data ) - and not _is_available_team( - team_id=complete_team_data.team_id, - user_api_key_dict=user_api_key_dict, - ) ): raise HTTPException( status_code=403, detail={ "error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format( - "/team/member_add", + "/team/permissions_update", complete_team_data.team_id, ) }, 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/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 410f636693..eb90d1b5ca 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -335,8 +335,9 @@ async def validate_key_mcp_servers_against_team( disallowed_servers = requested_servers - all_allowed_servers if disallowed_servers: if team_obj is not None: + team_id = team_obj.team_id detail = ( - f"Key requests MCP servers not allowed by team '{team_obj.team_id}': " + f"Key requests MCP servers not allowed by team '{team_id}': " f"{sorted(disallowed_servers)}. " f"Team allows: {sorted(team_allowed_servers)}. " f"Global (allow_all_keys) servers: {sorted(allow_all_keys_servers)}." @@ -365,8 +366,9 @@ async def validate_key_mcp_servers_against_team( disallowed_groups = requested_access_groups - team_access_groups if disallowed_groups: if team_obj is not None: + team_id = team_obj.team_id detail = ( - f"Key requests MCP access groups not allowed by team '{team_obj.team_id}': " + f"Key requests MCP access groups not allowed by team '{team_id}': " f"{sorted(disallowed_groups)}. " f"Team allows: {sorted(team_access_groups)}." ) @@ -390,13 +392,60 @@ async def validate_key_mcp_servers_against_team( if team_mcp_toolsets: disallowed_toolsets = requested_toolsets - set(team_mcp_toolsets) if disallowed_toolsets: + team_id = team_obj.team_id raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail={ "error": ( - f"Key requests MCP toolsets not allowed by team '{team_obj.team_id}': " + f"Key requests MCP toolsets not allowed by team '{team_id}': " f"{sorted(disallowed_toolsets)}. " f"Team allows: {sorted(team_mcp_toolsets)}." ) }, ) + + +def _extract_requested_search_tools(object_permission: Optional[dict]) -> List[str]: + """Return search_tool_name values from a key's object_permission dict.""" + if not object_permission or not isinstance(object_permission, dict): + return [] + raw = object_permission.get("search_tools") + if not isinstance(raw, list): + return [] + return [str(x) for x in raw if x] + + +async def validate_key_search_tools_against_team( + object_permission: Optional[dict], + team_obj: Optional["LiteLLM_TeamTableCachedObj"], +) -> None: + """ + Validate key object_permission.search_tools is a subset of the team's allowlist. + + Empty team allowlist means no restriction at team layer (skip). + """ + requested = _extract_requested_search_tools(object_permission) + if not requested: + return + + team_tools: List[str] = [] + if team_obj is not None and team_obj.object_permission is not None: + st = team_obj.object_permission.search_tools + if st: + team_tools = list(st) + + if not team_tools: + return + + disallowed = set(requested) - set(team_tools) + if disallowed: + team_id = team_obj.team_id if team_obj is not None else "unknown" + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ + "error": ( + f"Key requests search tools not allowed by team '{team_id}': " + f"{sorted(disallowed)}. Team allows: {sorted(team_tools)}." + ) + }, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 216eb61a9d..c42faa59cf 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -549,10 +549,16 @@ class AnthropicPassthroughLoggingHandler: # Create a mock user API key dict for the managed object storage from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + _request_metadata = (kwargs.get("litellm_params", {}) or {}).get( + "metadata", {} + ) or {} + user_api_key_dict = UserAPIKeyAuth( - user_id=kwargs.get("user_id", "default-user"), + user_id=_request_metadata.get( + "user_api_key_user_id", "default-user" + ), api_key="", - team_id=None, + team_id=_request_metadata.get("user_api_key_team_id"), team_alias=None, user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value user_email=None, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 86dd23e12c..6a13853261 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -849,10 +849,16 @@ class VertexPassthroughLoggingHandler: # Create a mock user API key dict for the managed object storage from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + _request_metadata = (kwargs.get("litellm_params", {}) or {}).get( + "metadata", {} + ) or {} + user_api_key_dict = UserAPIKeyAuth( - user_id=kwargs.get("user_id", "default-user"), + user_id=_request_metadata.get( + "user_api_key_user_id", "default-user" + ), api_key="", - team_id=None, + team_id=_request_metadata.get("user_api_key_team_id"), team_alias=None, user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value user_email=None, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 302d7e76ed..cbfcd34c43 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -36,21 +36,16 @@ class PassThroughStreamingHandler: passthrough_success_handler_obj: PassThroughEndpointLogging, url_route: str, ): - """ - - Yields chunks from the response - - Collect non-empty chunks for post-processing (logging) - - Inject cost into chunks if include_cost_in_streaming_usage is enabled - """ - try: - raw_bytes: List[bytes] = [] - # Extract model name for cost injection - model_name = PassThroughStreamingHandler._extract_model_for_cost_injection( - request_body=request_body, - url_route=url_route, - endpoint_type=endpoint_type, - litellm_logging_obj=litellm_logging_obj, - ) + raw_bytes: List[bytes] = [] + logging_scheduled = False + model_name = PassThroughStreamingHandler._extract_model_for_cost_injection( + request_body=request_body, + url_route=url_route, + endpoint_type=endpoint_type, + litellm_logging_obj=litellm_logging_obj, + ) + try: async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) if ( @@ -58,7 +53,6 @@ class PassThroughStreamingHandler: and model_name ): if endpoint_type == EndpointType.VERTEX_AI: - # Only handle streamRawPredict (uses Anthropic format) if "streamRawPredict" in url_route or "rawPredict" in url_route: modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( chunk, model_name @@ -73,25 +67,32 @@ class PassThroughStreamingHandler: chunk = modified_chunk yield chunk - - # After all chunks are processed, handle post-processing - end_time = datetime.now() - - asyncio.create_task( - PassThroughStreamingHandler._route_streaming_logging_to_handler( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body or {}, - endpoint_type=endpoint_type, - start_time=start_time, - raw_bytes=raw_bytes, - end_time=end_time, - ) - ) except Exception as e: verbose_proxy_logger.error(f"Error in chunk_processor: {str(e)}") raise + finally: + # GeneratorExit (raised on client disconnect) is not caught by + # `except Exception`; the finally block ensures partial usage + # still gets logged for spend tracking. See LIT-2642. + if not logging_scheduled and raw_bytes: + logging_scheduled = True + try: + asyncio.create_task( + PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body or {}, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=datetime.now(), + ) + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error scheduling chunk_processor logging: {str(e)}" + ) @staticmethod async def _route_streaming_logging_to_handler( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ce93f748b0..342e1c4ab3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -94,6 +94,7 @@ from litellm.proxy._types import ( TeamDefaultSettings, TokenCountRequest, TransformRequestBody, + UI_TEAM_ID, UserAPIKeyAuth, ) from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -431,6 +432,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, @@ -747,6 +751,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, @@ -757,6 +765,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) @@ -12065,7 +12076,7 @@ async def onboarding(invite_link: str, request: Request): """ - Get the invite link - Validate it's still 'valid' - - Invalidate the link (prevents abuse) + - Return a short-lived onboarding token - Get user from db - Pass in user_email if set """ @@ -12103,7 +12114,7 @@ async def onboarding(invite_link: str, request: Request): ) #### CHECK IF ALREADY USED - if invite_obj.is_accepted is True: + if invite_obj.is_accepted is True or invite_obj.accepted_at is not None: raise HTTPException( status_code=401, detail={"error": "Invitation link has already been used."}, @@ -12119,24 +12130,6 @@ async def onboarding(invite_link: str, request: Request): status_code=401, detail={"error": "User does not exist in db."} ) - user_email = user_obj.user_email - - response = await generate_key_helper_fn( - request_type="key", - **{ - "user_role": user_obj.user_role, - "duration": LITELLM_UI_SESSION_DURATION, - "key_max_budget": litellm.max_ui_session_budget, - "models": [], - "aliases": {}, - "config": {}, - "spend": 0, - "user_id": user_obj.user_id, - "team_id": "litellm-dashboard", - }, # type: ignore - ) - key = response["token"] # type: ignore - litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): litellm_dashboard_ui += "ui/onboarding" @@ -12144,13 +12137,24 @@ async def onboarding(invite_link: str, request: Request): litellm_dashboard_ui += "/ui/onboarding" import jwt + user_email = user_obj.user_email + onboarding_token = jwt.encode( # type: ignore + { + "token_type": "litellm_onboarding", + "invitation_link": invite_link, + "user_id": user_obj.user_id, + "exp": litellm.utils.get_utc_datetime() + timedelta(minutes=15), + }, + master_key, + algorithm="HS256", + ) disabled_non_admin_personal_key_creation = ( get_disabled_non_admin_personal_key_creation() ) returned_ui_token_object = ReturnedUITokenObject( user_id=user_obj.user_id, - key=key, + key=onboarding_token, user_email=user_obj.user_email, user_role=user_obj.user_role, login_method="username_password", @@ -12175,8 +12179,117 @@ async def onboarding(invite_link: str, request: Request): } +def _get_onboarding_claims_from_request(request: Request) -> dict: + global master_key, general_settings + + if master_key is None: + raise ProxyException( + message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", + type=ProxyErrorTypes.auth_error, + param="master_key", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + auth_header_name = general_settings.get("litellm_key_header_name", "Authorization") + onboarding_auth_header = request.headers.get(auth_header_name) + if onboarding_auth_header is None: + raise HTTPException( + status_code=401, + detail={"error": "Missing onboarding session for invitation link."}, + ) + onboarding_token = onboarding_auth_header + if onboarding_token.lower().startswith("bearer "): + onboarding_token = onboarding_token.split(" ", 1)[1] + + import jwt + + try: + return jwt.decode( + onboarding_token, + master_key, + algorithms=["HS256"], + ) + except Exception: + raise HTTPException( + status_code=401, + detail={"error": "Invalid onboarding session for invitation link."}, + ) + + +async def _rollback_onboarding_invite_claim( + invitation_link: str, + user_id: str, +) -> None: + global prisma_client + + if prisma_client is None: + return + + try: + await prisma_client.db.litellm_invitationlink.update_many( + where={"id": invitation_link, "is_accepted": True}, + data={ + "accepted_at": None, + "is_accepted": False, + "updated_at": litellm.utils.get_utc_datetime(), + "updated_by": user_id, + }, + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to roll back onboarding invitation after session key mint failed." + ) + + +async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: + global master_key, general_settings + + response = await generate_key_helper_fn( + request_type="key", + **{ + "user_role": user_obj.user_role, + "duration": LITELLM_UI_SESSION_DURATION, + "key_max_budget": litellm.max_ui_session_budget, + "models": [], + "aliases": {}, + "config": {}, + "spend": 0, + "user_id": user_obj.user_id, + "team_id": UI_TEAM_ID, + }, # type: ignore + ) + key = response["token"] # type: ignore + + from litellm.types.proxy.ui_sso import ReturnedUITokenObject + + import jwt + + disabled_non_admin_personal_key_creation = ( + get_disabled_non_admin_personal_key_creation() + ) + returned_ui_token_object = ReturnedUITokenObject( + user_id=user_obj.user_id, + key=key, + user_email=user_obj.user_email, + user_role=user_obj.user_role, + login_method="username_password", + premium_user=premium_user, + auth_header_name=general_settings.get( + "litellm_key_header_name", "Authorization" + ), + disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, + server_root_path=get_server_root_path(), + ) + assert master_key is not None + return jwt.encode( # type: ignore + cast(dict, returned_ui_token_object), + master_key, + algorithm="HS256", + ) + + @app.post("/onboarding/claim_token", include_in_schema=False) -async def claim_onboarding_link(data: InvitationClaim): +async def claim_onboarding_link(data: InvitationClaim, request: Request): """ Special route. Allows UI link share user to update their password. @@ -12188,7 +12301,7 @@ async def claim_onboarding_link(data: InvitationClaim): This route can only update user password. """ - global prisma_client + global prisma_client, master_key, general_settings ### VALIDATE INVITE LINK ### if prisma_client is None: raise HTTPException( @@ -12213,7 +12326,7 @@ async def claim_onboarding_link(data: InvitationClaim): ) #### CHECK IF ALREADY USED - if invite_obj.is_accepted is True: + if invite_obj.is_accepted is True or invite_obj.accepted_at is not None: raise HTTPException( status_code=401, detail={"error": "Invitation link has already been used."}, @@ -12229,32 +12342,87 @@ async def claim_onboarding_link(data: InvitationClaim): ) }, ) - ### UPDATE USER OBJECT ### - hashed_pw = hash_password(data.password) - user_obj = await prisma_client.db.litellm_usertable.update( - where={"user_id": invite_obj.user_id}, data={"password": hashed_pw} - ) - if user_obj is None: + onboarding_claims = _get_onboarding_claims_from_request(request=request) + if ( + onboarding_claims.get("token_type") != "litellm_onboarding" + or onboarding_claims.get("invitation_link") != data.invitation_link + or onboarding_claims.get("user_id") != data.user_id + ): raise HTTPException( - status_code=401, detail={"error": "User does not exist in db."} + status_code=401, + detail={"error": "Invalid onboarding session for invitation link."}, ) - #### MARK LINK AS USED + hashed_pw = hash_password(data.password) current_time = litellm.utils.get_utc_datetime() - await prisma_client.db.litellm_invitationlink.update( - where={"id": data.invitation_link}, - data={ - "accepted_at": current_time, - "updated_at": current_time, - "is_accepted": True, - "updated_by": invite_obj.user_id, # type: ignore - }, - ) + async with prisma_client.db.tx() as tx: + updated_count = await tx.litellm_invitationlink.update_many( + where={"id": data.invitation_link, "is_accepted": False}, + data={ + "is_accepted": True, + "updated_at": current_time, + "updated_by": invite_obj.user_id, # type: ignore + }, + ) + if updated_count == 0: + raise HTTPException( + status_code=401, + detail={"error": "Invitation link has already been used."}, + ) + + ### UPDATE USER OBJECT ### + user_obj = await tx.litellm_usertable.update( + where={"user_id": invite_obj.user_id}, data={"password": hashed_pw} + ) + + if user_obj is None: + raise HTTPException( + status_code=401, detail={"error": "User does not exist in db."} + ) + + #### MARK LINK AS USED + current_time = litellm.utils.get_utc_datetime() + await tx.litellm_invitationlink.update( + where={"id": data.invitation_link}, + data={ + "accepted_at": current_time, + "updated_at": current_time, + "updated_by": invite_obj.user_id, # type: ignore + }, + ) if user_obj and hasattr(user_obj, "__dict__"): user_obj.__dict__.pop("password", None) - return user_obj + + try: + jwt_token = await _generate_onboarding_ui_session_token(user_obj=user_obj) + except Exception as e: + await _rollback_onboarding_invite_claim( + invitation_link=data.invitation_link, + user_id=data.user_id, + ) + if isinstance(e, HTTPException): + raise e + raise HTTPException( + status_code=500, + detail={ + "error": "Failed to create onboarding session. Please retry the invitation link." + }, + ) from e + + litellm_dashboard_ui = get_custom_url(str(request.base_url)) + if litellm_dashboard_ui.endswith("/"): + litellm_dashboard_ui += "ui/" + else: + litellm_dashboard_ui += "/ui/" + litellm_dashboard_ui += "?login=success" + return { + "login_url": litellm_dashboard_ui, + "token": jwt_token, + "user_email": user_obj.user_email, + "user": user_obj, + } @app.get("/get_logo_url", include_in_schema=False) @@ -14319,6 +14487,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..a9d3911c07 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -277,6 +277,7 @@ model LiteLLM_ObjectPermissionTable { models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user + search_tools String[] @default([]) // search_tool_name values this key/team/user may call teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -1290,3 +1291,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/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index 8bed5b5407..15ed858b98 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -134,10 +134,48 @@ async def search( if "search_tool_name" in data and data["search_tool_name"]: data["model"] = data["search_tool_name"] + search_tool_name_value = data["search_tool_name"] + + # Authorization check: verify key can access this search tool + from litellm.proxy.auth.auth_checks import ( + can_key_call_search_tool, + can_team_call_search_tool, + get_team_object, + ) + + try: + # Check key-level access + await can_key_call_search_tool( + search_tool_name=search_tool_name_value, + valid_token=user_api_key_dict, + ) + + # Check team-level access if key is associated with a team + if user_api_key_dict.team_id: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + team_object = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + await can_team_call_search_tool( + search_tool_name=search_tool_name_value, + team_object=team_object, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Search tool authorization failed for {search_tool_name_value}: {str(e)}" + ) + raise if llm_router is not None and hasattr(llm_router, "search_tools"): - search_tool_name_value = data["search_tool_name"] - verbose_proxy_logger.debug( f"Search endpoint - Looking for search_tool_name: {search_tool_name_value}. " f"Available search tools in router: {[tool.get('search_tool_name') for tool in llm_router.search_tools]}. " @@ -163,6 +201,16 @@ async def search( data["metadata"] = {} data["metadata"]["model_group"] = search_tool_name_value + # Ensure team context is available to search router credential resolution. + # add_litellm_data_to_request() also injects these values, but this keeps + # search endpoint behavior explicit and resilient for direct router paths. + if "metadata" not in data or not isinstance(data.get("metadata"), dict): + data["metadata"] = {} + if getattr(user_api_key_dict, "team_metadata", None) is not None: + data["metadata"]["user_api_key_team_metadata"] = user_api_key_dict.team_metadata + if getattr(user_api_key_dict, "team_id", None) is not None: + data["metadata"]["user_api_key_team_id"] = user_api_key_dict.team_id + # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 889c781004..ec6245f47e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -53,20 +53,13 @@ def _get_max_string_length_prompt_in_db() -> int: def _is_master_key(api_key: Optional[str], _master_key: Optional[str]) -> bool: + """ + Raw-only constant-time master-key comparison. The hashed form is never + considered equivalent — only the raw master-key string matches. + """ if _master_key is None or api_key is None: return False - - ## string comparison - is_master_key = secrets.compare_digest(api_key, _master_key) - if is_master_key: - return True - - ## hash comparison - is_master_key = secrets.compare_digest(api_key, hash_token(_master_key)) - if is_master_key: - return True - - return False + return secrets.compare_digest(api_key, _master_key) def _get_spend_logs_metadata( @@ -235,8 +228,6 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d def get_logging_payload( # noqa: PLR0915 kwargs, response_obj, start_time, end_time ) -> SpendLogsPayload: - from litellm.proxy.proxy_server import general_settings, master_key - if kwargs is None: kwargs = {} @@ -295,11 +286,6 @@ def get_logging_payload( # noqa: PLR0915 if api_key.startswith("sk-"): # hash the api_key api_key = hash_token(api_key) - if ( - _is_master_key(api_key=api_key, _master_key=master_key) - and general_settings.get("disable_adding_master_key_hash_to_db") is True - ): - api_key = "litellm_proxy_master_key" # use a known alias, if the user disabled storing master key in db if ( standard_logging_payload is not None @@ -324,11 +310,6 @@ def get_logging_payload( # noqa: PLR0915 and standard_logging_payload.get("request_tags") is not None ): # use 'tags' from standard logging payload instead request_tags = json.dumps(standard_logging_payload["request_tags"]) - if ( - _is_master_key(api_key=api_key, _master_key=master_key) - and general_settings.get("disable_adding_master_key_hash_to_db") is True - ): - api_key = "litellm_proxy_master_key" # use a known alias, if the user disabled storing master key in db _model_id = metadata.get("model_info", {}).get("id", "") _model_group = metadata.get("model_group", "") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3a1184c434..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 ( @@ -2779,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)) @@ -4183,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 @@ -4204,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", "") @@ -4217,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/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index fefa6cb4e9..99a2085bfc 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -16,7 +16,9 @@ from fastapi import APIRouter, Depends, HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import ( LiteLLM_ManagedVectorStoresTable, ResponseLiteLLM_ManagedVectorStore, @@ -38,6 +40,81 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router = APIRouter() +_LITELLM_PARAMS_MASKER = SensitiveDataMasker() + + +_REDACT_LITELLM_PARAMS_MAX_DEPTH = 10 + + +def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> Any: + """ + Replace credential-bearing values in ``litellm_params`` with + ``REDACTED_BY_LITELM`` while preserving non-secret keys (``api_base``, + ``region``, ``model``, ``api_version``). + + Handles three input shapes: + + * ``dict`` — recurse into nested dicts (e.g. ``litellm_embedding_config`` + which itself carries ``api_key`` / ``aws_*`` / ``vertex_credentials``). + * ``str`` — the in-memory registry occasionally holds the params as a + JSON-serialized string. Parse, redact, re-serialize. If parsing + fails, return the redaction sentinel rather than echo the value + back verbatim. + * Anything else, or ``None`` — passed through. + + Recursion depth is bounded by ``_REDACT_LITELLM_PARAMS_MAX_DEPTH`` — + matching the convention of other allowlisted recursive helpers in the + repo (see ``tests/code_coverage_tests/recursive_detector.py``). + """ + if _depth >= _REDACT_LITELLM_PARAMS_MAX_DEPTH: + return REDACTED_BY_LITELM_STRING + if litellm_params is None: + return None + if isinstance(litellm_params, str): + try: + parsed = json.loads(litellm_params) + except (TypeError, ValueError): + return REDACTED_BY_LITELM_STRING + return json.dumps(_redact_sensitive_litellm_params(parsed, _depth + 1)) + if not isinstance(litellm_params, dict): + return litellm_params + out: Dict[str, Any] = {} + for k, v in litellm_params.items(): + if _LITELLM_PARAMS_MASKER.is_sensitive_key(k): + out[k] = REDACTED_BY_LITELM_STRING + elif isinstance(v, dict): + out[k] = _redact_sensitive_litellm_params(v, _depth + 1) + else: + out[k] = v + return out + + +async def _fetch_and_authorize_vector_store( + vector_store_id: str, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: Any, +) -> "LiteLLM_ManagedVectorStore": + """ + Look up a vector store by id and confirm the caller can access it. + Raises HTTPException(404) on miss and HTTPException(403) on access + denial. + """ + row = await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store_id} + ) + if row is None: + raise HTTPException( + status_code=404, + detail=f"Vector store with ID {vector_store_id} not found", + ) + typed = LiteLLM_ManagedVectorStore(**row.model_dump()) + if not await _check_vector_store_access(typed, user_api_key_dict): + raise HTTPException( + status_code=403, + detail="Access denied: You do not have permission to access this vector store", + ) + return typed + def _resolve_embedding_config_from_router( embedding_model: str, llm_router @@ -555,7 +632,11 @@ async def list_vector_stores( accessible_vector_stores = [] for vs in vector_store_map.values(): if await _check_vector_store_access(vs, user_api_key_dict): - accessible_vector_stores.append(vs) + redacted = LiteLLM_ManagedVectorStore(**vs) + redacted["litellm_params"] = _redact_sensitive_litellm_params( + vs.get("litellm_params") + ) + accessible_vector_stores.append(redacted) total_count = len(accessible_vector_stores) total_pages = (total_count + page_size - 1) // page_size @@ -716,33 +797,29 @@ async def get_vector_store_info( created_at=vector_store.get("created_at") or None, updated_at=vector_store.get("updated_at") or None, litellm_credential_name=vector_store.get("litellm_credential_name"), - litellm_params=vector_store.get("litellm_params") or None, + litellm_params=_redact_sensitive_litellm_params( + vector_store.get("litellm_params") + ), team_id=vector_store.get("team_id") or None, user_id=vector_store.get("user_id") or None, ) return {"vector_store": vector_store_pydantic_obj} - vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": data.vector_store_id} - ) + vector_store_typed = await _fetch_and_authorize_vector_store( + vector_store_id=data.vector_store_id, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, ) - if vector_store is None: - raise HTTPException( - status_code=404, - detail=f"Vector store with ID {data.vector_store_id} not found", + vector_store_dict = dict(vector_store_typed) + if "litellm_params" in vector_store_dict: + vector_store_dict["litellm_params"] = _redact_sensitive_litellm_params( + vector_store_dict["litellm_params"] ) - - # Check access control for DB vector store - vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined] - vector_store_typed = LiteLLM_ManagedVectorStore(**vector_store_dict) - if not await _check_vector_store_access(vector_store_typed, user_api_key_dict): - raise HTTPException( - status_code=403, - detail="Access denied: You do not have permission to access this vector store", - ) - return {"vector_store": vector_store_dict} + except HTTPException: + # Preserve 403/404 from the access-control / not-found checks above; + # the catch-all below would otherwise rewrite them as 500. + raise except Exception as e: verbose_proxy_logger.exception(f"Error getting vector store info: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) @@ -773,6 +850,15 @@ async def update_vector_store( update_data = data.model_dump(exclude_unset=True) vector_store_id = update_data.pop("vector_store_id") + # Per-store access control: anyone authenticated who passes the + # premium-feature gate could otherwise update *any* vector store — + # including stores belonging to other teams. + await _fetch_and_authorize_vector_store( + vector_store_id=vector_store_id, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + # Handle metadata serialization if update_data.get("vector_store_metadata") is not None: update_data["vector_store_metadata"] = safe_dumps( @@ -820,11 +906,24 @@ async def update_vector_store( f"Updated vector store {vector_store_id} in both database and in-memory registry" ) + # The DB row is returned in full, so the response would otherwise + # echo the persisted ``litellm_params`` (including provider + # credentials) back to the caller — even when the caller only + # changed unrelated fields like ``vector_store_description``. + response_vs = LiteLLM_ManagedVectorStore(**updated_vs) + response_vs["litellm_params"] = _redact_sensitive_litellm_params( + updated_vs.get("litellm_params") + ) return { "status": "success", "message": f"Vector store {vector_store_id} updated successfully", - "vector_store": updated_vs, + "vector_store": response_vs, } + except HTTPException: + # Preserve 403/404 responses from the access-control / not-found + # checks above; the catch-all below would otherwise rewrite them + # as 500 with the original status code embedded in the detail. + raise except Exception as e: verbose_proxy_logger.exception(f"Error updating vector store: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) 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/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index a26aa7e71e..9bcbcd8365 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -8,7 +8,7 @@ import asyncio import random import traceback from functools import partial -from typing import Any, Callable +from typing import Any, Callable, Dict, Optional, Tuple from litellm._logging import verbose_router_logger @@ -20,6 +20,28 @@ class SearchAPIRouter: Provides methods for search tool selection, load balancing, and fallback handling. """ + @staticmethod + def _resolve_search_provider_credentials( + *, + tool_litellm_params: Dict[str, Any], + ) -> Tuple[Optional[str], Optional[str]]: + """ + Resolve search provider credentials from tool configuration ONLY. + + Credentials are stored only in search_tool.litellm_params, never in team/key metadata. + This ensures secrets are not exposed in team/key API responses. + + Args: + tool_litellm_params: Search tool litellm_params with credentials + + Returns: + Tuple of (api_key, api_base) from tool configuration + """ + resolved_api_key: Optional[str] = tool_litellm_params.get("api_key") + resolved_api_base: Optional[str] = tool_litellm_params.get("api_base") + + return resolved_api_key, resolved_api_base + @staticmethod async def update_router_search_tools(router_instance: Any, search_tools: list): """ @@ -198,14 +220,15 @@ class SearchAPIRouter: # Extract search provider and other params from litellm_params litellm_params = selected_tool.get("litellm_params", {}) search_provider = litellm_params.get("search_provider") - api_key = litellm_params.get("api_key") - api_base = litellm_params.get("api_base") - if not search_provider: raise ValueError( f"search_provider not found in litellm_params for search tool '{search_tool_name}'" ) + api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( + tool_litellm_params=litellm_params, + ) + verbose_router_logger.debug( f"Selected search tool with provider: {search_provider}" ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index e3f63d0574..c376b8694a 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -330,7 +330,11 @@ class AnthropicMessagesToolResultParam(TypedDict, total=False): content: Union[ str, Iterable[ - Union[AnthropicMessagesToolResultContent, AnthropicMessagesImageParam] + Union[ + AnthropicMessagesToolResultContent, + AnthropicMessagesImageParam, + AnthropicMessagesDocumentParam, + ] ], ] cache_control: Optional[Union[dict, ChatCompletionCachedContent]] 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/utils.py b/litellm/utils.py index e63bf402bf..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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 114883f530..bbe13442d6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -712,6 +712,7 @@ }, "anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", @@ -735,6 +736,7 @@ }, "anthropic.claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", @@ -955,6 +957,7 @@ }, "anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -982,6 +985,7 @@ }, "anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1011,6 +1015,7 @@ }, "global.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1040,6 +1045,7 @@ }, "us.anthropic.claude-opus-4-6-v1": { "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", @@ -1127,6 +1133,7 @@ }, "anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1171,6 +1178,7 @@ }, "global.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", @@ -1201,6 +1209,7 @@ }, "us.anthropic.claude-opus-4-7": { "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", @@ -1291,6 +1300,7 @@ }, "anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", @@ -1319,6 +1329,7 @@ }, "global.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", @@ -1347,6 +1358,7 @@ }, "us.anthropic.claude-sonnet-4-6": { "cache_creation_input_token_cost": 4.125e-06, + "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", @@ -1461,11 +1473,13 @@ }, "anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, @@ -5117,6 +5131,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", @@ -17903,11 +17949,13 @@ }, "global.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "output_cost_per_token_above_200k_tokens": 2.25e-05, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, @@ -17964,6 +18012,7 @@ }, "global.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", @@ -19097,6 +19146,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", @@ -30106,6 +30187,7 @@ }, "us.anthropic.claude-haiku-4-5-20251001-v1:0": { "cache_creation_input_token_cost": 1.375e-06, + "cache_creation_input_token_cost_above_1hr": 2.2e-06, "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", @@ -30257,11 +30339,13 @@ }, "us.anthropic.claude-sonnet-4-5-20250929-v1:0": { "cache_creation_input_token_cost": 4.125e-06, + "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "input_cost_per_token_above_200k_tokens": 6.6e-06, "output_cost_per_token_above_200k_tokens": 2.475e-05, "cache_creation_input_token_cost_above_200k_tokens": 8.25e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", "max_input_tokens": 200000, @@ -30362,6 +30446,7 @@ }, "us.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.875e-06, + "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", @@ -30389,6 +30474,7 @@ }, "global.anthropic.claude-opus-4-5-20251101-v1:0": { "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", 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..a9d3911c07 100644 --- a/schema.prisma +++ b/schema.prisma @@ -277,6 +277,7 @@ model LiteLLM_ObjectPermissionTable { models String[] @default([]) blocked_tools String[] @default([]) // Tool names blocked for any key/team/user with this permission mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user + search_tools String[] @default([]) // search_tool_name values this key/team/user may call teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] @@ -1290,3 +1291,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/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 0aed224c25..a64f208b3f 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -557,6 +557,7 @@ async def test_avertex_batch_prediction(monkeypatch): mock_get_response = MagicMock() mock_get_response.json.return_value = mock_vertex_batch_response mock_get_response.status_code = 200 + mock_get_response.is_redirect = False mock_get_response.raise_for_status.return_value = None mock_get.return_value = mock_get_response diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index fc9c99f6af..07af2735df 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -46,6 +46,7 @@ IGNORE_FUNCTIONS = [ "dict", # max depth set. _LiteLLMParamsDictView.dict() calls builtin dict(), not itself. "_read_image_bytes", # max depth set. "_get_masked_values", # max depth set (default 20) to prevent infinite recursion while masking nested sensitive config dicts. + "_redact_sensitive_litellm_params", # max depth set (default 10). ] 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/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/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index 2d772c4a63..75061dda94 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -115,6 +115,13 @@ async def test_proxy_failure_metrics(): "litellm_llm_api_failed_requests_metric_total{", # Deprecated but may still be used ] + # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for + # hash_token(master_key) so the master key (or its hash) never + # propagates into metrics. See PR #26484. + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS + # Check if either pattern is in metrics and contains required fields found_metric = False for pattern in expected_patterns: @@ -125,8 +132,7 @@ async def test_proxy_failure_metrics(): 'api_key_alias="None"' in line and 'exception_class="Openai.RateLimitError"' in line and 'exception_status="429"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and 'route="/chat/completions"' in line ): @@ -135,8 +141,7 @@ async def test_proxy_failure_metrics(): # For deprecated llm_api metric, check llm-specific fields elif "litellm_llm_api_failed_requests_metric_total{" in line: if ( - 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + f'hashed_api_key="{expected_hashed_api_key}"' in line and 'model="429"' in line ): # The deprecated metric uses the actual model from the request found_metric = True @@ -156,8 +161,7 @@ async def test_proxy_failure_metrics(): for line in metrics.split("\n"): if ( total_requests_pattern in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and 'status_code="429"' in line ): @@ -195,6 +199,12 @@ async def test_proxy_success_metrics(): assert END_USER_ID not in metrics + # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for + # hash_token(master_key) (PR #26484). + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS + # Check if the success metric is present and correct - use flexible matching # Check for request_total_latency_metric with required fields # Note: The model can be "gpt-3.5-turbo-0301" or similar depending on what's returned @@ -203,8 +213,7 @@ async def test_proxy_success_metrics(): if ( "litellm_request_total_latency_metric_bucket{" in line and 'api_key_alias="None"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-openai-endpoint"' in line and 'le="0.005"' in line ): @@ -221,8 +230,7 @@ async def test_proxy_success_metrics(): if ( "litellm_llm_api_latency_metric_bucket{" in line and 'api_key_alias="None"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-openai-endpoint"' in line and 'le="0.005"' in line ): @@ -298,6 +306,12 @@ async def test_proxy_fallback_metrics(): print("/metrics", metrics) + # Master-key auth substitutes LITELLM_PROXY_MASTER_KEY_ALIAS for + # hash_token(master_key) (PR #26484). + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + expected_hashed_api_key = LITELLM_PROXY_MASTER_KEY_ALIAS + # Check if successful fallback metric is incremented - use flexible matching found_successful_fallback = False for line in metrics.split("\n"): @@ -307,8 +321,7 @@ async def test_proxy_fallback_metrics(): and 'exception_class="Openai.RateLimitError"' in line and 'exception_status="429"' in line and 'fallback_model="fake-openai-endpoint"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and "1.0" in line ): @@ -328,8 +341,7 @@ async def test_proxy_fallback_metrics(): and 'exception_class="Openai.RateLimitError"' in line and 'exception_status="429"' in line and 'fallback_model="unknown-model"' in line - and 'hashed_api_key="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"' - in line + and f'hashed_api_key="{expected_hashed_api_key}"' in line and 'requested_model="fake-azure-endpoint"' in line and "1.0" in line ): diff --git a/tests/proxy_admin_ui_tests/.npmrc b/tests/proxy_admin_ui_tests/.npmrc index 168e81a1c4..7999681cc3 100644 --- a/tests/proxy_admin_ui_tests/.npmrc +++ b/tests/proxy_admin_ui_tests/.npmrc @@ -2,4 +2,4 @@ # Packages needing lifecycle scripts: npm rebuild ignore-scripts=true # Protects local npm install only — npm ci (used in CI) ignores this -min-release-age=3d +min-release-age=3 diff --git a/tests/proxy_admin_ui_tests/ui_unit_tests/.npmrc b/tests/proxy_admin_ui_tests/ui_unit_tests/.npmrc index 168e81a1c4..7999681cc3 100644 --- a/tests/proxy_admin_ui_tests/ui_unit_tests/.npmrc +++ b/tests/proxy_admin_ui_tests/ui_unit_tests/.npmrc @@ -2,4 +2,4 @@ # Packages needing lifecycle scripts: npm rebuild ignore-scripts=true # Protects local npm install only — npm ci (used in CI) ignores this -min-release-age=3d +min-release-age=3 diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 19ee9f75d8..6a568d94f8 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -2715,7 +2715,12 @@ async def test_master_key_hashing(prisma_client): request=request, api_key=bearer_token ) - assert result.api_key == hash_token(master_key) + # Master-key auth substitutes a stable alias so the master key (or + # its hash) never propagates into spend logs / metrics / audit trails. + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + assert result.api_key != hash_token(master_key) except Exception as e: print("Got Exception", e) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 2c54244e25..543cabb6b4 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -1123,11 +1123,17 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): @pytest.mark.asyncio async def test_x_litellm_api_key(): """ - Check if auth can pick up x-litellm-api-key header, even if Bearer token is provided + Check if auth can pick up x-litellm-api-key header, even if Bearer token is provided. + + On a master-key match, ``UserAPIKeyAuth.api_key`` (and the derived + ``token``) are now the stable alias ``LITELLM_PROXY_MASTER_KEY_ALIAS`` + rather than ``hash_token(master_key)`` — the master key (or its hash) + must not propagate into spend logs / metrics / audit trails. """ from fastapi import Request from starlette.datastructures import URL + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, @@ -1153,7 +1159,8 @@ async def test_x_litellm_api_key(): api_key="Bearer " + ignored_key, custom_litellm_key_header=master_key, ) - assert valid_token.token == hash_token(master_key) + assert valid_token.token == LITELLM_PROXY_MASTER_KEY_ALIAS + assert valid_token.token != hash_token(master_key) @pytest.mark.asyncio 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 f8708dd2f7..d424cd8599 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 @@ -2476,3 +2476,186 @@ def test_bedrock_tools_pt_passes_ttl_for_claude_4_5(): cache_blocks_old = [b for b in result_old if "cachePoint" in b] assert len(cache_blocks_old) == 1 assert "ttl" not in cache_blocks_old[0]["cachePoint"] + + +def test_convert_to_anthropic_tool_result_openai_file_pdf_becomes_document(): + """ + OpenAI `{type: "file", file: {file_data: "data:application/pdf;..."}}` inside + a tool-message content list should translate to an Anthropic document block + inside the tool_result content. Reuses anthropic_process_openai_file_message, + which already handles this for user messages. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + pdf_b64 = "JVBERi0xLjQKJeLjz9MK" + message = { + "tool_call_id": "toolu_pdf_1", + "role": "tool", + "name": "fetch_document", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:application/pdf;base64,{pdf_b64}", + "filename": "summary.pdf", + }, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + assert result["type"] == "tool_result" + assert result["tool_use_id"] == "toolu_pdf_1" + content = result["content"] + assert isinstance(content, list) and len(content) == 1 + block = content[0] + assert block["type"] == "document" + assert block["source"]["type"] == "base64" + assert block["source"]["media_type"] == "application/pdf" + assert block["source"]["data"] == pdf_b64 + + +def test_convert_to_anthropic_tool_result_image_url_pdf_data_uri_becomes_document(): + """ + Regression: a PDF sent as an `image_url` data URI on the tool-result path + must translate to an Anthropic document block (not an image block — Anthropic + rejects image blocks whose media_type is a non-image like application/pdf). + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + pdf_b64 = "JVBERi0xLjQKJeLjz9MK" + message = { + "tool_call_id": "toolu_pdf_img_1", + "role": "tool", + "name": "fetch_document", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:application/pdf;base64,{pdf_b64}", + }, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + content = result["content"] + assert isinstance(content, list) and len(content) == 1 + block = content[0] + assert block["type"] == "document" + assert block["source"]["media_type"] == "application/pdf" + assert block["source"]["data"] == pdf_b64 + + +def test_convert_to_anthropic_tool_result_image_url_unsupported_mime_stays_image_path(): + """ + An `image_url` data URI whose mime is neither application/pdf nor text/plain + (e.g. application/json) must NOT be routed through the document path. Anthropic + only accepts application/pdf and text/plain as base64 document media_types — + anything else would produce a document block the API rejects. The old + (pre-fix) behavior was to wrap such data as an image block, which also + fails but stays on the image code path; preserve that failure mode rather + than switching to a document path that is equally broken. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + message = { + "tool_call_id": "toolu_json_1", + "role": "tool", + "name": "fetch_json", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:application/json;base64,eyJrIjoidiJ9", + }, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + content = result["content"] + assert isinstance(content, list) and len(content) == 1 + block = content[0] + assert block["type"] == "image", ( + f"unsupported mime {block.get('source', {}).get('media_type')!r} " + f"should not be routed to document path; got {block}" + ) + + +def test_convert_to_anthropic_tool_result_image_url_text_plain_data_uri_becomes_document(): + """ + text/plain is one of the two mimes Anthropic accepts as a base64 document + media_type. Confirm it routes through the document path so tightening the + gate to {application/pdf, text/plain} (not "application/*") covers both. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + txt_b64 = "aGVsbG8=" # "hello" + message = { + "tool_call_id": "toolu_txt_1", + "role": "tool", + "name": "fetch_text", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:text/plain;base64,{txt_b64}", + }, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + content = result["content"] + assert isinstance(content, list) and len(content) == 1 + block = content[0] + assert block["type"] == "document" + assert block["source"]["media_type"] == "text/plain" + assert block["source"]["data"] == txt_b64 + + +def test_convert_to_anthropic_tool_result_image_url_png_still_becomes_image(): + """ + Regression: image_url with a real image mime type must continue to translate + to an Anthropic image block. Locks in existing behavior after the + data-URI-mime-type branching for PDFs. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_anthropic_tool_result, + ) + + png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg==" + message = { + "tool_call_id": "toolu_png_1", + "role": "tool", + "name": "fetch_image", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{png_b64}", + }, + }, + ], + } + + result = convert_to_anthropic_tool_result(message) + + content = result["content"] + assert isinstance(content, list) and len(content) == 1 + block = content[0] + assert block["type"] == "image" + assert block["source"]["media_type"] == "image/png" 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/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 8e53e57f1e..589bcb4fc0 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -4146,3 +4146,257 @@ 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_bedrock_tool_message_openai_file_pdf_becomes_document(): + """ + OpenAI Chat Completions `{type: "file", file: {file_data: "data:application/pdf;...", filename}}` + inside a tool message content list should translate to a Bedrock + toolResult.content[].document block. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + pdf_b64 = "JVBERi0xLjQKJeLjz9MK" # tiny "%PDF-1.4\n" header + messages = [ + {"role": "user", "content": "Summarize the attached PDF."}, + { + "tool_call_id": "tooluse_pdf_1", + "role": "tool", + "name": "fetch_document", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:application/pdf;base64,{pdf_b64}", + "filename": "summary.pdf", + }, + }, + ], + }, + ] + + translated_msg = _bedrock_converse_messages_pt( + messages=messages, model="", llm_provider="" + ) + + tool_result = translated_msg[-1]["content"][-1]["toolResult"] + assert tool_result["toolUseId"] == "tooluse_pdf_1" + assert len(tool_result["content"]) == 1 + block = tool_result["content"][0] + assert "document" in block, f"expected document block, got {block}" + assert block["document"]["format"] == "pdf" + assert block["document"]["source"]["bytes"] == pdf_b64 + assert block["document"]["name"].startswith("DocumentPDFmessages_") + assert block["document"]["name"].endswith("_pdf") + + +def test_bedrock_tool_message_image_url_pdf_data_uri_becomes_document(): + """ + Regression for the processor-returns-document-but-wrapper-drops-it bug: + when a caller sends a PDF as an `image_url` data URI on the tool-result path, + BedrockImageProcessor correctly returns a {"document": ...} block, but the + tool-result wrapper only appended the "image" case, silently dropping documents. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + pdf_b64 = "JVBERi0xLjQKJeLjz9MK" + messages = [ + {"role": "user", "content": "Summarize the attached PDF."}, + { + "tool_call_id": "tooluse_pdf_img_1", + "role": "tool", + "name": "fetch_document", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:application/pdf;base64,{pdf_b64}", + }, + }, + ], + }, + ] + + translated_msg = _bedrock_converse_messages_pt( + messages=messages, model="", llm_provider="" + ) + + tool_result = translated_msg[-1]["content"][-1]["toolResult"] + assert tool_result["toolUseId"] == "tooluse_pdf_img_1" + assert len(tool_result["content"]) == 1 + block = tool_result["content"][0] + assert "document" in block, f"expected document block, got {block}" + assert block["document"]["format"] == "pdf" + assert block["document"]["source"]["bytes"] == pdf_b64 + + +def test_bedrock_tool_message_file_id_http_url_becomes_document(): + """ + OpenAI `file.file_id` is a server-side file reference. The Bedrock + user-message path (_process_file_message at factory.py:4796) accepts either + `file_data` or `file_id` and forwards to BedrockImageProcessor. The + tool-result path must match: when `file_id` is an http(s) PDF URL, it + should resolve to a Bedrock document block, not be silently dropped. + """ + from unittest.mock import patch + + from litellm.litellm_core_utils.prompt_templates.factory import ( + BedrockImageProcessor, + _bedrock_converse_messages_pt, + ) + + pdf_url = "https://example.com/whitepaper.pdf" + fake_document_block = { + "document": { + "format": "pdf", + "name": "fake_doc", + "source": {"bytes": "ZmFrZQ=="}, + } + } + messages = [ + {"role": "user", "content": "Summarize the attached PDF."}, + { + "tool_call_id": "tooluse_fid_1", + "role": "tool", + "name": "fetch_document", + "content": [ + { + "type": "file", + "file": { + "file_id": pdf_url, + "filename": "whitepaper.pdf", + }, + }, + ], + }, + ] + + with patch.object( + BedrockImageProcessor, + "process_image_sync", + return_value=fake_document_block, + ) as mock_proc: + translated_msg = _bedrock_converse_messages_pt( + messages=messages, model="", llm_provider="" + ) + + mock_proc.assert_called_once() + assert mock_proc.call_args.kwargs["image_url"] == pdf_url + + tool_result = translated_msg[-1]["content"][-1]["toolResult"] + assert len(tool_result["content"]) == 1 + block = tool_result["content"][0] + assert "document" in block, f"expected document block, got {block}" + assert block["document"]["source"]["bytes"] == "ZmFrZQ==" + + +def test_bedrock_tool_message_file_without_data_or_id_raises(): + """ + The user-message path raises BadRequestError when a `type: "file"` block + has neither `file_data` nor `file_id` (factory.py:4802-4809). The + tool-result path must match — silently dropping the block makes the model + see an empty tool result and obscures the caller bug. + """ + import litellm + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + messages = [ + {"role": "user", "content": "Summarize."}, + { + "tool_call_id": "tooluse_bad_1", + "role": "tool", + "name": "fetch_document", + "content": [ + { + "type": "file", + "file": {"filename": "nothing.pdf"}, + }, + ], + }, + ] + + with pytest.raises(litellm.BadRequestError): + _bedrock_converse_messages_pt(messages=messages, model="", llm_provider="") + + +def test_bedrock_tool_message_image_url_png_still_becomes_image(): + """ + Regression: image_url with an image mime type must continue to translate + to a Bedrock image block (not document). Locks in existing behavior after + the document-passthrough fix. + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABXvMqOgAAAABJRU5ErkJggg==" + messages = [ + {"role": "user", "content": "Describe the attached image."}, + { + "tool_call_id": "tooluse_png_1", + "role": "tool", + "name": "fetch_image", + "content": [ + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{png_b64}", + }, + }, + ], + }, + ] + + translated_msg = _bedrock_converse_messages_pt( + messages=messages, model="", llm_provider="" + ) + + tool_result = translated_msg[-1]["content"][-1]["toolResult"] + assert len(tool_result["content"]) == 1 + block = tool_result["content"][0] + assert "image" in block, f"expected image block, got {block}" + assert "document" not in block + assert block["image"]["format"] == "png" + assert block["image"]["source"]["bytes"] == png_b64 + + +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/passthrough/test_streaming_interrupt_spend_tracking.py b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py new file mode 100644 index 0000000000..f3fe3ae5c3 --- /dev/null +++ b/tests/test_litellm/passthrough/test_streaming_interrupt_spend_tracking.py @@ -0,0 +1,244 @@ +"""Regression tests for LIT-2642 — interrupted streams must still flush usage.""" + +import asyncio +from typing import List +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + + +def _make_streaming_response(chunks: List[bytes]): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"}) + mock.raise_for_status = MagicMock(return_value=None) + + async def _aiter_bytes(): + for chunk in chunks: + yield chunk + + mock.aiter_bytes = _aiter_bytes + mock.aclose = AsyncMock() + return mock + + +def _make_logging_obj(): + mock = MagicMock() + mock.async_flush_passthrough_collected_chunks = AsyncMock() + return mock + + +class _ImmediateExecutor: + def submit(self, fn, *args, **kwargs): + fn(*args, **kwargs) + + +@pytest.mark.asyncio +async def test_async_streaming_flushes_on_normal_completion(): + from litellm.passthrough.main import _async_streaming + + chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] + mock_response = _make_streaming_response(chunks) + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + provider_config = MagicMock() + + received = [] + async for chunk in _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ): + received.append(chunk) + + assert received == chunks + + await asyncio.sleep(0) + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = ( + mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs + ) + assert call_kwargs["raw_bytes"] == chunks + assert call_kwargs["provider_config"] is provider_config + + +@pytest.mark.asyncio +async def test_async_streaming_flushes_on_client_disconnect(): + from litellm.passthrough.main import _async_streaming + + chunks = [ + b'{"chunk": 1, "outputTokens": 10}', + b'{"chunk": 2, "outputTokens": 12}', + b'{"chunk": 3, "outputTokens": 8}', + ] + mock_response = _make_streaming_response(chunks) + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + provider_config = MagicMock() + + gen = _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + + received = [await gen.__anext__()] + await gen.aclose() + + assert received == [chunks[0]] + + await asyncio.sleep(0) + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = ( + mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs + ) + assert call_kwargs["raw_bytes"] == [chunks[0]] + + +@pytest.mark.asyncio +async def test_async_streaming_does_not_flush_on_4xx(): + from litellm.passthrough.main import _async_streaming + + err_response = MagicMock(spec=httpx.Response) + err_response.status_code = 429 + + def _raise(): + raise httpx.HTTPStatusError( + "429", + request=httpx.Request("POST", "https://example.com"), + response=httpx.Response( + 429, request=httpx.Request("POST", "https://example.com") + ), + ) + + err_response.raise_for_status = _raise + err_response.aclose = AsyncMock() + + async def response_coro(): + return err_response + + mock_logging_obj = _make_logging_obj() + + with pytest.raises(httpx.HTTPStatusError): + async for _ in _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=MagicMock(), + ): + pass + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_streaming_flushes_on_upstream_exception_with_partial_data(): + from litellm.passthrough.main import _async_streaming + + partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"] + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock(return_value=None) + mock_response.aclose = AsyncMock() + + async def _aiter_bytes_then_raise(): + for c in partial_chunks: + yield c + raise httpx.ReadError("upstream disconnected") + + mock_response.aiter_bytes = _aiter_bytes_then_raise + + async def response_coro(): + return mock_response + + mock_logging_obj = _make_logging_obj() + provider_config = MagicMock() + + received = [] + with pytest.raises(httpx.ReadError): + async for chunk in _async_streaming( + response=response_coro(), + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ): + received.append(chunk) + + assert received == partial_chunks + + await asyncio.sleep(0) + + mock_logging_obj.async_flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = ( + mock_logging_obj.async_flush_passthrough_collected_chunks.call_args.kwargs + ) + assert call_kwargs["raw_bytes"] == partial_chunks + + +def test_sync_streaming_flushes_on_normal_completion(): + from litellm.passthrough.main import _sync_streaming + + chunks = [b"a", b"b", b"c"] + + mock_response = MagicMock(spec=httpx.Response) + + def _iter_bytes(): + yield from chunks + + mock_response.iter_bytes = _iter_bytes + + mock_logging_obj = MagicMock() + mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() + provider_config = MagicMock() + + with patch("litellm.utils.executor", _ImmediateExecutor()): + received = list( + _sync_streaming( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + ) + + assert received == chunks + mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() + + +def test_sync_streaming_flushes_on_early_close(): + from litellm.passthrough.main import _sync_streaming + + chunks = [b"first", b"second", b"third"] + + mock_response = MagicMock(spec=httpx.Response) + + def _iter_bytes(): + yield from chunks + + mock_response.iter_bytes = _iter_bytes + + mock_logging_obj = MagicMock() + mock_logging_obj.flush_passthrough_collected_chunks = MagicMock() + provider_config = MagicMock() + + with patch("litellm.utils.executor", _ImmediateExecutor()): + gen = _sync_streaming( + response=mock_response, + litellm_logging_obj=mock_logging_obj, + provider_config=provider_config, + ) + + first = next(gen) + gen.close() + + assert first == chunks[0] + mock_logging_obj.flush_passthrough_collected_chunks.assert_called_once() + call_kwargs = mock_logging_obj.flush_passthrough_collected_chunks.call_args.kwargs + assert call_kwargs["raw_bytes"] == [chunks[0]] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index fd1a2b1236..dd352d0999 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -551,11 +551,14 @@ class TestMCPOAuth2AuthFlow: async def test_oauth2_token_in_authorization_header_fallback(self): """ - When only Authorization header is present with a non-LiteLLM OAuth2 token, + When only Authorization header is present with a non-LiteLLM OAuth2 token + AND the target server is operator-configured for ``auth_type=oauth2``, auth should fall back to permissive mode (OAuth2 passthrough). """ from fastapi import HTTPException + from litellm.types.mcp import MCPAuth + scope = { "type": "http", "method": "POST", @@ -568,10 +571,19 @@ class TestMCPOAuth2AuthFlow: async def mock_user_api_key_auth_fails(api_key, request): raise HTTPException(status_code=401, detail="Invalid API key") - with patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", - side_effect=mock_user_api_key_auth_fails, + oauth2_server = MagicMock() + oauth2_server.auth_type = MCPAuth.oauth2 + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, ): + mock_mgr.get_mcp_server_by_name.return_value = oauth2_server ( auth_result, mcp_auth_header, @@ -695,9 +707,11 @@ class TestMCPOAuth2AuthFlow: async def test_proxy_exception_oauth2_fallback(self): """ user_api_key_auth raises ProxyException (not HTTPException) in production. - The OAuth2 fallback must catch ProxyException with code 401/403 too. + The OAuth2 fallback must catch ProxyException with code 401/403 too, + but only when the target server is operator-configured for ``auth_type=oauth2``. """ from litellm.proxy._types import ProxyException + from litellm.types.mcp import MCPAuth scope = { "type": "http", @@ -716,10 +730,19 @@ class TestMCPOAuth2AuthFlow: code=401, ) - with patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", - side_effect=mock_user_api_key_auth_proxy_exception, + oauth2_server = MagicMock() + oauth2_server.auth_type = MCPAuth.oauth2 + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_proxy_exception, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, ): + mock_mgr.get_mcp_server_by_name.return_value = oauth2_server ( auth_result, mcp_auth_header, @@ -768,6 +791,290 @@ class TestMCPOAuth2AuthFlow: await MCPRequestHandler.process_mcp_request(scope) +@pytest.mark.asyncio +class TestMCPPublicRouteGuard: + """ + Regression tests for GHSA-7cwm-3279-qf3c / HW6xR21d: + the public-route bypass at the top of process_mcp_request must match + the exact `/.well-known/` path prefix, not a substring of the URL. + """ + + async def test_well_known_substring_in_query_does_not_bypass_auth(self): + """ + URL with `.well-known` smuggled into the query string must still + require valid LiteLLM auth. + """ + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/private_server", + "query_string": b"redirect=.well-known/oauth-protected-resource", + "headers": [(b"authorization", b"Bearer sk-bogus")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + # Explicit unresolvable target — proves auth still fails even + # when the registry has no info to fall back to. + mock_mgr.get_mcp_server_by_name.return_value = None + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 401 + + async def test_well_known_segment_in_middle_of_path_does_not_bypass_auth(self): + """ + Path containing `.well-known` as a non-prefix component (e.g. a server + name or sub-path) must still require auth. + """ + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/.well-known-fake/tools", + "headers": [(b"authorization", b"Bearer sk-bogus")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = None + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 401 + + async def test_legitimate_well_known_path_still_bypasses_auth(self): + """ + Real OAuth discovery routes registered under /.well-known/ must remain + public so unauthenticated clients can fetch them per RFC 8414/9728. + """ + scope = { + "type": "http", + "method": "GET", + "path": "/.well-known/oauth-protected-resource", + "headers": [], + } + + # No mock needed — public path should not call user_api_key_auth at all + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + ) as mock_auth: + (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + mock_auth.assert_not_called() + assert isinstance(auth_result, UserAPIKeyAuth) + + +@pytest.mark.asyncio +class TestMCPOAuth2FallbackTargetGating: + """ + Regression tests for GHSA-h8fm-g6wc-j228 / HW6xR21d: + The OAuth2 passthrough fallback must only fire when the target MCP server + is operator-configured for ``auth_type=oauth2``. A failed LiteLLM-auth + against a non-OAuth2 server (api_key, bearer_token, basic, etc.) must + propagate as a real auth error, not be exchanged for an anonymous session. + """ + + @staticmethod + def _make_server(auth_type): + server = MagicMock() + server.auth_type = auth_type + return server + + async def test_fallback_blocked_when_target_is_not_oauth2(self): + from fastapi import HTTPException + + from litellm.types.mcp import MCPAuth + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/api_key_server", + "headers": [(b"authorization", b"Bearer anything-at-all")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.api_key) + ) + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 401 + + async def test_fallback_blocked_when_target_unresolvable(self): + """ + If the target server cannot be resolved from path or x-mcp-servers, + we cannot prove it is OAuth2-mode, so we must fail closed. + """ + from fastapi import HTTPException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/never_registered_server", + "headers": [(b"authorization", b"Bearer anything")], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = None + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 401 + + async def test_fallback_allowed_when_target_is_oauth2_mode(self): + """ + Operator-configured OAuth2 passthrough still works: target server has + ``auth_type=oauth2`` → failed LiteLLM auth falls back to anonymous so + the bearer can be forwarded to upstream. + """ + from fastapi import HTTPException + + from litellm.types.mcp import MCPAuth + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/atlassian_mcp", + "headers": [ + (b"authorization", b"Bearer atlassian-oauth2-access-token-xyz"), + ], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.return_value = ( + TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.oauth2) + ) + (auth_result, *_rest) = await MCPRequestHandler.process_mcp_request(scope) + assert isinstance(auth_result, UserAPIKeyAuth) + + async def test_fallback_blocked_when_any_target_in_header_is_not_oauth2(self): + """ + x-mcp-servers can list multiple targets. If ANY of them is non-OAuth2, + the fallback must be blocked — otherwise an attacker can mix one + OAuth2-mode server in to enable bypass against the others. + """ + from fastapi import HTTPException + + from litellm.types.mcp import MCPAuth + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"authorization", b"Bearer anything"), + (b"x-mcp-servers", b"oauth2_server,api_key_server"), + ], + } + + async def mock_user_api_key_auth_fails(api_key, request): + raise HTTPException(status_code=401, detail="Invalid API key") + + def mock_lookup(name, client_ip=None): + if name == "oauth2_server": + return TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.oauth2) + return TestMCPOAuth2FallbackTargetGating._make_server(MCPAuth.api_key) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_fails, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_mgr, + ): + mock_mgr.get_mcp_server_by_name.side_effect = mock_lookup + with pytest.raises(HTTPException) as exc_info: + await MCPRequestHandler.process_mcp_request(scope) + assert exc_info.value.status_code == 401 + + async def test_proxy_exception_with_non_numeric_code_propagates(self): + """ + ``ProxyException`` normalises ``code`` via ``str()`` in its __init__, + so callers may produce ``"None"`` or any non-numeric string when no + explicit code was supplied. The exception handler must not coerce + with ``int(...)`` (which would raise ``ValueError`` and rewrite the + auth error as an unhandled 500); it must simply re-raise. + """ + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/atlassian_mcp", + "headers": [(b"authorization", b"Bearer anything")], + } + + async def mock_user_api_key_auth_no_code(api_key, request): + raise ProxyException( + message="Authentication Error", + type="auth_error", + param="api_key", + code=None, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth_no_code, + ): + with pytest.raises(ProxyException): + await MCPRequestHandler.process_mcp_request(scope) + + class TestMCPCustomHeaderName: """Test suite for custom MCP authentication header name functionality""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index ae3c16ce8b..558c677d2d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -230,7 +230,6 @@ async def test_authorize_endpoint_forwards_pkce_parameters(): async def test_token_endpoint_forwards_code_verifier(): """Test that token endpoint forwards code_verifier for PKCE flow""" try: - import httpx from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -632,8 +631,7 @@ async def test_token_endpoint_respects_x_forwarded_proto(): ) as mock_get_client: mock_get_client.return_value = mock_async_client - # Call token endpoint - response = await token_endpoint( + await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_code", @@ -933,8 +931,7 @@ async def test_token_endpoint_respects_x_forwarded_host(): ) as mock_get_client: mock_get_client.return_value = mock_async_client - # Call token endpoint - response = await token_endpoint( + await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_code", @@ -1240,6 +1237,7 @@ def _create_oauth2_server( alias="test_oauth", client_id="test_client_id", client_secret="test_client_secret", + available_on_public_internet=True, ): """Helper to create a mock OAuth2 MCPServer.""" from litellm.proxy._types import MCPTransport @@ -1258,6 +1256,7 @@ def _create_oauth2_server( authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], + available_on_public_internet=available_on_public_internet, ) @@ -1352,6 +1351,47 @@ async def test_authorize_root_fails_with_multiple_oauth2_servers(): global_mcp_server_manager.registry.clear() +@pytest.mark.asyncio +async def test_authorize_root_does_not_resolve_private_server_for_external_client(): + """Root /authorize must not auto-select an MCP server hidden from the caller IP.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + authorize, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server(available_on_public_internet=False) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", + return_value="198.51.100.10", + ): + with pytest.raises(HTTPException) as exc_info: + await authorize( + request=mock_request, + client_id="dummy_client", + mcp_server_name=None, + redirect_uri="http://localhost:62646/callback", + state="test_state", + ) + assert exc_info.value.status_code == 404 + finally: + global_mcp_server_manager.registry.clear() + + @pytest.mark.asyncio async def test_token_root_resolves_single_oauth2_server(): """When /token is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" @@ -1417,6 +1457,50 @@ async def test_token_root_resolves_single_oauth2_server(): global_mcp_server_manager.registry.clear() +@pytest.mark.asyncio +async def test_token_root_does_not_resolve_private_server_for_external_client(): + """Root /token must not exchange codes for a hidden MCP server.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + token_endpoint, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server(available_on_public_internet=False) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", + return_value="198.51.100.10", + ): + with pytest.raises(HTTPException) as exc_info: + await token_endpoint( + request=mock_request, + grant_type="authorization_code", + code="test_auth_code", + redirect_uri="http://localhost:62646/callback", + client_id="dummy_client", + mcp_server_name=None, + client_secret=None, + code_verifier="test_verifier", + ) + assert exc_info.value.status_code == 404 + finally: + global_mcp_server_manager.registry.clear() + + @pytest.mark.asyncio async def test_register_root_resolves_single_oauth2_server(): """When /register is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" @@ -1454,6 +1538,48 @@ async def test_register_root_resolves_single_oauth2_server(): global_mcp_server_manager.registry.clear() +@pytest.mark.asyncio +async def test_register_root_does_not_resolve_private_server_for_external_client(): + """Root /register must not reveal or use a hidden MCP server.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + register_client, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server(available_on_public_internet=False) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with ( + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", + new=AsyncMock(return_value={}), + ), + patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", + return_value="198.51.100.10", + ), + ): + result = await register_client(request=mock_request, mcp_server_name=None) + + assert result["client_id"] == "dummy_client" + assert result["redirect_uris"] == ["https://llm.example.com/callback"] + finally: + global_mcp_server_manager.registry.clear() + + @pytest.mark.asyncio async def test_discovery_root_includes_server_name_prefix(): """When root discovery is hit and exactly 1 OAuth2 server exists, include server name in URLs.""" @@ -1493,6 +1619,54 @@ async def test_discovery_root_includes_server_name_prefix(): global_mcp_server_manager.registry.clear() +@pytest.mark.asyncio +async def test_discovery_root_does_not_expose_private_server_for_external_client(): + """Root discovery must use caller visibility before adding server-specific metadata.""" + try: + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_authorization_server_response, + _build_oauth_protected_resource_response, + ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + global_mcp_server_manager.registry.clear() + oauth2_server = _create_oauth2_server(available_on_public_internet=False) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://llm.example.com/" + mock_request.headers = {} + + try: + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", + return_value="198.51.100.10", + ): + authorization_response = _build_oauth_authorization_server_response( + request=mock_request, + mcp_server_name=None, + ) + resource_response = _build_oauth_protected_resource_response( + request=mock_request, + mcp_server_name=None, + use_standard_pattern=False, + ) + + assert "/test_oauth/" not in authorization_response["authorization_endpoint"] + assert "/test_oauth/" not in authorization_response["token_endpoint"] + assert authorization_response["scopes_supported"] == [] + assert resource_response["authorization_servers"] == ["https://llm.example.com"] + assert resource_response["scopes_supported"] == [] + finally: + global_mcp_server_manager.registry.clear() + + @pytest.mark.asyncio async def test_oauth_callback_redirects_with_state(): """Test OAuth callback endpoint properly decodes state and redirects to client callback URL.""" @@ -1536,6 +1710,44 @@ async def test_oauth_callback_redirects_with_state(): mock_decode.assert_called_once_with("encrypted_state_value") +@pytest.mark.asyncio +async def test_oauth_callback_preserves_client_redirect_uri_query(): + """The callback should append code/state without dropping a client's existing query.""" + try: + from urllib.parse import parse_qs, urlparse + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + except ImportError: + pytest.skip("MCP discoverable endpoints not available") + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.return_value = { + "base_url": "http://localhost:3000/ui/mcp/oauth/callback", + "original_state": "test-uuid-state-123", + "code_challenge": "test_challenge", + "code_challenge_method": "S256", + "client_redirect_uri": ( + "http://localhost:3000/ui/mcp/oauth/callback?session=abc" + ), + } + + response = await callback( + code="test_authorization_code_12345", + state="encrypted_state_value", + ) + + assert response.status_code == 302 + parsed_location = urlparse(response.headers["location"]) + query_params = parse_qs(parsed_location.query) + assert query_params["session"] == ["abc"] + assert query_params["code"] == ["test_authorization_code_12345"] + assert query_params["state"] == ["test-uuid-state-123"] + + @pytest.mark.asyncio async def test_oauth_callback_handles_invalid_state(): """Test OAuth callback returns error page when state decryption fails.""" @@ -1948,6 +2160,48 @@ async def test_callback_revalidates_loopback_on_decoded_base_url(): assert exc_info.value.status_code == 400 +@pytest.mark.asyncio +async def test_callback_revalidates_loopback_on_decoded_client_redirect_uri(): + """If a state contains a full client_redirect_uri, validate that exact sink.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.return_value = { + "base_url": "http://localhost:3000/cb", + "original_state": "s", + "code_challenge": None, + "code_challenge_method": None, + "client_redirect_uri": "https://attacker.example.com/cb", + } + with pytest.raises(HTTPException) as exc_info: + await callback(code="stolen_code", state="encrypted_stale_state") + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_callback_rejects_state_missing_redirect_uri(): + """Malformed state without a redirect target should fail with a structured 400.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.return_value = { + "original_state": "s", + "code_challenge": None, + "code_challenge_method": None, + } + with pytest.raises(HTTPException) as exc_info: + await callback(code="code", state="encrypted_malformed_state") + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio async def test_token_endpoint_sets_no_store_cache_control(): """RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: the token response 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_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index dbee062e2b..4c21d0ec64 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2080,6 +2080,33 @@ class TestGuardrailModificationCheck: ) assert exc.value.status_code == 403 + @pytest.mark.parametrize( + "key", + [ + "guardrails", + "disable_global_guardrails", + "disable_global_guardrail", + "opted_out_global_guardrails", + ], + ) + @pytest.mark.parametrize("empty_value", [{}, [], "", 0, False]) + def test_rejects_empty_value_modification(self, key, empty_value): + """Regression: an explicitly-supplied empty/falsy value still expresses + intent to modify and must trigger the permission check. Truthiness-based + gating let callers bypass the check by sending e.g. + ``metadata={"guardrails": {}}``, which downstream evaluation interpreted + as "disable all guardrails" while the auth layer treated it as no-op. + """ + from fastapi import HTTPException + + with patch( + "litellm.proxy.guardrails.guardrail_helpers.can_modify_guardrails", + return_value=False, + ): + with pytest.raises(HTTPException) as exc: + self._call({"metadata": {key: empty_value}}) + assert exc.value.status_code == 403 + def test_rejects_injection_via_litellm_metadata_key(self): """Caller can populate the OTHER metadata key; that must also 403.""" from fastapi import HTTPException @@ -2338,3 +2365,134 @@ async def test_team_member_budget_check_per_member_override_wins_over_team_defau ) assert exc_info.value.current_cost == 250.0 assert exc_info.value.max_budget == 200.0 + + +@pytest.mark.asyncio +async def test_team_member_budget_check_null_clone_falls_back_to_team_default(): + """Per-member NULL max_budget falls through to the team default cap.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + team_object = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-default"}, + ) + user_object = LiteLLM_UserTable(user_id="test-user") + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user", + team_id="test-team", + ) + + # Per-member row exists with NULL max_budget (the cloned-from-incomplete-default case). + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + budget_id="budget-clone", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=None), + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + fake_default_row = MagicMock() + fake_default_row.max_budget = 65.0 + fake_default_row.dict = MagicMock( + return_value={"budget_id": "budget-default", "max_budget": 65.0} + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=fake_default_row + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:team_member:test-user:test-team": + return 500.0 + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.current_cost == 500.0 + assert exc_info.value.max_budget == 65.0 + prisma_client.db.litellm_budgettable.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_member_budget_check_null_clone_with_null_default_skips_enforcement(): + """When per-member and team default are both NULL, enforcement still skips.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + team_object = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-default"}, + ) + user_object = LiteLLM_UserTable(user_id="test-user") + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user", + team_id="test-team", + ) + + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + budget_id="budget-clone", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=None), + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + fake_default_row = MagicMock() + fake_default_row.max_budget = None + fake_default_row.dict = MagicMock( + return_value={"budget_id": "budget-default", "max_budget": None} + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=fake_default_row + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:team_member:test-user:test-team": + return 1000.0 + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), + ): + # No raise: both rows are NULL, so enforcement is correctly skipped. + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=proxy_logging_obj, + ) 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/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index 57fa871a35..c81f4cb7d6 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -2,14 +2,16 @@ Tests for the invite-link onboarding endpoints. Covers the security behavior of: - GET /onboarding/get_token – rejects already-used links before showing any user data - POST /onboarding/claim_token – rejects already-used links; marks is_accepted=True only - after the password is successfully written + GET /onboarding/get_token – rejects already-used links and returns only a + short-lived onboarding token, not a UI session key + POST /onboarding/claim_token – requires that onboarding token; mints the UI + session key only after the password is written """ from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch +import jwt import pytest from fastapi import HTTPException @@ -22,14 +24,27 @@ from litellm.proxy._types import InvitationClaim # --------------------------------------------------------------------------- -def _make_invite(*, is_accepted: bool, expired: bool = False) -> MagicMock: +class _AsyncTx: + def __init__(self, db: MagicMock): + self.db = db + + async def __aenter__(self) -> MagicMock: + return self.db + + async def __aexit__(self, exc_type, exc, tb): + return False + + +def _make_invite( + *, is_accepted: bool, expired: bool = False, claimed: bool = False +) -> MagicMock: now = litellm.utils.get_utc_datetime() invite = MagicMock() invite.id = "invite-abc" invite.user_id = "user-123" invite.is_accepted = is_accepted invite.expires_at = now - timedelta(days=1) if expired else now + timedelta(days=6) - invite.accepted_at = None + invite.accepted_at = now if claimed else None return invite @@ -45,11 +60,39 @@ def _make_prisma(invite: MagicMock, user: MagicMock | None = None) -> MagicMock: prisma = MagicMock() prisma.db.litellm_invitationlink.find_unique = AsyncMock(return_value=invite) prisma.db.litellm_invitationlink.update = AsyncMock() + prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=1) prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user) prisma.db.litellm_usertable.update = AsyncMock(return_value=user) + prisma.db.tx = MagicMock(return_value=_AsyncTx(prisma.db)) return prisma +def _make_onboarding_token( + *, + invitation_link: str = "invite-abc", + user_id: str = "user-123", + token_type: str = "litellm_onboarding", + master_key: str = "sk-test", +) -> str: + return jwt.encode( + { + "token_type": token_type, + "invitation_link": invitation_link, + "user_id": user_id, + "exp": litellm.utils.get_utc_datetime() + timedelta(minutes=15), + }, + master_key, + algorithm="HS256", + ) + + +def _make_claim_request(token: str | None = None) -> MagicMock: + request = MagicMock() + request.headers = {"Authorization": f"Bearer {token}"} if token is not None else {} + request.base_url = "http://localhost:4000/" + return request + + # --------------------------------------------------------------------------- # GET /onboarding/get_token # --------------------------------------------------------------------------- @@ -120,10 +163,10 @@ async def test_get_token_rejects_missing_link(): @pytest.mark.asyncio -async def test_get_token_does_not_set_is_accepted(): +async def test_get_token_returns_onboarding_token_without_minting_ui_key(): """ - A valid, unused link should succeed and must NOT flip is_accepted to True. - That flag is only written after the password is claimed. + A valid, unused link should return a short-lived onboarding token, but + must not reserve the invite or mint a usable UI/API key on GET. """ from litellm.proxy.proxy_server import onboarding @@ -133,6 +176,252 @@ async def test_get_token_does_not_set_is_accepted(): request = MagicMock() request.base_url = "http://localhost:4000/" + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key, + patch( + "litellm.proxy.proxy_server.get_custom_url", + return_value="http://localhost:4000/", + ), + patch( + "litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation", + return_value=False, + ), + patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), + ): + result = await onboarding(invite_link="invite-abc", request=request) + + # Endpoint succeeded + assert "token" in result + assert "login_url" in result + + outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"]) + onboarding_token = outer_claims["key"] + onboarding_claims = jwt.decode(onboarding_token, "sk-test", algorithms=["HS256"]) + assert onboarding_claims["token_type"] == "litellm_onboarding" + assert onboarding_claims["invitation_link"] == "invite-abc" + assert onboarding_claims["user_id"] == "user-123" + assert not onboarding_token.startswith("sk-") + + mock_generate_key.assert_not_called() + prisma.db.litellm_invitationlink.update_many.assert_not_called() + prisma.db.litellm_invitationlink.update.assert_not_called() + + +# --------------------------------------------------------------------------- +# POST /onboarding/claim_token +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_claim_token_rejects_already_used_link(): + """ + If is_accepted is True, the password has already been set. + A second claim attempt must be rejected with 401. + """ + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=True, claimed=True) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "already been used" in exc_info.value.detail["error"] + # Password must never have been written + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_expired_link(): + """An expired link must be rejected even if is_accepted is False.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False, expired=True) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "expired" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_claim_token_rejects_mismatched_user_id(): + """The user_id in the request must match the one on the invite.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="wrong-user", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "does not match" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_claim_token_rejects_missing_onboarding_token(): + """The password endpoint must require the onboarding token returned by get_token.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "Missing onboarding session" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_wrong_onboarding_session(): + """The onboarding token must be bound to the invite and user being claimed.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + request = _make_claim_request( + _make_onboarding_token(invitation_link="other-invite") + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 401 + assert "Invalid onboarding session" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_invalid_bearer_token(): + """A regular API key must not be accepted as an onboarding token.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + request = _make_claim_request("sk-regular-key") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 401 + assert "Invalid onboarding session" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_concurrent_reuse_before_password_write(): + """Only the first valid claim may reserve the invitation.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=0) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key, + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 401 + assert "already been used" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + mock_generate_key.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_sets_accepted_at_after_password_written(): + """ + A valid first-time claim must: + 1. Write the hashed password to the user table. + 2. Set accepted_at on the invitation link after the password write succeeds. + """ + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + user = _make_user() + prisma = _make_prisma(invite, user) + request = _make_claim_request(_make_onboarding_token()) + + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"} with ( @@ -155,113 +444,13 @@ async def test_get_token_does_not_set_is_accepted(): ), patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), ): - result = await onboarding(invite_link="invite-abc", request=request) - - # Endpoint succeeded - assert "token" in result - assert "login_url" in result - - # is_accepted must NOT have been updated here - prisma.db.litellm_invitationlink.update.assert_not_called() - - -# --------------------------------------------------------------------------- -# POST /onboarding/claim_token -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_claim_token_rejects_already_used_link(): - """ - If is_accepted is True, the password has already been set. - A second claim attempt must be rejected with 401. - """ - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=True) - prisma = _make_prisma(invite) - data = InvitationClaim( - invitation_link="invite-abc", - user_id="user-123", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - with pytest.raises(HTTPException) as exc_info: - await claim_onboarding_link(data=data) - - assert exc_info.value.status_code == 401 - assert "already been used" in exc_info.value.detail["error"] - # Password must never have been written - prisma.db.litellm_usertable.update.assert_not_called() - - -@pytest.mark.asyncio -async def test_claim_token_rejects_expired_link(): - """An expired link must be rejected even if is_accepted is False.""" - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=False, expired=True) - prisma = _make_prisma(invite) - data = InvitationClaim( - invitation_link="invite-abc", - user_id="user-123", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - with pytest.raises(HTTPException) as exc_info: - await claim_onboarding_link(data=data) - - assert exc_info.value.status_code == 401 - assert "expired" in exc_info.value.detail["error"] - - -@pytest.mark.asyncio -async def test_claim_token_rejects_mismatched_user_id(): - """The user_id in the request must match the one on the invite.""" - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=False) - prisma = _make_prisma(invite) - data = InvitationClaim( - invitation_link="invite-abc", - user_id="wrong-user", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - with pytest.raises(HTTPException) as exc_info: - await claim_onboarding_link(data=data) - - assert exc_info.value.status_code == 401 - assert "does not match" in exc_info.value.detail["error"] - - -@pytest.mark.asyncio -async def test_claim_token_sets_is_accepted_after_password_written(): - """ - A valid first-time claim must: - 1. Write the hashed password to the user table. - 2. Flip is_accepted to True on the invitation link — and only after the - password write succeeds. - """ - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=False) - user = _make_user() - prisma = _make_prisma(invite, user) - - data = InvitationClaim( - invitation_link="invite-abc", - user_id="user-123", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - result = await claim_onboarding_link(data=data) + result = await claim_onboarding_link(data=data, request=request) # Password was written + prisma.db.litellm_invitationlink.update_many.assert_called_once() + reserve_kwargs = prisma.db.litellm_invitationlink.update_many.call_args.kwargs + assert reserve_kwargs["where"] == {"id": "invite-abc", "is_accepted": False} + assert reserve_kwargs["data"]["is_accepted"] is True prisma.db.litellm_usertable.update.assert_called_once() call_kwargs = prisma.db.litellm_usertable.update.call_args assert call_kwargs.kwargs["where"] == {"user_id": "user-123"} @@ -270,5 +459,50 @@ async def test_claim_token_sets_is_accepted_after_password_written(): # is_accepted was flipped to True on the invitation link prisma.db.litellm_invitationlink.update.assert_called_once() link_update_data = prisma.db.litellm_invitationlink.update.call_args.kwargs["data"] - assert link_update_data["is_accepted"] is True + assert "is_accepted" not in link_update_data assert link_update_data["accepted_at"] is not None + outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"]) + assert outer_claims["key"] == "sk-generated-key" + + +@pytest.mark.asyncio +async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): + """A session key failure must not leave the invite permanently consumed.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + user = _make_user() + prisma = _make_prisma(invite, user) + request = _make_claim_request(_make_onboarding_token()) + + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + side_effect=Exception("key mint failed"), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 500 + assert "Failed to create onboarding session" in exc_info.value.detail["error"] + assert prisma.db.litellm_invitationlink.update_many.call_count == 2 + rollback_kwargs = prisma.db.litellm_invitationlink.update_many.call_args_list[ + 1 + ].kwargs + assert rollback_kwargs["where"] == { + "id": "invite-abc", + "is_accepted": True, + } + assert rollback_kwargs["data"]["accepted_at"] is None + assert rollback_kwargs["data"]["is_accepted"] is False diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 9c43ebcbe7..08f4bd0ebf 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2581,3 +2581,49 @@ async def test_centralized_common_checks_user_http_exception_isolates_to_user_on finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_master_key_auth_substitutes_alias_for_api_key(): + """ + When the master key authenticates a request, the resulting + ``UserAPIKeyAuth.api_key`` must be the stable alias + ``LITELLM_PROXY_MASTER_KEY_ALIAS`` — never the raw master key (which + would propagate downstream and be hashed into spend logs, Prometheus + ``/metrics`` labels, or audit trails) and never the master-key hash. + """ + from fastapi import Request + from starlette.datastructures import URL + + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + from litellm.proxy.utils import hash_token + + import litellm.proxy.proxy_server as _proxy_server_mod + + attrs = _proxy_server_attrs_for_custom_auth(user_custom_auth=None) + master_key = attrs["master_key"] + _orig = {k: getattr(_proxy_server_mod, k, None) for k in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + result = await _user_api_key_auth_builder( + request=request, + api_key=f"Bearer {master_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + + assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + assert result.api_key != master_key + assert result.api_key != hash_token(master_key) + finally: + for k, v in _orig.items(): + setattr(_proxy_server_mod, k, v) diff --git a/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py b/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py index 44288b027a..3e3e89f811 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_endpoint_utils.py @@ -104,7 +104,9 @@ def test_remove_sensitive_info_from_deployment_with_excluded_keys(): assert sanitized_config["litellm_params"]["access_token"] != "token-12345" assert "*" in sanitized_config["litellm_params"]["access_token"] - # With excluded_keys, litellm_credentials_name should NOT be masked (even if it would match patterns) + # With excluded_keys, litellm_credentials_name should NOT be masked. + # ``remove_sensitive_info_from_deployment`` mutates its input, so feed it + # a fresh copy rather than the already-sanitized one. sanitized_config = remove_sensitive_info_from_deployment( copy.deepcopy(base_config), excluded_keys={"litellm_credentials_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_team_callback_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py new file mode 100644 index 0000000000..6c4d1be783 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_callback_endpoints.py @@ -0,0 +1,294 @@ +""" +Audit-log emission for the team-callback admin endpoints. + +The endpoints in ``team_callback_endpoints.py`` mutate a team's logging +callbacks (``add_team_callbacks``) or zero them out entirely +(``disable_team_logging``). Both are admin-only mutations, and the +disable variant is itself a logging-control action, so when the operator +has Enterprise audit logging enabled (``litellm.store_audit_logs = True``) +each call must emit a row that captures who did it and what the metadata +looked like before/after. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import Request + +import litellm +from litellm.proxy._types import ( + AddTeamCallback, + LitellmTableNames, + UserAPIKeyAuth, +) +from litellm.proxy.management_endpoints.team_callback_endpoints import ( + add_team_callbacks, + disable_team_logging, +) + + +def _admin_auth() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="hashed", + user_id="admin-user", + user_role="proxy_admin", + ) + + +def _existing_team_row(metadata: dict) -> MagicMock: + row = MagicMock() + row.team_id = "team-1" + row.metadata = metadata + return row + + +def _patch_prisma(existing_metadata: dict): + """Build a context-manager that patches the proxy's ``prisma_client`` + to return ``existing_metadata`` from ``get_data`` and a stub team row + from ``litellm_teamtable.update``.""" + mock_prisma = MagicMock() + mock_prisma.get_data = AsyncMock(return_value=_existing_team_row(existing_metadata)) + + updated_row = MagicMock() + updated_row.team_id = "team-1" + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_row) + return mock_prisma + + +@pytest.mark.asyncio +async def test_disable_team_logging_emits_audit_log_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", True) + mock_prisma = _patch_prisma( + { + "callback_settings": { + "success_callback": ["langfuse"], + "failure_callback": [], + } + } + ) + + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + new=capture, + ), + ): + await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + # asyncio.create_task fires the coroutine eagerly; await one tick to let + # the audit-log emit run before the test exits. + import asyncio + + for _ in range(3): + await asyncio.sleep(0) + + assert len(audit_calls) == 1 + log = audit_calls[0] + assert log.table_name == LitellmTableNames.TEAM_TABLE_NAME + assert log.object_id == "team-1" + assert log.action == "updated" + assert log.changed_by == "admin-user" + + before = json.loads(log.before_value) + after = json.loads(log.updated_values) + # Before: the team's pre-existing success_callback survives in the snapshot. + assert before["metadata"]["callback_settings"]["success_callback"] == ["langfuse"] + # After: callbacks zeroed out by the endpoint. + assert after["metadata"]["callback_settings"]["success_callback"] == [] + assert after["metadata"]["callback_settings"]["failure_callback"] == [] + + +@pytest.mark.asyncio +async def test_disable_team_logging_no_audit_when_disabled(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", False) + mock_prisma = _patch_prisma( + { + "callback_settings": { + "success_callback": ["langfuse"], + "failure_callback": [], + } + } + ) + + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + new=capture, + ), + ): + await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + assert audit_calls == [] + + +@pytest.mark.asyncio +async def test_add_team_callbacks_emits_audit_log_when_enabled(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", True) + mock_prisma = _patch_prisma({"logging": []}) + + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + new=capture, + ), + ): + await add_team_callbacks( + data=AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + }, + ), + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by="ops-on-call", + ) + import asyncio + + for _ in range(3): + await asyncio.sleep(0) + + assert len(audit_calls) == 1 + log = audit_calls[0] + assert log.table_name == LitellmTableNames.TEAM_TABLE_NAME + assert log.object_id == "team-1" + assert log.action == "updated" + # ``litellm_changed_by`` header takes precedence over the auth user_id. + assert log.changed_by == "ops-on-call" + + before = json.loads(log.before_value) + after = json.loads(log.updated_values) + assert before["metadata"]["logging"] == [] + assert len(after["metadata"]["logging"]) == 1 + assert after["metadata"]["logging"][0]["callback_name"] == "langfuse" + + # Callback secrets MUST NOT leak into the audit log payload. + callback_vars = after["metadata"]["logging"][0]["callback_vars"] + assert callback_vars["langfuse_public_key"] != "pk" + assert callback_vars["langfuse_secret_key"] != "sk" + # Key names are preserved so the auditor can see which fields changed. + assert "langfuse_public_key" in callback_vars + assert "langfuse_secret_key" in callback_vars + # And no plaintext secret should appear anywhere in the serialized row. + assert "sk" not in log.updated_values.replace("sk-", "") # crude leak check + assert "pk" not in (log.updated_values.replace("pk-", "").replace("public_key", "")) + + +@pytest.mark.asyncio +async def test_disable_team_logging_redacts_existing_callback_secrets(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", True) + # Existing team has populated callback_vars containing secrets — redaction + # must apply to the BEFORE snapshot too. + mock_prisma = _patch_prisma( + { + "callback_settings": { + "success_callback": ["langfuse"], + "failure_callback": [], + "callback_vars": { + "langfuse_public_key": "pk-real", + "langfuse_secret_key": "sk-real-secret", + }, + } + } + ) + + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + new=capture, + ), + ): + await disable_team_logging( + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + import asyncio + + for _ in range(3): + await asyncio.sleep(0) + + assert len(audit_calls) == 1 + log = audit_calls[0] + # The pre-existing secret_key value must NOT appear in the serialized + # before_value or updated_values. + assert "sk-real-secret" not in log.before_value + assert "sk-real-secret" not in log.updated_values + assert "pk-real" not in log.before_value + assert "pk-real" not in log.updated_values + + +@pytest.mark.asyncio +async def test_add_team_callbacks_no_audit_when_disabled(monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", False) + mock_prisma = _patch_prisma({"logging": []}) + + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + new=capture, + ), + ): + await add_team_callbacks( + data=AddTeamCallback( + callback_name="langfuse", + callback_type="success", + callback_vars={ + "langfuse_public_key": "pk", + "langfuse_secret_key": "sk", + }, + ), + http_request=MagicMock(spec=Request), + team_id="team-1", + user_api_key_dict=_admin_auth(), + litellm_changed_by=None, + ) + + assert audit_calls == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index e4754e6d38..e668672dd2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -968,6 +968,20 @@ def test_add_new_models_to_team(): ) +def _make_team_member_add_request( + member_user_id: Optional[str] = "regular-user", + role: str = "user", + team_id: str = "test-team-123", +): + """Build a TeamMemberAddRequest with one Member entry for tests below.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + + return TeamMemberAddRequest( + team_id=team_id, + member=Member(role=role, user_id=member_user_id), + ) + + @pytest.mark.asyncio async def test_validate_team_member_add_permissions_admin(): """ @@ -977,17 +991,15 @@ async def test_validate_team_member_add_permissions_admin(): _validate_team_member_add_permissions, ) - # Create admin user admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) - # Create mock team team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = "test-team-123" - # Should not raise any exception for admin await _validate_team_member_add_permissions( user_api_key_dict=admin_user, complete_team_data=team, + data=_make_team_member_add_request(member_user_id="any-user", role="admin"), ) @@ -1000,20 +1012,17 @@ async def test_validate_team_member_add_permissions_non_admin(): _validate_team_member_add_permissions, ) - # Create non-admin user regular_user = UserAPIKeyAuth( user_id="regular-user", user_role=LitellmUserRoles.INTERNAL_USER, team_id="different-team", ) - # Create mock team team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = "test-team-123" team.members_with_roles = [] team.organization_id = None - # Mock the helper functions to return False with ( patch( "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", @@ -1024,17 +1033,303 @@ async def test_validate_team_member_add_permissions_non_admin(): return_value=False, ), ): - # Should raise HTTPException for non-admin with pytest.raises(HTTPException) as exc_info: await _validate_team_member_add_permissions( user_api_key_dict=regular_user, complete_team_data=team, + data=_make_team_member_add_request(), ) assert exc_info.value.status_code == 403 assert "not proxy admin OR team admin" in str(exc_info.value.detail) +# ── VERIA-56 regression tests for _is_available_team self-join enforcement ─── + + +@pytest.mark.asyncio +async def test_available_team_self_join_with_caller_user_id_allowed(): + """A standard user adding themselves to an available team with role=user + is the only legitimate use of the available-team bypass.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=_make_team_member_add_request(member_user_id="alice", role="user"), + ) + + +@pytest.mark.asyncio +async def test_available_team_self_join_blocks_admin_role(): + """Privesc shape from VERIA-56: caller adds themselves with role=admin + via the available-team bypass. Must be rejected.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=_make_team_member_add_request(member_user_id="alice", role="admin"), + ) + + assert exc_info.value.status_code == 403 + assert "admin" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_available_team_self_join_blocks_other_user_id(): + """Cross-user-injection shape from VERIA-56: caller adds someone else + via the available-team bypass. Must be rejected.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=_make_team_member_add_request( + member_user_id="bob-victim", role="user" + ), + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_available_team_self_join_blocks_when_caller_has_no_user_id(): + """If the auth context has no user_id we cannot prove self-join, so the + bypass must fail closed.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) # no user_id + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=_make_team_member_add_request(member_user_id="alice", role="user"), + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_available_team_self_join_blocks_email_only_member(): + """An email-only member entry can't be safely self-join-validated; the + caller must use their own user_id explicitly.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_email="alice@example.com"), + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_available_team_self_join_blocks_admin_role_in_member_list(): + """Bulk shape: list of members where one has role=admin must be rejected + even if the caller's own entry is correct.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=[ + Member(role="user", user_id="alice"), + Member(role="admin", user_id="alice"), + ], + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_update_team_member_permissions_blocks_non_admin_via_available_team( + mock_db_client, +): + """A non-admin caller invoking /team/permissions_update on an available + team must be rejected. The previous code path delegated to + ``_is_available_team`` and accepted the write; this PR removes that + bypass entirely so the result is 403 even with the bypass mocked True.""" + test_team_id = "public-team" + update_payload = { + "team_id": test_team_id, + "team_member_permissions": ["/key/generate"], + } + + existing_row = MagicMock(spec=LiteLLM_TeamTable) + existing_row.model_dump.return_value = { + "team_id": test_team_id, + "team_alias": "Public Team", + "team_member_permissions": [], + "spend": 0.0, + "models": [], + } + existing_row.team_id = test_team_id + existing_row.members_with_roles = [] + existing_row.organization_id = None + + non_admin_auth = UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=existing_row, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + # Even with the available-team bypass mocked True, the endpoint + # must NOT consult it any more — the gate should reject the + # non-admin caller outright. + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + ): + app.dependency_overrides[user_api_key_auth] = lambda: non_admin_auth + try: + response = client.post("/team/permissions_update", json=update_payload) + finally: + app.dependency_overrides = {} + + assert response.status_code == 403 + body = response.json() + assert "permissions_update" in str(body) or "not proxy admin" in str(body) + + @pytest.mark.asyncio async def test_process_team_members_single_member(): """ 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/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index 1b157a2f6b..b36383dfd9 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -16,6 +16,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _resolve_team_allowed_mcp_servers, _set_object_permission, validate_key_mcp_servers_against_team, + validate_key_search_tools_against_team, ) @@ -453,3 +454,52 @@ async def test_resolve_team_allowed_mcp_servers_dict_tool_permissions( result = await _resolve_team_allowed_mcp_servers(mock_perm) assert result == {"server-a"} + + +# ---- Tests for validate_key_search_tools_against_team ---- + + +def _make_team_obj_search(team_id="team-1", search_tools=None): + mock_team = MagicMock() + mock_team.team_id = team_id + if search_tools is not None: + mock_team.object_permission = MagicMock(spec=LiteLLM_ObjectPermissionTable) + mock_team.object_permission.search_tools = search_tools + else: + mock_team.object_permission = None + return mock_team + + +@pytest.mark.asyncio +async def test_validate_search_tools_no_key_request(): + await validate_key_search_tools_against_team( + object_permission=None, + team_obj=_make_team_obj_search(search_tools=["t1"]), + ) + + +@pytest.mark.asyncio +async def test_validate_search_tools_team_unrestricted(): + """Empty team search allowlist means unrestricted — key subset check skipped.""" + await validate_key_search_tools_against_team( + object_permission={"search_tools": ["any-tool"]}, + team_obj=_make_team_obj_search(search_tools=[]), + ) + + +@pytest.mark.asyncio +async def test_validate_search_tools_subset_ok(): + await validate_key_search_tools_against_team( + object_permission={"search_tools": ["t1"]}, + team_obj=_make_team_obj_search(search_tools=["t1", "t2"]), + ) + + +@pytest.mark.asyncio +async def test_validate_search_tools_raises_when_not_subset(): + with pytest.raises(HTTPException) as exc: + await validate_key_search_tools_against_team( + object_permission={"search_tools": ["bad"]}, + team_obj=_make_team_obj_search(search_tools=["t1"]), + ) + assert exc.value.status_code == 403 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 3def76b825..c16c42decc 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -573,6 +573,57 @@ class TestAnthropicBatchPassthroughCostTracking: or "claude-sonnet-4-5-20250929" in decoded ) + @pytest.mark.parametrize( + "kwargs,expected_user_id,expected_team_id", + [ + ( + { + "litellm_params": { + "metadata": { + "user_api_key_user_id": "real-user-123", + "user_api_key_team_id": "team-456", + } + } + }, + "real-user-123", + "team-456", + ), + ({}, "default-user", None), + ], + ) + def test_store_batch_managed_object_propagates_user_identity_from_metadata( + self, + mock_logging_obj, + kwargs, + expected_user_id, + expected_team_id, + ): + """The fabricated UserAPIKeyAuth must inherit user_id/team_id from the + request's litellm_params.metadata, not the (always-empty) top-level + kwargs lookup. Falls back to "default-user" only when metadata is + absent.""" + mock_managed_files_hook = MagicMock() + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler.verbose_proxy_logger" + ), + ): + mock_pl.get_proxy_hook.return_value = mock_managed_files_hook + + AnthropicPassthroughLoggingHandler._store_batch_managed_object( + unified_object_id="uoi", + batch_object={"id": "b1", "object": "batch", "status": "validating"}, + model_object_id="b1", + logging_obj=mock_logging_obj, + **kwargs, + ) + + mock_managed_files_hook.store_unified_object_id.assert_called_once() + call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1] + assert call_kwargs["user_api_key_dict"].user_id == expected_user_id + assert call_kwargs["user_api_key_dict"].team_id == expected_team_id + def test_batch_creation_handler_failure_status_code( self, mock_logging_obj, mock_request_body ): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py new file mode 100644 index 0000000000..f73aee77cc --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -0,0 +1,120 @@ +"""Regression tests for LIT-2642 — interrupted pass-through streams must still log usage.""" + +import asyncio +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, +) +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + +def _make_streaming_response(chunks): + mock = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes(): + for c in chunks: + yield c + + mock.aiter_bytes = _aiter_bytes + return mock + + +@pytest.mark.asyncio +async def test_chunk_processor_logs_on_normal_completion(): + chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] + response = _make_streaming_response(chunks) + + mock_logging_obj = MagicMock() + mock_passthrough_handler = MagicMock() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=mock_logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=mock_passthrough_handler, + url_route="/bedrock/model/claude/invoke-with-response-stream", + ): + received.append(chunk) + + await asyncio.sleep(0) + + assert received == chunks + mock_route.assert_called_once() + call_kwargs = mock_route.call_args.kwargs + assert call_kwargs["raw_bytes"] == chunks + + +@pytest.mark.asyncio +async def test_chunk_processor_logs_on_client_disconnect(): + chunks = [b"event-1", b"event-2", b"event-3"] + response = _make_streaming_response(chunks) + + mock_logging_obj = MagicMock() + mock_passthrough_handler = MagicMock() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=mock_logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=mock_passthrough_handler, + url_route="/bedrock/model/claude/invoke-with-response-stream", + ) + + first = await gen.__anext__() + await gen.aclose() + + await asyncio.sleep(0) + + assert first == chunks[0] + mock_route.assert_called_once() + call_kwargs = mock_route.call_args.kwargs + assert call_kwargs["raw_bytes"] == [chunks[0]] + + +@pytest.mark.asyncio +async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): + response = _make_streaming_response([]) + + mock_logging_obj = MagicMock() + mock_passthrough_handler = MagicMock() + + with patch.object( + PassThroughStreamingHandler, + "_route_streaming_logging_to_handler", + new=AsyncMock(), + ) as mock_route: + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=mock_logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=mock_passthrough_handler, + url_route="/bedrock/model/claude/invoke-with-response-stream", + ): + received.append(chunk) + + assert received == [] + mock_route.assert_not_called() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 756e5fa5bc..efa26a61bf 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -264,6 +264,57 @@ class TestVertexAIBatchPassthroughHandler: # Verify the managed files hook was called mock_managed_files_hook.store_unified_object_id.assert_called_once() + @pytest.mark.parametrize( + "kwargs,expected_user_id,expected_team_id", + [ + ( + { + "litellm_params": { + "metadata": { + "user_api_key_user_id": "real-user-123", + "user_api_key_team_id": "team-456", + } + } + }, + "real-user-123", + "team-456", + ), + ({}, "default-user", None), + ], + ) + def test_store_batch_managed_object_propagates_user_identity_from_metadata( + self, + mock_logging_obj, + mock_managed_files_hook, + kwargs, + expected_user_id, + expected_team_id, + ): + """The fabricated UserAPIKeyAuth must inherit user_id/team_id from the + request's litellm_params.metadata, not the (always-empty) top-level + kwargs lookup. Falls back to "default-user" only when metadata is + absent.""" + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger" + ), + ): + mock_pl.get_proxy_hook.return_value = mock_managed_files_hook + + VertexPassthroughLoggingHandler._store_batch_managed_object( + unified_object_id="uoi", + batch_object={"id": "b1", "object": "batch", "status": "validating"}, + model_object_id="b1", + logging_obj=mock_logging_obj, + **kwargs, + ) + + mock_managed_files_hook.store_unified_object_id.assert_called_once() + call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1] + assert call_kwargs["user_api_key_dict"].user_id == expected_user_id + assert call_kwargs["user_api_key_dict"].team_id == expected_team_id + def test_batch_cost_calculation_integration(self): """Single Vertex AI response → non-zero cost with correct token counts.""" from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index e532b948c7..185d337f90 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1513,9 +1513,13 @@ class TestIsMasterKey: def test_non_matching_key_returns_false(self): assert _is_master_key(api_key="sk-other", _master_key="sk-master") is False - def test_hashed_key_returns_true(self): + def test_master_key_hash_is_rejected(self): + """ + ``_is_master_key`` must not accept ``hash_token(master_key)`` as + equivalent to the raw master key — only the raw value matches. + """ from litellm.proxy.utils import hash_token master = "sk-master-key-123" hashed = hash_token(master) - assert _is_master_key(api_key=hashed, _master_key=master) is True + assert _is_master_key(api_key=hashed, _master_key=master) is False diff --git a/tests/test_litellm/proxy/test_batch_retrieve_bedrock.py b/tests/test_litellm/proxy/test_batch_retrieve_bedrock.py new file mode 100644 index 0000000000..1394575009 --- /dev/null +++ b/tests/test_litellm/proxy/test_batch_retrieve_bedrock.py @@ -0,0 +1,220 @@ +""" +Tests for the proxy /v1/batches/{batch_id} retrieve flow and the +/v1/files/{file_id}/content download flow with model-encoded IDs (Bedrock). + +Regression (retrieve): when the proxy decoded `model` from the encoded +batch_id, it did not forward `model` as a kwarg to `litellm.aretrieve_batch`. +That caused litellm to skip the `BedrockBatchesConfig` provider_config path +and fall into the legacy provider switch, which raises BadRequestError for +bedrock. + +The download path is included to lock in the end-to-end Bedrock batch flow: +retrieve returns an `output_file_id` re-encoded with model info, and that ID +must round-trip through `client.files.content(...)` back to bedrock with AWS +credentials and the raw S3 URI intact. +""" + +import os +import sys + +import httpx +import pytest +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, +) +from litellm.proxy.proxy_server import app +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.utils import LiteLLMBatch + +client = TestClient(app) + +BEDROCK_MODEL = "bedrock-claude-test" +BEDROCK_BATCH_ARN = ( + "arn:aws:bedrock:us-east-1:000000000000:model-invocation-job/test-job-id" +) +BEDROCK_OUTPUT_S3_URI = ( + "s3://test-bedrock-batch-output/job-output/test-job-id/output.jsonl.out" +) + + +@pytest.fixture +def bedrock_router() -> Router: + return Router( + model_list=[ + { + "model_name": BEDROCK_MODEL, + "litellm_params": { + "model": f"bedrock/{BEDROCK_MODEL}", + "aws_region_name": "us-east-1", + "aws_access_key_id": "test-access-key", + "aws_secret_access_key": "test-secret-key", + }, + "model_info": {"id": "bedrock-claude-test-id"}, + }, + ] + ) + + +def _setup_proxy(monkeypatch, llm_router: Router): + proxy_logging_obj = ProxyLogging( + user_api_key_cache=DualCache(default_in_memory_ttl=1) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + +def _encoded_bedrock_batch_id() -> str: + return encode_file_id_with_model( + file_id=BEDROCK_BATCH_ARN, model=BEDROCK_MODEL, id_type="batch" + ) + + +def _make_in_progress_batch_response(batch_id: str) -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-input", + object="batch", + status="in_progress", + ) + + +def test_retrieve_batch_passes_model_for_bedrock_encoded_id( + monkeypatch, bedrock_router +): + """Encoded batch_id → proxy must pass `model` to litellm.aretrieve_batch + so BedrockBatchesConfig is loaded. + + Without this, litellm falls into the legacy provider switch and raises + 'LiteLLM doesn't support bedrock for retrieve_batch'. + """ + _setup_proxy(monkeypatch, bedrock_router) + + user_key = UserAPIKeyAuth(api_key="test-key") + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + encoded_batch_id = _encoded_bedrock_batch_id() + captured_kwargs: dict = {} + + async def mock_aretrieve_batch(**kwargs): + captured_kwargs.update(kwargs) + return _make_in_progress_batch_response(BEDROCK_BATCH_ARN) + + monkeypatch.setattr(litellm, "aretrieve_batch", mock_aretrieve_batch) + + try: + response = client.get( + f"/v1/batches/{encoded_batch_id}", + headers={"Authorization": "Bearer test-key"}, + ) + assert response.status_code == 200, response.text + finally: + app.dependency_overrides.clear() + + assert captured_kwargs.get("custom_llm_provider") == "bedrock" + assert captured_kwargs.get("model") == BEDROCK_MODEL, ( + "model must be forwarded to litellm.aretrieve_batch so the bedrock " + "provider_config is loaded; got kwargs: " + repr(captured_kwargs) + ) + assert captured_kwargs.get("batch_id") == BEDROCK_BATCH_ARN + + +def test_retrieve_batch_response_id_is_re_encoded_with_model( + monkeypatch, bedrock_router +): + """After provider returns the raw ARN, the proxy must re-encode the + response id with the model so subsequent client calls keep routing to + bedrock.""" + _setup_proxy(monkeypatch, bedrock_router) + + user_key = UserAPIKeyAuth(api_key="test-key") + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + encoded_batch_id = _encoded_bedrock_batch_id() + + async def mock_aretrieve_batch(**kwargs): + return _make_in_progress_batch_response(BEDROCK_BATCH_ARN) + + monkeypatch.setattr(litellm, "aretrieve_batch", mock_aretrieve_batch) + + try: + response = client.get( + f"/v1/batches/{encoded_batch_id}", + headers={"Authorization": "Bearer test-key"}, + ) + assert response.status_code == 200, response.text + body = response.json() + finally: + app.dependency_overrides.clear() + + assert body["id"] == encoded_batch_id + + +def test_file_content_routes_to_bedrock_for_encoded_output_file_id( + monkeypatch, bedrock_router +): + """`client.files.content(output_file_id)` for a bedrock-encoded file ID + must reach `litellm.afile_content` with `custom_llm_provider="bedrock"`, + the raw S3 URI as `file_id`, and AWS credentials sourced from the router. + + This is the second half of the bedrock batch flow (the first being + retrieve). Without this round-trip, callers have to bypass the proxy and + call `litellm.file_content(...)` directly with hand-rolled AWS args. + """ + _setup_proxy(monkeypatch, bedrock_router) + + user_key = UserAPIKeyAuth(api_key="test-key") + app.dependency_overrides[user_api_key_auth] = lambda: user_key + + encoded_file_id = encode_file_id_with_model( + file_id=BEDROCK_OUTPUT_S3_URI, model=BEDROCK_MODEL, id_type="file" + ) + captured_kwargs: dict = {} + file_bytes = b'{"custom_id":"r1","response":{"body":{"choices":[]}}}\n' + + async def mock_afile_content(**kwargs): + captured_kwargs.update(kwargs) + return HttpxBinaryResponseContent( + response=httpx.Response( + status_code=200, + content=file_bytes, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url=BEDROCK_OUTPUT_S3_URI), + ) + ) + + monkeypatch.setattr(litellm, "afile_content", mock_afile_content) + + try: + response = client.get( + f"/v1/files/{encoded_file_id}/content", + headers={"Authorization": "Bearer test-key"}, + ) + assert response.status_code == 200, response.text + assert response.content == file_bytes + finally: + app.dependency_overrides.clear() + + assert captured_kwargs.get("custom_llm_provider") == "bedrock" + assert captured_kwargs.get("file_id") == BEDROCK_OUTPUT_S3_URI, ( + "file_id must be decoded back to the raw S3 URI before reaching " + "litellm.afile_content; got kwargs: " + repr(captured_kwargs) + ) + assert captured_kwargs.get("aws_region_name") == "us-east-1" + assert captured_kwargs.get("aws_access_key_id") == "test-access-key" + assert captured_kwargs.get("aws_secret_access_key") == "test-secret-key" 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/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 44cc5cc445..3d369ed224 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -1882,3 +1882,274 @@ async def test_create_vector_store_in_db_raises_when_no_db(): assert exc_info.value.status_code == 500 assert "database not connected" in exc_info.value.detail.lower() + + +class TestRedactSensitiveLitellmParams: + """ + ``litellm_params`` on a managed vector store carries the upstream + provider credential (OpenAI ``api_key``, AWS ``aws_secret_access_key``, + GCP ``vertex_credentials``, etc.). The list/info endpoints must redact + those values before returning them to any caller — including read-only + users and narrowly-scoped keys. + """ + + def test_redacts_well_known_credential_keys(self): + from litellm.constants import REDACTED_BY_LITELM_STRING + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + params = { + "api_key": "sk-real-openai-key-12345", + "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "vertex_credentials": ( + '{"type":"service_account","private_key":"-----BEGIN PRIVATE KEY-----..."}' + ), + "azure_authorization_token": "Bearer eyJhbGciOi...", + } + out = _redact_sensitive_litellm_params(params) + for k in params: + assert ( + out[k] == REDACTED_BY_LITELM_STRING + ), f"{k} should be redacted, got {out[k]!r}" + + def test_preserves_non_sensitive_keys(self): + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + params = { + "api_base": "https://api.openai.com/v1", + "model": "text-embedding-3-large", + "region": "us-east-1", + "vector_store_id": "vs_abc123", + "api_version": "2023-05-15", + } + out = _redact_sensitive_litellm_params(params) + for k, v in params.items(): + assert out[k] == v, f"{k} should be preserved verbatim" + + def test_handles_none_and_empty(self): + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + assert _redact_sensitive_litellm_params(None) is None + assert _redact_sensitive_litellm_params({}) == {} + + def test_redaction_does_not_mutate_input_litellm_params(self): + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + original = { + "api_key": "sk-real-openai-key-12345", + "api_base": "https://api.openai.com/v1", + } + snapshot = dict(original) + _redact_sensitive_litellm_params(original) + assert original == snapshot, "input dict must not be mutated" + + def test_redacts_nested_credentials_in_embedding_config(self): + """ + ``/vector_store/new`` and ``/vector_store/update`` auto-resolve + ``litellm_embedding_config`` from the model registry and store it + as a nested dict inside ``litellm_params``. The nested dict carries + its own ``api_key`` / ``aws_*`` / ``vertex_credentials``, and a + non-recursive redactor would leak them. + """ + from litellm.constants import REDACTED_BY_LITELM_STRING + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + params = { + "model": "openai/text-embedding-3-large", + "api_base": "https://api.openai.com/v1", + "litellm_embedding_config": { + "api_key": "sk-nested-secret", + "api_base": "https://nested.example.com", + "vertex_credentials": '{"private_key":"-----BEGIN..."}', + }, + } + out = _redact_sensitive_litellm_params(params) + nested = out["litellm_embedding_config"] + assert nested["api_key"] == REDACTED_BY_LITELM_STRING + assert nested["vertex_credentials"] == REDACTED_BY_LITELM_STRING + assert nested["api_base"] == "https://nested.example.com" + # Top-level non-secrets preserved + assert out["api_base"] == "https://api.openai.com/v1" + assert out["model"] == "openai/text-embedding-3-large" + + def test_redacts_json_string_litellm_params(self): + """ + The in-memory registry occasionally holds ``litellm_params`` as a + JSON-serialized string rather than a dict. The redactor must parse, + redact, and re-serialize so callers don't get the raw string back. + """ + import json as _json + + from litellm.constants import REDACTED_BY_LITELM_STRING + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + params_json = _json.dumps( + { + "api_key": "sk-secret-from-json-string", + "api_base": "https://api.openai.com/v1", + } + ) + out = _redact_sensitive_litellm_params(params_json) + assert isinstance(out, str) + parsed = _json.loads(out) + assert parsed["api_key"] == REDACTED_BY_LITELM_STRING + assert parsed["api_base"] == "https://api.openai.com/v1" + + def test_redacts_unparseable_string_litellm_params(self): + """ + If ``litellm_params`` is a string that isn't valid JSON, the + redactor must NOT echo the value back verbatim — it could contain + opaque credential material. + """ + from litellm.constants import REDACTED_BY_LITELM_STRING + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + _redact_sensitive_litellm_params, + ) + + out = _redact_sensitive_litellm_params( + "this is not json but might contain a secret" + ) + assert out == REDACTED_BY_LITELM_STRING + + +class TestUpdateVectorStoreAccessControlAndRedaction: + """ + ``/vector_store/update`` previously skipped per-store access control + (only the premium-feature gate ran), letting any authenticated + premium principal mutate *any* vector store. It also returned the + full DB row including ``litellm_params``, leaking provider + credentials to the caller. Both are fixed at the endpoint level. + """ + + @pytest.mark.asyncio + async def test_update_denied_when_caller_cannot_access_store(self): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + update_vector_store, + ) + from litellm.types.vector_stores import VectorStoreUpdateRequest + + existing_row = MagicMock() + existing_row.model_dump = MagicMock( + return_value={ + "vector_store_id": "vs_other_team", + "team_id": "team-A", + "litellm_params": {"api_key": "sk-team-A-secret"}, + } + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=existing_row + ) + + with ( + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access", + new_callable=AsyncMock, + return_value=False, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + ): + with pytest.raises(HTTPException) as exc_info: + await update_vector_store( + data=VectorStoreUpdateRequest( + vector_store_id="vs_other_team", + vector_store_description="hijacked", + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="attacker", team_id="team-B" + ), + ) + assert exc_info.value.status_code == 403 + # The attacker must NOT see the existing credential in the + # error message either. + assert "sk-team-A-secret" not in str(exc_info.value.detail) + # And the DB update must not have been called. + mock_prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called() + + @pytest.mark.asyncio + async def test_update_response_redacts_litellm_params(self): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.constants import REDACTED_BY_LITELM_STRING + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + update_vector_store, + ) + from litellm.types.vector_stores import VectorStoreUpdateRequest + + existing_row = MagicMock() + existing_row.model_dump = MagicMock( + return_value={ + "vector_store_id": "vs_owned", + "team_id": "team-A", + "litellm_params": { + "api_key": "sk-real-openai-key-123", + "api_base": "https://api.openai.com/v1", + }, + } + ) + updated_row = MagicMock() + updated_row.model_dump = MagicMock( + return_value={ + "vector_store_id": "vs_owned", + "team_id": "team-A", + "vector_store_description": "new desc", + "litellm_params": { + "api_key": "sk-real-openai-key-123", + "api_base": "https://api.openai.com/v1", + }, + } + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( + return_value=existing_row + ) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock( + return_value=updated_row + ) + + with ( + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.vector_store_endpoints.management_endpoints._check_vector_store_access", + new_callable=AsyncMock, + return_value=True, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", None), + ): + response = await update_vector_store( + data=VectorStoreUpdateRequest( + vector_store_id="vs_owned", + vector_store_description="new desc", + ), + user_api_key_dict=UserAPIKeyAuth(user_id="owner", team_id="team-A"), + ) + + params = response["vector_store"]["litellm_params"] + assert params["api_key"] == REDACTED_BY_LITELM_STRING + assert params["api_base"] == "https://api.openai.com/v1" diff --git a/tests/test_litellm/test_bedrock_anthropic_1hr_cache_pricing.py b/tests/test_litellm/test_bedrock_anthropic_1hr_cache_pricing.py new file mode 100644 index 0000000000..69af35dfea --- /dev/null +++ b/tests/test_litellm/test_bedrock_anthropic_1hr_cache_pricing.py @@ -0,0 +1,123 @@ +""" +Validate that Bedrock-hosted Anthropic Claude 4.5/4.6/4.7 entries carry the +1-hour prompt-cache write tier (`cache_creation_input_token_cost_above_1hr`) +in `model_prices_and_context_window.json`. + +AWS Bedrock pricing (https://aws.amazon.com/bedrock/pricing/) publishes a +separate 1-hour cache write column for the Claude 4.5 / 4.6 / 4.7 family. +Without these fields, cost tracking on Bedrock 1-hour-TTL prompt caching +falls back to the 5-minute write rate and undercounts spend by ~60%. + +Source values (per million tokens) for the 1-hour cache write column, +as published on the AWS Bedrock pricing page: + + Global pricing: + Opus 4.7 / Opus 4.6 / Opus 4.5 -> $10.00 + Sonnet 4.6 / Sonnet 4.5 (regular tier) -> $6.00 + Sonnet 4.5 long-context (>200K tier) -> $12.00 + Haiku 4.5 -> $2.00 + + US pricing (10% premium over Global): + Opus 4.7 / Opus 4.6 / Opus 4.5 -> $11.00 + Sonnet 4.6 / Sonnet 4.5 (regular tier) -> $6.60 + Sonnet 4.5 long-context (>200K tier) -> $13.20 + Haiku 4.5 -> $2.20 +""" + +import json +import os + +import pytest + + +@pytest.fixture(scope="module") +def model_data(): + json_path = os.path.join( + os.path.dirname(__file__), "../../model_prices_and_context_window.json" + ) + with open(json_path) as f: + return json.load(f) + + +# (model_key, expected 1hr cache write per token, expected 1hr LC tier or None) +GLOBAL_EXPECTED = [ + # Opus 4.7 - $10.00 / MTok + ("anthropic.claude-opus-4-7", 1e-05, None), + ("global.anthropic.claude-opus-4-7", 1e-05, None), + # Opus 4.6 - $10.00 / MTok + ("anthropic.claude-opus-4-6-v1", 1e-05, None), + ("global.anthropic.claude-opus-4-6-v1", 1e-05, None), + # Opus 4.5 - $10.00 / MTok + ("anthropic.claude-opus-4-5-20251101-v1:0", 1e-05, None), + ("global.anthropic.claude-opus-4-5-20251101-v1:0", 1e-05, None), + # Sonnet 4.6 - $6.00 / MTok (no separate LC tier per AWS) + ("anthropic.claude-sonnet-4-6", 6e-06, None), + ("global.anthropic.claude-sonnet-4-6", 6e-06, None), + # Sonnet 4.5 - $6.00 / MTok regular, $12.00 / MTok long-context (>200K) + ("anthropic.claude-sonnet-4-5-20250929-v1:0", 6e-06, 1.2e-05), + ("global.anthropic.claude-sonnet-4-5-20250929-v1:0", 6e-06, 1.2e-05), + # Haiku 4.5 - $2.00 / MTok + ("anthropic.claude-haiku-4-5-20251001-v1:0", 2e-06, None), + ("anthropic.claude-haiku-4-5@20251001", 2e-06, None), + ("global.anthropic.claude-haiku-4-5-20251001-v1:0", 2e-06, None), +] + +US_EXPECTED = [ + # US is +10% over Global. + ("us.anthropic.claude-opus-4-7", 1.1e-05, None), + ("us.anthropic.claude-opus-4-6-v1", 1.1e-05, None), + ("us.anthropic.claude-opus-4-5-20251101-v1:0", 1.1e-05, None), + ("us.anthropic.claude-sonnet-4-6", 6.6e-06, None), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", 6.6e-06, 1.32e-05), + ("us.anthropic.claude-haiku-4-5-20251001-v1:0", 2.2e-06, None), +] + + +@pytest.mark.parametrize( + "model_key, expected_1hr, expected_1hr_lc", GLOBAL_EXPECTED + US_EXPECTED +) +def test_bedrock_anthropic_1hr_cache_write_pricing( + model_data, model_key, expected_1hr, expected_1hr_lc +): + assert model_key in model_data, f"Missing model entry: {model_key}" + info = model_data[model_key] + + # 1hr cache write rate must be present and exact. + assert "cache_creation_input_token_cost_above_1hr" in info, ( + f"{model_key}: missing cache_creation_input_token_cost_above_1hr - " + "AWS Bedrock charges a separate 1-hour cache write rate for this model" + ) + assert info["cache_creation_input_token_cost_above_1hr"] == expected_1hr, ( + f"{model_key}: 1hr cache write rate " + f"{info['cache_creation_input_token_cost_above_1hr']} does not match " + f"expected {expected_1hr} from AWS Bedrock pricing" + ) + + # 1hr cache write rate must be 1.6x the 5-minute rate (AWS standard ratio). + five_min = info["cache_creation_input_token_cost"] + ratio = info["cache_creation_input_token_cost_above_1hr"] / five_min + assert ( + abs(ratio - 1.6) < 1e-9 + ), f"{model_key}: 1hr/5min ratio is {ratio}, expected 1.6" + + # Long-context (>200K) tier, where AWS publishes one. + if expected_1hr_lc is not None: + assert ( + "cache_creation_input_token_cost_above_1hr_above_200k_tokens" in info + ), f"{model_key}: missing 1hr cache write tier for >200K context" + assert ( + info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"] + == expected_1hr_lc + ), ( + f"{model_key}: long-context 1hr cache write rate " + f"{info['cache_creation_input_token_cost_above_1hr_above_200k_tokens']} " + f"does not match expected {expected_1hr_lc}" + ) + five_min_lc = info["cache_creation_input_token_cost_above_200k_tokens"] + ratio_lc = ( + info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"] + / five_min_lc + ) + assert ( + abs(ratio_lc - 1.6) < 1e-9 + ), f"{model_key}: long-context 1hr/5min ratio is {ratio_lc}, expected 1.6" 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 b8a4220c67..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, 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/.npmrc b/ui/litellm-dashboard/.npmrc index 168e81a1c4..7999681cc3 100644 --- a/ui/litellm-dashboard/.npmrc +++ b/ui/litellm-dashboard/.npmrc @@ -2,4 +2,4 @@ # Packages needing lifecycle scripts: npm rebuild ignore-scripts=true # Protects local npm install only — npm ci (used in CI) ignores this -min-release-age=3d +min-release-age=3 diff --git a/ui/litellm-dashboard/next.config.mjs b/ui/litellm-dashboard/next.config.mjs index bdf492de33..cfaeb24dc5 100644 --- a/ui/litellm-dashboard/next.config.mjs +++ b/ui/litellm-dashboard/next.config.mjs @@ -7,6 +7,11 @@ const __dirname = path.dirname(__filename); const nextConfig = { output: "export", + // Required with output: "export" — default image optimizer runs only in server mode. + // See https://nextjs.org/docs/messages/export-image-api + images: { + unoptimized: true, + }, basePath: "", assetPrefix: "/litellm-asset-prefix", turbopack: { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx index 1859970091..1a8c6632a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx @@ -15,10 +15,18 @@ import ModelAliasManager from "@/components/common_components/ModelAliasManager" import React, { useEffect, useState } from "react"; import { useQueryClient } from "@tanstack/react-query"; import NotificationsManager from "@/components/molecules/notifications_manager"; -import { fetchMCPAccessGroups, getGuardrailsList, getPoliciesList, Organization, Team, teamCreateCall } from "@/components/networking"; +import { + fetchMCPAccessGroups, + getGuardrailsList, + getPoliciesList, + Organization, + Team, + teamCreateCall, +} from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { organizationKeys } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; +import SearchToolSelector from "@/components/SearchTools/SearchToolSelector"; interface ModelAliases { [key: string]: string; @@ -212,15 +220,27 @@ const CreateTeamModal = ({ } } - // Transform allowed_vector_store_ids and allowed_mcp_servers_and_groups into object_permission + // Transform integrations into object_permission (vector stores, MCP, agents, search tools) + const hasAgents = + formValues.allowed_agents_and_groups && + ((formValues.allowed_agents_and_groups.agents?.length ?? 0) > 0 || + (formValues.allowed_agents_and_groups.accessGroups?.length ?? 0) > 0); + const hasSearchTools = + Array.isArray(formValues.object_permission_search_tools) && + formValues.object_permission_search_tools.length > 0; + if ( (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) || (formValues.allowed_mcp_servers_and_groups && (formValues.allowed_mcp_servers_and_groups.servers?.length > 0 || formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 || - formValues.allowed_mcp_servers_and_groups.toolPermissions)) + formValues.allowed_mcp_servers_and_groups.toolPermissions)) || + hasAgents || + hasSearchTools ) { - formValues.object_permission = {}; + if (!formValues.object_permission) { + formValues.object_permission = {}; + } if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { formValues.object_permission.vector_stores = formValues.allowed_vector_store_ids; delete formValues.allowed_vector_store_ids; @@ -238,9 +258,6 @@ const CreateTeamModal = ({ // Add tool permissions separately if (formValues.mcp_tool_permissions && Object.keys(formValues.mcp_tool_permissions).length > 0) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } formValues.object_permission.mcp_tool_permissions = formValues.mcp_tool_permissions; delete formValues.mcp_tool_permissions; } @@ -248,9 +265,6 @@ const CreateTeamModal = ({ // Handle agent permissions if (formValues.allowed_agents_and_groups) { const { agents, accessGroups } = formValues.allowed_agents_and_groups; - if (!formValues.object_permission) { - formValues.object_permission = {}; - } if (agents && agents.length > 0) { formValues.object_permission.agents = agents; } @@ -259,6 +273,11 @@ const CreateTeamModal = ({ } delete formValues.allowed_agents_and_groups; } + + if (hasSearchTools) { + formValues.object_permission.search_tools = formValues.object_permission_search_tools; + delete formValues.object_permission_search_tools; + } } // Transform allowed_mcp_access_groups into object_permission @@ -729,6 +748,34 @@ const CreateTeamModal = ({ + + + Search Tool Settings + + + + Allowed Search Tools{" "} + + + + + } + name="object_permission_search_tools" + className="mt-4" + help="Restrict which configured search tools keys on this team may call." + > + form.setFieldValue("object_permission_search_tools", vals)} + value={form.getFieldValue("object_permission_search_tools")} + accessToken={accessToken || ""} + placeholder="Select search tools (optional, empty = all allowed)" + /> + + + + Logging Settings diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx index d9e75388d1..bf73e802cd 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx @@ -1,6 +1,6 @@ -import { render, screen } from "@testing-library/react"; +import { act, fireEvent, render, screen } from "@testing-library/react"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { OnboardingForm } from "./OnboardingForm"; const mockUseOnboardingCredentials = vi.fn(); @@ -36,14 +36,33 @@ vi.mock("./OnboardingErrorView", () => ({ })); vi.mock("./OnboardingFormBody", () => ({ - OnboardingFormBody: ({ variant, userEmail }: { variant: string; userEmail: string }) => ( + OnboardingFormBody: ({ + variant, + userEmail, + claimError, + onSubmit, + }: { + variant: string; + userEmail: string; + claimError: string | null; + onSubmit: (formValues: { password: string }) => void; + }) => (
Form Body + + {claimError ?
{claimError}
: null}
), })); describe("OnboardingForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 GMT; path=/"; + }); + it("should render loading view when credentials are loading", () => { mockUseOnboardingCredentials.mockReturnValue({ data: undefined, @@ -92,4 +111,24 @@ describe("OnboardingForm", () => { expect(screen.getByTestId("form-body")).toHaveAttribute("data-variant", "reset_password"); }); + + it("should show claim error when claim response is missing final token", async () => { + mockUseOnboardingCredentials.mockReturnValue({ + data: { token: "fake-jwt-token" }, + isLoading: false, + isError: false, + }); + mockClaimToken.mockImplementation((_params, options) => { + options.onSuccess({}); + }); + + render(); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: "Submit" })); + }); + + expect(screen.getByTestId("claim-error")).toHaveTextContent("Failed to start session"); + expect(document.cookie).not.toContain("fake-jwt-token"); + }); }); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx index 9a9d9d7e72..1e8afe7645 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx @@ -31,18 +31,21 @@ export function OnboardingForm({ variant }: OnboardingFormProps) { const userEmail: string = decoded?.user_email ?? ""; const userId: string | null = decoded?.user_id ?? null; const accessToken: string | null = decoded?.key ?? null; - const jwtToken: string | null = credentialsData?.token ?? null; const handleSubmit = (formValues: { password: string }) => { - if (!accessToken || !jwtToken || !userId || !inviteId) return; + if (!accessToken || !userId || !inviteId) return; setClaimError(null); claimToken( { accessToken, inviteId, userId, password: formValues.password }, { - onSuccess: () => { - document.cookie = `token=${jwtToken}; path=/; SameSite=Lax`; + onSuccess: (data: { token?: string }) => { + if (!data?.token) { + setClaimError("Failed to start session. Please try again."); + return; + } + document.cookie = `token=${data.token}; path=/; SameSite=Lax`; const proxyBaseUrl = getProxyBaseUrl(); window.location.href = proxyBaseUrl ? `${proxyBaseUrl}/ui/?login=success` 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" ? ( = ({ } } - // Transform allowed_vector_store_ids and allowed_mcp_servers_and_groups into object_permission + const hasSearchTools = + Array.isArray(formValues.object_permission_search_tools) && + formValues.object_permission_search_tools.length > 0; + if ( (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) || (formValues.allowed_mcp_servers_and_groups && @@ -513,7 +523,9 @@ const Teams: React.FC = ({ formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 || formValues.allowed_mcp_servers_and_groups.toolPermissions)) ) { - formValues.object_permission = {}; + if (!formValues.object_permission) { + formValues.object_permission = {}; + } if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { formValues.object_permission.vector_stores = formValues.allowed_vector_store_ids; delete formValues.allowed_vector_store_ids; @@ -529,11 +541,7 @@ const Teams: React.FC = ({ delete formValues.allowed_mcp_servers_and_groups; } - // Add tool permissions separately if (formValues.mcp_tool_permissions && Object.keys(formValues.mcp_tool_permissions).length > 0) { - if (!formValues.object_permission) { - formValues.object_permission = {}; - } formValues.object_permission.mcp_tool_permissions = formValues.mcp_tool_permissions; delete formValues.mcp_tool_permissions; } @@ -563,6 +571,14 @@ const Teams: React.FC = ({ delete formValues.allowed_agents_and_groups; } + if (hasSearchTools) { + if (!formValues.object_permission) { + formValues.object_permission = {}; + } + formValues.object_permission.search_tools = formValues.object_permission_search_tools; + delete formValues.object_permission_search_tools; + } + // Add model_aliases if any are defined if (Object.keys(modelAliases).length > 0) { formValues.model_aliases = modelAliases; @@ -1516,6 +1532,34 @@ const Teams: React.FC = ({ + + + Search Tool Settings + + + + Allowed Search Tools{" "} + + + + + } + name="object_permission_search_tools" + className="mt-4" + help="Restrict which configured search tools keys on this team may call." + > + form.setFieldValue("object_permission_search_tools", vals)} + value={form.getFieldValue("object_permission_search_tools")} + accessToken={accessToken || ""} + placeholder="Select search tools (optional, empty = all allowed)" + /> + + + + Logging Settings diff --git a/ui/litellm-dashboard/src/components/SearchTools/SearchToolSelector.tsx b/ui/litellm-dashboard/src/components/SearchTools/SearchToolSelector.tsx new file mode 100644 index 0000000000..c93ff35e7d --- /dev/null +++ b/ui/litellm-dashboard/src/components/SearchTools/SearchToolSelector.tsx @@ -0,0 +1,69 @@ +import React, { useEffect, useState } from "react"; +import { Select } from "antd"; +import { fetchSearchTools } from "../networking"; + +export interface SearchToolSelectorProps { + onChange: (selected: string[]) => void; + value?: string[]; + className?: string; + accessToken: string; + placeholder?: string; + disabled?: boolean; +} + +const SearchToolSelector: React.FC = ({ + onChange, + value, + className, + accessToken, + placeholder = "Select search tools (optional)", + disabled = false, +}) => { + const [options, setOptions] = useState<{ label: string; value: string }[]>([]); + const [loading, setLoading] = useState(false); + + useEffect(() => { + const load = async () => { + if (!accessToken) return; + setLoading(true); + try { + const data = await fetchSearchTools(accessToken); + const tools = Array.isArray(data?.search_tools) + ? data.search_tools + : Array.isArray(data?.data) + ? data.data + : []; + setOptions( + tools + .map((tool: { search_tool_name?: string }) => tool?.search_tool_name) + .filter((name: unknown): name is string => typeof name === "string" && name.length > 0) + .map((name: string) => ({ label: name, value: name })), + ); + } catch (e) { + console.error("Failed to load search tools:", e); + } finally { + setLoading(false); + } + }; + load(); + }, [accessToken]); + + return ( + = ({ )} + )} diff --git a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx index 427c55c92b..937844e1c1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx @@ -7,15 +7,25 @@ import type { Row } from "@tanstack/react-table"; import { renderWithProviders } from "../../../tests/test-utils"; const mockHandleFilterResetFromHook = vi.fn(); -vi.mock("./log_filter_logic", () => ({ - useLogFilterLogic: vi.fn(() => ({ - filters: {}, - filteredLogs: { data: [], total: 0, page: 1, page_size: 50, total_pages: 1 }, - allTeams: [], - handleFilterChange: vi.fn(), - handleFilterReset: mockHandleFilterResetFromHook, - })), -})); +vi.mock("./log_filter_logic", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + useLogFilterLogic: vi.fn(() => ({ + filters: {}, + filteredLogs: { + data: [], + total: 0, + page: 1, + page_size: 50, + total_pages: 1, + }, + allTeams: [], + handleFilterChange: vi.fn(), + handleFilterReset: mockHandleFilterResetFromHook, + })), + }; +}); vi.mock("../networking", async (importOriginal) => { const actual = await importOriginal(); diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 8a015d8305..2f9e8fe878 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -24,7 +24,7 @@ import { ConfigInfoMessage } from "./ConfigInfoMessage"; import { AGENT_CALL_TYPES, ERROR_CODE_OPTIONS, MCP_CALL_TYPES, QUICK_SELECT_OPTIONS } from "./constants"; import { CostBreakdownViewer } from "./CostBreakdownViewer"; import { ErrorViewer } from "./ErrorViewer"; -import { useLogFilterLogic } from "./log_filter_logic"; +import { FILTER_KEYS, useLogFilterLogic } from "./log_filter_logic"; import { LogDetailsDrawer } from "./LogDetailsDrawer"; import { getTimeRangeDisplay } from "./logs_utils"; import { RequestResponsePanel } from "./RequestResponsePanel"; @@ -415,6 +415,11 @@ export default function SpendLogsTable({ label: "Model", customComponent: PaginatedModelSelect, }, + { + name: FILTER_KEYS.PUBLIC_MODEL_OR_SEARCH_TOOL, + label: "Public model / search tool", + isSearchable: false, + }, { name: "Key Alias", label: "Key Alias", diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx index 8ccf0a9e1b..17c5077152 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx @@ -112,6 +112,7 @@ describe("useLogFilterLogic", () => { expect(filters["Key Alias"]).toBe(""); expect(filters["Error Code"]).toBe(""); expect(filters["Error Message"]).toBe(""); + expect(filters["Public model / search tool"]).toBe(""); }); it("should return all logs when no filters are applied", () => { @@ -200,6 +201,51 @@ describe("useLogFilterLogic", () => { ); }); + it("should pass model param and filter search-tool rows by spend log model column", async () => { + const searchRows = [ + createLogEntry({ + request_id: "s1", + call_type: "asearch", + model: "tavily-marketing", + model_id: "", + team_id: "team-x", + }), + ]; + vi.mocked(uiSpendLogsCall).mockResolvedValue(createPaginatedResponse(searchRows)); + const logs = createPaginatedResponse([ + ...searchRows, + createLogEntry({ + request_id: "c1", + call_type: "chat", + model: "gpt-4o", + model_id: "mid-1", + team_id: "team-x", + }), + ]); + const { result } = renderHook(() => useLogFilterLogic({ ...defaultProps, logs }), { wrapper }); + + act(() => { + result.current.handleFilterChange({ "Public model / search tool": "tavily-marketing" }); + }); + + await waitFor( + () => { + expect(result.current.filteredLogs.data).toHaveLength(1); + expect(result.current.filteredLogs.data[0].model).toBe("tavily-marketing"); + expect(result.current.filteredLogs.data[0].call_type).toBe("asearch"); + }, + { timeout: 500 }, + ); + + expect(vi.mocked(uiSpendLogsCall)).toHaveBeenCalledWith( + expect.objectContaining({ + params: expect.objectContaining({ + model: "tavily-marketing", + }), + }), + ); + }); + it("should filter logs by api_key when Key Hash filter is set", async () => { const filteredLog = createLogEntry({ request_id: "req-1", api_key: "key-x" }); vi.mocked(uiSpendLogsCall).mockResolvedValue( diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx index a538872bfd..8f916999c1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx @@ -9,11 +9,14 @@ import { defaultPageSize } from "../constants"; import { PaginatedResponse } from "."; import type { LogsSortField } from "./columns"; -const FILTER_KEYS = { +/** Spend log `model` column (LLM public model name or `search_tool_name` for /search). */ +export const FILTER_KEYS = { TEAM_ID: "Team ID", KEY_HASH: "Key Hash", REQUEST_ID: "Request ID", MODEL: "Model", + /** Exact match on LiteLLM_SpendLogs.model — use for search tools and public model names. */ + PUBLIC_MODEL_OR_SEARCH_TOOL: "Public model / search tool", USER_ID: "User ID", END_USER: "End User", STATUS: "Status", @@ -58,6 +61,7 @@ export function useLogFilterLogic({ [FILTER_KEYS.KEY_HASH]: "", [FILTER_KEYS.REQUEST_ID]: "", [FILTER_KEYS.MODEL]: "", + [FILTER_KEYS.PUBLIC_MODEL_OR_SEARCH_TOOL]: "", [FILTER_KEYS.USER_ID]: "", [FILTER_KEYS.END_USER]: "", [FILTER_KEYS.STATUS]: "", @@ -107,6 +111,7 @@ export function useLogFilterLogic({ end_user: filters[FILTER_KEYS.END_USER] || undefined, status_filter: filters[FILTER_KEYS.STATUS] || undefined, model_id: filters[FILTER_KEYS.MODEL] || undefined, + model: filters[FILTER_KEYS.PUBLIC_MODEL_OR_SEARCH_TOOL] || undefined, key_alias: filters[FILTER_KEYS.KEY_ALIAS] || undefined, error_code: filters[FILTER_KEYS.ERROR_CODE] || undefined, error_message: filters[FILTER_KEYS.ERROR_MESSAGE] || undefined, @@ -155,7 +160,8 @@ export function useLogFilterLogic({ filters[FILTER_KEYS.END_USER] || filters[FILTER_KEYS.ERROR_CODE] || filters[FILTER_KEYS.ERROR_MESSAGE] || - filters[FILTER_KEYS.MODEL] + filters[FILTER_KEYS.MODEL] || + filters[FILTER_KEYS.PUBLIC_MODEL_OR_SEARCH_TOOL] ), [filters], ); @@ -218,6 +224,11 @@ export function useLogFilterLogic({ filteredData = filteredData.filter((log) => log.model_id === filters[FILTER_KEYS.MODEL]); } + if (filters[FILTER_KEYS.PUBLIC_MODEL_OR_SEARCH_TOOL]) { + const m = filters[FILTER_KEYS.PUBLIC_MODEL_OR_SEARCH_TOOL]; + filteredData = filteredData.filter((log) => log.model === m); + } + if (filters[FILTER_KEYS.KEY_HASH]) { filteredData = filteredData.filter((log) => log.api_key === filters[FILTER_KEYS.KEY_HASH]); } 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" },