Merge remote-tracking branch 'origin/main' into litellm_websearch_activation_status

This commit is contained in:
Yuneng Jiang 2026-09-19 16:16:25 -07:00
commit dd6a3558d5
No known key found for this signature in database
108 changed files with 6952 additions and 1620 deletions

View file

@ -1,3 +1,5 @@
import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
@ -15,6 +17,12 @@ def main() -> int:
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
if expected_count is not None and (
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
):
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
return 1
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
@ -38,6 +46,13 @@ def main() -> int:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")):
continue
_ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n")
for prop in case.findall("./properties/property"):
name = prop.get("name", "")
value = prop.get("value", "")
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
r"[A-Za-z0-9_.:<>-]{1,240}", value
):
_ = sys.stdout.write(f" {name}: {value}\n")
if (
selected
and not missing

View file

@ -5,6 +5,7 @@ from typing import Final
SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
UNSUPPORTED: Final = re.compile(
r"^tests/e2e/(ui|claude_code|load)/"
r"|^tests/e2e/mcp/test_mcp_oauth_happy_path_e2e\.py$"
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"

177
.github/workflows/test-mcp-oauth-e2e.yml vendored Normal file
View file

@ -0,0 +1,177 @@
name: MCP OAuth happy path
on:
pull_request:
paths:
- '.github/workflows/test-mcp-oauth-e2e.yml'
- '.github/e2e-stack/**'
- 'tests/e2e/*.py'
- 'tests/e2e/pytest.ini'
- 'tests/e2e/idp_realm.json'
- 'tests/e2e/mcp/**'
- 'litellm/experimental_mcp_client/**'
- 'litellm/proxy/_experimental/mcp_server/**'
- 'litellm/proxy/auth/**'
- 'litellm/proxy/management_endpoints/*sso*.py'
- 'litellm/proxy/management_endpoints/sso/**'
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
- 'litellm/proxy/proxy_server.py'
- 'litellm/proxy/schema.prisma'
- 'ui/litellm-dashboard/src/app/connect/**'
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
- 'pyproject.toml'
- 'uv.lock'
workflow_dispatch:
permissions: {}
concurrency:
group: mcp-oauth-${{ github.ref }}
cancel-in-progress: true
jobs:
oauth:
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
environment: e2e-changed
timeout-minutes: 45
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_HOST: 127.0.0.1
DATABASE_PORT: '5432'
DATABASE_USER: litellm
DATABASE_PASSWORD: dbpassword9090
DATABASE_NAME: litellm
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
E2E_KEYCLOAK_ADMIN_USER: admin
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
E2E_FIXTURE_MODE: live
E2E_PROVIDER_CACHE: '0'
E2E_MCP_OAUTH_LIVE: '1'
E2E_REQUIRED_TEST_COUNT: '4'
steps:
- name: Checkout the tested source
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
persist-credentials: false
- name: Require and materialize the upstream login
env:
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
run: |
umask 077
python3 - <<'PY'
import base64
import json
import os
import secrets
from pathlib import Path
encoded = os.environ.get("STORAGE_STATE", "")
if not encoded:
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
state = json.loads(base64.b64decode(encoded, validate=True))
if not isinstance(state, dict) or not state.get("cookies"):
raise SystemExit("The captured login must contain browser cookies")
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
directory.mkdir(mode=0o700)
path = directory / "linear-state.json"
path.write_text(json.dumps(state))
with open(os.environ["GITHUB_ENV"], "a") as output:
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
value = "sk-e2e-" + secrets.token_hex(24)
print(f"::add-mask::{value}")
output.write(f"{name}={value}\n")
PY
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: '3.13'
- uses: ./.github/actions/setup-uv-with-retries
with:
version: '0.10.9'
- uses: ./.github/actions/cache-cargo-build
- name: Install the frozen E2E environment
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
uv run --no-sync python scripts/prisma_generate_if_needed.py
uv run --no-sync playwright install --with-deps chromium
- name: Configure license access
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: mcp-oauth-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Load the E2E license
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
test -n "${license}"
echo "::add-mask::${license}"
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
- name: Build the gateway consent UI at the tested commit
run: |
cd ui/litellm-dashboard
../../scripts/with_dashboard_node.sh npm ci
../../scripts/with_dashboard_node.sh npm run build
mkdir -p ../../litellm/proxy/_experimental/out
cp -r out/. ../../litellm/proxy/_experimental/out/
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
mkdir -p "${page%.html}"
mv "${page}" "${page%.html}/index.html"
done
- name: Prepare the isolated database and IdP
run: |
umask 077
bash .github/e2e-stack/start-idp.sh
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
- name: Run every required OAuth variant without retries
run: |
umask 077
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
- name: Report JUnit results and reject skipped or missing cases
if: always()
run: |
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
- name: Remove private login and logs
if: always()
run: |
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"

View file

@ -60,6 +60,11 @@ async def _get_email_settings(prisma_client) -> Dict[str, bool]:
async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
"""Helper function to save email settings to general_settings in db"""
from litellm.proxy.proxy_server import proxy_config
proxy_config.reject_config_owned_writes(
section_name="general_settings", changed_keys={"email_settings": settings}
)
try:
verbose_proxy_logger.debug(
f"Saving email settings to general_settings: {settings}"
@ -168,6 +173,8 @@ async def update_event_settings(
await _save_email_settings(prisma_client, settings_dict)
return {"message": "Email event settings updated successfully"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error updating email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@ -197,6 +204,8 @@ async def reset_event_settings(
await _save_email_settings(prisma_client, default_settings)
return {"message": "Email event settings reset to defaults"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error resetting email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))

View file

@ -20,6 +20,7 @@ from typing import (
)
from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
@ -34,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from openai.types.file_deleted import FileDeleted
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
@ -59,6 +61,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_content_type_from_file_object,
get_model_id_from_unified_batch_id,
get_original_file_id,
is_litellm_executed_batch,
map_raw_file_ids_to_unified,
normalize_mime_type_for_provider,
resolve_managed_output_file_model_name,
@ -75,6 +78,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
CreateFileRequest,
FileListPage,
FileObject,
HttpxBinaryResponseContent,
OpenAIFileObject,
ResponsesAPIResponse,
)
@ -86,10 +90,6 @@ from litellm.types.utils import (
SpecialEnums,
)
if TYPE_CHECKING:
from litellm.types.llms.openai import HttpxBinaryResponseContent
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prisma.models import (
@ -204,6 +204,19 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
return prisma_client.db.litellm_managedobjecttable
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
)
return MappingProxyType(
{
key: value
for key in ("storage_backend", "storage_url")
if isinstance(value := hidden_params.get(key), str)
}
)
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
@ -226,6 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
) -> None:
verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
storage_metadata: Final = _storage_metadata_of(file_object)
if file_object is not None:
litellm_managed_file_object = LiteLLM_ManagedFileTable(
unified_file_id=file_id,
@ -235,6 +249,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
created_by=resolve_resource_owner_id(user_api_key_dict),
team_id=user_api_key_dict.team_id,
updated_by=user_api_key_dict.user_id,
storage_backend=storage_metadata.get("storage_backend"),
storage_url=storage_metadata.get("storage_url"),
)
await self.internal_usage_cache.async_set_cache(
key=file_id,
@ -262,14 +278,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object_json = file_object.model_dump_json()
db_data["file_object"] = file_object_json
update_data["file_object"] = file_object_json
# Extract storage metadata from hidden params if present
hidden_params = getattr(file_object, "_hidden_params", {}) or {}
if "storage_backend" in hidden_params:
db_data["storage_backend"] = hidden_params["storage_backend"]
update_data["storage_backend"] = hidden_params["storage_backend"]
if "storage_url" in hidden_params:
db_data["storage_url"] = hidden_params["storage_url"]
update_data["storage_url"] = hidden_params["storage_url"]
db_data.update(storage_metadata)
update_data.update(storage_metadata)
verbose_logger.debug(
f"Storage metadata: storage_backend={db_data.get('storage_backend')}, "
@ -314,6 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
create_if_missing: bool = True,
batch_processed: bool = False,
) -> None:
"""Persist a managed object row, caching it and upserting it in the DB.
@ -328,6 +339,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
row absent from the table is left absent rather than created with the
observer as its creator, because created_by and team_id are written from
whoever calls the create branch.
batch_processed is set by callers that have already billed the batch
themselves, so CheckBatchCost skips the row instead of billing it twice.
It is written only in the upsert create branch.
"""
verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
litellm_managed_object = LiteLLM_ManagedObjectTable(
@ -379,6 +394,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
**attribution_columns,
"batch_processed": batch_processed,
},
"update": update_columns,
},
@ -1343,6 +1359,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> LLMResponseTypes:
if isinstance(response, LiteLLMBatch):
decoded_batch_id: Final = _is_base64_encoded_unified_file_id(response.id)
if decoded_batch_id and is_litellm_executed_batch(decoded_batch_id):
return response
## Check if unified_file_id is in the response
unified_file_id = response._hidden_params.get("unified_file_id") # managed file id
unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id
@ -1794,24 +1813,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Check if file deletion should be blocked due to batch references
await self._check_file_deletion_allowed(file_id)
# file_id = convert_b64_uid_to_unified_uid(file_id)
model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
if specific_model_file_id_mapping:
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
await self._delete_storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
else:
await self._delete_provider_files(file_id, litellm_parent_otel_span, llm_router, data)
await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1820,16 +1826,53 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
return FileDeleted(id=file_id, object="file", deleted=True)
async def _delete_storage_backend_content(self, storage_backend_name: str, storage_url: str) -> None:
try:
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
raise HTTPException(status_code=400, detail=f"Cannot delete the stored file content: {e}") from e
await storage_backend.delete_file(storage_url)
async def _delete_provider_files(
self,
file_id: str,
litellm_parent_otel_span: Span | None,
llm_router: Router,
data: Mapping[str, object],
) -> None:
model_file_id_mapping: Final = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
if not specific_model_file_id_mapping:
return
filtered_data: Final = {
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**filtered_data,
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
async def afile_content(
self,
file_id: str,
litellm_parent_otel_span: Optional[Span],
llm_router: Router,
**data: Dict,
) -> "HttpxBinaryResponseContent":
) -> HttpxBinaryResponseContent:
"""
Get the content of a file from first model that has it
"""
managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
return await self._storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
model_file_id_mapping = data.pop("model_file_id_mapping", None)
model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping(
[file_id], litellm_parent_otel_span
@ -1859,6 +1902,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
else:
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
async def _storage_backend_content(self, storage_backend_name: str, storage_url: str) -> HttpxBinaryResponseContent:
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
content: Final = await storage_backend.download_file(storage_url)
return HttpxBinaryResponseContent(response=httpx.Response(status_code=httpx.codes.OK, content=content))
async def _convert_storage_files_to_base64(
self,
messages: List[AllMessageValues],
@ -1889,16 +1937,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# File is stored in a storage backend, download and convert to base64
try:
from litellm.llms.base_llm.files.storage_backend_factory import (
get_storage_backend,
)
storage_backend_name = db_file.storage_backend
storage_url = db_file.storage_url
# Get storage backend (uses same env vars as callback)
try:
storage_backend = get_storage_backend(storage_backend_name)
storage_backend = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
verbose_logger.warning(
f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.68"
version = "0.1.69"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.68"
version = "0.1.69"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -0,0 +1,8 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedFileContentTable" (
"id" TEXT NOT NULL,
"content" BYTEA NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_ManagedFileContentTable_pkey" PRIMARY KEY ("id")
);

View file

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.99"
version = "0.4.100"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.99"
version = "0.4.100"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -3046,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2 0.4.15",

View file

@ -34,7 +34,7 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }

10
litellm-rust/clippy.toml Normal file
View file

@ -0,0 +1,10 @@
# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate
# must see every entry. Going around it makes a fork-after-use hang instead of raising.
disallowed-methods = [
{ path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" },
]

View file

@ -4,6 +4,7 @@ use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use std::time::Duration;
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
use crate::{Pythonized, panic_to_pyerr, release_gil};
use futures_util::FutureExt;
use pyo3::exceptions::PyRuntimeError;
@ -12,6 +13,67 @@ use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pyo3::create_exception!(
_native,
ForkedAfterNativeRuntimeStarted,
PyRuntimeError,
"This process was forked after the native runtime started. Runtime threads do not survive fork(), so native routes cannot run here."
);
pyo3::create_exception!(
_native,
ProcessReservedForForking,
PyRuntimeError,
"This process was reserved for forking workers, so native routes cannot run here."
);
static FORK_GATE: ForkGate = ForkGate::new();
/// Whether this process has started the Tokio runtime.
pub fn runtime_started() -> bool {
FORK_GATE.started(std::process::id())
}
/// Declares that this process exists to fork workers, so it must never start the runtime.
/// Fails if it already has. Workers are unaffected: the reservation is keyed by pid.
pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> {
FORK_GATE.reserve(std::process::id())
}
/// The only door to the Tokio runtime: every route reaches it through this module, which is
/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it.
fn enter_runtime() -> PyResult<()> {
FORK_GATE
.enter(std::process::id())
.map_err(|refused| match refused {
Refused::ReservedForForking => ProcessReservedForForking::new_err(
"this process is reserved for forking workers and cannot run native routes; \
move the call into a worker, after the fork",
),
Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(
"this process was forked after the native runtime started, and runtime threads \
do not survive fork(); start workers with spawn or forkserver, or fork before \
the first native call",
),
})
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn runtime() -> PyResult<&'static Runtime> {
enter_runtime()?;
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where
F: Future<Output = PyResult<T>> + Send + 'static,
T: for<'py> IntoPyObject<'py> + Send + 'static,
{
enter_runtime()?;
pyo3_async_runtimes::tokio::future_into_py(py, future)
}
pub fn run_sync<T, E, F>(
py: Python<'_>,
future: F,
@ -22,12 +84,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
run_sync_on(py, runtime()?, future, map_error)
}
pub fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
@ -35,7 +92,7 @@ where
T: Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future)
run_sync_value_on(py, runtime()?, future)
}
fn run_sync_value_on<T, F>(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult<T>
@ -83,7 +140,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
future_into_py(py, async move {
let result = catch_future_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
@ -95,7 +152,7 @@ where
T: for<'py> IntoPyObject<'py> + Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
future_into_py(py, async move { catch_future_panic(future).await? })
}
pub fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
@ -103,8 +160,9 @@ where
T: Send,
F: Future<Output = PyResult<T>> + Send,
{
let runtime = runtime()?;
let result = release_gil(py, || {
let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter();
let _runtime = runtime.enter();
std::panic::catch_unwind(AssertUnwindSafe(|| {
future.poll(&mut Context::from_waker(Waker::noop()))
}))
@ -286,27 +344,25 @@ mod tests {
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
fn runtime_worker_count() -> PyResult<usize> {
Ok(runtime()?.metrics().num_workers())
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> PyResult<bool> {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
return Ok(false);
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
runtime()?.spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
Ok(heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok())
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
@ -317,6 +373,16 @@ mod tests {
.expect("result should convert")
}
#[rstest]
fn reaching_the_runtime_marks_the_process_as_started(
#[from(initialized_python)] python: &InitializedPython,
) {
python.attach(|py| {
run_sync_value(py, async { Ok(()) }).unwrap();
assert!(runtime_started());
});
}
#[rstest]
fn inline_poll_releases_gil_and_enters_runtime(
#[from(initialized_python)] python: &InitializedPython,

View file

@ -0,0 +1,139 @@
use std::sync::atomic::{AtomicU32, Ordering};
const UNSET: u32 = 0;
/// Decides which process may use the Tokio runtime. Its worker threads do not survive
/// `fork()`: a child forked after they started hangs on its first native call. The gate turns
/// both halves of that hazard into errors, keyed by pid so a fork needs no hook to be seen:
/// a process reserved for forking can never start the runtime, and a child of a process that
/// did start it is refused instead of hanging.
pub(crate) struct ForkGate {
runtime_pid: AtomicU32,
fork_only_pid: AtomicU32,
}
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Refused {
ReservedForForking,
ForkedAfterStart,
}
#[derive(Debug, PartialEq, Eq)]
pub struct RuntimeAlreadyStarted;
impl ForkGate {
pub(crate) const fn new() -> Self {
Self {
runtime_pid: AtomicU32::new(UNSET),
fork_only_pid: AtomicU32::new(UNSET),
}
}
/// Claims the runtime for `pid`. Claim first, then look for a reservation: `reserve` does
/// the mirror image, so when the two race at least one of them sees the other.
pub(crate) fn enter(&self, pid: u32) -> Result<(), Refused> {
match self
.runtime_pid
.compare_exchange(UNSET, pid, Ordering::SeqCst, Ordering::SeqCst)
{
Err(owner) if owner != pid => return Err(Refused::ForkedAfterStart),
_ => {}
}
if self.fork_only_pid.load(Ordering::SeqCst) == pid {
// Nothing was started, so the workers forked from here must still find it unclaimed.
let _ =
self.runtime_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(Refused::ReservedForForking);
}
Ok(())
}
/// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does
/// the mirror image, so when the two race at least one of them sees the other. A refused
/// reservation leaves the gate exactly as it was, so a process already running the runtime
/// keeps refusing the children it forks.
pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> {
self.fork_only_pid.store(pid, Ordering::SeqCst);
if self.runtime_pid.load(Ordering::SeqCst) == pid {
let _ =
self.fork_only_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(RuntimeAlreadyStarted);
}
Ok(())
}
pub(crate) fn started(&self, pid: u32) -> bool {
self.runtime_pid.load(Ordering::SeqCst) == pid
}
}
#[cfg(test)]
mod tests {
use super::*;
const MASTER: u32 = 100;
const WORKER: u32 = 101;
#[test]
fn unreserved_process_starts_the_runtime_and_stays_started() {
let gate = ForkGate::new();
assert!(!gate.started(MASTER));
assert_eq!(gate.enter(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
}
#[test]
fn reserved_process_can_never_start_the_runtime() {
let gate = ForkGate::new();
assert_eq!(gate.reserve(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert!(!gate.started(MASTER));
}
#[test]
fn workers_forked_from_a_reserved_process_start_their_own_runtime() {
let gate = ForkGate::new();
gate.reserve(MASTER).unwrap();
gate.enter(MASTER).unwrap_err();
assert_eq!(gate.enter(WORKER), Ok(()));
assert!(gate.started(WORKER));
}
#[test]
fn reserving_after_the_runtime_started_is_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
}
#[test]
fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
}
#[test]
fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
assert!(!gate.started(WORKER));
assert_eq!(gate.enter(MASTER), Ok(()));
}
}

View file

@ -8,6 +8,7 @@ mod argument;
mod callable;
mod driver;
mod execution;
mod fork_gate;
mod gil;
mod handle;
mod marshal;
@ -18,7 +19,12 @@ pub use adapter::{
pub use argument::lookup;
pub use callable::wrap_failure;
pub use driver::run_call;
pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value};
pub use execution::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value,
reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value,
runtime_started,
};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep};
pub use marshal::{Pythonized, from_py, from_py_argument, panic_to_pyerr, to_py};

View file

@ -1,5 +1,5 @@
use litellm_host_python::release_count;
use pyo3::{prelude::*, types::PyDict};
use litellm_host_python::{release_count, runtime_started};
use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
#[pyfunction]
pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
@ -8,6 +8,20 @@ pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
Ok(stats.into_any().unbind())
}
/// True once this process has started the native runtime, which does not survive `fork()`.
#[pyfunction]
pub(crate) fn process_state_started() -> bool {
runtime_started()
}
/// Declares that this process only forks workers: from now on every native route raises here,
/// so the runtime can never start. Raises if it already has. Forked workers are unaffected.
#[pyfunction]
pub(crate) fn reserve_process_for_forking() -> PyResult<()> {
litellm_host_python::reserve_process_for_forking()
.map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process"))
}
#[cfg(feature = "panic-test")]
#[pyfunction]
pub(crate) fn _panic_for_test() {

View file

@ -13,7 +13,7 @@ mod _native {
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
#[pymodule_export]
use crate::diagnostics::gil_stats;
use crate::diagnostics::{gil_stats, process_state_started, reserve_process_for_forking};
#[pymodule_export]
use crate::errors::{RustBridgeDeclined, RustUpstreamError};
#[pymodule_export]
@ -30,6 +30,8 @@ mod _native {
use crate::routes::responses::ResponsesWebSocketConnection;
#[pymodule_export]
use crate::token_counter::TokenCounter;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
}
use pyo3::prelude::*;
@ -50,6 +52,8 @@ mod tests {
let mut expected = vec![
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ocr",
"aocr",
"transcription",
@ -62,6 +66,8 @@ mod tests {
"ResponsesWebSocketConnection",
"TokenCounter",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",
];
expected.sort_unstable();

View file

@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection {
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(responses_error_to_pyerr)?;
@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection {
fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner
.send_text(text)
.await
@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection {
fn recv_text<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.recv_text().await.map_err(responses_error_to_pyerr)
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.close().await.map_err(responses_error_to_pyerr)
})
}
@ -68,6 +68,10 @@ mod tests {
use tokio_tungstenite::{accept_async, tungstenite::Message};
#[test]
#[expect(
clippy::disallowed_methods,
reason = "the test server shares the routes' runtime"
)]
fn responses_websocket_connection_round_trips_through_python() {
Python::initialize();
let runtime = pyo3_async_runtimes::tokio::get_runtime();

View file

@ -1694,6 +1694,7 @@ LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, int(os.getenv("LITELLM_EXECUTED_BATCH_CONCURRENCY", "4")))
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"

View file

@ -7,9 +7,11 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from enum import Enum
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final, cast
from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast
from urllib.parse import urlsplit
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity
from litellm.integrations.otel.model.semconv import (
GenAIOperation,
@ -25,6 +27,7 @@ from litellm.integrations.otel.model.utils import (
as_float,
as_int,
as_str,
as_str_mapping,
as_str_tuple,
)
@ -424,7 +427,7 @@ class LLMCallSpanData:
# plain ``.get`` — no repeated ``isinstance`` guards.
raw_response: Final = payload.get("response")
response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {})
choices_out: Final = _dicts(response.get("choices"))
choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response)
# ``finish_reasons`` is metadata, not content, so derive it from
# ``choices_out`` before gating. The raw message/choice bodies are only
# retained when content capture is enabled (see ``capture_span_content``);
@ -703,6 +706,84 @@ def _finish_reasons(choices: tuple[Mapping[str, object], ...]) -> tuple[str, ...
return tuple(r for c in choices if (r := as_str(c.get("finish_reason"))))
class _ToolFunction(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class _ToolCall(TypedDict):
id: ReadOnly[str]
type: ReadOnly[Literal["function"]]
function: ReadOnly[_ToolFunction]
class _AssistantMessage(TypedDict):
role: ReadOnly[str]
content: ReadOnly[str | None]
refusal: ReadOnly[str | None]
tool_calls: ReadOnly[tuple[_ToolCall, ...] | None]
class _Choice(TypedDict):
message: ReadOnly[_AssistantMessage]
finish_reason: ReadOnly[str | None]
_RESPONSES_TOOL_CALL_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]:
"""A Responses API ``output`` folded into one chat-shaped assistant choice."""
items: Final = _dicts(response.get("output"))
messages: Final = tuple(item for item in items if item.get("type") == "message")
parts: Final = tuple(part for item in messages for part in _dicts(item.get("content")))
tool_calls: Final = tuple(
_responses_tool_call(item) for item in items if item.get("type") in _RESPONSES_TOOL_CALL_TYPES
)
if not messages and not tool_calls:
return ()
message: Final[_AssistantMessage] = {
"role": next((role for item in messages if (role := as_str(item.get("role")))), "assistant"),
"content": _responses_parts_text(parts, "output_text", "text"),
"refusal": _responses_parts_text(parts, "refusal", "refusal"),
"tool_calls": tool_calls or None,
}
choice: Final[_Choice] = {"message": message, "finish_reason": _responses_finish_reason(response, bool(tool_calls))}
return (choice,)
def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None:
texts: Final = tuple(
text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None
)
return "".join(texts) if texts else None
def _responses_tool_call(item: Mapping[str, object]) -> _ToolCall:
custom: Final = item.get("type") == "custom_tool_call"
function: Final[_ToolFunction] = {
"name": as_str(item.get("name")) or "",
"arguments": as_str(item.get("input" if custom else "arguments")) or "",
}
tool_call: Final[_ToolCall] = {
"id": as_str(item.get("call_id")) or as_str(item.get("id")) or "",
"type": "function",
"function": function,
}
return tool_call
def _responses_finish_reason(response: Mapping[str, object], has_tool_calls: bool) -> str | None:
status: Final = as_str(response.get("status"))
if status == "completed":
return "tool_calls" if has_tool_calls else "stop"
if status != "incomplete":
return None
details: Final = as_str_mapping(response.get("incomplete_details"))
reason: Final = details.get("reason") if details is not None else None
return "content_filter" if reason == "content_filter" else "length"
def _parse_error(payload: StandardLoggingPayload) -> SpanError | None:
"""A ``SpanError`` for a failed request, or ``None`` on success."""
if payload.get("status") != "failure":

View file

@ -5528,6 +5528,10 @@ class StandardLoggingPayloadSetup:
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
clean_metadata[key] = metadata[key]
recorded_guardrails: Final = metadata.get("applied_guardrails")
if applied_guardrails and isinstance(recorded_guardrails, list):
clean_metadata["applied_guardrails"] = list(dict.fromkeys([*applied_guardrails, *recorded_guardrails]))
user_api_key: Final = metadata.get("user_api_key")
if user_api_key and isinstance(user_api_key, str) and is_valid_sha256_hash(user_api_key):
clean_metadata["user_api_key_hash"] = user_api_key

View file

@ -128,6 +128,8 @@ def _redact_responses_api_output(output_items):
for content_part in output_item.content:
if getattr(content_part, "text", None) is not None:
content_part.text = REDACTED_BY_LITELLM
if getattr(content_part, "refusal", None) is not None:
content_part.refusal = REDACTED_BY_LITELLM
# Redact reasoning items in output array
if hasattr(output_item, "type") and output_item.type == "reasoning":
@ -138,6 +140,8 @@ def _redact_responses_api_output(output_items):
if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"):
output_item.arguments = REDACTED_BY_LITELLM
if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"):
output_item.input = REDACTED_BY_LITELLM
def _redact_responses_api_output_dict(output_items, redacted_str: str):
@ -153,6 +157,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
for content_item in output_item["content"]:
if isinstance(content_item, dict) and content_item.get("text") is not None:
content_item["text"] = redacted_str
if isinstance(content_item, dict) and content_item.get("refusal") is not None:
content_item["refusal"] = redacted_str
if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list):
for summary_item in output_item["summary"]:
@ -161,6 +167,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
if output_item.get("type") == "function_call" and "arguments" in output_item:
output_item["arguments"] = redacted_str
if output_item.get("type") == "custom_tool_call" and "input" in output_item:
output_item["input"] = redacted_str
def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:

View file

@ -280,10 +280,17 @@ class AzureOpenAIConfig(BaseConfig):
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
request_params: Final = MappingProxyType(
{
key: value
for key, value in optional_params.items()
if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions")
}
)
return {
"model": model,
"messages": azure_messages,
**optional_params,
**request_params,
**sanitized_tools_update(optional_params),
}

View file

@ -0,0 +1,40 @@
from typing import TYPE_CHECKING, Final
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.repositories.managed_file_content_repository import ManagedFileContentRepository
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
LITELLM_DB_STORAGE_BACKEND_NAME: Final = "litellm_db"
LITELLM_DB_STORAGE_URL_PREFIX: Final = f"{LITELLM_DB_STORAGE_BACKEND_NAME}://"
def storage_url_to_row_id(storage_url: str) -> str:
if not storage_url.startswith(LITELLM_DB_STORAGE_URL_PREFIX):
raise ValueError(f"Not a {LITELLM_DB_STORAGE_BACKEND_NAME} storage url: {storage_url}")
return storage_url.removeprefix(LITELLM_DB_STORAGE_URL_PREFIX)
class LiteLLMDbStorageBackend(BaseFileStorageBackend):
def __init__(self, prisma_client: "PrismaClient") -> None:
self._contents = ManagedFileContentRepository(prisma_client)
async def upload_file(
self,
file_content: bytes,
filename: str,
content_type: str,
path_prefix: str | None = None,
file_naming_strategy: str = "uuid",
) -> str:
return f"{LITELLM_DB_STORAGE_URL_PREFIX}{await self._contents.store(file_content)}"
async def download_file(self, storage_url: str) -> bytes:
content: Final = await self._contents.load(storage_url_to_row_id(storage_url))
if content is None:
raise ValueError(f"No stored file content for {storage_url}")
return content
async def delete_file(self, storage_url: str) -> None:
await self._contents.delete(storage_url_to_row_id(storage_url))

View file

@ -6,32 +6,46 @@ based on the backend type. Backends use the same configuration as their correspo
callbacks (e.g., azure_storage uses the same env vars as AzureBlobStorageLogger).
"""
from typing import TYPE_CHECKING
from litellm._logging import verbose_logger
from .azure_blob_storage_backend import AzureBlobStorageBackend
from .litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME, LiteLLMDbStorageBackend
from .storage_backend import BaseFileStorageBackend
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
def get_storage_backend(backend_type: str) -> BaseFileStorageBackend:
def get_storage_backend(backend_type: str, prisma_client: "PrismaClient | None" = None) -> BaseFileStorageBackend:
"""
Factory function to create a storage backend instance.
Backends are configured using the same environment variables as their
corresponding callbacks. For example, "azure_storage" uses the same
env vars as AzureBlobStorageLogger.
env vars as AzureBlobStorageLogger. "litellm_db" stores file bytes in the
proxy's own database and needs the connected Prisma client.
Args:
backend_type: Backend type identifier (e.g., "azure_storage")
backend_type: Backend type identifier (e.g., "azure_storage", "litellm_db")
prisma_client: The proxy's database client, required by "litellm_db"
Returns:
BaseFileStorageBackend: Instance of the appropriate storage backend
Raises:
ValueError: If backend_type is not supported
ValueError: If backend_type is not supported, or "litellm_db" is asked for without a database
"""
verbose_logger.debug("Creating storage backend: type=%s", backend_type)
if backend_type == "azure_storage":
return AzureBlobStorageBackend()
else:
raise ValueError(f"Unsupported storage backend type: {backend_type}. Supported types: azure_storage")
if backend_type == LITELLM_DB_STORAGE_BACKEND_NAME:
if prisma_client is None:
raise ValueError(f"Storage backend {LITELLM_DB_STORAGE_BACKEND_NAME} requires a database-connected proxy")
return LiteLLMDbStorageBackend(prisma_client)
raise ValueError(
f"Unsupported storage backend type: {backend_type}. "
f"Supported types: azure_storage, {LITELLM_DB_STORAGE_BACKEND_NAME}"
)

View file

@ -4012,6 +4012,64 @@ async def get_org_object(
return _org_obj
def _last_known_org_cache_key(org_id: str) -> str:
return f"org_id:{org_id}:with_budget:last_known"
async def _keep_last_known_org(
org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache
) -> None:
cache_key: Final = _last_known_org_cache_key(org_id)
held_locally: Final = await user_api_key_cache.async_get_cache(
key=cache_key, local_only=True, model_type=LiteLLM_OrganizationTable
)
if held_locally is not None:
return
await user_api_key_cache.async_set_cache(
key=cache_key,
value=org,
model_type=LiteLLM_OrganizationTable,
ttl=get_management_object_ttl(user_api_key_cache),
)
async def get_org_object_for_request(
org_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_OrganizationTable | None:
try:
org: Final = await get_org_object(
org_id=org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
include_budget_table=True,
)
except OrganizationNotFoundError:
return None
except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits
if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True)
return None
last_known_org: Final = await user_api_key_cache.async_get_cache(
key=_last_known_org_cache_key(org_id),
model_type=LiteLLM_OrganizationTable,
)
if last_known_org is not None:
return last_known_org
if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable():
return None
raise
if org is None:
return None
await _keep_last_known_org(org, org_id, user_api_key_cache)
return org
async def _get_resources_from_access_groups(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
@ -5680,7 +5738,7 @@ async def _project_max_budget_check(
if project_object.litellm_budget_table is not None:
max_budget = project_object.litellm_budget_table.max_budget
if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget):
if max_budget is None or not math.isfinite(max_budget):
return
from litellm.proxy.proxy_server import get_current_spend

View file

@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_checks import (
get_jwt_key_mapping_object,
get_key_end_user_budget_id,
get_object_permission,
get_org_object_for_request,
get_project_object,
get_team_membership,
get_team_object,
@ -2611,6 +2612,47 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc
return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
async def _inherit_org_identity(
user_api_key_auth_obj: UserAPIKeyAuth,
team_object: LiteLLM_TeamTableCachedObj | None,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> None:
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
user_api_key_auth_obj.org_id = team_object.organization_id
already_populated: Final = any(
value is not None
for value in (
user_api_key_auth_obj.organization_alias,
user_api_key_auth_obj.organization_max_budget,
user_api_key_auth_obj.organization_tpm_limit,
user_api_key_auth_obj.organization_rpm_limit,
user_api_key_auth_obj.organization_metadata,
)
)
if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None:
return
org_object: Final = await get_org_object_for_request(
org_id=user_api_key_auth_obj.org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if org_object is None:
return
user_api_key_auth_obj.organization_alias = org_object.organization_alias
user_api_key_auth_obj.organization_metadata = org_object.metadata
budget: Final = org_object.litellm_budget_table
if budget is None:
return
user_api_key_auth_obj.organization_max_budget = budget.max_budget
user_api_key_auth_obj.organization_tpm_limit = budget.tpm_limit
user_api_key_auth_obj.organization_rpm_limit = budget.rpm_limit
def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, object]) -> bool:
return master_key is None and not any(
general_settings.get(flag, False)
@ -2849,8 +2891,14 @@ async def _run_centralized_common_checks(
user_object=user_object,
)
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
user_api_key_auth_obj.org_id = team_object.organization_id
await _inherit_org_identity(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
# common_checks identifies admin via user_object, not the token
# (non_proxy_admin_allowed_routes_check). JWT admin shortcut and

View file

@ -7,10 +7,12 @@
import asyncio
import os
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -18,6 +20,15 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE,
LiteLLMExecutedBatchRunner,
ManagedBatchStore,
batch_error,
executed_batch_runner_lost,
litellm_executed_provider_for,
resolve_litellm_executed_provider,
)
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
@ -47,16 +58,87 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_model_id_from_unified_batch_id,
get_models_from_unified_file_id,
get_original_file_id,
is_litellm_executed_batch,
prepare_data_with_credentials,
update_batch_in_database,
validate_managed_id_requirement,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.router import Router
from litellm.types.llms.openai import LiteLLMBatchCreateRequest
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from prisma.models import LiteLLM_ManagedObjectTable
router: Final = APIRouter()
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None:
metadata: Final = data.get("litellm_metadata")
if metadata is None:
return None
return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata))
def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyLogging) -> LiteLLMExecutedBatchRunner:
from litellm.proxy.proxy_server import general_settings, prisma_client
managed_files: Final = proxy_logging_obj.get_proxy_hook("managed_files")
if prisma_client is None or not isinstance(managed_files, ManagedBatchStore):
raise batch_error(
400,
"LiteLLM-executed batches need a database: set DATABASE_URL so LiteLLM can keep the batch and its files",
)
return LiteLLMExecutedBatchRunner(
llm_router=llm_router,
prisma_client=prisma_client,
managed_files=managed_files,
batches=ManagedBatchRepository(prisma_client),
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
)
async def _batch_from_database(
batch_id: str,
unified_batch_id: str | Literal[False],
executed_batch: bool,
managed_files_obj: object,
prisma_client: PrismaClient | None,
llm_router: Router | None,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple["LiteLLM_ManagedObjectTable | None", LiteLLMBatch | None]:
row, batch = await get_batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
)
updated_at: Final[object] = getattr(row, "updated_at", None)
if not executed_batch or batch is None or llm_router is None or not isinstance(updated_at, datetime):
return row, batch
if not executed_batch_runner_lost(batch.status, updated_at):
return row, batch
runner: Final = _litellm_executed_batch_runner(llm_router, proxy_logging_obj)
return row, await runner.fail_abandoned(batch, user_api_key_dict)
async def _raise_when_input_file_must_be_managed(model: str, credentials: Mapping[str, object]) -> None:
if await litellm_executed_provider_for(credentials) is None:
return
raise batch_error(
400,
f"Batches for {model} run inside LiteLLM, so the input file must be a LiteLLM managed file: "
f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}",
)
def _raise_not_found_when_openai_fallback_unservable(
@ -101,6 +183,24 @@ async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str |
return db_file.storage_url or None
async def _create_provider_batch_for_managed_file(
llm_router: Router,
create_batch_data: LiteLLMBatchCreateRequest,
input_file_id: str,
unified_file_id: str,
) -> LiteLLMBatch:
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
request: Final[LiteLLMBatchCreateRequest] = {
**create_batch_data,
"input_file_id": resolved_storage_url or input_file_id,
"disable_fallbacks": True,
}
response: Final = await llm_router.acreate_batch(**request)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
return response
@router.post(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -296,24 +396,35 @@ async def create_batch(
await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
_create_batch_data["model"] = model
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
if resolved_storage_url is not None:
_create_batch_data["input_file_id"] = resolved_storage_url
if llm_router is None:
raise HTTPException(
status_code=500,
detail={"error": "LLM Router not initialized. Ensure models added to proxy."},
)
_create_batch_data.update(disable_fallbacks=True) # pyright: ignore[reportCallIssue] # router flag
response = await llm_router.acreate_batch(**_create_batch_data)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
executed_provider: Final = await resolve_litellm_executed_provider(
llm_router, model, user_api_key_dict.team_id
)
response = (
await _litellm_executed_batch_runner(llm_router, proxy_logging_obj).create(
create_request=_create_batch_data,
unified_input_file_id=input_file_id,
model=model,
provider=executed_provider,
user_api_key_dict=user_api_key_dict,
request_tags=_request_tags(_create_batch_data),
)
if executed_provider is not None
else await _create_provider_batch_for_managed_file(
llm_router, _create_batch_data, input_file_id, unified_file_id
)
)
else:
# Check if model specified via header/query/body param
model_param: Final = (
data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model")
_create_batch_data.get("model")
or request.query_params.get("model")
or request.headers.get("x-litellm-model")
)
# SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback
@ -325,6 +436,7 @@ async def create_batch(
user_api_key_dict=user_api_key_dict,
operation_context="batch creation",
)
await _raise_when_input_file_must_be_managed(model_param, credentials)
prepare_data_with_credentials(
data=_create_batch_data,
@ -486,23 +598,26 @@ async def retrieve_batch(
managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
from litellm.proxy.proxy_server import prisma_client
db_batch_object, response = await get_batch_from_database(
executed_batch: Final = isinstance(unified_batch_id, str) and is_litellm_executed_batch(unified_batch_id)
db_batch_object, response = await _batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
executed_batch=executed_batch,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
)
if executed_batch and response is None:
raise batch_error(404, f"No batch found with id '{batch_id}'.")
# If batch is in a terminal state, return immediately.
# Include "complete" (DB-normalized form of "completed").
if response is not None and response.status in [
"completed",
"complete",
"failed",
"cancelled",
"expired",
]:
if response is not None and (
response.status in ("completed", "complete", "failed", "cancelled", "expired") or executed_batch
):
# Call hooks and return
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
@ -978,6 +1093,17 @@ async def cancel_batch(
proxy_config=proxy_config,
)
unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) if unified_batch_id else None
if unified_model_id is not None:
resolved_unified_model: Final = (
llm_router.resolve_model_name_from_model_id(unified_model_id) if llm_router is not None else None
)
await authorize_model_for_key(
model_id=resolved_unified_model or unified_model_id,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
# SCENARIO 1: Batch ID is encoded with model info
if model_from_id is not None:
credentials: Final = await get_authorized_credentials_for_model(
@ -1009,6 +1135,12 @@ async def cancel_batch(
)
# SCENARIO 2: target_model_names based routing
elif unified_batch_id and is_litellm_executed_batch(unified_batch_id):
if llm_router is None:
raise batch_error(500, "LLM Router not initialized. Ensure models added to proxy.")
response = await _litellm_executed_batch_runner( # rebind-ok: each cancel path sets the route's response
llm_router, proxy_logging_obj
).cancel(batch_id, user_api_key_dict)
elif unified_batch_id:
if llm_router is None:
raise HTTPException(
@ -1022,11 +1154,6 @@ async def cancel_batch(
status_code=400,
detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."},
)
await authorize_model_for_key(
model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
data["model"] = model_id_from_batch
data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id)
response = await llm_router.acancel_batch(**data)

View file

@ -0,0 +1,716 @@
import asyncio
import json
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import pairwise
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable
import httpx
from openai.types.batch import Errors
from openai.types.batch_error import BatchError
from openai.types.batch_request_counts import BatchRequestCounts
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid as uuid_module
from litellm.constants import LITELLM_EXECUTED_BATCH_CONCURRENCY
from litellm.integrations.prometheus import PrometheusLogger
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import is_request_body_safe
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.openai_files_endpoints.common_utils import LITELLM_EXECUTED_BATCH_ID_PREFIX
from litellm.proxy.openai_files_endpoints.storage_backend_service import StorageBackendFileService
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.types.llms.openai import LiteLLMBatchCreateRequest, OpenAIFileObject, OpenAIFilesPurpose
from litellm.types.utils import LITELLM_EXECUTED_BATCH_PROVIDERS, ExtractedFileData, LiteLLMBatch, LlmProviders
if TYPE_CHECKING:
from prisma import types as prisma_types
from litellm.router import Router
BatchEndpoint: TypeAlias = Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"]
BatchStatus: TypeAlias = Literal[
"in_progress", "finalizing", "completed", "failed", "cancelling", "cancelled", "expired"
]
TERMINAL_BATCH_STATUSES: Final[frozenset[str]] = frozenset({"completed", "failed", "cancelled", "expired"})
_STOP_STATUSES: Final[frozenset[str]] = TERMINAL_BATCH_STATUSES | frozenset({"cancelling"})
_BATCH_ENDPOINT_ADAPTER: Final[TypeAdapter[BatchEndpoint]] = TypeAdapter(BatchEndpoint)
_CANCEL_POLL_SECONDS: Final = 1.0
_HEARTBEAT_SECONDS: Final = 30.0
_STALE_AFTER_SECONDS: Final = 180.0
_FILES_API_PROBE_TIMEOUT_SECONDS: Final = 5.0
_COMPLETION_WINDOW_SECONDS: Final = 24 * 60 * 60
_RUNNER_LOST_MESSAGE: Final = "the proxy replica running this batch stopped before it finished; resubmit the batch"
_EXPIRED_MESSAGE: Final = "This request could not be executed before the completion window expired."
_ROUTER_METHODS: Final[Mapping[BatchEndpoint, str]] = MappingProxyType(
{
"/v1/chat/completions": "acompletion",
"/v1/completions": "atext_completion",
"/v1/embeddings": "aembedding",
"/v1/responses": "aresponses",
}
)
_CANCELLING_TRANSITIONS: Final[Mapping[BatchStatus, BatchStatus]] = MappingProxyType(
{"completed": "cancelled", "expired": "cancelled", "in_progress": "cancelling", "finalizing": "cancelling"}
)
LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE: Final = (
"upload it through POST /v1/files with purpose=batch and either the x-litellm-model header or the "
"target_model_names form field naming the model, so LiteLLM keeps the file and runs the batch itself"
)
_RUNNING_BATCHES: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong references keep running batch tasks alive
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
class _ErrorDetail(TypedDict):
message: ReadOnly[str]
type: ReadOnly[str]
param: ReadOnly[None]
code: ReadOnly[None]
class _ErrorBody(TypedDict):
error: ReadOnly[_ErrorDetail]
class _ResultResponse(TypedDict):
status_code: ReadOnly[int]
request_id: ReadOnly[str]
body: ReadOnly[Mapping[str, object]]
class _LineError(TypedDict):
code: ReadOnly[str]
message: ReadOnly[str]
class _ResultLine(TypedDict):
id: ReadOnly[str]
custom_id: ReadOnly[str]
response: ReadOnly[_ResultResponse | None]
error: ReadOnly[_LineError | None]
class BatchInputLine(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
custom_id: str
method: Literal["POST"]
url: str
body: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class InvalidBatchInput:
line_number: int | None
reason: str
def describe(self) -> str:
return f"line {self.line_number}: {self.reason}" if self.line_number is not None else self.reason
@dataclass(frozen=True, slots=True)
class RowOutcome:
custom_id: str
status_code: int
body: Mapping[str, object]
succeeded: bool
@dataclass(frozen=True, slots=True)
class ExpiredRow:
custom_id: str
@dataclass(frozen=True, slots=True)
class _BatchRun:
unified_batch_id: str
llm_batch_id: str
model: str
endpoint: BatchEndpoint
lines: tuple[BatchInputLine, ...]
user_api_key_dict: UserAPIKeyAuth
request_tags: tuple[str, ...]
deadline: float
@runtime_checkable
class ManagedBatchStore(Protocol):
def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: ...
async def get_unified_file_id(
self, file_id: str, litellm_parent_otel_span: object | None = None
) -> LiteLLM_ManagedFileTable | None: ...
async def store_unified_object_id(
self,
unified_object_id: str,
file_object: LiteLLMBatch,
litellm_parent_otel_span: object | None,
model_object_id: str,
file_purpose: Literal["batch", "fine-tune", "response"],
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
batch_processed: bool = False,
) -> None: ...
class _StorageBackendFactory(Protocol):
def __call__(self, backend_type: str, prisma_client: PrismaClient | None = None) -> BaseFileStorageBackend: ...
class _ResultFileUploader(Protocol):
def __call__(
self,
file_data: Mapping[str, object],
target_storage: str,
target_model_names: Sequence[str],
purpose: OpenAIFilesPurpose,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None = None,
) -> Awaitable[OpenAIFileObject]: ...
@runtime_checkable
class _RouterCall(Protocol):
def __call__(self, **params: object) -> Awaitable[object]: ... # kwargs-ok: the request body is passed as keywords
def litellm_executed_provider_of(credentials: Mapping[str, object]) -> str | None:
explicit_provider: Final = credentials.get("custom_llm_provider")
provider: Final = (
explicit_provider if isinstance(explicit_provider, str) else _provider_of(credentials.get("model"))
)
return provider if provider in LITELLM_EXECUTED_BATCH_PROVIDERS else None
class _HttpGetter(Protocol):
async def get(
self, url: str, *, headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None
) -> httpx.Response: ...
class FilesApiProbe(Protocol):
async def __call__(self, api_base: str, api_key: str | None) -> bool: ...
class BodyRejection(Protocol):
def __call__(self, body: Mapping[str, object], /) -> str | None: ...
async def upstream_lacks_files_api(api_base: str, api_key: str | None, http_client: _HttpGetter | None = None) -> bool:
client: Final = http_client or get_async_httpx_client(llm_provider=LlmProviders.HOSTED_VLLM)
try:
response: Final = await client.get(
f"{api_base.rstrip('/')}/files",
headers=(
{"Authorization": f"Bearer {api_key}"} # mutable-ok: AsyncHTTPHandler.get wants a plain dict
if api_key
else None
),
timeout=_FILES_API_PROBE_TIMEOUT_SECONDS,
)
except httpx.HTTPError:
return False
return response.status_code == httpx.codes.NOT_FOUND
def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None:
model: Final = credentials.get("model")
api_base: Final = credentials.get("api_base")
api_key: Final = credentials.get("api_key")
if not isinstance(model, str):
return None
try:
_, _, resolved_api_key, resolved_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=provider,
api_base=api_base if isinstance(api_base, str) else None,
api_key=api_key if isinstance(api_key, str) else None,
)
except Exception: # noqa: BLE001 # get_llm_provider raises on a model it cannot map, which means nothing to probe
return None
return None if resolved_api_base is None else (resolved_api_base, resolved_api_key)
async def litellm_executed_provider_for(
credentials: Mapping[str, object], lacks_files_api: FilesApiProbe = upstream_lacks_files_api
) -> str | None:
provider: Final = litellm_executed_provider_of(credentials)
if provider is None:
return None
upstream: Final = _upstream_of(credentials, provider)
if upstream is None:
return None
return provider if await lacks_files_api(*upstream) else None
async def resolve_litellm_executed_provider(
llm_router: "Router",
model: str,
team_id: str | None,
lacks_files_api: FilesApiProbe = upstream_lacks_files_api,
) -> str | None:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model, team_id=team_id)
return None if credentials is None else await litellm_executed_provider_for(credentials, lacks_files_api)
def _provider_of(model: object) -> str | None:
if not isinstance(model, str):
return None
try:
return litellm.get_llm_provider(model=model)[1]
except Exception: # noqa: BLE001 # get_llm_provider raises on an unknown model, which means no provider
return None
def _validation_reason(error: ValidationError) -> str:
return "; ".join(
f"{'.'.join(str(part) for part in item['loc'])}: {item['msg']}" if item["loc"] else item["msg"]
for item in error.errors()
)
def _accept_every_body(_body: Mapping[str, object]) -> str | None:
return None
def _parse_line(
line_number: int, raw: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection
) -> BatchInputLine | InvalidBatchInput:
try:
line: Final = BatchInputLine.model_validate_json(raw)
except ValidationError as e:
return InvalidBatchInput(line_number, _validation_reason(e))
if line.url != endpoint:
return InvalidBatchInput(line_number, f"url {line.url!r} does not match the batch endpoint {endpoint!r}")
if line.body.get("stream"):
return InvalidBatchInput(line_number, "streaming requests are not supported in a batch")
rejection: Final = reject_body(line.body)
if rejection is not None:
return InvalidBatchInput(line_number, rejection)
return line
def parse_batch_input(
content: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection = _accept_every_body
) -> tuple[BatchInputLine, ...] | InvalidBatchInput:
raw_lines: Final = tuple((number, raw) for number, raw in enumerate(content.splitlines(), start=1) if raw.strip())
if not raw_lines:
return InvalidBatchInput(None, "the input file has no requests")
parsed: Final = tuple(_parse_line(number, raw, endpoint, reject_body) for number, raw in raw_lines)
first_invalid: Final = next((item for item in parsed if isinstance(item, InvalidBatchInput)), None)
if first_invalid is not None:
return first_invalid
lines: Final = tuple(item for item in parsed if isinstance(item, BatchInputLine))
custom_ids: Final = sorted(line.custom_id for line in lines)
duplicate: Final = next((first for first, second in pairwise(custom_ids) if first == second), None)
if duplicate is not None:
return InvalidBatchInput(None, f"custom_id {duplicate!r} is used more than once")
return lines
def batch_error(status_code: int, message: str) -> ProxyException:
error_type: Final = "invalid_request_error" if status_code < 500 else ProxyErrorTypes.internal_server_error.value
return ProxyException(message=message, type=error_type, param=None, code=status_code)
def _validate_endpoint(endpoint: object) -> BatchEndpoint:
try:
return _BATCH_ENDPOINT_ADAPTER.validate_python(endpoint)
except ValidationError:
raise batch_error(400, f"endpoint {endpoint!r} is not supported for a LiteLLM-executed batch")
def _status_code_of(error: Exception) -> int:
status_code: Final[object] = getattr(error, "status_code", None)
return status_code if isinstance(status_code, int) else 500
def _error_body(error: Exception) -> _ErrorBody:
body: Final[_ErrorBody] = {
"error": {"message": str(error), "type": type(error).__name__, "param": None, "code": None}
}
return body
def _line_response(outcome: RowOutcome | ExpiredRow) -> _ResultResponse | None:
if isinstance(outcome, ExpiredRow):
return None
response: Final[_ResultResponse] = {
"status_code": outcome.status_code,
"request_id": f"req_{uuid_module.uuid4().hex[:24]}",
"body": outcome.body,
}
return response
def _line_error(outcome: RowOutcome | ExpiredRow) -> _LineError | None:
if isinstance(outcome, RowOutcome):
return None
error: Final[_LineError] = {"code": "batch_expired", "message": _EXPIRED_MESSAGE}
return error
def _result_line(outcome: RowOutcome | ExpiredRow) -> _ResultLine:
line: Final[_ResultLine] = {
"id": f"batch_req_{uuid_module.uuid4().hex[:24]}",
"custom_id": outcome.custom_id,
"response": _line_response(outcome),
"error": _line_error(outcome),
}
return line
def _dump(response: object) -> Mapping[str, object]:
if isinstance(response, BaseModel):
return response.model_dump(mode="json")
raise TypeError(f"Batch rows must return a single response object, got {type(response).__name__}")
def _resolve_transition(current_status: str, requested: BatchStatus) -> BatchStatus:
if current_status != "cancelling":
return requested
return _CANCELLING_TRANSITIONS.get(requested, requested)
def executed_batch_runner_lost(status: str, updated_at: datetime) -> bool:
if status in TERMINAL_BATCH_STATUSES:
return False
return (datetime.now(timezone.utc) - updated_at).total_seconds() > _STALE_AFTER_SECONDS
class _StopWatch:
def __init__(self, load_status: Callable[[], Awaitable[str | None]], interval_seconds: float) -> None:
self._load_status = load_status
self._interval_seconds = interval_seconds
self._checked_at = float("-inf")
self._stopped = False
async def stopped(self) -> bool:
if self._stopped:
return True
now: Final = time.monotonic()
if now - self._checked_at < self._interval_seconds:
return False
self._checked_at = now
self._stopped = await self._load_status() in _STOP_STATUSES
return self._stopped
class LiteLLMExecutedBatchRunner:
def __init__(
self,
llm_router: "Router",
prisma_client: PrismaClient,
managed_files: ManagedBatchStore,
batches: ManagedBatchRepository,
proxy_logging_obj: ProxyLogging,
general_settings: Mapping[str, object],
concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY,
heartbeat_seconds: float = _HEARTBEAT_SECONDS,
completion_window_seconds: float = _COMPLETION_WINDOW_SECONDS,
storage_backend_factory: _StorageBackendFactory = get_storage_backend,
upload_result_file: _ResultFileUploader = StorageBackendFileService.upload_file_to_storage_backend,
) -> None:
self.llm_router = llm_router
self.prisma_client = prisma_client
self.managed_files = managed_files
self.batches = batches
self.proxy_logging_obj = proxy_logging_obj
self.general_settings = general_settings
self.concurrency = concurrency
self.heartbeat_seconds = heartbeat_seconds
self.completion_window_seconds = completion_window_seconds
self.storage_backend_factory = storage_backend_factory
self.upload_result_file = upload_result_file
async def create(
self,
create_request: LiteLLMBatchCreateRequest,
unified_input_file_id: str,
model: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None,
) -> LiteLLMBatch:
endpoint: Final = _validate_endpoint(create_request.get("endpoint"))
content: Final = await self._download_input(unified_input_file_id, user_api_key_dict)
parsed: Final = parse_batch_input(content, endpoint, self._body_rejection(model))
if isinstance(parsed, InvalidBatchInput):
raise batch_error(400, f"Invalid batch input file: {parsed.describe()}")
llm_batch_id: Final = f"{LITELLM_EXECUTED_BATCH_ID_PREFIX}{uuid_module.uuid4().hex}"
model_id: Final = next(iter(self.llm_router.get_model_ids(model_name=model)), model)
unified_batch_id: Final = self.managed_files.get_unified_batch_id(batch_id=llm_batch_id, model_id=model_id)
now: Final = time.time()
created_at: Final = int(now)
batch: Final = LiteLLMBatch(
id=unified_batch_id,
object="batch",
endpoint=endpoint,
input_file_id=unified_input_file_id,
completion_window="24h",
status="validating",
created_at=created_at,
expires_at=created_at + int(self.completion_window_seconds),
metadata=create_request.get("metadata"),
model=model,
request_counts=BatchRequestCounts(completed=0, failed=0, total=len(parsed)),
)
await self.managed_files.store_unified_object_id(
unified_object_id=unified_batch_id,
file_object=batch,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
model_object_id=llm_batch_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=request_tags,
persist_attribution=True,
batch_processed=True,
)
_record_batch_created(model, provider, user_api_key_dict)
run: Final = _BatchRun(
unified_batch_id=unified_batch_id,
llm_batch_id=llm_batch_id,
model=model,
endpoint=endpoint,
lines=parsed,
user_api_key_dict=user_api_key_dict,
request_tags=tuple(request_tags or ()),
deadline=now + self.completion_window_seconds,
)
task: Final = asyncio.create_task(self._run(run))
_RUNNING_BATCHES.add(task)
task.add_done_callback(_RUNNING_BATCHES.discard)
return batch
async def cancel(self, unified_batch_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch:
current: Final = await self.batches.load_batch(unified_batch_id)
if current is None:
raise batch_error(404, f"Batch {unified_batch_id} not found")
if current.status in TERMINAL_BATCH_STATUSES:
raise batch_error(400, f"Cannot cancel a batch with status '{current.status}'")
if current.status == "cancelling":
return current
cancelling: Final = current.model_copy(
update=MappingProxyType({"status": "cancelling", "cancelling_at": int(time.time())})
)
unchanged: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {"status": current.status}
if await self.batches.compare_and_set(cancelling, unchanged, user_api_key_dict.user_id):
return cancelling
return await self.cancel(unified_batch_id, user_api_key_dict)
async def fail_abandoned(self, batch: LiteLLMBatch, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch:
error: Final = BatchError(message=_RUNNER_LOST_MESSAGE, code="runner_lost")
errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list
failed: Final = batch.model_copy(
update=MappingProxyType({"status": "failed", "failed_at": int(time.time()), "errors": errors})
)
untouched: Final[prisma_types.DateTimeFilter] = {
"lt": datetime.now(timezone.utc) - timedelta(seconds=_STALE_AFTER_SECONDS)
}
still_abandoned: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {
"status": batch.status,
"updated_at": untouched,
}
if await self.batches.compare_and_set(failed, still_abandoned, user_api_key_dict.user_id):
return failed
return await self.batches.load_batch(batch.id) or batch
def _body_rejection(self, model: str) -> BodyRejection:
def reject(body: Mapping[str, object]) -> str | None:
try:
is_request_body_safe(
request_body=dict(body), # mutable-ok: is_request_body_safe takes a dict
general_settings=dict(self.general_settings), # mutable-ok: is_request_body_safe takes a dict
llm_router=self.llm_router,
model=model,
)
except ValueError as e:
return str(e)
return None
return reject
async def _download_input(self, unified_input_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bytes:
stored: Final = await self.managed_files.get_unified_file_id(
unified_input_file_id, litellm_parent_otel_span=user_api_key_dict.parent_otel_span
)
if stored is None or not stored.storage_backend or not stored.storage_url:
raise batch_error(
400,
f"LiteLLM does not hold the content of input file {unified_input_file_id}: "
f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}",
)
try:
backend: Final = self.storage_backend_factory(stored.storage_backend, prisma_client=self.prisma_client)
return await backend.download_file(stored.storage_url)
except ValueError as e:
raise batch_error(400, str(e))
async def _run(self, run: _BatchRun) -> None:
heartbeat: Final = asyncio.create_task(self._heartbeat(run))
try:
await self._execute(run)
except Exception as e: # noqa: BLE001 # whatever fails, the batch must end up marked failed
verbose_proxy_logger.exception("LiteLLM-executed batch %s failed: %s", run.unified_batch_id, e)
error: Final = BatchError(message=str(e), code="internal_error")
errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list
try:
await self._advance(run, "failed", MappingProxyType({"errors": errors}))
except Exception as advance_error: # noqa: BLE001 # a failed status write is logged, never raised
verbose_proxy_logger.exception(
"LiteLLM-executed batch %s could not be marked failed: %s", run.unified_batch_id, advance_error
)
finally:
heartbeat.cancel()
async def _heartbeat(self, run: _BatchRun) -> None:
while True:
await asyncio.sleep(self.heartbeat_seconds)
try:
await self._touch(run)
except Exception as e: # noqa: BLE001 # a missed beat is logged and the next one retries
verbose_proxy_logger.warning("LiteLLM-executed batch %s heartbeat failed: %s", run.unified_batch_id, e)
async def _touch(self, run: _BatchRun) -> None:
await self.batches.touch(run.unified_batch_id, run.user_api_key_dict.user_id)
async def _execute(self, run: _BatchRun) -> None:
await self._advance(run, "in_progress")
watch: Final = _StopWatch(lambda: self.batches.load_status(run.unified_batch_id), _CANCEL_POLL_SECONDS)
semaphore: Final = asyncio.Semaphore(self.concurrency)
results: Final = await asyncio.gather(*(self._run_row(run, line, watch, semaphore) for line in run.lines))
outcomes: Final = tuple(outcome for outcome in results if outcome is not None)
if await self._advance(run, "finalizing") is None:
return
succeeded: Final = tuple(
outcome for outcome in outcomes if isinstance(outcome, RowOutcome) and outcome.succeeded
)
failed: Final = tuple(
outcome for outcome in outcomes if isinstance(outcome, ExpiredRow) or not outcome.succeeded
)
output_file_id: Final = await self._upload_results(run, "output", succeeded)
error_file_id: Final = await self._upload_results(run, "error", failed)
request_counts: Final = BatchRequestCounts(completed=len(succeeded), failed=len(failed), total=len(run.lines))
final_status: Final[BatchStatus] = (
"expired" if any(isinstance(outcome, ExpiredRow) for outcome in outcomes) else "completed"
)
await self._advance(
run,
final_status,
MappingProxyType(
{"output_file_id": output_file_id, "error_file_id": error_file_id, "request_counts": request_counts}
),
)
async def _run_row(
self, run: _BatchRun, line: BatchInputLine, watch: _StopWatch, semaphore: asyncio.Semaphore
) -> RowOutcome | ExpiredRow | None:
async with semaphore:
if await watch.stopped():
return None
remaining: Final = run.deadline - time.time()
if remaining <= 0:
return ExpiredRow(custom_id=line.custom_id)
try:
return await asyncio.wait_for(self._row_outcome(run, line), timeout=remaining)
except asyncio.TimeoutError:
return ExpiredRow(custom_id=line.custom_id)
async def _row_outcome(self, run: _BatchRun, line: BatchInputLine) -> RowOutcome:
try:
body: Final = await self._dispatch(run, line)
except Exception as e: # noqa: BLE001 # a provider error becomes the row's error line, never a crashed batch
return RowOutcome(
custom_id=line.custom_id, status_code=_status_code_of(e), body=_error_body(e), succeeded=False
)
return RowOutcome(custom_id=line.custom_id, status_code=200, body=body, succeeded=True)
async def _dispatch(self, run: _BatchRun, line: BatchInputLine) -> Mapping[str, object]:
params: Final = MappingProxyType(
{**line.body, "model": run.model, "metadata": self._row_metadata(run), "disable_fallbacks": True}
)
return _dump(await self._router_call(run.endpoint)(**params))
def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall:
method: Final[object] = getattr(self.llm_router, _ROUTER_METHODS[endpoint], None)
if not isinstance(method, _RouterCall):
raise TypeError(f"the router has no callable for {endpoint}")
return method
def _row_metadata(self, run: _BatchRun) -> dict[str, object]: # mutable-ok: router updates metadata in place
return { # mutable-ok: the router updates request metadata in place
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(run.user_api_key_dict),
"user_api_key": run.user_api_key_dict.api_key,
"user_api_end_user_max_budget": run.user_api_key_dict.end_user_max_budget,
"tags": list(run.request_tags), # mutable-ok: litellm types request tags as a list
"batch_id": run.unified_batch_id,
}
async def _upload_results(
self, run: _BatchRun, kind: Literal["output", "error"], outcomes: Sequence[RowOutcome | ExpiredRow]
) -> str | None:
if not outcomes:
return None
content: Final = "".join(f"{json.dumps(_result_line(outcome))}\n" for outcome in outcomes).encode()
file_data: Final[ExtractedFileData] = {
"filename": f"{run.llm_batch_id}_{kind}.jsonl",
"content": content,
"content_type": "application/jsonl",
"headers": _NO_HEADERS,
}
file_object: Final = await self.upload_result_file(
file_data=file_data,
target_storage=LITELLM_DB_STORAGE_BACKEND_NAME,
target_model_names=(run.model,),
purpose="batch_output",
proxy_logging_obj=self.proxy_logging_obj,
user_api_key_dict=run.user_api_key_dict,
prisma_client=self.prisma_client,
)
return file_object.id
async def _advance(
self, run: _BatchRun, requested: BatchStatus, fields: Mapping[str, object] = _NO_FIELDS
) -> BatchStatus | None:
current: Final = await self.batches.load_batch(run.unified_batch_id)
if current is None:
raise RuntimeError(f"Batch {run.unified_batch_id} is no longer stored")
if current.status in TERMINAL_BATCH_STATUSES:
return None
status: Final = _resolve_transition(current.status, requested)
updated: Final = current.model_copy(
update=MappingProxyType({**fields, "status": status, f"{status}_at": int(time.time())})
)
unchanged: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {"status": current.status}
if await self.batches.compare_and_set(updated, unchanged, run.user_api_key_dict.user_id):
return status
return await self._advance(run, requested, fields)
def _record_batch_created(model: str, provider: str, user_api_key_dict: UserAPIKeyAuth) -> None:
prometheus_logger: Final = PrometheusLogger.get_instance()
if prometheus_logger is None:
return
prometheus_logger.record_managed_batch_created(
model=model,
api_provider=provider,
user=user_api_key_dict.user_id or "",
user_email=user_api_key_dict.user_email or "",
api_key_alias=user_api_key_dict.key_alias or "",
)

View file

@ -60,9 +60,7 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def rejected_writes(self, incoming: Mapping[str, JsonValue]) -> tuple[str, ...]:
return tuple(
sorted(
key for key, value in incoming.items() if self.owned_by_config(key) and value != self._yaml_values[key]
)
sorted(key for key, value in incoming.items() if self.owned_by_config(key) and value != self.get(key))
)
def shadowed_db_keys(self) -> tuple[str, ...]:
@ -74,8 +72,13 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def apply_db_row(self, row: DbRow, db_row: Mapping[str, JsonValue]) -> None:
previous_row: Final = self._database_rows.get(row, _EMPTY_VALUES)
changed: Final = frozenset(
key
for key in (*previous_row, *db_row)
if previous_row.get(key, ABSENT) != db_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare
)
self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))})
self._clear_runtime_keys(frozenset((*previous_row, *db_row)))
self._clear_runtime_keys(changed)
def resolved(self) -> Mapping[str, JsonValue]:
return MappingProxyType(dict(self))
@ -130,6 +133,9 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def __len__(self) -> int:
return sum(1 for _ in self)
def __bool__(self) -> bool:
return any(True for _ in self)
def _clear_runtime(self) -> None:
self._runtime_values = _EMPTY_VALUES
self._deleted_runtime_keys = frozenset()
@ -160,7 +166,12 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def _db_value_is_shadowed(self, key: str) -> bool:
db_value: Final = self._db_value(key)
return not isinstance(db_value, Absent) and db_value is not None and db_value != self.get(key)
return (
not isinstance(db_value, Absent)
and db_value is not None
and db_value != self.get(key)
and db_value != self.config_value(key)
)
def _resolution_for(self, key: str) -> Resolved:
yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT)

View file

@ -364,6 +364,11 @@ async def update_coordination_redis_settings(
settings: Final = _merge_over_saved(request.settings, saved_settings or {})
_validated_params(settings)
from litellm.proxy.proxy_server import proxy_config
proxy_config.reject_config_owned_writes(
section_name=_GENERAL_SETTINGS_PARAM_NAME, changed_keys={_COORDINATION_REDIS_KEY: settings}
)
general_settings: Final = await _read_general_settings()
before_settings: Final = general_settings.get(_COORDINATION_REDIS_KEY)
action: Final[AUDIT_ACTIONS] = "updated" if isinstance(before_settings, dict) else "created"

View file

@ -38,6 +38,7 @@ if TYPE_CHECKING:
FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500
BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create"
LITELLM_EXECUTED_BATCH_ID_PREFIX: Final = "litellm_batch_"
def validate_file_list_limit(limit: int | None) -> None:
@ -179,6 +180,11 @@ def get_batch_id_from_unified_batch_id(file_id: str) -> str:
return re.split(r"[;,]", batch_id, maxsplit=1)[0]
def is_litellm_executed_batch(decoded_unified_batch_id: str) -> bool:
_, marker, batch_id = decoded_unified_batch_id.partition("llm_batch_id:")
return bool(marker) and batch_id.startswith(LITELLM_EXECUTED_BATCH_ID_PREFIX)
def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", "batch"] = "file") -> str:
"""
Encode a file/batch ID with model routing information.

View file

@ -7,7 +7,7 @@
import asyncio
import traceback
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Any, BinaryIO, Final, TypedDict, cast, get_args
import httpx
@ -32,10 +32,15 @@ from litellm.litellm_core_utils.cloud_storage_security import (
is_managed_cloud_storage_uri,
)
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import build_list_page
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
litellm_executed_provider_of,
resolve_litellm_executed_provider,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
@ -67,6 +72,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
add_internal_model_credentials,
apply_team_provider_credentials,
authorize_model_for_key,
encode_file_id_with_model,
extract_file_creation_params,
get_authorized_credentials_for_model,
@ -86,7 +92,7 @@ from litellm.proxy.openai_files_endpoints.general_upload_validation import (
coerce_optional_str_list_setting,
raise_upload_validation_failure,
)
from litellm.proxy.utils import ProxyLogging, is_known_model
from litellm.proxy.utils import PrismaClient, ProxyLogging, is_known_model
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.router import Router
from litellm.types.llms.openai import (
@ -99,6 +105,64 @@ from litellm.types.llms.openai import (
router: Final = APIRouter()
def _names_a_litellm_executed_provider(llm_router: Router, candidate: str, team_id: str | None) -> bool:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=candidate, team_id=team_id)
return credentials is not None and litellm_executed_provider_of(credentials) is not None
async def _litellm_executed_batch_input_model(
llm_router: Router | None,
purpose: OpenAIFilesPurpose,
model: str | None,
target_model_names_list: Sequence[str],
user_api_key_dict: UserAPIKeyAuth,
explicit_storage: str | None,
) -> str | None:
if llm_router is None:
return None
candidates: Final = (model,) if model is not None else tuple(target_model_names_list)
team_id: Final = user_api_key_dict.team_id
await asyncio.gather(
*(
authorize_model_for_key(model_id=candidate, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
for candidate in candidates
if _names_a_litellm_executed_provider(llm_router, candidate, team_id)
)
)
if explicit_storage is not None:
return None
providers: Final = await asyncio.gather(
*(resolve_litellm_executed_provider(llm_router, candidate, team_id) for candidate in candidates)
)
executed: Final = tuple(
candidate for candidate, provider in zip(candidates, providers, strict=True) if provider is not None
)
if not executed:
return None
if purpose != "batch":
raise ProxyException(
message=(
f"The server behind {', '.join(executed)} has no Files API, so LiteLLM keeps only batch input "
f"files for it and runs the batch itself: upload with purpose=batch; got purpose={purpose}"
),
type="invalid_request_error",
param="purpose",
code=400,
)
if len(candidates) == 1:
return executed[0]
raise ProxyException(
message=(
f"LiteLLM runs batches for {', '.join(executed)} itself and keeps their input files, so a batch "
f"input file can target only that one model; got target_model_names={', '.join(candidates)}"
),
type="invalid_request_error",
param="target_model_names",
code=400,
)
_MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None)
_LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject])
@ -244,30 +308,41 @@ async def route_create_file(
5. Else -> use custom_llm_provider with files_settings
"""
# Handle custom storage backend
if target_storage and target_storage != "default":
explicit_storage: Final = target_storage if target_storage and target_storage != "default" else None
if explicit_storage == LITELLM_DB_STORAGE_BACKEND_NAME:
raise ProxyException(
message=(
f"target_storage={LITELLM_DB_STORAGE_BACKEND_NAME} is not a storage a caller can pick: LiteLLM "
"chooses it on its own for the batch input files of a model whose batches it runs itself, so "
"upload with purpose=batch and name that model instead of target_storage"
),
type="invalid_request_error",
param="target_storage",
code=400,
)
executed_model: Final = await _litellm_executed_batch_input_model(
llm_router, purpose, model, target_model_names_list, user_api_key_dict, explicit_storage
)
storage: Final = explicit_storage or (LITELLM_DB_STORAGE_BACKEND_NAME if executed_model is not None else None)
if storage is not None:
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
from litellm.proxy.openai_files_endpoints.storage_backend_service import (
StorageBackendFileService,
)
from litellm.proxy.proxy_server import prisma_client
# Extract file data
file_data: Final = extract_file_data(cast(Any, _create_file_request.get("file")))
# Use storage backend service to handle upload
file_object: Final = await StorageBackendFileService.upload_file_to_storage_backend(
file_data=file_data,
target_storage=target_storage,
target_model_names=target_model_names_list,
return await StorageBackendFileService.upload_file_to_storage_backend(
file_data=extract_file_data(cast(Any, _create_file_request.get("file"))),
target_storage=storage,
target_model_names=(executed_model,) if executed_model is not None else target_model_names_list,
purpose=purpose,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
return file_object
# NEW: Handle model-based routing (no DB required)
if model is not None:
# Get credentials from model_list via router
@ -848,7 +923,7 @@ async def get_file_content(
# Check if file is stored in a storage backend (check DB)
if hasattr(managed_files_obj, "prisma_client") and getattr(managed_files_obj, "prisma_client", None):
prisma_client: Final = getattr(managed_files_obj, "prisma_client")
prisma_client: Final[PrismaClient] = getattr(managed_files_obj, "prisma_client")
db_file: Final = await ManagedFileRepository(prisma_client).table.find_first(
where={"unified_file_id": file_id}
)
@ -863,7 +938,7 @@ async def get_file_content(
try:
# Get storage backend (uses same env vars as callback)
storage_backend: Final = get_storage_backend(storage_backend_name)
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=prisma_client)
file_content: Final = await storage_backend.download_file(storage_url)
# Return file content

View file

@ -7,15 +7,16 @@ storage backends (e.g., Azure Blob Storage) and managing associated metadata.
import base64
import time
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Any, Final, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid as uuid_module
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.types.llms.openai import OpenAIFileObject, OpenAIFilesPurpose
from litellm.types.utils import SpecialEnums
@ -35,21 +36,23 @@ class StorageBackendFileService:
async def upload_file_to_storage_backend(
file_data: Mapping[str, Any],
target_storage: str,
target_model_names: list[str],
target_model_names: Sequence[str],
purpose: OpenAIFilesPurpose,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None = None,
) -> OpenAIFileObject:
"""
Upload a file to a storage backend and create a file object.
Args:
file_data: File data dictionary from extract_file_data()
target_storage: Storage backend name (e.g., "azure_storage")
target_storage: Storage backend name (e.g., "azure_storage", "litellm_db")
target_model_names: List of model names for managed files
purpose: File purpose (e.g., "user_data", "batch")
proxy_logging_obj: Proxy logging object for accessing hooks
user_api_key_dict: User API key authentication data
prisma_client: The proxy's database client, required by the "litellm_db" backend
Returns:
OpenAIFileObject: Created file object with storage metadata
@ -59,7 +62,7 @@ class StorageBackendFileService:
"""
# Get storage backend instance
try:
storage_backend: Final = get_storage_backend(target_storage)
storage_backend: Final = get_storage_backend(target_storage, prisma_client=prisma_client)
except ValueError as e:
raise ProxyException(
message=str(e),
@ -103,8 +106,9 @@ class StorageBackendFileService:
storage_url=storage_url,
)
# Store in managed files if target_model_names provided
if target_model_names:
if not target_model_names:
return file_object
try:
await StorageBackendFileService._store_in_managed_files(
file_object=file_object,
file_data=file_data,
@ -114,9 +118,25 @@ class StorageBackendFileService:
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
)
except Exception:
await StorageBackendFileService._discard_orphaned_content(storage_backend, storage_url, target_storage)
raise
return file_object
@staticmethod
async def _discard_orphaned_content(
storage_backend: BaseFileStorageBackend, storage_url: str, target_storage: str
) -> None:
try:
await storage_backend.delete_file(storage_url)
except Exception as e: # noqa: BLE001 # the metadata failure is what surfaces; a failed cleanup is only logged
verbose_proxy_logger.warning(
"Could not delete orphaned content at %s on %s after its metadata write failed: %s",
storage_url,
target_storage,
e,
)
@staticmethod
def _create_file_object_with_storage_metadata(
file_content: bytes,
@ -164,7 +184,7 @@ class StorageBackendFileService:
@staticmethod
def _create_unified_file_id(
file_type: str,
target_model_names: list[str],
target_model_names: Sequence[str],
file_id: str,
) -> str:
"""
@ -194,7 +214,7 @@ class StorageBackendFileService:
async def _store_in_managed_files(
file_object: OpenAIFileObject,
file_data: Mapping[str, Any],
target_model_names: list[str],
target_model_names: Sequence[str],
target_storage: str,
storage_url: str,
proxy_logging_obj: ProxyLogging,

View file

@ -589,6 +589,11 @@ class ProxyInitializationHelpers:
gunicorn_options["certfile"] = ssl_certfile_path
gunicorn_options["keyfile"] = ssl_keyfile_path
# The master preloads the app and then forks every worker, so native routes are
# forbidden in it: their runtime threads would not survive the fork.
from litellm.rust_bridge.fork_guard import reserve_process_for_forking
reserve_process_for_forking("the gunicorn master")
start_query_engine_reaper()
StandaloneApplication(app=app, options=gunicorn_options).run() # Run gunicorn

View file

@ -5134,6 +5134,15 @@ class ProxyConfig:
store.apply_db_row(cast(DbRow, section_name), wrote_section)
await invalidate_config_param(section_name)
def reject_config_owned_deletes(self, *, section_name: str, keys: tuple[str, ...]) -> None:
"""Refuse a delete of a setting the config file owns; unlike a write, the value never makes it allowed."""
store: Final = self._settings_stores.get(cast(Section, section_name))
if store is None:
return
owned: Final = tuple(sorted(key for key in keys if store.owned_by_config(key)))
if owned:
self._raise_config_owned(section_name=section_name, rejected=owned, store=store)
def reject_config_owned_writes(self, *, section_name: str, changed_keys: Mapping[str, JsonValue]) -> None:
"""Refuse a write to a setting the config file owns, rather than storing a value that never applies."""
store: Final = self._settings_stores.get(cast(Section, section_name))
@ -5142,6 +5151,9 @@ class ProxyConfig:
rejected: Final = store.rejected_writes(changed_keys)
if not rejected:
return
self._raise_config_owned(section_name=section_name, rejected=rejected, store=store)
def _raise_config_owned(self, *, section_name: str, rejected: tuple[str, ...], store: SettingsStore) -> None:
subject: Final = (
f"key '{rejected[0]}' is" if len(rejected) == 1 else f"keys {', '.join(repr(key) for key in rejected)} are"
)
@ -9764,10 +9776,12 @@ class ProxyStartupEvent:
user_api_key_cache: UserApiKeyCache,
):
"""Initialize JWT auth on startup"""
if general_settings.get("litellm_jwtauth", None) is not None:
for k, v in general_settings["litellm_jwtauth"].items():
if isinstance(v, str) and v.startswith("os.environ/"):
general_settings["litellm_jwtauth"][k] = get_secret(v)
declared_jwtauth: Final = general_settings.get("litellm_jwtauth", None)
if declared_jwtauth is not None:
resolved_jwtauth: Final = {
key: (get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value)
for key, value in declared_jwtauth.items()
}
# ``user_config_file_path`` is set by ``ProxyConfig._get_config_from_file``
# during startup. Threading it through lets an operator-
# configured ``custom_validate: s3://...`` resolve through
@ -9775,7 +9789,7 @@ class ProxyStartupEvent:
# file context) hit the gate and refuse remote loads.
litellm_jwtauth = LiteLLM_JWTAuth(
config_file_path=user_config_file_path,
**general_settings["litellm_jwtauth"],
**resolved_jwtauth,
)
else:
litellm_jwtauth = LiteLLM_JWTAuth()
@ -17745,9 +17759,12 @@ async def get_config_general_settings(
detail={"error": f"Field name={field_name} is not set"},
)
declared: Final = (
settings.config_value(field_name) if settings.owned_by_config(field_name) else settings[field_name]
)
field_value = _redact_general_setting_value(
field_name,
settings[field_name],
declared,
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN,
)
if field_name == "plugins" and isinstance(field_value, list):
@ -18121,6 +18138,8 @@ async def delete_config_general_settings(
detail={"error": f"Invalid field={data.field_name} passed in."},
)
proxy_config.reject_config_owned_deletes(section_name="general_settings", keys=(data.field_name,))
## get general settings from db
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}

View file

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID

View file

@ -470,12 +470,16 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje
class _UpstreamStreamBoundary(Generic[_T]):
__slots__ = ("_upstream", "failure")
__slots__ = ("_source", "_upstream", "failure")
def __init__(self, upstream: AsyncIterable[_T]) -> None:
self._source: Final = upstream
self._upstream: Final = upstream.__aiter__()
self.failure: BaseException | None = None
def __getattr__(self, name: str) -> object:
return getattr(self._source, name)
def __aiter__(self) -> "_UpstreamStreamBoundary[_T]":
return self

View file

@ -0,0 +1,48 @@
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from litellm.repositories.table_repositories import PrismaTableRepository
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from prisma import models as prisma_models
def _batch_of(blob: object) -> LiteLLMBatch:
return LiteLLMBatch.model_validate_json(blob) if isinstance(blob, str) else LiteLLMBatch.model_validate(blob)
class ManagedBatchRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedObjectTable"]):
table_name = "litellm_managedobjecttable"
async def load_batch(self, unified_batch_id: str) -> LiteLLMBatch | None:
row: Final = await self._find_row(unified_batch_id)
return None if row is None or not row.file_object else _batch_of(row.file_object)
async def load_status(self, unified_batch_id: str) -> str | None:
row: Final = await self._find_row(unified_batch_id)
return row.status if row is not None else None
async def compare_and_set(
self, batch: LiteLLMBatch, unchanged: Mapping[str, object], updated_by: str | None
) -> bool:
updated_rows: Final = await self.table.update_many(
where={"unified_object_id": batch.id, **unchanged}, # mutable-ok: prisma filters are plain dicts
data={ # mutable-ok: prisma payloads are plain dicts
"file_object": batch.model_dump_json(),
"status": batch.status,
"updated_by": updated_by,
},
)
return updated_rows > 0
async def touch(self, unified_batch_id: str, updated_by: str | None) -> None:
await self.table.update_many(
where={"unified_object_id": unified_batch_id}, # mutable-ok: prisma filters are plain dicts
data={"updated_by": updated_by}, # mutable-ok: prisma payloads are plain dicts
)
async def _find_row(self, unified_batch_id: str) -> "prisma_models.LiteLLM_ManagedObjectTable | None":
return await self.table.find_first(
where={"unified_object_id": unified_batch_id} # mutable-ok: prisma filters are plain dicts
)

View file

@ -0,0 +1,32 @@
from typing import TYPE_CHECKING, Final
from litellm.repositories.table_repositories import PrismaTableRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
class ManagedFileContentRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedFileContentTable"]):
table_name = "litellm_managedfilecontenttable"
async def store(self, content: bytes) -> str:
from prisma import Base64
row: Final = await self.table.create(
data={"content": Base64.encode(content)} # mutable-ok: prisma payloads are plain dicts
)
return row.id
async def load(self, row_id: str) -> bytes | None:
row: Final[prisma_models.LiteLLM_ManagedFileContentTable | None] = await self.table.find_unique(
where={"id": row_id} # mutable-ok: prisma filters are plain dicts
)
return None if row is None else row.content.decode()
async def delete(self, row_id: str) -> None:
from prisma.errors import RecordNotFoundError
try:
await self.table.delete(where={"id": row_id}) # mutable-ok: prisma filters are plain dicts
except RecordNotFoundError:
return

View file

@ -9,6 +9,8 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMe
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
class ForkedAfterNativeRuntimeStarted(RuntimeError): ...
class ProcessReservedForForking(RuntimeError): ...
def ocr(
request: LiteLLMOcrRequest,
@ -101,8 +103,12 @@ class TokenCounter:
def acount_request(self, body: bytes) -> Future[dict[str, object]]: ...
def gil_stats() -> dict[str, int]: ...
def process_state_started() -> bool: ...
def reserve_process_for_forking() -> None: ...
__all__ = [
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",
"RustUpstreamError",
@ -116,5 +122,7 @@ __all__ = [
"gil_stats",
"messages",
"ocr",
"process_state_started",
"reserve_process_for_forking",
"transcription",
]

View file

@ -0,0 +1,47 @@
"""Fork safety of the Rust extension.
Its runtime threads do not survive ``fork``, so a child forked after the first native call
cannot run native routes: it raises ``ForkedAfterNativeRuntimeStarted`` instead of hanging.
Fork before the first native call, or start workers with ``spawn`` / ``forkserver``.
A process whose job is to fork workers (the gunicorn master under ``preload``) reserves itself:
from then on any native route called in it raises ``ProcessReservedForForking`` at the call
site, so the runtime can never start there. Workers forked from it are unaffected.
"""
from __future__ import annotations
from typing import Final
from litellm.rust_bridge.loader import get_native_bridge
class NativeStateStartedBeforeFork(RuntimeError):
pass
class _NeverRaised(RuntimeError):
"""Stands in for a native exception when the extension is unavailable or predates it."""
_native: Final = get_native_bridge()
ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr(
_native, "ForkedAfterNativeRuntimeStarted", _NeverRaised
)
ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised)
def reserve_process_for_forking(where: str) -> None:
"""Forbid native routes in this process. Raises if one already ran here."""
native: Final = get_native_bridge()
reserve: Final = getattr(native, "reserve_process_for_forking", None)
if not callable(reserve):
return
try:
reserve()
except RuntimeError as error:
raise NativeStateStartedBeforeFork(
f"The LiteLLM Rust extension already ran a native route in {where}, and its runtime "
"threads do not survive fork(). Move the native call (warm-up, health check, "
"import-time initialization) into the worker, after the fork."
) from error

View file

@ -512,6 +512,7 @@ class CreateBatchRequest(TypedDict, total=False):
class LiteLLMBatchCreateRequest(CreateBatchRequest, total=False):
model: str
disable_fallbacks: ReadOnly[bool]
class RetrieveBatchRequest(TypedDict, total=False):

View file

@ -4178,6 +4178,8 @@ FILE_CONTENT_STREAMING_PROVIDERS: Final[frozenset[str]] = frozenset(
{*OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, LlmProviders.VERTEX_AI.value}
)
LITELLM_EXECUTED_BATCH_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.HOSTED_VLLM.value})
ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "litellm_proxy", "vertex_ai"]
LIST_BATCHES_SUPPORTED_PROVIDERS: Final[frozenset[str]] = frozenset(get_args(ListBatchesSupportedProvider))

View file

@ -71,8 +71,8 @@ proxy = [
"mcp>=2.2.0,<3",
"httpx2>=2.5.0,<3",
"pydantic>=2.12.0,<3",
"litellm-proxy-extras==0.4.99",
"litellm-enterprise==0.1.68",
"litellm-proxy-extras==0.4.100",
"litellm-enterprise==0.1.69",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",

View file

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID

View file

@ -60,6 +60,13 @@ TQ007 A module global that a conftest saves before every test and restores aft
names are read from the keys the conftest assigns directly and from whatever the
save loop iterates, including a module-level tuple or dict it names rather than
spells out.
TQ009 A child interpreter spawned as `subprocess.run([sys.executable, ...])` without
`-I`/`-P` as its first flag. Without isolation the child's sys.path leads with
the working directory, so a source checkout shadows the installed package and
the child tests a different `litellm` than the parent imported -- TQ003 is the
same working-directory hazard seen from the child's side. Use
tests.test_litellm_rust.support.child_interpreter.run_child_interpreter, which
also asserts the child resolved the same `litellm.__file__` as the parent.
Every rule is suppressible with `# test-quality-ok: <reason>` on the reported
line, following the repo's `*-ok: <reason>` convention. A suppression without a
@ -140,6 +147,9 @@ SKIP_CALLS: Final = frozenset(("pytest.skip", "skip"))
CONFTEST_NAME: Final = "conftest.py"
SDK_MODULE: Final = "litellm"
SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_call", "call"))
INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P"))
CREDENTIAL_NAME_RE: Final = re.compile(
r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$"
)
@ -709,6 +719,35 @@ def _snapshotted_names(tree: ast.Module) -> Iterator[tuple[str, int]]:
yield from _string_members(iterable)
def iter_child_interpreter_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
for node in ast.walk(tree):
if not (isinstance(node, ast.Call) and node.args):
continue
if _dotted_name(node.func).rsplit(".", 1)[-1] not in SUBPROCESS_SPAWNS:
continue
argv: Final = node.args[0]
if not isinstance(argv, (ast.List, ast.Tuple)) or not argv.elts:
continue
if _dotted_name(argv.elts[0]) != "sys.executable":
continue
isolated: Final = (
len(argv.elts) > 1
and isinstance(argv.elts[1], ast.Constant)
and argv.elts[1].value in INTERPRETER_ISOLATION_FLAGS
)
if isolated:
continue
yield Violation(
path,
node.lineno,
"TQ009",
"child interpreter spawned without -I/-P; the working directory lands on sys.path "
"and a source checkout can shadow the installed package, use "
"tests.test_litellm_rust.support.child_interpreter.run_child_interpreter or pass -I "
f"(suppress: `# {SUPPRESSION_TOKEN}: <reason>`)",
)
def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
if path.name != CONFTEST_NAME:
return
@ -746,6 +785,7 @@ def check_file(path: Path) -> tuple[Violation, ...]:
*iter_credential_skip_violations(path, tree),
*iter_conftest_inventory_violations(path, tree),
*iter_internal_patch_violations(path, tree),
*iter_child_interpreter_violations(path, tree),
)
if violation.line not in skip
)

View file

@ -43,22 +43,22 @@ resource "litellm_team" "dev_team" {
The LiteLLM provider supports the following resources:
* [`litellm_model`](./resources/model) - Manage LiteLLM model configurations
* [`litellm_team`](./resources/team) - Manage teams and their permissions
* [`litellm_team_member`](./resources/team_member) - Manage team member configurations
* [`litellm_team_member_add`](./resources/team_member_add) - Add members to teams
* [`litellm_key`](./resources/key) - Manage API keys
* [`litellm_mcp_server`](./resources/mcp_server) - Manage MCP (Model Context Protocol) servers
* [`litellm_credential`](./resources/credential) - Manage credentials for various providers
* [`litellm_vector_store`](./resources/vector_store) - Manage vector stores
* [`litellm_jwt_key_mapping`](./resources/jwt_key_mapping) - Map JWT claim values to virtual keys
* `litellm_model` - Manage LiteLLM model configurations
* `litellm_team` - Manage teams and their permissions
* `litellm_team_member` - Manage team member configurations
* `litellm_team_member_add` - Add members to teams
* `litellm_key` - Manage API keys
* `litellm_mcp_server` - Manage MCP (Model Context Protocol) servers
* `litellm_credential` - Manage credentials for various providers
* `litellm_vector_store` - Manage vector stores
* `litellm_jwt_key_mapping` - Map JWT claim values to virtual keys
## Available Data Sources
The LiteLLM provider supports the following data sources:
* [`litellm_credential`](./data-sources/credential) - Retrieve credential information
* [`litellm_vector_store`](./data-sources/vector_store) - Retrieve vector store information
* `litellm_credential` - Retrieve credential information
* `litellm_vector_store` - Retrieve vector store information
## Authentication

View file

@ -22,5 +22,8 @@
},
"TQ008": {
"limit": 10993
},
"TQ009": {
"limit": 59
}
}

View file

@ -226,3 +226,31 @@ def test_an_unusable_secret_is_named_without_printing_its_value(
assert unprintable not in result.stderr
assert result.stdout == ""
assert not env_path.exists()
@pytest.mark.parametrize("phase", ("setup", "call", "teardown"))
def test_oauth_failure_diagnostics_do_not_publish_private_payloads(tmp_path: Path, phase: str) -> None:
suite: Final = ET.Element("testsuite")
case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
private: Final = "private-token-in-exception-message"
failure: Final = ET.SubElement(case, "failure", message=private)
failure.text = private
properties: Final = ET.SubElement(case, "properties")
for name, value in (
("oauth_failure_phase", phase),
("oauth_exception_type", "AssertionError"),
("oauth_frame", "oauth_gateway.py:120:start"),
("oauth_frame", f"injected\\n{private}"),
("unrelated_property", private),
):
_ = ET.SubElement(properties, "property", name=name, value=value)
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run(
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]], capture_output=True, text=True
)
assert result.returncode == 1
assert f"oauth_failure_phase: {phase}" in result.stdout
assert "oauth_exception_type: AssertionError" in result.stdout
assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout
assert private not in result.stdout + result.stderr

View file

@ -14,7 +14,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright)
- `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0)
- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token
- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token
- `logging/` - logging-integration delivery (datadog and friends)
- `security/` - secret handling and log-leak protection
- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory regression test (`test_reliability_memory_e2e.py`: a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak)
@ -26,14 +26,14 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
## MCP suite: real Datadog only
Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite
Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server, except the two Linear OAuth tests `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite
- Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env
- Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`); never skip for a missing fake upstream
- Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters
- Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down
- If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog
- The one standing exception is `test_mcp_chat_completion_oauth_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so it cannot exercise gateway-managed OAuth or per-user token seeding in any form. That test drives a real Linear MCP server instead; it is still a real remote upstream, so the no-mock, no-fixture rule above holds unchanged
- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead; they are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts
## Lay the pattern down in a class
@ -152,7 +152,7 @@ MCPs - endpoint features with the protocol op as the variant
mcp.<operation>.<auth_family>.<assertion>
operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt
auth_family : none | api_key | bearer | oauth
assertion : succeeds | denied_without_permission
assertion : succeeds | denied_without_permission | persists_across_processes
e.g. mcp.call_tool.oauth.succeeds
```

View file

@ -248,3 +248,56 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr
Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the existing legacy harness files with `--noconftest -o pythonpath=tests/e2e`; they need only synthetic HTTP providers and temporary fixture storage
## MCP OAuth happy path
`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants:
aggregate gateway SSO and explicitly configured per-server JWT, each directly
against Linear and through the live provider edge. The edge forwards to real
Linear without replay and compares the forwarded bearer to the encrypted
canonical user/server credential. This observes the forwarding boundary, not
Linear's internal logs. Direct variants independently exercise discovery
Use the existing database preparation, Prisma generation and Keycloak setup.
Build and stage the dashboard from the tested checkout as in the UI runner.
Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`,
and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using
`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private
file. The test workspace must contain a team. Do not publish browser state or
raw test/proxy output
```bash
E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \
uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0
```
The test starts and restarts its own source-built proxy on a free loopback port,
retaining its database and SSO client but no Redis or process-local cache. It
does not restart an existing proxy or clear shared databases. Gateway login,
consent, immediate list/call and post-restart reconnect must all succeed. The
aggregate client never injects a gateway header; the explicitly labeled JWT
variant configures `x-litellm-api-key` for the first consent and reconnects with
only its gateway JWT after restart
`.github/workflows/test-mcp-oauth-e2e.yml` automatically requests a run for
same-repository pull requests changing MCP, gateway authentication/SSO, consent
UI, dependencies or the relevant E2E harness/workflow paths. It retains manual
`workflow_dispatch` for targeted verification. The four cases run in the
protected `e2e-changed` environment after its normal deployment approval;
reviewers should approve and inspect this separate OAuth check when it appears.
Fork pull requests do not run this credentialed job; use a reviewed
same-repository branch for their verification. The workflow's path-filtered
check is not configured here as a globally required branch-protection check.
Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret there and retain the existing E2E license/AWS role configuration. A missing or
expired session fails the job; collection, deselection and skips are not passes.
The generic changed-test job excludes this file because it requires an owned
proxy and consent UI. No LLM call is needed
Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO,
PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this
scenario; consult the registry and LIT-3559 for their existing coverage and gaps.
LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership
of dependency/Python compatibility and its matrix; this test reuses its delivered
environment and does not change dependency constraints or compatibility gates

View file

@ -1160,62 +1160,151 @@ def _vllm_params(api_base: str, api_key: str | None, model_id: str) -> LiteLLMPa
)
class TestHostedVllmBatch:
"""hosted_vllm file upload + batch create (OpenAI-compatible path, LIT-3266).
HOSTED_VLLM_DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
HOSTED_VLLM_BAD_LINE_CUSTOM_ID = "req-bad"
hosted_vllm is in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, so /v1/files
and /v1/batches route through the OpenAI handler against the deployment's
api_base. Skipped for now: it needs a live vLLM (or OpenAI-compatible) server
exposing the files/batches APIs (HOSTED_VLLM_API_BASE), which the e2e
environment does not currently provision.
def _hosted_vllm_deployment(client: BatchClient, resources: ResourceManager) -> str:
api_base = os.environ.get("HOSTED_VLLM_API_BASE")
if api_base is None:
pytest.skip("set HOSTED_VLLM_API_BASE (the live vLLM server this deployment targets)")
api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None
model_id = (os.environ.get("HOSTED_VLLM_MODEL") or HOSTED_VLLM_DEFAULT_MODEL).strip()
proxy_name = batch_model_name("hosted-vllm-batch")
model_row_id = client.create_model(proxy_name, _vllm_params(api_base, api_key, model_id))
resources.defer(lambda: client.delete_model(model_row_id))
return proxy_name
def _upload_hosted_vllm_input(
client: BatchClient, content: bytes, *, proxy_name: str, key: str, upload_route: str
) -> Result[FileObject]:
if upload_route == "model_query":
return client.upload_file(content=content, form=FileUploadForm(purpose="batch"), model=proxy_name, key=key)
return client.upload_file(
content=content, form=FileUploadForm(purpose="batch", target_model_names=proxy_name), key=key
)
def _jsonl_with_a_failing_line(model: str) -> bytes:
bad_line = {
"custom_id": HOSTED_VLLM_BAD_LINE_CUSTOM_ID,
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": -1},
}
return render_jsonl(model) + (json.dumps(bad_line) + "\n").encode()
def _download_managed_file(client: BatchClient, file_id: str, *, key: str) -> list[str]:
downloaded = client.proxy.transport.download(
f"/v1/files/{file_id}/content", headers=client.proxy.transport.bearer(key)
)
assert downloaded.status_code == 200, (
f"file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}"
)
return downloaded.body.strip().splitlines()
class TestHostedVllmBatch:
"""hosted_vllm file upload + batch execution (LIT-5739).
vLLM implements neither /v1/files nor /v1/batches, so LiteLLM keeps the batch
input in its own database, runs every line through the deployment's
/v1/chat/completions itself, and serves the batch plus its output and error
files from that database under the creating key. Needs a live vLLM server
(HOSTED_VLLM_API_BASE), which the default e2e stack does not provision, so
the cases skip without it.
"""
@pytest.mark.skip(
reason="hosted_vllm batch/files needs a live vLLM server (HOSTED_VLLM_API_BASE) "
"not provisioned in the e2e environment; re-enable when available (LIT-3266)"
)
@pytest.mark.parametrize("upload_route", ["target_model_names", "model_query"])
@pytest.mark.covers(
"llm.batches.hosted_vllm.basic.nonstream.works",
"llm.files.hosted_vllm.upload.nonstream.works",
exercised_on=["batches", "files"],
)
def test_unified_file_and_batch_create(
self, client: BatchClient, resources: ResourceManager
def test_batch_runs_to_completion_with_a_downloadable_output(
self, client: BatchClient, resources: ResourceManager, upload_route: str
) -> None:
api_base = os.environ["HOSTED_VLLM_API_BASE"]
api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None
model_id = (
os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct"
).strip()
proxy_name = batch_model_name("hosted-vllm-batch")
model_row_id = client.create_model(
proxy_name, _vllm_params(api_base, api_key, model_id)
)
resources.defer(lambda: client.delete_model(model_row_id))
proxy_name = _hosted_vllm_deployment(client, resources)
key = resources.key()
file = unwrap(
client.upload_file(
content=render_jsonl(model_id),
form=FileUploadForm(purpose="batch", target_model_names=proxy_name),
key=key,
_upload_hosted_vllm_input(
client, render_jsonl(proxy_name), proxy_name=proxy_name, key=key, upload_route=upload_route
)
)
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider="hosted_vllm")
assert is_managed_id(file.id), f"hosted_vllm batch input must stay in LiteLLM, got file id {file.id!r}"
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}"
assert batch.status in CREATED_BATCH_STATUSES, (
f"hosted_vllm batch has non-transitional status {batch.status!r}"
)
resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True))
assert is_managed_id(batch.id), f"hosted_vllm batch must be LiteLLM-managed, got {batch.id!r}"
assert batch.status in CREATED_BATCH_STATUSES, f"hosted_vllm batch has non-transitional status {batch.status!r}"
assert_batch_object(batch)
finished = _poll_until_terminal(client, batch.id, key)
assert finished.status == "completed", f"hosted_vllm batch ended {finished.status!r}: {finished.errors!r}"
assert finished.output_file_id, "completed hosted_vllm batch has no output_file_id"
assert finished.error_file_id is None, f"all lines succeeded but error_file_id={finished.error_file_id!r}"
output_lines = _download_managed_file(client, finished.output_file_id, key=key)
assert len(output_lines) == 1, f"one input line must yield one output line, got {output_lines!r}"
first_line = BatchOutputLine.model_validate_json(output_lines[0])
assert first_line.custom_id == "req-1", f"output line lost its custom_id: {output_lines[0][:300]}"
assert first_line.response.status_code == 200, f"batch output line reports failure: {output_lines[0][:400]}"
assert first_line.response.body is not None and first_line.response.body.choices, (
"batch output line has no choices"
)
rows = client.proxy.poll_logs_for_key(
key, predicate=lambda found: any(row.call_type == "acompletion" for row in found)
)
line_rows = [row for row in rows if row.call_type == "acompletion"]
assert line_rows, f"the batch line's chat call was not logged under the creating key: {rows!r}"
assert all(row.custom_llm_provider == "hosted_vllm" for row in line_rows), (
f"batch line rows must be attributed to hosted_vllm: {line_rows!r}"
)
@pytest.mark.covers("llm.batches.hosted_vllm.basic.nonstream.works", exercised_on=["batches", "files"])
def test_failing_line_lands_in_the_error_file_not_the_batch_status(
self, client: BatchClient, resources: ResourceManager
) -> None:
proxy_name = _hosted_vllm_deployment(client, resources)
key = resources.key()
file = unwrap(
_upload_hosted_vllm_input(
client,
_jsonl_with_a_failing_line(proxy_name),
proxy_name=proxy_name,
key=key,
upload_route="target_model_names",
)
)
resources.defer(lambda: cleanup_file(client, file.id, key=key))
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True))
finished = _poll_until_terminal(client, batch.id, key)
assert finished.status == "completed", f"a failing line must not fail the batch, got {finished.status!r}"
assert finished.output_file_id, "the good line must still produce an output file"
assert finished.error_file_id, "the failing line must produce an error file"
output_lines = _download_managed_file(client, finished.output_file_id, key=key)
error_lines = _download_managed_file(client, finished.error_file_id, key=key)
assert [BatchOutputLine.model_validate_json(line).custom_id for line in output_lines] == ["req-1"]
assert len(error_lines) == 1, f"one failing line must yield one error line, got {error_lines!r}"
error_line = BatchOutputLine.model_validate_json(error_lines[0])
assert error_line.custom_id == HOSTED_VLLM_BAD_LINE_CUSTOM_ID
assert error_line.response.status_code == 400, f"error line must carry the provider's 4xx: {error_lines[0][:400]}"
BATCH_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"})
FAILED_BATCH_POLL_SECONDS = 120.0
@ -1443,6 +1532,7 @@ class BatchOutputResponse(BaseModel):
class BatchOutputLine(BaseModel):
custom_id: str | None = None
response: BatchOutputResponse

View file

@ -17,6 +17,7 @@ import functools
import os
from collections.abc import Generator, Iterator
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -28,6 +29,7 @@ from e2e_config import (
FIXTURE_DIR,
FIXTURE_MODE_RAW,
MANAGED_FILES_OPT_IN_ENV,
MCP_OAUTH_LIVE_OPT_IN_ENV,
PROMPT_CACHING_OPT_IN_ENV,
PROXY_BASE_URL,
REDIS_CHAOS_OPT_IN_ENV,
@ -56,6 +58,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV,
"redis_chaos": REDIS_CHAOS_OPT_IN_ENV,
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
"mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV,
}
)
@ -132,6 +135,11 @@ def pytest_configure(config: pytest.Config) -> None:
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set",
)
config.addinivalue_line(
"markers",
"mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless "
"E2E_MCP_OAUTH_LIVE is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:
@ -213,6 +221,8 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
if item.get_closest_marker("e2e") is None:
return
if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames:
return
reason = _proxy_fail_reason()
if reason is not None:
pytest.fail(reason)
@ -236,6 +246,13 @@ def pytest_runtest_makereport(
"""Stash the call-phase outcome so teardown can tell a passed test from a
failed one without re-deriving it."""
report = yield
if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None:
# Publish code locations only, never exception messages, source text or locals.
item.user_properties.append(("oauth_failure_phase", report.when))
item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__))
for entry in call.excinfo.traceback:
item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}"))
report.user_properties = list(item.user_properties)
if report.when == "call":
item.stash[_CALL_PASSED] = report.passed
return report

View file

@ -71,6 +71,14 @@
assertions: [succeeds]
source: "db.py user_oauth_credential lookup"
rationale: OAuth2 token passthrough; per-user credential storage
- id: mcp.call_tool.oauth.persists_across_processes
module: mcp
tier: P1
operation: call_tool
auth_family: oauth
assertions: [persists_across_processes]
source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore"
rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache
- id: mcp.list_tools.none.succeeds
module: mcp
tier: P1

View file

@ -72,7 +72,6 @@
- {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"}
- {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"}
- {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"}
- {id: mgmt.config.allowed_ip.changed_key_only, module: mgmt, tier: P2, surface: api, assertions: [persists], source: "proxy_setting_endpoints.py:496", rationale: "An allowed-IP change leaves unrelated file settings out of the DB row"}
- {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"}
- {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"}
- {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"}

View file

@ -52,6 +52,7 @@ CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp")
LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "")
LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed
# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger`
# service in docker-compose.yml maps it to host 16686). Trace-completeness tests
@ -144,6 +145,7 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS"
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))

View file

@ -1,5 +1,4 @@
general_settings:
max_parallel_requests: 100
proxy_batch_write_at: 5
enable_jwt_auth: true
litellm_jwtauth:

View file

@ -435,7 +435,7 @@ def _signal_process_group(process_id: int, signum: int) -> bool:
return True
def _stop_process_group(child: subprocess.Popen[bytes]) -> None:
def stop_process_group(child: subprocess.Popen[bytes]) -> None:
_signal_process_group(child.pid, signal.SIGTERM)
deadline: Final = time.monotonic() + 5
while _process_group_exists(child.pid):
@ -476,7 +476,7 @@ def run_oidc_profile(proxy_url: str, command: list[str]) -> int:
try:
return child.wait()
finally:
_stop_process_group(child)
stop_process_group(child)
if __name__ == "__main__":

View file

@ -7,13 +7,18 @@ so a read-back reflects the change. Router settings, which mutate global proxy
state, are exercised with a benign, self-restoring change so a shared proxy is left
as it was found.
Cache settings and the Vault config override are deliberately not covered here.
Both routes reconfigure the whole proxy: /cache/settings persists what it receives
into a row that outranks the YAML cache_params and is re-applied on a timer, and
/config_overrides/hashicorp_vault swaps the process-wide secret manager. Neither can
be exercised safely against the shared proxy the suites run on, so they need an
isolated proxy before a test lands. Do not add a read-then-write-back test for
either one.
Cache settings, the Vault config override and the allowed-IP routes are deliberately
not covered here. All three reconfigure the whole proxy: /cache/settings persists what
it receives into a row that outranks the YAML cache_params and is re-applied on a timer,
/config_overrides/hashicorp_vault swaps the process-wide secret manager, and
/add/allowed_ip mutates the live general_settings["allowed_ips"] that
auth_utils._check_valid_ip reads, so the first call locks every other client out of the
shared proxy. The allowlist is an exact string match with no CIDR support, and no route
reports the caller's address as the proxy sees it, so a test cannot allowlist itself
first; /delete/allowed_ip sits behind the same auth dependency, so the cleanup is locked
out too and the proxy stays poisoned for the rest of the build. None of the three can be
exercised safely against the shared proxy the suites run on, so they need an isolated
proxy before a test lands. Do not add a read-then-write-back test for any of them.
"""
from __future__ import annotations
@ -21,10 +26,9 @@ from __future__ import annotations
import math
import time
from collections.abc import Callable
from typing import Final
import pytest
from pydantic import BaseModel, JsonValue, RootModel
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_http import NoBody, Success, unwrap, unwrap_status
@ -188,7 +192,7 @@ class JwtKeyMappingResponse(BaseModel):
class RouterSettingsPatch(BaseModel):
num_retries: int
retry_after: int
class ConfigUpdateBody(BaseModel):
@ -199,39 +203,8 @@ class ConfigUpdateResponse(BaseModel):
message: str
class AllowedIpBody(BaseModel):
ip: str
class ConfigFieldInfoParams(BaseModel):
field_name: str
class ConfigFieldInfoResponse(BaseModel):
field_name: str
field_value: JsonValue
source: str
editable: bool
class ConfigListParams(BaseModel):
config_type: str
class ConfigListEntry(BaseModel):
field_name: str
field_value: JsonValue
stored_in_db: bool | None
source: str
editable: bool
class ConfigListResponse(RootModel[list[ConfigListEntry]]):
pass
class RouterCurrentValues(BaseModel):
num_retries: int | None = None
retry_after: int | None = None
class RouterSettingsResponse(BaseModel):
@ -493,17 +466,25 @@ class TestRouterSettings:
) -> None:
"""/config/update is the only write path for router_settings (there is no
dedicated router-settings write route). The change is restored on teardown so
the shared proxy keeps its original retry policy."""
original = self._read_num_retries(client)
assert original is not None, "GET /router/settings did not report num_retries; cannot prove a change"
resources.defer(lambda: self._write_num_retries(client, original))
the shared proxy keeps its original retry policy.
target = original + 5
retry_after is the subject because it satisfies all three constraints at once:
no lane's config file declares it, so the database owns it and the write is not
refused as config-owned; it is in RUNTIME_UPDATABLE_ROUTER_SETTINGS, so
/config/update accepts it; and it is in ROUTER_SETTINGS_FIELDS backed by an
always-set Router attribute, so GET /router/settings reports it for the
read-back. Bumping it by one second is the smallest change that proves the
round-trip without slowing a concurrent test that hits a retry."""
original = self._read_retry_after(client)
assert original is not None, "GET /router/settings did not report retry_after; cannot prove a change"
resources.defer(lambda: self._write_retry_after(client, original))
target = original + 1
response = unwrap(
client.proxy.transport.post(
"/config/update",
headers=client.proxy.transport.master,
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(num_retries=target)),
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(retry_after=target)),
response_type=ConfigUpdateResponse,
)
)
@ -513,20 +494,20 @@ class TestRouterSettings:
_ = _poll(
client,
lambda: True if self._read_num_retries(client) == target else None,
f"GET /router/settings never reported num_retries {target} after /config/update",
lambda: True if self._read_retry_after(client) == target else None,
f"GET /router/settings never reported retry_after {target} after /config/update",
)
self._write_num_retries(client, original)
self._write_retry_after(client, original)
restored = _poll(
client,
lambda: original if self._read_num_retries(client) == original else None,
f"GET /router/settings never returned to the original num_retries {original} after the restore",
lambda: original if self._read_retry_after(client) == original else None,
f"GET /router/settings never returned to the original retry_after {original} after the restore",
)
assert restored == original, f"router num_retries left at {restored}, expected the original {original}"
assert restored == original, f"router retry_after left at {restored}, expected the original {original}"
@staticmethod
def _read_num_retries(client: ManagementClient) -> int | None:
def _read_retry_after(client: ManagementClient) -> int | None:
return unwrap(
client.proxy.transport.get(
"/router/settings",
@ -534,72 +515,20 @@ class TestRouterSettings:
params=NoBody(),
response_type=RouterSettingsResponse,
)
).current_values.num_retries
).current_values.retry_after
@staticmethod
def _write_num_retries(client: ManagementClient, value: int) -> None:
def _write_retry_after(client: ManagementClient, value: int) -> None:
_ = unwrap(
client.proxy.transport.post(
"/config/update",
headers=client.proxy.transport.master,
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(num_retries=value)),
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(retry_after=value)),
response_type=ConfigUpdateResponse,
)
)
class TestConfigPersistence:
@pytest.mark.covers("mgmt.config.allowed_ip.changed_key_only")
def test_add_allowed_ip_does_not_store_unrelated_config_value(
self, client: ManagementClient, resources: ResourceManager
) -> None:
allowed_ip: Final = "127.0.0.1"
added: Final = unwrap(
client.proxy.transport.post(
"/add/allowed_ip",
headers=client.proxy.transport.master,
json=AllowedIpBody(ip=allowed_ip),
response_type=ConfigUpdateResponse,
)
)
resources.defer(
lambda: unwrap(
client.proxy.transport.post(
"/delete/allowed_ip",
headers=client.proxy.transport.master,
json=AllowedIpBody(ip=allowed_ip),
response_type=ConfigUpdateResponse,
)
)
)
assert added.message == f"IP {allowed_ip} address added successfully"
listed: Final = unwrap(
client.proxy.transport.get(
"/config/list",
headers=client.proxy.transport.master,
params=ConfigListParams(config_type="general_settings"),
response_type=ConfigListResponse,
)
)
unrelated: Final = next(entry for entry in listed.root if entry.field_name == "max_parallel_requests")
assert unrelated.stored_in_db is not True
assert unrelated.source == "config"
assert unrelated.editable is False
field_info: Final = unwrap(
client.proxy.transport.get(
"/config/field/info",
headers=client.proxy.transport.master,
params=ConfigFieldInfoParams(field_name="max_parallel_requests"),
response_type=ConfigFieldInfoResponse,
)
)
assert field_info.source == "config"
assert field_info.editable is False
assert field_info.field_value == unrelated.field_value
class TestMcpServerSubmission:
@pytest.mark.covers("mgmt.mcp_server.register.happy_path")
def test_register_submits_pending_server(self, client: ManagementClient, resources: ResourceManager) -> None:

View file

@ -18,7 +18,7 @@ import asyncio
import re
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl
import httpx
@ -26,11 +26,21 @@ import httpx2
import pytest
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
from e2e_http import AuthHeaders, NoBody, unwrap
from idp import Identity
from mcp import ClientSession
from mcp.client.auth import OAuthClientProvider
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo
from mcp.types import TextContent
from models import (
ChatBody,
ChatResponse,
McpOauthUserCredentialStatus,
McpServerCreateBody,
McpServerInfo,
McpServerUserCredentialListResponse,
McpServerUserCredentialRow,
)
from proxy_client import ProxyClient
if TYPE_CHECKING:
@ -44,8 +54,8 @@ OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
BROWSER_CONSENT_TIMEOUT = 60.0
def _mcp_url(alias: str) -> str:
return f"{PROXY_BASE_URL}/{alias}/mcp"
def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str:
return f"{base_url.rstrip('/')}/{alias}/mcp"
class InMemoryTokenStorage:
@ -69,7 +79,13 @@ class InMemoryTokenStorage:
self._client_info = client_info
async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> tuple[str, str | None]:
async def _browser_follow_authorize(
start_url: str,
storage_state_path: str,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> tuple[str, str | None]:
"""Play the browser's role for a real upstream whose authorize endpoint
serves an interactive consent page (Linear). A headless Chromium primed
with a human's saved Linear session opens the gateway authorize URL and
@ -85,6 +101,9 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
def _note_request(request: object) -> None:
url = getattr(request, "url", "")
host = httpx.URL(url).host
if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")):
captured["upstream_consent"] = "seen"
if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured:
captured["url"] = url
@ -96,7 +115,7 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
context = await browser.new_context(storage_state=storage_state_path)
await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect)
page = await context.new_page()
page.on("request", _note_request)
context.on("request", _note_request)
page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
await page.goto(start_url, wait_until="domcontentloaded")
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
@ -105,8 +124,29 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
await page.wait_for_load_state("networkidle", timeout=8000)
except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it
pass
if "url" in captured:
if "upstream_consent" in captured or "url" in captured:
break
if await page.locator("#username").count() and identity is not None:
await page.locator("#username").fill(identity.username)
await page.locator("#password").fill(identity.password)
await page.locator("#kc-login").click()
continue
if "/ui/connect" in page.url and server_alias is not None:
card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True))
if await card.count() != 1:
await asyncio.sleep(0.5)
continue
connect = card.get_by_text("Connect", exact=True)
if await connect.count():
await connect.click()
continue
if not await card.locator("svg.text-success").count():
await asyncio.sleep(0.5)
continue
finish = page.get_by_role("button", name="Finish connecting", exact=True)
if await finish.count() and await finish.is_enabled():
await finish.click()
continue
control = page.locator(
'button[name="action"][value="approve"], button:has-text("Authorize"), '
'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
@ -118,27 +158,44 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
final_url = page.url
await browser.close()
# A redirect chain can finish inside goto/networkidle before the loop checks the page.
assert "upstream_consent" not in captured, "cold reconnect required upstream consent"
landing = captured.get("url")
assert landing is not None, (
f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; "
f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
)
params = dict(parse_qsl(httpx.URL(landing).query.decode()))
assert "code" in params, f"client redirect_uri carried no code: {landing}"
assert "code" in params, "client redirect_uri carried no authorization code"
return params["code"], params.get("state")
def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str) -> OAuthClientProvider:
def _oauth_provider(
url: str,
storage: InMemoryTokenStorage,
storage_state_path: str | None,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> OAuthClientProvider:
"""The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR,
PKCE, token exchange) with the browser leg driven by Playwright against the
upstream's consent screen."""
code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks
async def redirect_handler(authorize_url: str) -> None:
code, state = await _browser_follow_authorize(authorize_url, storage_state_path)
async def _reject_redirect(_: str) -> None:
raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused")
async def _follow_redirect(authorize_url: str) -> None:
assert storage_state_path is not None
code, state = await _browser_follow_authorize(
authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent
)
code_holder["code"] = code
code_holder["state"] = state
redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect
async def callback_handler() -> AuthorizationCodeResult:
code = code_holder.get("code")
assert code is not None, "callback_handler ran before the authorize redirect completed"
@ -167,24 +224,45 @@ class _HeaderInjectingTransport(httpx2.AsyncBaseTransport):
store the upstream token for from the key on the token exchange, exactly
like a production MCP host configured with a LiteLLM key header."""
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str]) -> None:
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None:
self._inner = inner
self._headers = headers
self._gateway_url = httpx2.URL(gateway_url)
@staticmethod
def _port(url: httpx2.URL) -> int | None:
if url.port is not None:
return url.port
return {"http": 80, "https": 443}.get(url.scheme)
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
for name, value in self._headers.items():
if name not in request.headers:
request.headers[name] = value
same_origin: Final = (
request.url.scheme == self._gateway_url.scheme
and request.url.host == self._gateway_url.host
and self._port(request.url) == self._port(self._gateway_url)
)
if same_origin:
for name, value in self._headers.items():
if name not in request.headers:
request.headers[name] = value
else:
for name, value in self._headers.items():
if request.headers.get(name) == value:
del request.headers[name]
return await self._inner.handle_async_request(request)
async def aclose(self) -> None:
await self._inner.aclose()
def _oauth_http_client(headers: dict[str, str], auth: OAuthClientProvider) -> httpx2.AsyncClient:
def _oauth_http_client(
headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL
) -> httpx2.AsyncClient:
return httpx2.AsyncClient(
headers=headers,
auth=auth,
timeout=httpx2.Timeout(REQUEST_TIMEOUT),
follow_redirects=True,
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers),
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url),
)
@ -199,6 +277,43 @@ async def _seed_via_dance(
return tuple(sorted(tool.name for tool in listed.tools))
@dataclass(frozen=True, slots=True)
class OauthToolRun:
tools: tuple[str, ...]
is_error: bool
text: str
async def _list_and_call(
url: str,
headers: dict[str, str],
storage: InMemoryTokenStorage,
storage_state_path: str | None,
tool: str,
arguments: dict[str, str],
gateway_url: str = PROXY_BASE_URL,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> OauthToolRun:
async with _oauth_http_client(
headers,
_oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent),
gateway_url,
) as http_client:
async with streamable_http_client(url, http_client=http_client) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
listed: Final = await session.list_tools()
result: Final = await session.call_tool(tool, arguments)
text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent))
return OauthToolRun(
tools=tuple(sorted(tool_item.name for tool_item in listed.tools)),
is_error=result.is_error,
text=text,
)
@dataclass(frozen=True, slots=True)
class ChatMcpClient:
proxy: ProxyClient
@ -252,6 +367,53 @@ class ChatMcpClient:
f"last error: {last_error!r}"
)
def list_and_call(
self,
alias: str,
headers: dict[str, str],
storage: InMemoryTokenStorage,
storage_state_path: str | None,
tool: str,
arguments: dict[str, str],
base_url: str = PROXY_BASE_URL,
identity: Identity | None = None,
allow_upstream_consent: bool = True,
) -> OauthToolRun:
return asyncio.run(
_list_and_call(
f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url),
headers,
storage,
storage_state_path,
tool,
arguments,
base_url,
identity,
alias,
allow_upstream_consent,
)
)
def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]:
return unwrap(
self.proxy.transport.get(
f"/v1/mcp/server/{server_id}/user-credentials",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServerUserCredentialListResponse,
)
).root
def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None:
_ = unwrap(
self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}/oauth-user-credential",
headers=headers,
json=NoBody(),
response_type=McpOauthUserCredentialStatus,
)
)
def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse:
"""POST /chat/completions carrying the LiteLLM key in `headers` (either
ingress form) with an MCP server attached in `body.tools`. The gateway

View file

@ -0,0 +1,198 @@
"""An owned, source-built OAuth gateway with cold restarts and credential observations.
Only this child process is restarted. Its database and SSO client survive while
its process-local caches do not; Redis is deliberately absent from its config.
The optional live edge measures headers without recording credentials or bodies.
"""
from __future__ import annotations
import os
import socket
import subprocess
import sys
import threading
import time
from collections.abc import Callable, Mapping
from contextlib import ExitStack
from dataclasses import dataclass, field
from pathlib import Path
from typing import Final
import psycopg
from e2e_http import NoBody
from idp import Keycloak, stop_process_group
from proxy_client import ProxyClient, build_proxy_client
from psycopg.rows import class_row
from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError
INHERITED_ENV_PREFIXES: Final = ("REDIS_", "MICROSOFT_", "GOOGLE_", "GENERIC_", "PROXY_")
class StoredOAuth(BaseModel):
type: str
access_token: SecretStr
@dataclass(frozen=True, slots=True)
class CredentialRow:
credential_b64: str = field(repr=False)
def stored_oauth(user_id: str, server_id: str) -> StoredOAuth:
"""Read the encrypted credential because management APIs omit the plaintext token."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
with psycopg.Connection[CredentialRow].connect(
os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow)
) as conn:
row: Final = conn.execute(
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
(user_id, server_id),
).fetchone()
assert row is not None, "canonical user/server has no persisted credential"
plaintext: Final = decrypt_value_helper(
row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False
)
assert plaintext is not None, "persisted credential must decrypt with the gateway salt"
assert plaintext != row.credential_b64, "persisted credential must be encrypted"
try:
credential: Final = StoredOAuth.model_validate_json(plaintext)
except ValidationError:
raise AssertionError("decrypted credential is not an OAuth payload") from None
assert credential.type == "oauth2"
assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty"
return credential
class RpcMethod(BaseModel):
method: str = ""
@dataclass(slots=True)
class OAuthObservation:
gateway_token: str = field(default="", repr=False)
_seen: tuple[tuple[str, str, bool], ...] = field(default=(), init=False, repr=False)
_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None:
if body is None or not url.endswith("/mcp"):
return
try:
operation: Final = RpcMethod.model_validate_json(body).method
except ValidationError:
return
if operation not in ("tools/list", "tools/call"):
return
received: Final = headers.get("authorization", "")
gateway_leaked: Final = any(
value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values()
)
with self._lock:
self._seen = (*self._seen, (operation, received, gateway_leaked))
def assert_forwarded(self, expected: StoredOAuth) -> None:
with self._lock:
snapshot: Final = self._seen
self._seen = ()
assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations"
expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}"
assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token"
assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream"
def available_port() -> int:
with socket.socket() as listener:
listener.bind(("127.0.0.1", 0))
return TypeAdapter(tuple[str, int]).validate_python(listener.getsockname())[1]
@dataclass(slots=True)
class OAuthGateway:
base_url: str
proxy: ProxyClient
_environment: Mapping[str, str] = field(repr=False)
_command: tuple[str, ...] = field(repr=False)
_log_path: Path
_child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
def start(self) -> None:
with self._log_path.open("ab") as log:
self._child = subprocess.Popen(
self._command,
env=self._environment,
stdout=log,
stderr=log,
start_new_session=True,
)
deadline: Final = time.monotonic() + 120
while time.monotonic() < deadline:
assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log"
result = self.proxy.transport.probe("/health/liveliness", params=NoBody())
if result.status_code == 200:
return
time.sleep(0.5)
raise AssertionError("owned OAuth gateway did not become ready")
def stop(self) -> None:
if self._child is not None:
stop_process_group(self._child)
assert self._child.poll() is not None, "old gateway process is still alive"
def restart(self) -> None:
assert self._child is not None
previous: Final = self._child.pid
self.stop()
self.start()
assert self._child.pid != previous, "gateway restart did not create a new process"
def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway:
for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
assert os.environ.get(name), f"{name} is required for the owned OAuth gateway"
port: Final = available_port()
base_url: Final = f"http://127.0.0.1:{port}"
def defer(callback: Callable[[], object]) -> None:
cleanup.callback(callback)
browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer)
config: Final = directory / "oauth-gateway.yaml"
config.write_text(
"model_list: []\n"
"general_settings:\n"
" master_key: os.environ/LITELLM_MASTER_KEY\n"
" database_url: os.environ/DATABASE_URL\n"
" enable_jwt_auth: true\n"
" litellm_jwtauth:\n"
" user_id_jwt_field: sub\n"
" user_email_jwt_field: email\n"
" team_ids_jwt_field: groups\n"
" user_id_upsert: true\n"
)
environment: Final = {
**{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)},
**browser.environment(idp.discovery()),
"PROXY_BASE_URL": base_url,
"JWT_PUBLIC_KEY_URL": idp.jwks_url,
"JWT_ISSUER": idp.issuer,
"JWT_AUDIENCE": "litellm-e2e",
"DISABLE_SCHEMA_UPDATE": "true",
"STORE_MODEL_IN_DB": "True",
"PYTHONPATH": str(Path(__file__).resolve().parents[3]),
}
gateway: Final = OAuthGateway(
base_url=base_url,
proxy=build_proxy_client(
base_url=base_url,
control_plane_base_url=base_url,
replica_urls=(base_url,),
master_key=os.environ["LITELLM_MASTER_KEY"],
),
_environment=environment,
_command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)),
_log_path=directory / "oauth-gateway.log",
)
cleanup.callback(gateway.stop)
gateway.start()
return gateway

View file

@ -0,0 +1,207 @@
"""Real OAuth consent, immediate MCP operations and cold-restart persistence.
Aggregate SSO uses the SDK's normal authentication. The per-server variant is
explicitly a configured two-header client, not an Authorization-only OAuth host.
The observed variants forward to the same real Linear upstream and compare its
bearer at the forwarding boundary; direct variants retain unmodified discovery.
"""
from __future__ import annotations
import os
from collections.abc import Iterator
from contextlib import ExitStack
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal
import pytest
from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker
from e2e_http import AuthHeaders, NoBody, get_external, unwrap
from idp import Identity, Keycloak
from lifecycle import ResourceManager
from models import (
McpOauthCredentials,
McpServerCreateBody,
ObjectPermission,
TeamMemberAddBody,
TeamMemberEntry,
TeamUpdateBody,
)
from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client
from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth
from provider_edge import LiveEdge, start_provider_edge
from proxy_client import ProxyClient
from pydantic import BaseModel, ValidationError
pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live]
class OAuthMetadata(BaseModel):
authorization_endpoint: str
token_endpoint: str
registration_endpoint: str
class LinearTeam(BaseModel):
id: str
name: str
class LinearTeams(BaseModel):
teams: tuple[LinearTeam, ...]
def assert_tool_result(run: OauthToolRun, tool: str) -> None:
assert tool in run.tools
assert run.is_error is False
try:
result: Final = LinearTeams.model_validate_json(run.text)
except ValidationError:
raise AssertionError("list_teams did not return the expected teams payload") from None
assert result.teams, "the test workspace must contain at least one team"
assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names"
@pytest.fixture(scope="module")
def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]:
assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), (
"E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py"
)
assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay"
with ExitStack() as cleanup:
yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup)
@pytest.fixture(scope="module")
def proxy(oauth_gateway: OAuthGateway) -> ProxyClient:
return oauth_gateway.proxy
@pytest.fixture(scope="module")
def client(proxy: ProxyClient) -> ChatMcpClient:
return build_chat_client(proxy)
class TestMcpOauthHappyPath:
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes")
@pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt"))
@pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed"))
def test_consent_list_call_and_cold_restart(
self,
client: ChatMcpClient,
resources: ResourceManager,
jwt_identity: Identity,
idp: Keycloak,
oauth_gateway: OAuthGateway,
route: Literal["aggregate_sso", "explicit_header_jwt"],
observed: bool,
) -> None:
alias: Final = f"e2elinear{unique_marker()}"
tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}"
token: Final = idp.access_token(jwt_identity)
observation: Final = OAuthObservation(gateway_token=token)
edge: Final = (
start_provider_edge(
LiveEdge(observe_request=observation.observe),
mounts=MappingProxyType(
{"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"}
),
)
if observed
else None
)
if edge is not None:
resources.defer(edge.shutdown)
metadata: Final = (
unwrap(
get_external(
"https://mcp.linear.app/.well-known/oauth-authorization-server",
response_type=OAuthMetadata,
)
)
if observed
else None
)
created: Final = client.create_server(
McpServerCreateBody(
alias=alias,
server_name=alias,
url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL,
transport="http",
allow_all_keys=False,
auth_type="oauth2",
oauth2_flow="authorization_code",
per_server_oauth_discovery=route == "explicit_header_jwt",
authorization_url=metadata.authorization_endpoint if metadata else None,
token_url=metadata.token_endpoint if metadata else None,
registration_url=metadata.registration_endpoint if metadata else None,
credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None,
)
)
resources.defer(lambda: client.delete_server(created.server_id))
assert client.server_user_credentials(created.server_id) == (), (
"scenario must start without upstream credentials"
)
unwrap(
client.proxy.transport.post(
"/team/member_add",
headers=client.proxy.transport.master,
json=TeamMemberAddBody(
team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user")
),
response_type=NoBody,
)
)
client.proxy.update_team(
TeamUpdateBody(
team_id=jwt_identity.group,
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
)
)
headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {}
resources.defer(
lambda: client.revoke_user_token(
created.server_id,
AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"),
)
)
identity: Final = jwt_identity if route == "aggregate_sso" else None
first: Final = client.list_and_call(
alias,
headers,
InMemoryTokenStorage(),
LINEAR_STORAGE_STATE,
tool,
{},
base_url=oauth_gateway.base_url,
identity=identity,
)
assert_tool_result(first, tool)
credentials: Final = client.server_user_credentials(created.server_id)
assert len(credentials) == 1
assert credentials[0].user_id == jwt_identity.user_id
assert credentials[0].credential_type == "oauth2"
first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
if observed:
observation.assert_forwarded(first_stored_oauth)
oauth_gateway.restart()
fresh_token: Final = idp.access_token(jwt_identity)
observation.gateway_token = fresh_token
second: Final = client.list_and_call(
alias,
{"Authorization": f"Bearer {fresh_token}"} if identity is None else {},
InMemoryTokenStorage(),
LINEAR_STORAGE_STATE if identity is not None else None,
tool,
{},
base_url=oauth_gateway.base_url,
identity=identity,
allow_upstream_consent=False,
)
assert_tool_result(second, tool)
second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
if observed:
observation.assert_forwarded(second_stored_oauth)

View file

@ -192,7 +192,7 @@ class ImageUrl(BaseModel):
class TextContentPart(BaseModel):
type: str = "text"
text: str
cache_control: "CacheControl | None" = None
cache_control: CacheControl | None = None
class ImageContentPart(BaseModel):
@ -572,6 +572,10 @@ class McpInfo(BaseModel):
logo_url: str | None = None
class McpOauthCredentials(BaseModel):
upstream_resource: str
class McpServerCreateBody(BaseModel):
"""POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
`oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
@ -584,8 +588,11 @@ class McpServerCreateBody(BaseModel):
allow_all_keys: bool = True
auth_type: str | None = None
oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
per_server_oauth_discovery: bool | None = None
authorization_url: str | None = None
token_url: str | None = None
registration_url: str | None = None
credentials: McpOauthCredentials | None = None
server_name: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
@ -625,6 +632,26 @@ class McpServerListResponse(RootModel[list[McpServerRow]]):
"""GET /v1/mcp/server answers with a bare array of servers."""
class McpServerUserCredentialRow(BaseModel):
user_id: str
credential_type: Literal["oauth2", "byok"]
expires_at: str | None = None
connected_at: str | None = None
updated_at: str
class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]):
"""GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array."""
class McpOauthUserCredentialStatus(BaseModel):
server_id: str
has_credential: bool
expires_at: str | None = None
is_expired: bool = False
connected_at: str | None = None
class ToolsetTool(BaseModel):
server_id: str
tool_name: str
@ -1172,8 +1199,9 @@ class TeamNewResponse(BaseModel):
class TeamUpdateBody(BaseModel):
team_id: str
team_alias: str
team_alias: str | None = None
models: list[str] | None = None
object_permission: ObjectPermission | None = None
class TeamInfoParams(BaseModel):

View file

@ -46,7 +46,7 @@ import os
import re
import threading
from collections import deque
from collections.abc import Generator, Mapping, Sequence
from collections.abc import Callable, Generator, Mapping, Sequence
from contextlib import closing, contextmanager
from dataclasses import dataclass, field, replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@ -538,7 +538,7 @@ class ReplayEdge:
@dataclass(frozen=True, slots=True)
class LiveEdge:
pass
observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
@ -787,10 +787,13 @@ def _handle_record(
def _handle_live(
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None,
observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None,
) -> EdgeOutcome:
forwarded: Final = {
name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
}
if observe_request is not None:
observe_request(url, forwarded, body)
head: Final = (
forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
if cache is None else cache.forward(mount, method, url, forwarded, body, timeout, test_key=test_key)
@ -868,9 +871,10 @@ def handle_edge_request(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
backend, mount, test_key,
)
case LiveEdge():
case LiveEdge(observe_request=observe_request):
return _handle_live(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
observe_request=observe_request,
)
case RecordEdge():
return _handle_record(

View file

@ -89,6 +89,7 @@ from models import (
TeamDeleteBody,
TeamNewBody,
TeamNewResponse,
TeamUpdateBody,
ToolsetCreateBody,
ToolsetRow,
ToolsetUpdateBody,
@ -871,6 +872,16 @@ class ProxyClient:
)
).team_id
def update_team(self, body: TeamUpdateBody) -> None:
unwrap(
self.transport.post(
"/team/update",
headers=self.transport.master,
json=body,
response_type=NoBody,
)
)
def delete_team(self, team_id: str) -> None:
result = self.transport.post(
"/team/delete",

View file

@ -12,3 +12,4 @@ markers =
prompt_caching_stack: needs a proxy running with router_settings.optional_pre_call_checks including prompt_caching; deselected unless E2E_PROMPT_CACHING_STACK is set
cli_determinism: drives the real claude CLI for several seconds; deselected unless E2E_CLI_DETERMINISM is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set
mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless E2E_MCP_OAUTH_LIVE is set

View file

@ -9,7 +9,7 @@ import httpx
from integration._support.asgi import asgi_server
from integration._support.client import Gateway, Scenario
from integration._support.database import read_rows
from mcp.server.fastmcp import FastMCP
from mcp.server.mcpserver import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
from mcp_tests.mcp_e2e_upstream_server import add, multiply
from starlette.requests import Request
@ -27,12 +27,7 @@ class McpPeer:
@contextmanager
def mcp_peer() -> Iterator[McpPeer]:
service: Final = FastMCP(
"integration-math",
stateless_http=True,
json_response=True,
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
service: Final = MCPServer("integration-math")
service.add_tool(add)
service.add_tool(multiply)
@ -40,7 +35,11 @@ def mcp_peer() -> Iterator[McpPeer]:
def fail() -> str:
raise ValueError("synthetic tool failure")
app: Final = service.streamable_http_app()
app: Final = service.streamable_http_app(
stateless_http=True,
json_response=True,
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
observed: Final[queue.Queue[dict[str, object]]] = queue.Queue()
async def capture(scope: Scope, receive: Receive, send: Send) -> None:
@ -94,9 +93,7 @@ def tool_names(gateway: Gateway, key: str, identity: str) -> dict[str, str]:
}
def call_tool(
gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]
) -> httpx.Response:
def call_tool(gateway: Gateway, key: str, identity: str, name: str, arguments: dict[str, object]) -> httpx.Response:
return gateway.client.post(
"/mcp-rest/tools/call",
headers={"x-litellm-api-key": key},

View file

@ -1311,6 +1311,27 @@
],
"tests/integration/cost_calculation/test_cost_tracking.py::test_case_bills_expected_cost[us.anthropic.claude-opus-5-v1:0-stream_full_usage]": [
"quota_management.spend_tracking.scripted_wire.logs_cost"
],
"tests/integration/mcp/test_mcp_lifecycle.py::test_health_intersects_route_restricted_key_grants_in_both_management_modes": [
"other.mcp.health.restricted_keys_intersect_grants_in_both_modes"
],
"tests/integration/mcp/test_mcp_lifecycle.py::test_warm_credential_removal_rejects_without_upstream_traffic": [
"other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic"
],
"tests/integration/observability/test_guardrail_effects.py::test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls": [
"other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution"
],
"tests/integration/mcp/test_oauth_configuration.py::test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server[revoke]": [
"other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server"
],
"tests/integration/mcp/test_oauth_configuration.py::test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server[expire]": [
"other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server"
],
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[anonymous]": [
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
],
"tests/integration/mcp/test_mcp_lifecycle.py::test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution[bearer]": [
"other.mcp.permissions.same_url_servers_enforce_discovery_and_execution"
]
},
"browser": {

View file

@ -0,0 +1,28 @@
# MCP security regression coverage
[LIT-4506](https://linear.app/litellm-ai/issue/LIT-4506) tracks ten gateway guards and the later JWT/OAuth acceptance. This inventory distinguishes executable assertions from unresolved coverage. A listed test counts as verified only when its exact commit has an executed, passing result
Run the controlled gateway cases through `python tests/integration/run.py extensions`. They use real HTTP, PostgreSQL, scoped non-master keys and an SDK upstream. The existing runner supplies test entitlement; these tests do not validate licenses or external-provider consent. Canonical nodes and contract IDs live in `../contracts.json`
| Requested guard | Existing or added coverage | Remaining limitation and owner |
| --- | --- | --- |
| 1. Discovery scoped by org/team/user/key | `test_mcp_lifecycle.py` checks the exact key-granted catalog and health visibility in both management modes. [PR #38680](https://github.com/BerriAI/litellm/pull/38680) adds team/org/user toolset E2E assertions | Per-principal native MCP coverage is not established by REST results; reuse #38680 rather than duplicate it |
| 2. Users cannot attach unauthorized servers to their own keys | Existing live probes are recorded on LIT-4506; they are not durable endpoint regression tests | Own-key create/update escalation and its permission-validator boundary remain on existing management security tickets, including [LIT-4502](https://linear.app/litellm-ai/issue/LIT-4502). A generic route denial does not prove that validator ran |
| 3. UI/API permission parity | Existing dashboard tests cover admin operations | The same non-admin actor must be tested through browser and API; admin UI tests do not establish parity. Retained with [LIT-3644](https://linear.app/litellm-ai/issue/LIT-3644) |
| 4. Server ID determines identity | `test_mcp_lifecycle.py` grants one of two servers sharing a URL and denies calls to the other, using explicit server IDs for direct REST calls and server-qualified search results for virtual calls, with and without bearer credentials | Virtual calls identify the target by the searched tool name, not the REST `server_id` field. Bare names such as `add` are ambiguous across servers; duplicate aliases/names and unprefixed protocol routing remain with [LIT-4500](https://linear.app/litellm-ai/issue/LIT-4500) |
| 5. Same-URL servers do not share credentials | `test_oauth_configuration.py` crosses two gateway users with two server IDs and four distinct stored OAuth tokens. It checks actual upstream headers and successful results, then invalidates only one tuple | Controlled stored-token tests do not prove separate external-provider accounts or consent flows |
| 6. OAuth never falls back to anonymous | OAuth isolation variants remove a stored token or expire it without refresh, require separate list/call auth failures and no upstream requests, and preserve all other valid tuples. `test_mcp_lifecycle.py` also covers warm static-header removal and OBO without a caller JWT | External upstream revocation, refresh/reauthorization and aggregate challenges remain with [LIT-4501](https://linear.app/litellm-ai/issue/LIT-4501), [LIT-3433](https://linear.app/litellm-ai/issue/LIT-3433), [LIT-4422](https://linear.app/litellm-ai/issue/LIT-4422) and [LIT-4436](https://linear.app/litellm-ai/issue/LIT-4436) |
| 7. Stateful HTTP/session continuity | Legacy public-client tests exercise initialized sessions | No claim here proves upstream session state continuity; retained under [LIT-3143](https://linear.app/litellm-ai/issue/LIT-3143) |
| 8. Production guardrails/hooks run | `../observability/test_guardrail_effects.py` checks selected pre-call guards on direct and virtual execution, key/team/request selection, allowed results and zero denied executions. The two legacy test-owned dispatcher files are removed | This does not establish every post-call/output-scanning or concurrent hook contract |
| 9. Permissions enforced at discovery and execution | Exact key catalog and virtual search results plus forbidden direct/virtual calls in `test_mcp_lifecycle.py`; existing `../compatibility/test_persisted_toolsets.py` checks tool-level ceiling, denied sibling and allowed control | All principal/transport combinations are not established; link #38680's evidence for its additional principal cases |
| 10. Stateless/stateful matrix | These controlled peers use stateless HTTP upstreams | Stateful combinations depend on LIT-3143 and shared conformance runs. Modern-agent/legacy-upstream interaction remains deferred; legacy passes do not establish modern conformance |
## Additional JWT/OAuth acceptance
[LIT-3467 / PR #41909](https://github.com/BerriAI/litellm/pull/41909) owns one shared real login/consent, immediate list/call and cold-restart implementation, with aggregate SSO and explicitly configured per-server JWT variants. Reuse that implementation and its protected login secret; do not create another browser bootstrap here. Credit its exact-commit evidence separately from these controlled credential tests
The two-user/two-server cases here create non-admin users and scoped API keys through management APIs. They store synthetic upstream OAuth credentials through the real credential endpoint and assert the actual bearer at the owned upstream. This deliberately isolates credential lookup, expiry and revocation from consent. No gateway API key may replace the expected upstream token
Gateway JWT precedence, invalid/expired gateway JWTs, inactive-user denial, and their MCP-specific interaction with isolated credential lookup remain unverified by these API-key cases. General JWT unit/API tests are useful existing coverage but do not substitute for those MCP outcomes. Real-provider auth failures should extend LIT-3467's settled helpers; its explicit-header case must not be described as an uninterrupted Authorization-only OAuth flow
[PR #41718 / LIT-7737](https://github.com/BerriAI/litellm/pull/41718) owns dependency and public-client compatibility checks. This suite consumes the merged SDK2 API and keeps the existing dependency constraints. Its result must be reported independently of an installation-matrix pass

View file

@ -1,14 +1,18 @@
import json
import uuid
from contextlib import ExitStack
from pathlib import Path
from typing import Final
import pytest
import yaml
from hypothesis import strategies as st
from hypothesis.stateful import RuleBasedStateMachine, invariant, rule, run_state_machine_as_test
from integration._support.client import Gateway
from integration._support.database import read_rows
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
from integration._support.process import owned_proxy
from integration._support.mcp import call_tool, mcp_peer, register_mcp, tool_names
@ -53,7 +57,7 @@ def test_tool_error_remains_error_and_healthy_sibling_returns_value(gateway: Gat
failure: Final = call_tool(gateway, key, identity, names["fail"], {})
assert failure.status_code == 200, failure.text
assert failure.json()["isError"] is True
assert "synthetic tool failure" in failure.json()["content"][0]["text"]
assert failure.json()["content"][0]["text"] == "Error executing tool fail"
healthy: Final = call_tool(gateway, key, identity, names["multiply"], {"a": 3, "b": 5})
assert healthy.status_code == 200, healthy.text
assert healthy.json()["isError"] is False
@ -121,3 +125,182 @@ def test_generated_mcp_edits_preserve_actual_headers_and_tool_results(gateway: G
self.resources.close()
run_state_machine_as_test(Servers, settings=LIFECYCLE_SETTINGS)
@pytest.mark.covers("other.mcp.health.restricted_keys_intersect_grants_in_both_modes")
def test_health_intersects_route_restricted_key_grants_in_both_management_modes(
gateway: Gateway, tmp_path: Path
) -> None:
for mode in ("restricted", "view_all"):
config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"]["user_mcp_management_mode"] = mode
path = tmp_path / f"health-{mode}.yaml"
path.write_text(yaml.safe_dump(config))
with (
owned_proxy(gateway, tmp_path, {}, config=path) as candidate,
mcp_peer() as peer,
candidate.scenario() as scenario,
):
first = register_mcp(scenario, peer, "health" + uuid.uuid4().hex)
second = register_mcp(scenario, peer, "health" + uuid.uuid4().hex)
control = scenario.key(object_permission={"mcp_servers": [first]})
names = tool_names(candidate, control, first)
healthy = call_tool(candidate, control, first, names["add"], {"a": 3, "b": 5})
assert healthy.status_code == 200 and healthy.json()["content"][0]["text"] == "8", healthy.text
for grants in ([first], [second], []):
key = scenario.key(
allowed_routes=["/v1/mcp/server", "/v1/mcp/server/health"],
object_permission={"mcp_servers": grants},
)
listed = candidate.request("GET", "/v1/mcp/server", key=key)
assert listed.status_code == 200, listed.text
assert {row["server_id"] for row in listed.json()} == set(grants), listed.text
for requested in (None, [second], [first, second]):
response = candidate.client.get(
"/v1/mcp/server/health",
headers={"Authorization": f"Bearer {key}"},
params=[] if requested is None else [("server_ids", identity) for identity in requested],
)
assert response.status_code == 200, response.text
expected = set(grants) if requested is None else set(grants).intersection(requested)
assert {row["server_id"] for row in response.json()} == expected, response.text
assert all(row["status"] == "healthy" for row in response.json())
@pytest.mark.covers("other.mcp.credentials.warm_removal_fails_closed_without_upstream_traffic")
def test_warm_credential_removal_rejects_without_upstream_traffic(gateway: Gateway) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
identity = register_mcp(
scenario,
peer,
"credentials" + uuid.uuid4().hex,
auth_type="bearer_token",
static_headers={"Authorization": "Bearer synthetic-upstream-credential"},
)
key = scenario.key(object_permission={"mcp_servers": [identity]})
names = tool_names(gateway, key, identity)
warm = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
assert warm.status_code == 200 and warm.json()["content"][0]["text"] == "8", warm.text
calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
assert len(calls) == 1
assert calls[0]["headers"][b"authorization"] == b"Bearer synthetic-upstream-credential"
removed = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "static_headers": {}})
assert removed.status_code == 202, removed.text
stored = gateway.request("GET", f"/v1/mcp/server/{identity}")
assert stored.status_code == 200, stored.text
assert stored.json()["auth_type"] == "bearer_token"
assert not stored.json().get("static_headers"), stored.text
peer.drain()
for operation in ("list", "call"):
rejected = (
gateway.client.get(
"/mcp-rest/tools/list", params={"server_id": identity}, headers={"x-litellm-api-key": key}
)
if operation == "list"
else call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
)
assert rejected.status_code == 500, rejected.text
if operation == "list":
assert rejected.json()["detail"]["error"] == "internal", rejected.text
assert "Failed to list tools from server" in rejected.json()["detail"]["message"], rejected.text
else:
assert "requires a usable upstream credential" in rejected.text, rejected.text
assert peer.drain() == (), "missing static credential escaped to upstream"
changed = gateway.request(
"PUT",
"/v1/mcp/server",
{
"server_id": identity,
"auth_type": "oauth2_token_exchange",
"token_exchange_endpoint": peer.url + "/token",
"credentials": {"client_id": "synthetic-client"},
},
)
assert changed.status_code == 202, changed.text
peer.drain()
rejected_subject = call_tool(gateway, key, identity, names["add"], {"a": 3, "b": 5})
assert rejected_subject.status_code == 401, rejected_subject.text
assert peer.drain() == (), "virtual key cannot supply an OBO subject token"
control_id = register_mcp(scenario, peer, "control" + uuid.uuid4().hex, auth_type="none")
control_key = scenario.key(object_permission={"mcp_servers": [control_id]})
control_names = tool_names(gateway, control_key, control_id)
control = call_tool(gateway, control_key, control_id, control_names["multiply"], {"a": 3, "b": 5})
assert control.status_code == 200 and control.json()["content"][0]["text"] == "15", control.text
@pytest.mark.parametrize("authenticated", (False, True), ids=("anonymous", "bearer"))
@pytest.mark.covers("other.mcp.permissions.same_url_servers_enforce_discovery_and_execution")
def test_same_url_server_grants_scope_discovery_and_direct_or_virtual_execution(
gateway: Gateway, authenticated: bool
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
aliases: Final = tuple("scope" + uuid.uuid4().hex for _ in range(2))
servers: Final = tuple(
register_mcp(
scenario,
peer,
alias,
auth_type="bearer_token" if authenticated else "none",
static_headers={
"X-Integration-Server": alias,
**({"Authorization": f"Bearer synthetic-{alias}"} if authenticated else {}),
},
)
for alias in aliases
)
for virtual in (False, True):
keys: Final = tuple(
scenario.key(object_permission={"mcp_servers": [server], "mcp_tool_search_enabled": virtual})
for server in servers
)
for server, alias, key in zip(servers, aliases, keys):
catalog: Final = gateway.request("GET", "/mcp-rest/tools/list", key=key)
assert catalog.status_code == 200, catalog.text
if virtual:
assert {tool["name"] for tool in catalog.json()["tools"]} == {
"mcp_tool_search",
"mcp_tool_call",
"agent_search",
"skill_search",
}, catalog.text
search: Final = gateway.request(
"POST",
"/mcp-rest/tools/call",
{"name": "mcp_tool_search", "arguments": {"query": "add", "top_k": 10}},
key=key,
)
assert search.status_code == 200 and search.json()["isError"] is False, search.text
assert [tool["name"] for tool in json.loads(search.json()["content"][0]["text"])] == [
f"{alias}-add"
], search.text
else:
assert {tool["mcp_info"]["server_id"] for tool in catalog.json()["tools"]} == {server}
assert {tool["name"] for tool in catalog.json()["tools"]} == {"add", "multiply", "fail"}
for server_index, caller_index in ((0, 0), (1, 0), (1, 1)):
peer.drain()
response: Final = gateway.request(
"POST",
"/mcp-rest/tools/call",
{
"name": "mcp_tool_call" if virtual else "add",
**({} if virtual else {"server_id": servers[server_index]}),
"arguments": (
{"tool_name": f"{aliases[server_index]}-add", "arguments": {"a": 3, "b": 5}}
if virtual
else {"a": 3, "b": 5}
),
},
key=keys[caller_index],
)
observed: Final = peer.drain()
if server_index != caller_index:
assert response.status_code == 403 and "not allowed" in response.text, response.text
assert observed == (), "forbidden server reached the upstream"
continue
assert response.status_code == 200 and response.json()["isError"] is False, response.text
assert response.json()["content"][0]["text"] == "8", response.text
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
assert len(calls) == 1
assert calls[0]["headers"][b"x-integration-server"] == aliases[server_index].encode()
expected_auth: Final = f"Bearer synthetic-{aliases[server_index]}".encode() if authenticated else None
assert all(item["headers"].get(b"authorization") == expected_auth for item in observed)

View file

@ -2,14 +2,14 @@ import json
import queue
import uuid
from urllib.parse import parse_qs, urlsplit
from typing import Final
from typing import Final, Literal
from pathlib import Path
import pytest
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.mcp import McpPeer, register_mcp
from integration._support.mcp import McpPeer, call_tool, mcp_peer, register_mcp, tool_names
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
@ -102,3 +102,87 @@ def test_partial_discovery_and_unrelated_edit_keep_actual_authorization_destinat
"PUT", "/v1/mcp/server", {"server_id": identity, "server_name": alias + "renamed"}
)
assert updated.status_code == 202, updated.text
@pytest.mark.covers("other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server")
@pytest.mark.parametrize("transition", ("revoke", "expire"))
def test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server(
gateway: Gateway,
transition: Literal["revoke", "expire"],
) -> None:
with mcp_peer() as peer, gateway.scenario() as scenario:
servers: Final = tuple(
register_mcp(
scenario,
peer,
"oauth" + uuid.uuid4().hex,
auth_type="oauth2",
oauth2_flow="authorization_code",
authorization_url=peer.url + "/authorize",
token_url=peer.url + "/token",
credentials={"client_id": "synthetic-oauth-client"},
)
for _ in range(2)
)
users: Final = tuple(scenario.user(user_role="internal_user") for _ in range(2))
keys: Final = tuple(
scenario.key(user_id=user, object_permission={"mcp_servers": list(servers)}) for user in users
)
for user_index, key in enumerate(keys):
for server_index, server_id in enumerate(servers):
stored: Final = gateway.request(
"POST",
f"/v1/mcp/server/{server_id}/oauth-user-credential",
{"access_token": f"synthetic-user-{user_index}-server-{server_index}", "expires_in": 3600},
key=key,
)
assert stored.status_code == 200 and stored.json()["has_credential"] is True, stored.text
scenario.cleanups.callback(
gateway.request,
"DELETE",
f"/v1/mcp/server/{server_id}/oauth-user-credential",
key=key,
)
names: Final = tuple(tool_names(gateway, keys[0], server) for server in servers)
for generation in range(2):
for user_index, key in enumerate(keys):
for server_index, server_id in enumerate(servers):
peer.drain()
discovery: Final = gateway.request(
"GET",
"/mcp-rest/tools/list",
key=key,
params={"server_id": server_id},
)
call: Final = call_tool(gateway, key, server_id, names[server_index]["add"], {"a": 3, "b": 5})
observed: Final = peer.drain()
if generation == 1 and user_index == 0 and server_index == 0:
for rejected in (discovery, call):
assert rejected.status_code == 401, rejected.text
assert rejected.json() == {"detail": "Unauthorized"}, rejected.text
assert "resource_metadata=" in rejected.headers["www-authenticate"]
assert observed == (), "unusable credentials must not fall back to another user or server"
else:
assert discovery.status_code == 200, discovery.text
assert {tool["name"] for tool in discovery.json()["tools"]} == set(names[server_index].values())
assert call.status_code == 200 and call.json()["isError"] is False, call.text
assert call.json()["content"][0]["text"] == "8", call.text
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
assert len(calls) == 1
expected: Final = f"Bearer synthetic-user-{user_index}-server-{server_index}".encode()
assert calls[0]["headers"][b"authorization"] == expected
assert all(item["headers"].get(b"authorization") == expected for item in observed)
if generation == 0:
changed: Final = gateway.request(
"DELETE" if transition == "revoke" else "POST",
f"/v1/mcp/server/{servers[0]}/oauth-user-credential",
None
if transition == "revoke"
else {
"access_token": "synthetic-expired-user-0-server-0",
"expires_in": -60,
},
key=keys[0],
)
assert changed.status_code == 200, changed.text
assert changed.json()["has_credential"] is (transition == "expire"), changed.text

View file

@ -8,6 +8,7 @@ import yaml
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.mcp import mcp_peer, register_mcp, tool_names
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
@ -143,3 +144,73 @@ def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gatewa
)
assert len(observed.get("/__observations").json()["requests"]) == 1
assert len(policy.drain()) == 2
@pytest.mark.covers("other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution")
def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway: Gateway, tmp_path: Path) -> None:
guardrail = "mcp-policy-" + uuid.uuid4().hex
config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["guardrails"] = [
{
"guardrail_name": guardrail,
"litellm_params": {
"guardrail": "custom_code",
"mode": "pre_mcp_call",
"default_on": False,
"custom_code": (
"def apply_guardrail(inputs, request_data, input_type):\n"
' if inputs.get("tools", [{}])[0].get("function", {}).get("name") == "add":\n'
' return block("integration resolved add denied")\n'
" return allow()\n"
),
},
}
]
path = tmp_path / "mcp-guardrail.yaml"
path.write_text(yaml.safe_dump(config))
with (
owned_proxy(gateway, tmp_path, {}, config=path) as candidate,
mcp_peer() as peer,
candidate.scenario() as scenario,
):
identity = register_mcp(scenario, peer, "guardrail" + uuid.uuid4().hex)
permission = {"mcp_servers": [identity], "mcp_tool_search_enabled": True}
key = scenario.key(object_permission=permission)
key_selected = scenario.key(object_permission=permission, guardrails=[guardrail])
team = scenario.team(guardrails=[guardrail], object_permission={"mcp_servers": [identity]})
team_selected = scenario.key(team_id=team, object_permission=permission)
catalog_key = scenario.key(object_permission={"mcp_servers": [identity]})
names = tool_names(candidate, catalog_key, identity)
assert set(names) == {"add", "multiply", "fail"}
for virtual in (False, True):
for caller, selected, tool, expected in (
(key, [], "add", 8),
(key, [guardrail], "add", None),
(key_selected, [], "add", None),
(team_selected, [], "add", None),
(key, [guardrail], "multiply", 15),
):
arguments = {"a": 3, "b": 5}
peer.drain()
response = candidate.client.post(
"/mcp-rest/tools/call",
headers={"x-litellm-api-key": caller},
json={
"server_id": identity,
"name": "mcp_tool_call" if virtual else names[tool],
"arguments": {"tool_name": names[tool], "arguments": arguments} if virtual else arguments,
"guardrails": selected,
},
)
calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call")
if expected is None:
assert response.status_code == 400, response.text
assert "integration resolved add denied" in response.text, response.text
assert calls == (), "pre-call denial must prevent upstream execution"
else:
assert response.status_code == 200, response.text
assert response.json()["isError"] is False
assert response.json()["content"][0]["text"] == str(expected), response.text
assert len(calls) == 1
assert calls[0]["body"]["params"]["name"] == tool
assert calls[0]["body"]["params"]["arguments"] == arguments

View file

@ -1,6 +1,6 @@
"""Deterministic upstream MCP server for the mcp e2e suite.
A tiny FastMCP server exposing `add` and `multiply` over streamable-http so the
A tiny MCP server exposing `add` and `multiply` over streamable-http so the
suite has a self-hosted, offline upstream to register and exercise. DNS-rebinding
protection is turned off because the litellm container reaches this over the
compose network by service name (`mcp-upstream:8090`), not localhost, and the
@ -9,15 +9,10 @@ stack is an isolated throwaway. Bind host/port come from MCP_HOST/MCP_PORT.
import os
from mcp.server.fastmcp import FastMCP
from mcp.server.mcpserver import MCPServer
from mcp.server.transport_security import TransportSecuritySettings
mcp: FastMCP = FastMCP(
"e2e-math",
host=os.getenv("MCP_HOST", "0.0.0.0"),
port=int(os.getenv("MCP_PORT", "8090")),
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
mcp: MCPServer = MCPServer("e2e-math")
@mcp.tool()
@ -33,7 +28,12 @@ def multiply(a: int, b: int) -> int:
def main() -> None:
mcp.run(transport="streamable-http")
mcp.run(
transport="streamable-http",
host=os.getenv("MCP_HOST", "0.0.0.0"),
port=int(os.getenv("MCP_PORT", "8090")),
transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False),
)
if __name__ == "__main__":

View file

@ -1,770 +0,0 @@
"""
Test file for MCP Guardrails Feature
This file tests the MCP guardrails functionality for both pre and during MCP call hooks,
including various guardrail types and proper exception handling.
"""
import asyncio
import pytest
from datetime import datetime
from typing import Optional, Dict, Any
from unittest.mock import MagicMock, AsyncMock, patch
# Add the project root to the path
import litellm
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching.caching import DualCache
from litellm.types.mcp import (
MCPPreCallRequestObject,
MCPPreCallResponseObject,
MCPDuringCallRequestObject,
MCPDuringCallResponseObject,
)
from litellm.types.llms.base import HiddenParams
from litellm.types.guardrails import GuardrailEventHooks
from fastapi import HTTPException
class MockPiiGuardrail(CustomGuardrail):
"""Mock PII guardrail that raises BlockedPiiEntityError"""
def __init__(self, should_block: bool = True, entity_type: str = "EMAIL_ADDRESS"):
super().__init__()
self.should_block = should_block
self.entity_type = entity_type
self.guardrail_name = "mock-pii-guardrail"
self.call_count = 0
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
"""Always run for testing"""
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
"""Mock pre-call hook that raises BlockedPiiEntityError"""
self.call_count += 1
if self.should_block:
raise BlockedPiiEntityError(
entity_type=self.entity_type,
guardrail_name=self.guardrail_name,
)
return None
class MockContentGuardrail(CustomGuardrail):
"""Mock content guardrail that raises GuardrailRaisedException"""
def __init__(self, should_block: bool = True):
super().__init__()
self.should_block = should_block
self.guardrail_name = "mock-content-guardrail"
self.call_count = 0
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
"""Always run for testing"""
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
"""Mock pre-call hook that raises GuardrailRaisedException"""
self.call_count += 1
if self.should_block:
raise GuardrailRaisedException(
guardrail_name=self.guardrail_name, message="Content violates policy"
)
return None
class MockHttpGuardrail(CustomGuardrail):
"""Mock HTTP guardrail that raises HTTPException"""
def __init__(self, should_block: bool = True):
super().__init__()
self.should_block = should_block
self.guardrail_name = "mock-http-guardrail"
self.call_count = 0
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
"""Always run for testing"""
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
"""Mock pre-call hook that raises HTTPException"""
self.call_count += 1
if self.should_block:
raise HTTPException(
status_code=400, detail={"error": "Violated guardrail policy"}
)
return None
class MockDuringCallGuardrail(CustomGuardrail):
"""Mock guardrail for during-call testing"""
def __init__(self, should_block: bool = True):
super().__init__()
self.should_block = should_block
self.guardrail_name = "mock-during-guardrail"
self.call_count = 0
def should_run_guardrail(self, data: dict, event_type: GuardrailEventHooks) -> bool:
"""Always run for testing"""
return True
async def async_moderation_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
call_type: str,
):
"""Mock during-call hook that raises exceptions"""
self.call_count += 1
if self.should_block:
raise BlockedPiiEntityError(
entity_type="PHONE_NUMBER",
guardrail_name=self.guardrail_name,
)
return None
class MockProxyLogging:
"""Mock proxy logging object for testing MCP guardrails"""
def __init__(self, guardrails: Optional[list] = None):
self.guardrails = guardrails if guardrails is not None else []
self.call_details = {"user_api_key_cache": DualCache()}
self.dynamic_success_callbacks = []
self.call_count = 0
def get_combined_callback_list(self, dynamic_success_callbacks, global_callbacks):
"""Return the guardrails for testing"""
return self.guardrails
def _convert_mcp_to_llm_format(self, request_obj, kwargs: dict) -> dict:
"""Convert MCP tool call to LLM message format"""
tool_call_content = (
f"Tool: {request_obj.tool_name}\nArguments: {request_obj.arguments}"
)
return {
"messages": [{"role": "user", "content": tool_call_content}],
"model": kwargs.get("model", "mcp-tool-call"),
"user_api_key_user_id": kwargs.get("user_api_key_user_id"),
"user_api_key_team_id": kwargs.get("user_api_key_team_id"),
}
def _convert_llm_result_to_mcp_response(self, llm_result, request_obj):
"""Convert LLM result back to MCP response format"""
return None # For testing, we don't need to convert back
def _parse_pre_mcp_call_hook_response(self, response, original_request):
"""Parse pre MCP call hook response"""
return response
async def async_pre_mcp_tool_call_hook(
self,
kwargs: dict,
request_obj: Any,
start_time: datetime,
end_time: datetime,
) -> Optional[Any]:
"""Mock pre MCP tool call hook"""
self.call_count += 1
# Simulate the actual hook logic
for guardrail in self.guardrails:
if isinstance(guardrail, CustomGuardrail):
try:
synthetic_data = self._convert_mcp_to_llm_format(
request_obj, kwargs
)
# Check if guardrail should run
if not guardrail.should_run_guardrail(
synthetic_data, GuardrailEventHooks.pre_mcp_call
):
continue
result = await guardrail.async_pre_call_hook(
user_api_key_dict=kwargs.get("user_api_key_auth"),
cache=self.call_details["user_api_key_cache"],
data=synthetic_data,
call_type="mcp_call",
)
if result is not None:
return self._parse_pre_mcp_call_hook_response(
result, request_obj
)
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions
raise e
except Exception as e:
# Log non-guardrail exceptions as non-blocking
print(
f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {str(e)}"
)
return None
async def async_during_mcp_tool_call_hook(
self,
kwargs: dict,
request_obj: Any,
start_time: datetime,
end_time: datetime,
) -> Optional[Any]:
"""Mock during MCP tool call hook"""
self.call_count += 1
# Simulate the actual hook logic
for guardrail in self.guardrails:
if isinstance(guardrail, CustomGuardrail):
try:
synthetic_data = self._convert_mcp_to_llm_format(
request_obj, kwargs
)
result = await guardrail.async_moderation_hook(
data=synthetic_data,
user_api_key_dict=kwargs.get("user_api_key_auth"),
call_type="mcp_call",
)
if result is not None:
return result
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions
raise e
except Exception as e:
# Log non-guardrail exceptions as non-blocking
print(
f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging {str(e)}"
)
return None
@pytest.fixture
def mock_user_api_key():
"""Mock user API key for testing"""
return UserAPIKeyAuth(api_key="test_key", user_id="test_user")
@pytest.fixture
def mock_cache():
"""Mock cache for testing"""
return DualCache()
@pytest.fixture
def mock_pii_guardrail():
"""Mock PII guardrail that blocks"""
return MockPiiGuardrail(should_block=True)
@pytest.fixture
def mock_pii_guardrail_allow():
"""Mock PII guardrail that allows"""
return MockPiiGuardrail(should_block=False)
@pytest.fixture
def mock_content_guardrail():
"""Mock content guardrail that blocks"""
return MockContentGuardrail(should_block=True)
@pytest.fixture
def mock_http_guardrail():
"""Mock HTTP guardrail that blocks"""
return MockHttpGuardrail(should_block=True)
@pytest.fixture
def mock_during_guardrail():
"""Mock during-call guardrail that blocks"""
return MockDuringCallGuardrail(should_block=True)
@pytest.fixture
def mock_proxy_logging():
"""Mock proxy logging object"""
return MockProxyLogging()
class TestMCPGuardrailsPreCall:
"""Test MCP guardrails for pre-call hooks"""
@pytest.mark.asyncio
async def test_pii_guardrail_blocks_pre_call(
self, mock_pii_guardrail, mock_user_api_key, mock_cache
):
"""Test that PII guardrail properly blocks pre-call"""
proxy_logging = MockProxyLogging([mock_pii_guardrail])
# Create MCP request
request_obj = MCPPreCallRequestObject(
tool_name="email_tool",
arguments={"email": "test@example.com"},
server_name="email_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "email_tool",
"arguments": {"email": "test@example.com"},
"server_name": "email_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that BlockedPiiEntityError is raised
with pytest.raises(BlockedPiiEntityError) as excinfo:
await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Verify the error details
assert excinfo.value.entity_type == "EMAIL_ADDRESS"
assert excinfo.value.guardrail_name == "mock-pii-guardrail"
assert mock_pii_guardrail.call_count == 1
@pytest.mark.asyncio
async def test_pii_guardrail_allows_pre_call(
self, mock_pii_guardrail_allow, mock_user_api_key, mock_cache
):
"""Test that PII guardrail allows pre-call when configured to allow"""
proxy_logging = MockProxyLogging([mock_pii_guardrail_allow])
request_obj = MCPPreCallRequestObject(
tool_name="email_tool",
arguments={"email": "test@example.com"},
server_name="email_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "email_tool",
"arguments": {"email": "test@example.com"},
"server_name": "email_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that no exception is raised
result = await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is None
assert mock_pii_guardrail_allow.call_count == 1
@pytest.mark.asyncio
async def test_content_guardrail_blocks_pre_call(
self, mock_content_guardrail, mock_user_api_key, mock_cache
):
"""Test that content guardrail properly blocks pre-call"""
proxy_logging = MockProxyLogging([mock_content_guardrail])
request_obj = MCPPreCallRequestObject(
tool_name="content_tool",
arguments={"content": "sensitive content"},
server_name="content_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "content_tool",
"arguments": {"content": "sensitive content"},
"server_name": "content_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that GuardrailRaisedException is raised
with pytest.raises(GuardrailRaisedException) as excinfo:
await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Verify the error details
assert "Content violates policy" in str(excinfo.value)
assert excinfo.value.guardrail_name == "mock-content-guardrail"
assert mock_content_guardrail.call_count == 1
@pytest.mark.asyncio
async def test_http_guardrail_blocks_pre_call(
self, mock_http_guardrail, mock_user_api_key, mock_cache
):
"""Test that HTTP guardrail properly blocks pre-call"""
proxy_logging = MockProxyLogging([mock_http_guardrail])
request_obj = MCPPreCallRequestObject(
tool_name="http_tool",
arguments={"url": "http://example.com"},
server_name="http_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "http_tool",
"arguments": {"url": "http://example.com"},
"server_name": "http_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that HTTPException is raised
with pytest.raises(HTTPException) as excinfo:
await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Verify the error details
assert excinfo.value.status_code == 400
assert "Violated guardrail policy" in str(excinfo.value.detail)
assert mock_http_guardrail.call_count == 1
@pytest.mark.asyncio
async def test_multiple_guardrails_pre_call(
self, mock_pii_guardrail, mock_content_guardrail, mock_user_api_key, mock_cache
):
"""Test multiple guardrails - first one should block"""
proxy_logging = MockProxyLogging([mock_pii_guardrail, mock_content_guardrail])
request_obj = MCPPreCallRequestObject(
tool_name="test_tool",
arguments={"email": "test@example.com"},
server_name="test_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "test_tool",
"arguments": {"email": "test@example.com"},
"server_name": "test_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that first guardrail blocks
with pytest.raises(BlockedPiiEntityError):
await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Verify only first guardrail was called
assert mock_pii_guardrail.call_count == 1
assert mock_content_guardrail.call_count == 0
class TestMCPGuardrailsDuringCall:
"""Test MCP guardrails for during-call hooks"""
@pytest.mark.asyncio
async def test_during_call_guardrail_blocks(
self, mock_during_guardrail, mock_user_api_key, mock_cache
):
"""Test that during-call guardrail properly blocks execution"""
proxy_logging = MockProxyLogging([mock_during_guardrail])
request_obj = MCPDuringCallRequestObject(
tool_name="phone_tool",
arguments={"phone": "555-123-4567"},
server_name="phone_server",
start_time=datetime.now().timestamp(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "phone_tool",
"arguments": {"phone": "555-123-4567"},
"server_name": "phone_server",
}
# Test that BlockedPiiEntityError is raised
with pytest.raises(BlockedPiiEntityError) as excinfo:
await proxy_logging.async_during_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Verify the error details
assert excinfo.value.entity_type == "PHONE_NUMBER"
assert excinfo.value.guardrail_name == "mock-during-guardrail"
assert mock_during_guardrail.call_count == 1
class TestMCPGuardrailsIntegration:
"""Test MCP guardrails integration with MCP server manager"""
@pytest.mark.asyncio
async def test_mcp_server_manager_with_guardrails(self):
"""Test MCP server manager with guardrail integration"""
mock_proxy_logging = MockProxyLogging([MockPiiGuardrail(should_block=True)])
# Test that guardrail exception is properly raised in the hook
with pytest.raises(BlockedPiiEntityError):
await mock_proxy_logging.async_pre_mcp_tool_call_hook(
kwargs={
"name": "email_tool",
"arguments": {"email": "test@example.com"},
},
request_obj=MagicMock(),
start_time=datetime.now(),
end_time=datetime.now(),
)
@pytest.mark.asyncio
async def test_guardrail_exception_propagation(self):
"""Test that guardrail exceptions properly propagate through the system"""
# Test BlockedPiiEntityError
with pytest.raises(BlockedPiiEntityError):
raise BlockedPiiEntityError(
entity_type="EMAIL_ADDRESS", guardrail_name="test-guardrail"
)
# Test GuardrailRaisedException
with pytest.raises(GuardrailRaisedException):
raise GuardrailRaisedException(
guardrail_name="test-guardrail", message="Test message"
)
# Test HTTPException
with pytest.raises(HTTPException):
raise HTTPException(status_code=400, detail={"error": "Test error"})
class TestMCPGuardrailsErrorHandling:
"""Test MCP guardrails error handling scenarios"""
@pytest.mark.asyncio
async def test_non_guardrail_exception_logging(self, mock_user_api_key, mock_cache):
"""Test that non-guardrail exceptions are logged as non-blocking"""
class MockFailingGuardrail(CustomGuardrail):
def should_run_guardrail(
self, data: dict, event_type: GuardrailEventHooks
) -> bool:
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
raise Exception("Non-guardrail error")
proxy_logging = MockProxyLogging([MockFailingGuardrail()])
request_obj = MCPPreCallRequestObject(
tool_name="test_tool",
arguments={"test": "data"},
server_name="test_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "test_tool",
"arguments": {"test": "data"},
"server_name": "test_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that non-guardrail exceptions are handled gracefully
result = await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Should return None (not raise exception)
assert result is None
@pytest.mark.asyncio
async def test_guardrail_should_not_run(self, mock_user_api_key, mock_cache):
"""Test that guardrails don't run when should_run_guardrail returns False"""
class MockConditionalGuardrail(CustomGuardrail):
def should_run_guardrail(
self, data: dict, event_type: GuardrailEventHooks
) -> bool:
return False # Don't run
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
raise BlockedPiiEntityError("EMAIL_ADDRESS", "test-guardrail")
proxy_logging = MockProxyLogging([MockConditionalGuardrail()])
request_obj = MCPPreCallRequestObject(
tool_name="test_tool",
arguments={"test": "data"},
server_name="test_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "test_tool",
"arguments": {"test": "data"},
"server_name": "test_server",
"user_api_key_auth": mock_user_api_key,
}
# Test that guardrail doesn't run and no exception is raised
result = await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Should return None (guardrail didn't run)
assert result is None
class TestMCPGuardrailsEdgeCases:
"""Test MCP guardrails edge cases and error conditions"""
@pytest.mark.asyncio
async def test_empty_guardrails_list(self, mock_user_api_key, mock_cache):
"""Test behavior with empty guardrails list"""
proxy_logging = MockProxyLogging([]) # No guardrails
request_obj = MCPPreCallRequestObject(
tool_name="test_tool",
arguments={"test": "data"},
server_name="test_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "test_tool",
"arguments": {"test": "data"},
"server_name": "test_server",
"user_api_key_auth": mock_user_api_key,
}
# Should return None without any issues
result = await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is None
@pytest.mark.asyncio
async def test_guardrail_with_invalid_data(self, mock_user_api_key, mock_cache):
"""Test guardrail behavior with invalid data"""
class MockInvalidDataGuardrail(CustomGuardrail):
def should_run_guardrail(
self, data: dict, event_type: GuardrailEventHooks
) -> bool:
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
# Try to access invalid data
invalid_data = data.get("invalid_key", {})
if invalid_data.get("should_fail"):
raise BlockedPiiEntityError("EMAIL_ADDRESS", "test-guardrail")
return None
proxy_logging = MockProxyLogging([MockInvalidDataGuardrail()])
request_obj = MCPPreCallRequestObject(
tool_name="test_tool",
arguments={"test": "data"},
server_name="test_server",
user_api_key_auth=mock_user_api_key.model_dump(),
hidden_params=HiddenParams(),
)
kwargs = {
"name": "test_tool",
"arguments": {"test": "data"},
"server_name": "test_server",
"user_api_key_auth": mock_user_api_key,
}
# Should handle invalid data gracefully
result = await proxy_logging.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is None
if __name__ == "__main__":
pytest.main([__file__])

View file

@ -1,475 +0,0 @@
"""
Test file for MCP Hook Architecture
This file demonstrates the new MCP hook system with comprehensive examples
and validation tests.
"""
import asyncio
import pytest
from datetime import datetime
from typing import Optional
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.mcp import (
MCPPreCallRequestObject,
MCPPreCallResponseObject,
MCPDuringCallRequestObject,
MCPDuringCallResponseObject,
MCPPostCallResponseObject,
)
from litellm.types.llms.base import HiddenParams
class TestMCPAccessControlHook(CustomLogger):
"""Test hook for access control functionality"""
def __init__(self):
self.allowed_tools = {"github/create_issue", "zapier/send_email"}
self.blocked_users = {"user123", "user456"}
self.call_count = 0
async def async_pre_mcp_tool_call_hook(
self, kwargs, request_obj: MCPPreCallRequestObject, start_time, end_time
) -> Optional[MCPPreCallResponseObject]:
"""Test access control validation"""
self.call_count += 1
tool_name = request_obj.tool_name
user_id = kwargs.get("user_api_key_auth", {}).get("user_id")
# Check if user is blocked
if user_id in self.blocked_users:
return MCPPreCallResponseObject(
should_proceed=False,
error_message=f"User {user_id} is not authorized to use MCP tools",
)
# Check if tool is allowed
if tool_name not in self.allowed_tools:
return MCPPreCallResponseObject(
should_proceed=False,
error_message=f"Tool {tool_name} is not authorized",
)
return None # Allow execution to proceed
class TestMCPCostTrackingHook(CustomLogger):
"""Test hook for cost tracking functionality"""
def __init__(self):
self.cost_map = {
"github/create_issue": 0.10,
"zapier/send_email": 0.05,
"default": 0.01,
}
self.call_count = 0
async def async_post_mcp_tool_call_hook(
self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time
) -> Optional[MCPPostCallResponseObject]:
"""Test cost calculation after tool execution"""
self.call_count += 1
tool_name = kwargs.get("name", "")
cost = self.cost_map.get(tool_name, self.cost_map["default"])
# Set the response cost
response_obj.hidden_params.response_cost = cost
return response_obj
class TestMCPMonitoringHook(CustomLogger):
"""Test hook for real-time monitoring functionality"""
def __init__(self):
self.max_execution_time = 30.0 # seconds
self.call_count = 0
async def async_during_mcp_tool_call_hook(
self, kwargs, request_obj: MCPDuringCallRequestObject, start_time, end_time
) -> Optional[MCPDuringCallResponseObject]:
"""Test execution time monitoring"""
self.call_count += 1
tool_name = request_obj.tool_name
execution_time = (datetime.now() - start_time).total_seconds()
# Check if execution is taking too long
if execution_time > self.max_execution_time:
return MCPDuringCallResponseObject(
should_continue=False,
error_message=f"Tool {tool_name} execution timeout after {execution_time}s",
)
return None # Allow execution to continue
class TestMCPArgumentValidationHook(CustomLogger):
"""Test hook for argument validation functionality"""
def __init__(self):
self.call_count = 0
async def async_pre_mcp_tool_call_hook(
self, kwargs, request_obj: MCPPreCallRequestObject, start_time, end_time
) -> Optional[MCPPreCallResponseObject]:
"""Test argument validation and sanitization"""
self.call_count += 1
tool_name = request_obj.tool_name
arguments = request_obj.arguments.copy() # Create a copy to modify
# Example: Validate GitHub issue creation
if tool_name == "github/create_issue":
if not arguments.get("title"):
return MCPPreCallResponseObject(
should_proceed=False, error_message="GitHub issue title is required"
)
# Sanitize the title
title = arguments["title"]
if len(title) > 100:
title = title[:97] + "..."
arguments["title"] = title
# Example: Validate email sending
elif tool_name == "zapier/send_email":
if not arguments.get("to"):
return MCPPreCallResponseObject(
should_proceed=False, error_message="Email recipient is required"
)
return MCPPreCallResponseObject(
should_proceed=True, modified_arguments=arguments
)
# Test fixtures
@pytest.fixture
def access_control_hook():
return TestMCPAccessControlHook()
@pytest.fixture
def cost_tracking_hook():
return TestMCPCostTrackingHook()
@pytest.fixture
def monitoring_hook():
return TestMCPMonitoringHook()
@pytest.fixture
def argument_validation_hook():
return TestMCPArgumentValidationHook()
# Test cases
class TestMCPHooks:
"""Test cases for MCP hook functionality"""
@pytest.mark.asyncio
async def test_access_control_hook_allowed_tool(self, access_control_hook):
"""Test that allowed tools pass validation"""
kwargs = {
"user_api_key_auth": {"user_id": "user789"},
"name": "github/create_issue",
}
request_obj = MCPPreCallRequestObject(
tool_name="github/create_issue",
arguments={"title": "Test issue"},
user_api_key_auth={"user_id": "user789"},
)
result = await access_control_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is None # Should allow execution
assert access_control_hook.call_count == 1
@pytest.mark.asyncio
async def test_access_control_hook_blocked_user(self, access_control_hook):
"""Test that blocked users are rejected"""
kwargs = {
"user_api_key_auth": {"user_id": "user123"},
"name": "github/create_issue",
}
request_obj = MCPPreCallRequestObject(
tool_name="github/create_issue",
arguments={"title": "Test issue"},
user_api_key_auth={"user_id": "user123"},
)
result = await access_control_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is False
assert "not authorized" in result.error_message
@pytest.mark.asyncio
async def test_access_control_hook_unauthorized_tool(self, access_control_hook):
"""Test that unauthorized tools are rejected"""
kwargs = {
"user_api_key_auth": {"user_id": "user789"},
"name": "unauthorized_tool",
}
request_obj = MCPPreCallRequestObject(
tool_name="unauthorized_tool",
arguments={"param": "value"},
user_api_key_auth={"user_id": "user789"},
)
result = await access_control_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is False
assert "not authorized" in result.error_message
@pytest.mark.asyncio
async def test_cost_tracking_hook(self, cost_tracking_hook):
"""Test cost tracking functionality"""
kwargs = {"name": "github/create_issue"}
response_obj = MCPPostCallResponseObject(
mcp_tool_call_response=[], hidden_params=HiddenParams()
)
result = await cost_tracking_hook.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=response_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.hidden_params.response_cost == 0.10
assert cost_tracking_hook.call_count == 1
@pytest.mark.asyncio
async def test_cost_tracking_hook_default_cost(self, cost_tracking_hook):
"""Test default cost assignment"""
kwargs = {"name": "unknown_tool"}
response_obj = MCPPostCallResponseObject(
mcp_tool_call_response=[], hidden_params=HiddenParams()
)
result = await cost_tracking_hook.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=response_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.hidden_params.response_cost == 0.01 # Default cost
@pytest.mark.asyncio
async def test_monitoring_hook_normal_execution(self, monitoring_hook):
"""Test monitoring hook with normal execution time"""
kwargs = {"name": "test_tool"}
request_obj = MCPDuringCallRequestObject(
tool_name="test_tool", arguments={}, start_time=datetime.now().timestamp()
)
result = await monitoring_hook.async_during_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is None # Should allow execution to continue
assert monitoring_hook.call_count == 1
@pytest.mark.asyncio
async def test_argument_validation_hook_valid_github_issue(
self, argument_validation_hook
):
"""Test argument validation for valid GitHub issue"""
kwargs = {"name": "github/create_issue"}
request_obj = MCPPreCallRequestObject(
tool_name="github/create_issue", arguments={"title": "Valid issue title"}
)
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is True
assert result.modified_arguments == {"title": "Valid issue title"}
assert argument_validation_hook.call_count == 1
@pytest.mark.asyncio
async def test_argument_validation_hook_missing_title(
self, argument_validation_hook
):
"""Test argument validation for missing GitHub issue title"""
kwargs = {"name": "github/create_issue"}
request_obj = MCPPreCallRequestObject(
tool_name="github/create_issue", arguments={} # Missing title
)
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is False
assert "title is required" in result.error_message
@pytest.mark.asyncio
async def test_argument_validation_hook_long_title_sanitization(
self, argument_validation_hook
):
"""Test argument validation with title sanitization"""
kwargs = {"name": "github/create_issue"}
long_title = "A" * 150 # Very long title
request_obj = MCPPreCallRequestObject(
tool_name="github/create_issue", arguments={"title": long_title}
)
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is True
assert len(result.modified_arguments["title"]) == 100 # Truncated
assert result.modified_arguments["title"].endswith("...")
@pytest.mark.asyncio
async def test_argument_validation_hook_email_validation(
self, argument_validation_hook
):
"""Test argument validation for email sending"""
kwargs = {"name": "zapier/send_email"}
request_obj = MCPPreCallRequestObject(
tool_name="zapier/send_email",
arguments={"to": "test@example.com", "subject": "Test"},
)
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is True
assert result.modified_arguments == {
"to": "test@example.com",
"subject": "Test",
}
@pytest.mark.asyncio
async def test_argument_validation_hook_missing_email_recipient(
self, argument_validation_hook
):
"""Test argument validation for missing email recipient"""
kwargs = {"name": "zapier/send_email"}
request_obj = MCPPreCallRequestObject(
tool_name="zapier/send_email",
arguments={"subject": "Test"}, # Missing 'to' field
)
result = await argument_validation_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert result is not None
assert result.should_proceed is False
assert "recipient is required" in result.error_message
# Integration test
class TestMCPHookIntegration:
"""Integration tests for MCP hook system"""
@pytest.mark.asyncio
async def test_hook_chain_execution(self):
"""Test that multiple hooks can work together"""
access_hook = TestMCPAccessControlHook()
cost_hook = TestMCPCostTrackingHook()
validation_hook = TestMCPArgumentValidationHook()
# Test data
kwargs = {
"user_api_key_auth": {"user_id": "user789"},
"name": "github/create_issue",
}
request_obj = MCPPreCallRequestObject(
tool_name="github/create_issue",
arguments={"title": "Integration test issue"},
user_api_key_auth={"user_id": "user789"},
)
# Execute pre-hooks
access_result = await access_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
validation_result = await validation_hook.async_pre_mcp_tool_call_hook(
kwargs=kwargs,
request_obj=request_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Both hooks should allow execution
assert access_result is None
assert validation_result is not None
assert validation_result.should_proceed is True
# Simulate post-hook execution
response_obj = MCPPostCallResponseObject(
mcp_tool_call_response=[], hidden_params=HiddenParams()
)
cost_result = await cost_hook.async_post_mcp_tool_call_hook(
kwargs=kwargs,
response_obj=response_obj,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert cost_result is not None
assert cost_result.hidden_params.response_cost == 0.10
if __name__ == "__main__":
# Run the tests
pytest.main([__file__, "-v"])

View file

@ -260,3 +260,90 @@ async def test_endpoint_with_no_prisma_client(mock_user_api_key_auth):
with pytest.raises(HTTPException) as exc_info:
await reset_event_settings(user_api_key_dict=mock_user_api_key_auth)
assert exc_info.value.status_code == 500
def _prisma_recording_upserts(upserts):
client = mock.MagicMock()
async def find_unique(*args, **kwargs):
return None
async def upsert(*args, **kwargs):
upserts.append(kwargs)
return None
client.db.litellm_config.find_unique = find_unique
client.db.litellm_config.upsert = upsert
return client
def _proxy_config_owning(general_settings):
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config._load_yaml_settings_stores({"general_settings": general_settings})
return proxy_config
@pytest.mark.asyncio
async def test_save_email_settings_refuses_a_config_owned_email_settings():
upserts = []
client = _prisma_recording_upserts(upserts)
proxy_config = _proxy_config_owning({"email_settings": {EmailEvent.new_user_invitation.value: True}})
with mock.patch("litellm.proxy.proxy_server.proxy_config", proxy_config): # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
with pytest.raises(HTTPException) as refused:
await _save_email_settings(client, {EmailEvent.new_user_invitation.value: False})
assert refused.value.status_code == 400
assert refused.value.detail["keys"] == ["email_settings"]
assert upserts == []
@pytest.mark.asyncio
async def test_update_event_settings_surfaces_the_config_owned_refusal(mock_user_api_key_auth):
upserts = []
client = _prisma_recording_upserts(upserts)
proxy_config = _proxy_config_owning({"email_settings": {EmailEvent.virtual_key_created.value: False}})
request = EmailEventSettingsUpdateRequest(
settings=[EmailEventSettings(event=EmailEvent.virtual_key_created, enabled=True)]
)
with mock.patch("litellm.proxy.proxy_server.prisma_client", client): # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
with mock.patch("litellm.proxy.proxy_server.proxy_config", proxy_config): # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
with pytest.raises(HTTPException) as refused:
await update_event_settings(request=request, user_api_key_dict=mock_user_api_key_auth)
assert refused.value.status_code == 400
assert refused.value.detail["keys"] == ["email_settings"]
assert upserts == []
@pytest.mark.asyncio
async def test_save_email_settings_still_writes_when_the_config_file_is_silent():
upserts = []
client = _prisma_recording_upserts(upserts)
proxy_config = _proxy_config_owning({})
with mock.patch("litellm.proxy.proxy_server.proxy_config", proxy_config): # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
await _save_email_settings(client, {EmailEvent.new_user_invitation.value: False})
assert len(upserts) == 1
written = json.loads(upserts[0]["data"]["create"]["param_value"])
assert written["email_settings"] == {EmailEvent.new_user_invitation.value: False}
@pytest.mark.asyncio
async def test_reset_event_settings_surfaces_the_config_owned_refusal(mock_user_api_key_auth):
upserts = []
client = _prisma_recording_upserts(upserts)
proxy_config = _proxy_config_owning({"email_settings": {EmailEvent.new_user_invitation.value: True}})
with mock.patch("litellm.proxy.proxy_server.prisma_client", client): # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
with mock.patch("litellm.proxy.proxy_server.proxy_config", proxy_config): # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
with pytest.raises(HTTPException) as refused:
await reset_event_settings(user_api_key_dict=mock_user_api_key_auth)
assert refused.value.status_code == 400
assert refused.value.detail["keys"] == ["email_settings"]
assert upserts == []

View file

@ -1067,6 +1067,7 @@ async def test_afile_content_passes_trusted_model_credentials_to_router():
managed_files = _make_managed_files_instance()
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
@ -1238,6 +1239,7 @@ async def test_afile_content_bedrock_unified_id_end_to_end(monkeypatch):
managed_files = _make_managed_files_instance()
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
@ -1268,6 +1270,7 @@ async def test_afile_content_error_reports_unified_id_not_provider_uri():
managed_files = _make_managed_files_instance()
unified_file_id = "litellm_proxy_unified_id_abc"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
@ -1732,6 +1735,40 @@ async def test_batch_retrieve_hook_does_not_claim_attribution():
assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False
def _unified_batch_id(llm_batch_id: str) -> str:
decoded = f"litellm_proxy;model_id:my-vllm;llm_batch_id:{llm_batch_id}"
return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=")
@pytest.mark.asyncio
@pytest.mark.parametrize(
"llm_batch_id, stores",
[("litellm_batch_abc", False), ("batch_abc", True)],
ids=["litellm-executed batch is left alone", "provider batch is still stored"],
)
async def test_post_call_hook_leaves_litellm_executed_batches_untouched(llm_batch_id: str, stores: bool):
managed_files = _make_managed_files_instance()
response = _make_batch_response(status="in_progress", output_file_id=None)
response.id = _unified_batch_id(llm_batch_id)
response._hidden_params = {
"unified_batch_id": response.id,
"model_id": "my-vllm",
"model_name": "hosted_vllm/qwen",
}
original_id = response.id
returned = await managed_files.async_post_call_success_hook(
data={},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-the-poller", user_id="bob", parent_otel_span=None),
response=response,
)
assert returned is response
assert managed_files.store_unified_object_id.await_count == (1 if stores else 0)
if not stores:
assert response.id == original_id
@pytest.mark.asyncio
async def test_afile_delete_passes_trusted_model_credentials_to_router():
"""
@ -1743,6 +1780,7 @@ async def test_afile_delete_passes_trusted_model_credentials_to_router():
managed_files = _make_managed_files_instance()
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
@ -1809,6 +1847,7 @@ async def test_afile_delete_bedrock_unified_id_end_to_end(monkeypatch):
managed_files = _make_managed_files_instance()
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
managed_files.get_unified_file_id = AsyncMock(return_value=None)
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
@ -1827,3 +1866,147 @@ async def test_afile_delete_bedrock_unified_id_end_to_end(monkeypatch):
assert response.id == unified_file_id
assert response.model_dump() == {"id": unified_file_id, "object": "file", "deleted": True}
managed_files.delete_unified_file_id.assert_awaited_once_with(unified_file_id, None)
@pytest.mark.asyncio
async def test_afile_delete_storage_backed_row_deletes_stored_content_not_provider_files():
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
from openai.types import FileDeleted
from litellm.caching import DualCache
from litellm.models.managed_files import LiteLLM_ManagedFileTable
storage_url = "litellm_db://content-row-1"
unified_file_id = _managed_deletion_file_id(storage_url)
row = LiteLLM_ManagedFileTable(
unified_file_id=unified_file_id,
model_mappings={"vllm-batch": storage_url},
flat_model_file_ids=[storage_url],
file_object=_make_file_object(unified_file_id),
storage_backend="litellm_db",
storage_url=storage_url,
)
file_table = MagicMock(find_first=AsyncMock(return_value=row), delete=AsyncMock())
content_table = MagicMock(delete=AsyncMock())
managed_files = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
),
)
router = MagicMock(
get_deployment_credentials_with_provider=MagicMock(return_value=None),
afile_delete=AsyncMock(),
)
response = await managed_files.afile_delete(
file_id=unified_file_id,
litellm_parent_otel_span=None,
llm_router=router,
)
content_table.delete.assert_awaited_once_with(where={"id": "content-row-1"})
router.afile_delete.assert_not_awaited()
file_table.delete.assert_awaited_once_with(where={"unified_file_id": unified_file_id})
assert response == FileDeleted(id=unified_file_id, object="file", deleted=True)
@pytest.mark.asyncio
async def test_afile_content_storage_backed_row_returns_stored_bytes_not_provider_content():
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
from prisma import Base64
from litellm.caching import DualCache
from litellm.models.managed_files import LiteLLM_ManagedFileTable
storage_url = "litellm_db://content-row-1"
unified_file_id = _managed_deletion_file_id(storage_url)
stored_bytes = b'{"custom_id": "line-1", "method": "POST", "url": "/v1/chat/completions", "body": {}}\n'
row = LiteLLM_ManagedFileTable(
unified_file_id=unified_file_id,
model_mappings={"vllm-batch": storage_url},
flat_model_file_ids=[storage_url],
file_object=_make_file_object(unified_file_id),
storage_backend="litellm_db",
storage_url=storage_url,
)
file_table = MagicMock(find_first=AsyncMock(return_value=row))
content_table = MagicMock(find_unique=AsyncMock(return_value=MagicMock(content=Base64.encode(stored_bytes))))
managed_files = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(
db=MagicMock(litellm_managedfiletable=file_table, litellm_managedfilecontenttable=content_table)
),
)
router = MagicMock(
get_deployment_credentials_with_provider=MagicMock(return_value=None),
afile_content=AsyncMock(),
)
response = await managed_files.afile_content(
file_id=unified_file_id,
litellm_parent_otel_span=None,
llm_router=router,
)
assert response.content == stored_bytes
content_table.find_unique.assert_awaited_once_with(where={"id": "content-row-1"})
router.afile_content.assert_not_awaited()
@pytest.mark.asyncio
async def test_store_unified_object_id_batch_processed_is_written_only_when_asked():
managed_files, mock_prisma = _make_object_store_instance()
upsert = mock_prisma.db.litellm_managedobjecttable.upsert
creator = UserAPIKeyAuth(api_key="sk-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None)
await managed_files.store_unified_object_id(
unified_object_id="uoi-processed",
file_object=_make_batch_response(status="completed"),
litellm_parent_otel_span=None,
model_object_id="batch-processed",
file_purpose="batch",
user_api_key_dict=creator,
batch_processed=True,
)
await managed_files.store_unified_object_id(
unified_object_id="uoi-default",
file_object=_make_batch_response(status="completed"),
litellm_parent_otel_span=None,
model_object_id="batch-default",
file_purpose="batch",
user_api_key_dict=creator,
)
processed_create, default_create = (call.kwargs["data"]["create"] for call in upsert.await_args_list)
assert processed_create["batch_processed"] is True
assert default_create["batch_processed"] is False
@pytest.mark.asyncio
async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets():
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
from litellm.caching import DualCache
file_table = MagicMock(upsert=AsyncMock(), find_first=AsyncMock(side_effect=AssertionError("cache miss")))
managed_files = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=MagicMock(db=MagicMock(litellm_managedfiletable=file_table)),
)
stored = _make_file_object("file-kept").model_copy(update={"purpose": "batch"})
stored._hidden_params = {"storage_backend": "litellm_db", "storage_url": "litellm_db://content-row-1"}
await managed_files.store_unified_file_id(
file_id="unified-kept",
file_object=stored,
litellm_parent_otel_span=None,
model_mappings={"vllm-batch": "litellm_db://content-row-1"},
user_api_key_dict=_make_user_api_key_dict(),
)
cached = await managed_files.get_unified_file_id("unified-kept")
assert cached is not None
assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1")
create_data = file_table.upsert.await_args.kwargs["data"]["create"]
assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1")

View file

@ -738,6 +738,140 @@ def test_embedding_summary_is_absent_without_vectors_and_for_chat_data_lists():
assert chat.embedding_output is None
def _responses_payload(output: list[object], status: str = "completed", **response_fields: object):
return _sample_payload(
call_type="aresponses",
model="gpt-5.4-nano",
response={"id": "resp_1", "object": "response", "status": status, "output": output, **response_fields},
)
_RESPONSES_TEXT_ITEM = {
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "po", "annotations": []}, {"type": "output_text", "text": "ng"}],
}
def test_responses_output_text_becomes_one_assistant_choice_with_stop():
data = LLMCallSpanData.from_standard_logging_payload(
_responses_payload([{"type": "reasoning", "summary": []}, _RESPONSES_TEXT_ITEM]), capture_content=True
)
assert json.loads(json.dumps(data.choices_out)) == [
{
"message": {"role": "assistant", "content": "pong", "refusal": None, "tool_calls": None},
"finish_reason": "stop",
}
]
assert data.finish_reasons == ("stop",)
assert data.response_id == "resp_1"
def test_responses_tool_calls_fold_into_the_assistant_message_with_tool_calls_finish_reason():
data = LLMCallSpanData.from_standard_logging_payload(
_responses_payload(
[
_RESPONSES_TEXT_ITEM,
{"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city": "sf"}'},
{"type": "custom_tool_call", "call_id": "call_2", "name": "grep", "input": "-r TODO"},
]
),
capture_content=True,
)
assert len(data.choices_out) == 1
message = data.choices_out[0]["message"]
assert message["content"] == "pong"
assert json.loads(json.dumps(message["tool_calls"])) == [
{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}},
{"id": "call_2", "type": "function", "function": {"name": "grep", "arguments": "-r TODO"}},
]
assert data.finish_reasons == ("tool_calls",)
def test_responses_tool_call_only_output_has_no_content():
data = LLMCallSpanData.from_standard_logging_payload(
_responses_payload([{"type": "function_call", "id": "fc_1", "name": "get_weather", "arguments": "{}"}]),
capture_content=True,
)
assert data.choices_out[0]["message"]["content"] is None
assert data.choices_out[0]["message"]["tool_calls"][0]["id"] == "fc_1"
@pytest.mark.parametrize(
("status", "response_fields", "expected"),
[
("incomplete", {"incomplete_details": {"reason": "max_output_tokens"}}, ("length",)),
("incomplete", {"incomplete_details": {"reason": "content_filter"}}, ("content_filter",)),
("incomplete", {}, ("length",)),
("failed", {}, ()),
],
)
def test_responses_status_maps_to_finish_reasons(status, response_fields, expected):
data = LLMCallSpanData.from_standard_logging_payload(
_responses_payload([_RESPONSES_TEXT_ITEM], status=status, **response_fields), capture_content=True
)
assert data.finish_reasons == expected
assert data.choices_out[0]["message"]["content"] == "pong"
def test_responses_output_follows_the_content_capture_gate_but_finish_reasons_do_not():
data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([_RESPONSES_TEXT_ITEM]))
assert data.choices_out == ()
assert data.finish_reasons == ("stop",)
def test_responses_content_only_reads_output_text_parts():
item = {
"type": "message",
"role": "assistant",
"content": [{"type": "refusal", "refusal": "no", "text": "not output"}, {"type": "output_text", "text": "ok"}],
}
data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([item]), capture_content=True)
assert data.choices_out[0]["message"]["content"] == "ok"
assert data.choices_out[0]["message"]["refusal"] == "no"
def test_responses_refusal_only_output_keeps_the_refusal_text():
item = {
"type": "message",
"role": "assistant",
"content": [{"type": "refusal", "refusal": "I can't "}, {"type": "refusal", "refusal": "help with that."}],
}
data = LLMCallSpanData.from_standard_logging_payload(_responses_payload([item]), capture_content=True)
assert json.loads(json.dumps(data.choices_out)) == [
{
"message": {"role": "assistant", "content": None, "refusal": "I can't help with that.", "tool_calls": None},
"finish_reason": "stop",
}
]
def test_responses_output_without_messages_or_tool_calls_stays_empty():
data = LLMCallSpanData.from_standard_logging_payload(
_responses_payload([{"type": "reasoning", "summary": []}]), capture_content=True
)
assert data.choices_out == ()
assert data.finish_reasons == ()
def test_chat_choices_win_over_a_responses_output_list():
payload = _sample_payload(response={"choices": [{"finish_reason": "stop", "message": {"content": "chat"}}]})
payload["response"]["output"] = [_RESPONSES_TEXT_ITEM]
data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True)
assert data.choices_out[0]["message"]["content"] == "chat"
assert data.finish_reasons == ("stop",)
def test_request_identity_prefers_canonical_team_keys():
from litellm.integrations.otel.model.payloads import RequestIdentity

View file

@ -196,6 +196,37 @@ def test_langfuse_mapper_keeps_chat_output_when_no_embedding_summary():
assert json.loads(attrs["langfuse.observation.output"]) == [{"role": "assistant", "content": "Sunny."}]
def test_langfuse_mapper_renders_a_responses_api_call_from_the_standard_logging_payload():
payload = {
"call_type": "aresponses",
"custom_llm_provider": "openai",
"model": "gpt-5.4-nano",
"messages": [{"role": "user", "content": "weather in sf?"}],
"response": {
"id": "resp_1",
"status": "completed",
"output": [
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Checking."}]},
{"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city": "sf"}'},
],
},
}
data = LLMCallSpanData.from_standard_logging_payload(payload, capture_content=True)
attrs = LangfuseMapper().map(data)
assert json.loads(attrs["langfuse.observation.output"]) == [
{
"role": "assistant",
"content": "Checking.",
"refusal": None,
"tool_calls": [
{"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "sf"}'}}
],
}
]
assert attrs["langfuse.observation.type"] == "generation"
# --------------------------------------------------------------------------- #
# Weave
# --------------------------------------------------------------------------- #

View file

@ -3532,6 +3532,24 @@ def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
assert merged.get("applied_guardrails") == ["pam-ethical-request"]
def test_get_standard_logging_metadata_merges_recorded_applied_guardrails():
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
result = StandardLoggingPayloadSetup.get_standard_logging_metadata(
metadata={"applied_guardrails": ["blocker"]},
litellm_params={},
applied_guardrails=["guard-a", "blocker", "guard-b"],
)
assert result["applied_guardrails"] == ["guard-a", "blocker", "guard-b"]
result = StandardLoggingPayloadSetup.get_standard_logging_metadata(
metadata={"applied_guardrails": ["blocker"]},
litellm_params={},
applied_guardrails=["guard-a"],
)
assert result["applied_guardrails"] == ["guard-a", "blocker"]
def test_function_setup_metadata_takes_precedence_over_litellm_metadata():
"""
Test that when BOTH metadata and litellm_metadata are present (e.g., user sets

View file

@ -493,6 +493,40 @@ class TestPerformRedaction:
assert redacted["output"][0]["arguments"] == "redacted-by-litellm"
assert redacted["output"][0]["name"] == "get_weather"
def test_redacts_responses_api_custom_tool_call_input_dict(self):
result = {
"output": [
{"type": "custom_tool_call", "name": "grep", "input": "-r secret-token", "call_id": "call_1"},
{"type": "function_call", "name": "get_weather", "input": "not-a-custom-input", "call_id": "call_2"},
]
}
redacted = perform_redaction({}, result)
assert redacted["output"][0]["input"] == "redacted-by-litellm"
assert redacted["output"][0]["name"] == "grep"
assert redacted["output"][1]["input"] == "not-a-custom-input"
def test_redacts_responses_api_refusal_parts_dict(self):
result = {
"output": [
{
"type": "message",
"role": "assistant",
"content": [
{"type": "refusal", "refusal": "I cannot share the secret"},
{"type": "output_text", "text": "ok"},
],
}
]
}
redacted = perform_redaction({}, result)
assert redacted["output"][0]["content"][0]["refusal"] == "redacted-by-litellm"
assert redacted["output"][0]["content"][0]["type"] == "refusal"
assert redacted["output"][0]["content"][1]["text"] == "redacted-by-litellm"
def test_redacts_every_tool_call_in_multi_element_list(self):
result = litellm.ModelResponse(
id="resp-multi",
@ -563,6 +597,23 @@ class TestPerformRedaction:
assert output_item.arguments == "redacted-by-litellm"
assert output_item.name == "get_weather"
def test_redacts_responses_api_custom_tool_call_input_object(self):
output_item = SimpleNamespace(type="custom_tool_call", name="grep", input="-r secret-token", call_id="call_1")
_redact_responses_api_output([output_item])
assert output_item.input == "redacted-by-litellm"
assert output_item.name == "grep"
def test_redacts_responses_api_refusal_parts_object(self):
refusal = SimpleNamespace(type="refusal", refusal="I cannot share the secret")
output_item = SimpleNamespace(type="message", role="assistant", content=[refusal])
_redact_responses_api_output([output_item])
assert refusal.refusal == "redacted-by-litellm"
assert refusal.type == "refusal"
def test_redacts_response_output_objects_with_top_level_text(self):
output_items = [
SimpleNamespace(text="top-level output"),

View file

@ -335,6 +335,129 @@ class TestAzureToolSchemaCombinatorFlattening:
assert request["temperature"] == 0.2
@pytest.mark.parametrize("tool_choice", ["none", "auto"])
def test_azure_drops_tool_choice_without_tools_or_functions(tool_choice: str) -> None:
optional_params = {"tool_choice": tool_choice, "temperature": 0.2}
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params=optional_params,
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert "tool_choice" not in request
assert request["temperature"] == 0.2
assert optional_params["tool_choice"] == tool_choice
def test_azure_tools_empty_drops_tool_choice() -> None:
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"tools": [], "tool_choice": "auto"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["tools"] == []
assert "tool_choice" not in request
def test_azure_functions_empty_drops_tool_choice() -> None:
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"functions": [], "tool_choice": "none"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["functions"] == []
assert "tool_choice" not in request
def test_azure_preserves_tool_choice_with_tools() -> None:
tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}]
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"tools": tools, "tool_choice": "auto"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["tools"] == tools
assert request["tool_choice"] == "auto"
def test_azure_preserves_tool_choice_with_legacy_functions() -> None:
functions = [{"name": "get_weather", "parameters": {}}]
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"functions": functions, "tool_choice": "auto"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["functions"] == functions
assert request["tool_choice"] == "auto"
def test_azure_preserves_function_call_without_tools() -> None:
request = AzureOpenAIConfig().transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"function_call": "none", "tool_choice": "auto"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["function_call"] == "none"
assert "tool_choice" not in request
def test_azure_gpt5_drops_tool_choice_without_tools() -> None:
request = AzureOpenAIGPT5Config().transform_request(
model="gpt5_series/gpt-5.6-sol",
messages=[{"role": "user", "content": "hi"}],
optional_params={"tool_choice": "none"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["model"] == "gpt-5.6-sol"
assert "tool_choice" not in request
@pytest.mark.asyncio
async def test_azure_async_transform_drops_tool_choice_without_tools() -> None:
request = await AzureOpenAIConfig().async_transform_request(
model="gpt-4o",
messages=[{"role": "user", "content": "hi"}],
optional_params={"tool_choice": "none"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert "tool_choice" not in request
@pytest.mark.asyncio
async def test_azure_gpt5_async_transform_drops_tool_choice_without_tools() -> None:
request = await AzureOpenAIGPT5Config().async_transform_request(
model="gpt5_series/gpt-5.6-sol",
messages=[{"role": "user", "content": "hi"}],
optional_params={"tool_choice": "auto"},
litellm_params={"custom_llm_provider": "azure"},
headers={},
)
assert request["model"] == "gpt-5.6-sol"
assert "tool_choice" not in request
def test_transform_request_strips_litellm_format_from_managed_file_id():
import base64

View file

@ -0,0 +1,92 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from prisma import Base64
from prisma.errors import RecordNotFoundError
from litellm.llms.base_llm.files.litellm_db_storage_backend import (
LITELLM_DB_STORAGE_URL_PREFIX,
LiteLLMDbStorageBackend,
storage_url_to_row_id,
)
def _backend_with_table():
table = MagicMock(create=AsyncMock(), find_unique=AsyncMock(), delete=AsyncMock())
prisma_client = MagicMock(db=MagicMock(litellm_managedfilecontenttable=table))
return LiteLLMDbStorageBackend(prisma_client), table
@pytest.mark.asyncio
async def test_upload_stores_bytes_and_returns_prefixed_row_id():
backend, table = _backend_with_table()
table.create.return_value = SimpleNamespace(id="row-1")
content = b"\x00\x01binary jsonl\n"
storage_url = await backend.upload_file(file_content=content, filename="input.jsonl", content_type="text/plain")
assert storage_url == f"{LITELLM_DB_STORAGE_URL_PREFIX}row-1"
stored = table.create.await_args.kwargs["data"]["content"]
assert isinstance(stored, Base64)
assert stored.decode() == content
@pytest.mark.asyncio
async def test_download_returns_exact_bytes_of_the_row():
backend, table = _backend_with_table()
content = b'{"custom_id": "1"}\n'
table.find_unique.return_value = SimpleNamespace(id="row-1", content=Base64.encode(content))
downloaded = await backend.download_file(f"{LITELLM_DB_STORAGE_URL_PREFIX}row-1")
assert downloaded == content
table.find_unique.assert_awaited_once_with(where={"id": "row-1"})
@pytest.mark.asyncio
async def test_download_missing_row_raises_value_error_naming_the_url():
backend, table = _backend_with_table()
table.find_unique.return_value = None
storage_url = f"{LITELLM_DB_STORAGE_URL_PREFIX}missing"
with pytest.raises(ValueError, match="missing"):
await backend.download_file(storage_url)
@pytest.mark.asyncio
async def test_download_rejects_url_without_prefix_before_touching_the_db():
backend, table = _backend_with_table()
with pytest.raises(ValueError, match="https://elsewhere/blob"):
await backend.download_file("https://elsewhere/blob")
table.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_delete_removes_the_parsed_row():
backend, table = _backend_with_table()
await backend.delete_file(f"{LITELLM_DB_STORAGE_URL_PREFIX}row-1")
table.delete.assert_awaited_once_with(where={"id": "row-1"})
@pytest.mark.asyncio
async def test_delete_tolerates_a_row_that_is_already_gone():
backend, table = _backend_with_table()
table.delete.side_effect = RecordNotFoundError({"user_facing_error": {"message": "gone"}})
await backend.delete_file(f"{LITELLM_DB_STORAGE_URL_PREFIX}row-1")
table.delete.assert_awaited_once_with(where={"id": "row-1"})
def test_storage_url_to_row_id_round_trips():
assert storage_url_to_row_id(f"{LITELLM_DB_STORAGE_URL_PREFIX}abc-123") == "abc-123"
def test_storage_url_to_row_id_rejects_foreign_urls():
with pytest.raises(ValueError, match="s3://bucket/key"):
storage_url_to_row_id("s3://bucket/key")

View file

@ -0,0 +1,34 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.llms.base_llm.files.litellm_db_storage_backend import (
LITELLM_DB_STORAGE_BACKEND_NAME,
LITELLM_DB_STORAGE_URL_PREFIX,
LiteLLMDbStorageBackend,
)
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
@pytest.mark.asyncio
async def test_litellm_db_backend_stores_through_the_given_prisma_client():
table = MagicMock(create=AsyncMock(return_value=SimpleNamespace(id="row-1")))
prisma_client = MagicMock(db=MagicMock(litellm_managedfilecontenttable=table))
backend = get_storage_backend(LITELLM_DB_STORAGE_BACKEND_NAME, prisma_client=prisma_client)
assert isinstance(backend, LiteLLMDbStorageBackend)
stored_at = await backend.upload_file(file_content=b"line\n", filename="input.jsonl", content_type="text/plain")
assert stored_at == f"{LITELLM_DB_STORAGE_URL_PREFIX}row-1"
table.create.assert_awaited_once()
def test_litellm_db_backend_without_a_database_is_rejected():
with pytest.raises(ValueError, match="database-connected proxy"):
get_storage_backend(LITELLM_DB_STORAGE_BACKEND_NAME)
def test_unknown_backend_is_still_rejected():
with pytest.raises(ValueError, match="Unsupported storage backend type: nope"):
get_storage_backend("nope", prisma_client=MagicMock())

View file

@ -5739,9 +5739,15 @@ class TestMCPDcrBridgeDelegateAdmission:
prisma and are swallowed (``_safe_fetch`` / the SCIM gate's fail-open), so their checks
skip. Yields the ``get_key_object`` mock so callers can assert the sealed ``key_hash`` was
the reload key."""
from litellm.proxy.auth.auth_checks import OrganizationNotFoundError
get_key_object = AsyncMock(return_value=return_value, side_effect=side_effect)
get_org_object = AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db."))
patchers = [
patch("litellm.proxy.auth.auth_checks.get_key_object", get_key_object),
patch( # test-quality-ok: central auth now resolves org limits; this fixture models a missing org row
"litellm.proxy.auth.auth_checks.get_org_object", get_org_object
),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
]

View file

@ -11974,9 +11974,13 @@ async def test_oauth_credential_write_keeps_virtual_key_permissions(
from litellm.proxy._experimental.mcp_server import mcp_server_manager
from litellm.proxy._experimental.mcp_server.bridge_token_flow import authorize_oauth_credential_request
from litellm.proxy._types import UserAPIKeyAuth, hash_token
from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key
from litellm.proxy.auth.auth_checks import OrganizationNotFoundError, jwt_key_mapping_cache_key
handler, signing_key = jwt_oauth_identity
monkeypatch.setattr(
"litellm.proxy.auth.auth_checks.get_org_object",
AsyncMock(side_effect=OrganizationNotFoundError("Organization doesn't exist in db.")),
)
key: Final = "sk-oauth-permission-test"
hashed: Final = hash_token(key)
credential: Final = UserAPIKeyAuth(

View file

@ -6087,6 +6087,107 @@ async def test_organization_budget_check_carries_org_state_on_the_token():
assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0)
@pytest.mark.parametrize("warmed_by_auth_prefetch", [False, True])
@pytest.mark.asyncio
async def test_get_org_object_for_request_serves_last_known_org_through_db_outage(warmed_by_auth_prefetch):
"""A JWT whose team sits in an org resolves the org on every request, and the org row
is cached for only DEFAULT_IN_MEMORY_TTL seconds while the team and user rows ride the
60s management-object TTL. Without a last-known copy, a DB outage a few seconds old
turned that traffic into 503s while the same request through a virtual key kept
succeeding on its cached team. The copy must exist whoever filled the short-lived entry:
this lookup's own DB read, or the virtual-key auth prefetch warming it for the same org."""
from litellm.proxy._types import LiteLLM_OrganizationTable
from litellm.proxy.auth.auth_checks import get_org_object_for_request
org_columns = {
"organization_id": "org-1",
"organization_alias": "platform-org",
"budget_id": "b1",
"created_by": "admin",
"updated_by": "admin",
"litellm_budget_table": {"budget_id": "b1", "max_budget": 50.0, "tpm_limit": 700, "rpm_limit": 7},
}
org_row = MagicMock()
org_row.model_dump = lambda: org_columns
db_outage = ConnectionRefusedError("db unavailable")
prisma_client = MagicMock()
prisma_client.db.litellm_organizationtable.find_unique = AsyncMock(
side_effect=[db_outage] if warmed_by_auth_prefetch else [org_row, db_outage]
)
user_api_key_cache = UserApiKeyCache()
if warmed_by_auth_prefetch:
await user_api_key_cache.async_set_cache(
key="org_id:org-1:with_budget",
value=LiteLLM_OrganizationTable.model_validate(org_columns),
model_type=LiteLLM_OrganizationTable,
)
async def _lookup():
return await get_org_object_for_request(
org_id="org-1",
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists
warm = await _lookup()
assert warm is not None and warm.organization_alias == "platform-org"
await user_api_key_cache.async_delete_cache("org_id:org-1:with_budget")
during_outage = await _lookup()
assert prisma_client.db.litellm_organizationtable.find_unique.await_count == (1 if warmed_by_auth_prefetch else 2)
assert during_outage is not None
assert during_outage.organization_alias == "platform-org"
assert during_outage.litellm_budget_table is not None
assert during_outage.litellm_budget_table.rpm_limit == 7
assert during_outage.litellm_budget_table.max_budget == 50.0
@pytest.mark.asyncio
async def test_get_org_object_for_request_writes_the_last_known_org_only_when_absent():
"""The last-known copy is written when this worker holds none, never per request:
with Redis attached, a write on every cached org hit would cost one SET per JWT request."""
from litellm.proxy._types import LiteLLM_OrganizationTable
from litellm.proxy.auth.auth_checks import get_org_object_for_request
class _WriteRecordingCache(UserApiKeyCache):
def __init__(self):
super().__init__()
self.written_keys = []
async def async_set_cache(self, key, value, local_only=False, **kwargs):
self.written_keys.append(key)
return await super().async_set_cache(key=key, value=value, local_only=local_only, **kwargs)
user_api_key_cache = _WriteRecordingCache()
await user_api_key_cache.async_set_cache(
key="org_id:org-1:with_budget",
value=LiteLLM_OrganizationTable(
organization_id="org-1",
organization_alias="platform-org",
budget_id="b1",
created_by="admin",
updated_by="admin",
),
model_type=LiteLLM_OrganizationTable,
)
for _ in range(3):
org = await get_org_object_for_request(
org_id="org-1",
prisma_client=MagicMock(),
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=None,
)
assert org is not None and org.organization_alias == "platform-org"
assert user_api_key_cache.written_keys.count("org_id:org-1:with_budget:last_known") == 1
@pytest.mark.parametrize(
"max_budget, spend, expect_blocked",
[
@ -7572,7 +7673,7 @@ async def test_project_allowlist_enforced_when_key_models_empty():
assert exc_info.value.code == "403"
def _project_with_budget(spend: float, max_budget: float):
def _project_with_budget(spend: float, max_budget: float | None):
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_ProjectTableCachedObj
return LiteLLM_ProjectTableCachedObj(
@ -7592,11 +7693,12 @@ def _project_with_budget(spend: float, max_budget: float):
pytest.param(4.99, 0.0, 5.0, False, id="counter-under-budget-admits"),
pytest.param(None, 5.0, 5.0, True, id="no-counter-falls-back-to-persisted-spend"),
pytest.param(None, 0.0, 5.0, False, id="no-counter-and-no-persisted-spend-admits"),
pytest.param(12.5, 12.5, 0.0, False, id="zero-budget-is-unbudgeted"),
pytest.param(12.5, 12.5, -1.0, False, id="negative-budget-is-unbudgeted"),
pytest.param(None, 0.0, 0.0, True, id="zero-budget-blocks-before-any-spend"),
pytest.param(12.5, 12.5, 0.0, True, id="zero-budget-blocks-with-spend"),
pytest.param(12.5, 12.5, None, False, id="null-budget-is-unlimited"),
],
)
async def test_project_max_budget_check_blocks_only_when_live_spend_reaches_a_positive_budget(
async def test_project_max_budget_check_blocks_when_live_spend_reaches_the_budget(
counter_spend, db_spend, max_budget, blocks
):
from litellm.caching.dual_cache import DualCache
@ -7631,7 +7733,7 @@ async def test_project_max_budget_check_blocks_only_when_live_spend_reaches_a_po
assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value
assert exc_info.value.entity_id == "p-budget"
assert exc_info.value.current_cost == 5.0
assert exc_info.value.current_cost == (db_spend if counter_spend is None else counter_spend)
proxy_logging_obj.budget_alerts.assert_awaited_once()
assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget"

View file

@ -25,6 +25,7 @@ from litellm.proxy._types import (
LiteLLM_JWTAuth,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
LiteLLM_OrganizationTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LitellmUserRoles,
@ -35,6 +36,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.auth_checks import (
OrganizationNotFoundError,
TeamNotFoundError,
UserNotFoundError,
get_key_object,
@ -5986,6 +5988,152 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id,
setattr(_proxy_server_mod, k, v)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"key_org_id,team_id,team_org_id,existing_alias,existing_rpm,lookup_mode,allow_db_unavailable,expect_lookup_error,expected_org_id,expected_alias,expected_limits",
[
(None, "t1", "org-from-team", None, None, "success", False, False, "org-from-team", "acme-org", (12.5, 700, 7)),
("org-jwt", None, None, None, None, "success", False, False, "org-jwt", "acme-org", (12.5, 700, 7)),
("org-pinned", None, None, "preset", None, "success", False, False, "org-pinned", "preset", (None, None, None)),
("org-view", None, None, None, 3, "success", False, False, "org-view", None, (None, None, 3)),
("org-missing", None, None, None, None, "missing", False, False, "org-missing", None, (None, None, None)),
("org-db-failure-allowed", None, None, None, None, "db_failure", True, False, "org-db-failure-allowed", None, (None, None, None)),
("org-db-failure-denied", None, None, None, None, "db_failure", False, True, "org-db-failure-denied", None, (None, None, None)),
("org-bad-row", None, None, None, None, "bad_row", False, False, "org-bad-row", None, (None, None, None)),
("org-nobudget", None, None, None, None, "no_budget", False, False, "org-nobudget", "acme-org", (None, None, None)),
],
)
async def test_centralized_common_checks_inherits_org_identity(
key_org_id: str | None,
team_id: str | None,
team_org_id: str | None,
existing_alias: str | None,
existing_rpm: int | None,
lookup_mode: str,
allow_db_unavailable: bool,
expect_lookup_error: bool,
expected_org_id: str | None,
expected_alias: str | None,
expected_limits: tuple[float | None, int | None, int | None],
) -> None:
import litellm.proxy.proxy_server as _proxy_server_mod
from fastapi import Request
from starlette.datastructures import URL
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
token = UserAPIKeyAuth(
api_key="sk-test",
user_id="u",
team_id=team_id,
org_id=key_org_id,
organization_alias=existing_alias,
organization_rpm_limit=existing_rpm,
)
request = Request(scope={"type": "http"})
request._url = URL(url="/chat/completions")
fetched_team = (
LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) if team_id is not None else None
)
organization = LiteLLM_OrganizationTable(
organization_id=expected_org_id,
organization_alias="acme-org",
budget_id="budget-id",
metadata={"model_rpm_limit": {"gpt-4o": 2}},
models=[],
created_by="test",
updated_by="test",
litellm_budget_table=(
None
if lookup_mode == "no_budget"
else LiteLLM_BudgetTable(budget_id="budget-id", max_budget=12.5, tpm_limit=700, rpm_limit=7)
),
)
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
attrs["prisma_client"] = MagicMock()
attrs["general_settings"] = {"allow_requests_on_db_unavailable": allow_db_unavailable}
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
try:
for k, v in attrs.items():
setattr(_proxy_server_mod, k, v)
with (
patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists
"litellm.proxy.auth.user_api_key_auth.get_team_object",
new_callable=AsyncMock,
return_value=fetched_team,
) as mock_get_team_object,
patch( # test-quality-ok: centralized auth calls this module helper directly; no dependency injection seam exists
"litellm.proxy.auth.auth_checks.get_org_object",
new_callable=AsyncMock,
return_value=organization,
) as mock_get_org_object,
patch( # test-quality-ok: capture downstream token state without invoking unrelated common checks
"litellm.proxy.auth.user_api_key_auth.common_checks",
new_callable=AsyncMock,
) as mock_checks,
):
if lookup_mode == "missing":
mock_get_org_object.side_effect = OrganizationNotFoundError("x")
elif lookup_mode == "db_failure":
mock_get_org_object.side_effect = ConnectionRefusedError("db unavailable")
elif lookup_mode == "bad_row":
mock_get_org_object.side_effect = ValueError("row failed validation")
if expect_lookup_error:
with pytest.raises(ConnectionRefusedError, match="db unavailable"):
await _run_centralized_common_checks(
user_api_key_auth_obj=token,
request=request,
request_data={"model": "gpt-4o"},
route="/chat/completions",
)
else:
await _run_centralized_common_checks(
user_api_key_auth_obj=token,
request=request,
request_data={"model": "gpt-4o"},
route="/chat/completions",
)
assert token.org_id == expected_org_id
if expect_lookup_error:
mock_checks.assert_not_awaited()
assert token.organization_alias is None
assert token.organization_max_budget is None
assert token.organization_tpm_limit is None
assert token.organization_rpm_limit is None
return
mock_checks.assert_awaited_once()
assert token.organization_alias == expected_alias
assert (
token.organization_max_budget,
token.organization_tpm_limit,
token.organization_rpm_limit,
) == expected_limits
checked_token = mock_checks.await_args.kwargs["valid_token"]
assert checked_token.org_id == expected_org_id
assert checked_token.organization_alias == expected_alias
if team_id is None:
mock_get_team_object.assert_not_awaited()
else:
mock_get_team_object.assert_awaited_once()
if existing_alias is not None or existing_rpm is not None:
mock_get_org_object.assert_not_awaited()
assert token.organization_metadata is None
else:
mock_get_org_object.assert_awaited_once()
assert mock_get_org_object.await_args.kwargs["org_id"] == expected_org_id
assert mock_get_org_object.await_args.kwargs["include_budget_table"] is True
if lookup_mode not in {"missing", "db_failure", "bad_row"}:
assert token.organization_metadata == {"model_rpm_limit": {"gpt-4o": 2}}
finally:
for k, v in originals.items():
setattr(_proxy_server_mod, k, v)
@pytest.mark.asyncio
async def test_cli_session_token_org_backfilled_from_team(monkeypatch):
"""LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted

View file

@ -34,10 +34,13 @@ import json
import logging
from contextlib import ExitStack
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
import litellm
@ -52,7 +55,7 @@ from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.llms.openai import BatchJobStatus
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import CredentialItem, LiteLLMBatch
from litellm.types.utils import CredentialItem, LiteLLMBatch, SpecialEnums
from fastapi import Request, Response
@ -74,6 +77,12 @@ CREDS: Dict[str, Dict[str, str]] = {
"api_base": "https://vertex.test",
"model": "vertex_ai/gemini-2.0",
},
"my-vllm": {
"custom_llm_provider": "hosted_vllm",
"api_key": "sk-vllm",
"api_base": "http://vllm.test/v1",
"model": "hosted_vllm/qwen",
},
}
# A real model-encoded file id: decodes to "azure/gpt-4o", strips to "file-original123".
@ -147,6 +156,7 @@ class Harness:
router: MagicMock
logging: MagicMock
creds_resolver: MagicMock
upstream_files_route: respx.Route
@property
def router_acreate(self) -> AsyncMock:
@ -162,13 +172,14 @@ class Harness:
return dict(self.router_acreate.call_args.kwargs)
def _creds_lookup(*, model_id: str) -> Dict[str, str]:
# KeyError on an unknown/hardcoded model_id - the bug cannot hide.
return dict(CREDS[model_id])
def _creds_lookup(*, model_id: str, team_id: str | None = None) -> dict[str, str] | None:
# An unknown/hardcoded model_id resolves to None exactly like the real router,
# which the endpoint turns into a 400 and a missing dispatch - the bug cannot hide.
return dict(CREDS[model_id]) if model_id in CREDS else None
@pytest.fixture
def harness():
def harness(monkeypatch: pytest.MonkeyPatch):
"""Seam harness. Patches only true I/O boundaries; pure encode/decode/merge
helpers run for real. Object mocks are spec'd so unknown method calls raise."""
body_holder: Dict[str, Any] = {}
@ -192,6 +203,7 @@ def harness():
provider_from_headers = MagicMock(return_value=None)
is_known_model = MagicMock(return_value=False)
litellm_acreate = AsyncMock(return_value=make_batch())
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
with ExitStack() as stack:
stack.enter_context(patch.object(endpoints, "_read_request_body", read_body))
@ -213,6 +225,10 @@ def harness():
stack.enter_context(patch.object(endpoints, "is_known_model", is_known_model))
stack.enter_context(patch.object(litellm, "acreate_batch", litellm_acreate))
stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
upstream = stack.enter_context(respx.mock(assert_all_called=False))
upstream_files_route = upstream.get(f"{CREDS['my-vllm']['api_base']}/files").mock(
return_value=httpx.Response(404, json={"detail": "Not Found"})
)
stack.enter_context(patch.object(proxy_server, "llm_router", router))
stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
stack.enter_context(patch.object(proxy_server, "general_settings", {}))
@ -231,6 +247,7 @@ def harness():
router=router,
logging=logging,
creds_resolver=router.get_deployment_credentials_with_provider,
upstream_files_route=upstream_files_route,
)
yield h
@ -255,6 +272,25 @@ async def call_create(
)
@pytest.fixture
def executed_runner():
runner = MagicMock(spec=endpoints.LiteLLMExecutedBatchRunner)
runner.create = AsyncMock(return_value=make_batch(id="litellm-executed-batch"))
runner.cancel = AsyncMock(return_value=make_batch(id="litellm-executed-batch", status="cancelling"))
factory = MagicMock(return_value=runner)
with patch.object( # test-quality-ok: the route builds its runner from proxy_server globals; the factory is the only seam
endpoints, "_litellm_executed_batch_runner", factory
):
yield runner, factory
def _managed_input_file_id(model: str) -> str:
unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
"application/jsonl", "managed-id", model, "file-id", "file-model-id"
)
return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
# =========================================================================== #
# SCENARIO 1 - input_file_id encoded with model. The full showcase: every
# assertion type from the design lives here.
@ -766,6 +802,137 @@ async def test_create__unified_file_id_legacy_row_without_storage_url_dispatches
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
# --------------------------------------------------------------------------- #
# LiteLLM-executed batches: a unified file targeting a provider whose API has
# no /v1/batches (hosted_vllm) runs inside LiteLLM instead of being forwarded.
# --------------------------------------------------------------------------- #
@pytest.mark.asyncio
async def test_create__unified_executed_provider_runs_inside_litellm(harness, executed_runner):
runner, factory = executed_runner
caller = UserAPIKeyAuth(api_key="sk-test", team_id="team-vllm")
input_file_id = _managed_input_file_id("my-vllm")
set_body(
harness,
{
"input_file_id": input_file_id,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"litellm_metadata": {"tags": ["batch-tag"]},
},
)
resp = await call_create(harness, user=caller)
harness.router_acreate.assert_not_called()
harness.litellm_acreate.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="my-vllm", team_id="team-vllm")
factory.assert_called_once_with(harness.router, harness.logging)
runner.create.assert_awaited_once()
create_kwargs = runner.create.call_args.kwargs
assert create_kwargs["unified_input_file_id"] == input_file_id
assert create_kwargs["model"] == "my-vllm"
assert create_kwargs["provider"] == "hosted_vllm"
assert create_kwargs["request_tags"] == ("batch-tag",)
assert create_kwargs["user_api_key_dict"] is caller
assert create_kwargs["create_request"]["model"] == "my-vllm"
assert resp.id == "litellm-executed-batch"
@pytest.mark.asyncio
async def test_create__unified_executed_provider_without_database_400(harness):
set_body(
harness,
{
"input_file_id": _managed_input_file_id("my-vllm"),
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
)
with pytest.raises(ProxyException) as exc:
await call_create(harness)
assert exc.value.code == "400"
assert "need a database" in exc.value.message
harness.router_acreate.assert_not_called()
harness.litellm_acreate.assert_not_called()
@pytest.mark.asyncio
async def test_create__unified_executed_provider_with_its_own_files_api_goes_to_the_provider(harness, executed_runner):
runner, factory = executed_runner
harness.upstream_files_route.mock(return_value=httpx.Response(200, json={"object": "list", "data": []}))
set_body(
harness,
{
"input_file_id": _managed_input_file_id("my-vllm"),
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
)
await call_create(harness)
factory.assert_not_called()
runner.create.assert_not_called()
assert harness.router_kwargs()["model"] == "my-vllm"
@pytest.mark.asyncio
async def test_create__unified_provider_model_never_touches_executed_runner(harness, executed_runner):
runner, factory = executed_runner
set_body(
harness,
{
"input_file_id": _managed_input_file_id("azure/gpt-4o"),
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
)
await call_create(harness)
factory.assert_not_called()
runner.create.assert_not_called()
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
assert harness.router_kwargs()["model"] == "azure/gpt-4o"
@pytest.mark.asyncio
@pytest.mark.parametrize("via", ["body", "header"])
async def test_create__raw_file_with_executed_model_400_with_upload_guidance(harness, via):
body = {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}
set_body(harness, {**body, "model": "my-vllm"} if via == "body" else body)
headers = {"x-litellm-model": "my-vllm"} if via == "header" else None
with pytest.raises(ProxyException) as exc:
await call_create(harness, headers=headers)
assert exc.value.code == "400"
assert "POST /v1/files" in exc.value.message
assert "x-litellm-model" in exc.value.message
harness.litellm_acreate.assert_not_called()
harness.router_acreate.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"upstream_answer",
[httpx.Response(200, json={"object": "list", "data": []}), httpx.Response(405), httpx.ConnectError("refused")],
ids=["lists files", "files route without list", "unreachable"],
)
async def test_create__raw_file_with_executed_model_is_forwarded_unless_the_server_lacks_a_files_api(
harness, upstream_answer
):
harness.upstream_files_route.mock(side_effect=[upstream_answer])
set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"})
await call_create(harness, headers={"x-litellm-model": "my-vllm"})
forwarded = harness.acreate_kwargs()
assert forwarded["input_file_id"] == "file-plain"
assert forwarded["custom_llm_provider"] == "hosted_vllm"
assert forwarded["api_base"] == CREDS["my-vllm"]["api_base"]
@pytest.mark.asyncio
async def test_create__model_encoded_beats_unified(harness):
"""Precedence row: a file id that is BOTH model-encoded and (pretend) unified
@ -1146,6 +1313,11 @@ AZURE_BATCH_ID = encode_file_id_with_model("batch_orig123", "azure/gpt-4o", id_t
# returns). model_id / llm_batch_id are parsed out of this by the real helpers.
UNIFIED_BATCH_ID = "litellm_proxy;model_id:gpt-4o-mini;llm_batch_id:batch-raw-xyz"
# A decoded unified id of a batch LiteLLM runs itself: the llm_batch_id carries
# the litellm_batch_ prefix, so no provider holds a batch to sync with.
EXECUTED_BATCH_ID = "litellm_proxy;model_id:my-vllm;llm_batch_id:litellm_batch_abc"
EXECUTED_BATCH_B64 = base64.urlsafe_b64encode(EXECUTED_BATCH_ID.encode()).decode().rstrip("=")
@dataclass
class RetrieveHarness:
@ -1580,6 +1752,69 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn
assert retrieve_harness.update_batch_in_db.call_count == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status):
db_response = make_batch(id="litellm-executed-batch", status=status)
db_batch_object = MagicMock()
db_batch_object.updated_at = datetime.now(timezone.utc)
retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)
resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64)
assert resp is db_response
retrieve_harness.litellm_aretrieve.assert_not_called()
retrieve_harness.router_aretrieve.assert_not_called()
retrieve_harness.update_batch_in_db.assert_not_called()
retrieve_harness.ensure_managed_files.assert_called_once()
assert retrieve_harness.ensure_managed_files.call_args.kwargs["unified_batch_id"] == EXECUTED_BATCH_ID
@pytest.mark.asyncio
async def test_retrieve__executed_batch_abandoned_by_its_runner_is_served_failed(retrieve_harness, executed_runner):
runner, _ = executed_runner
failed = make_batch(id="litellm-executed-batch", status="failed")
runner.fail_abandoned = AsyncMock(return_value=failed)
db_response = make_batch(id="litellm-executed-batch", status="in_progress")
db_batch_object = MagicMock()
db_batch_object.updated_at = datetime.now(timezone.utc) - timedelta(minutes=10)
retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)
user = UserAPIKeyAuth(api_key="sk-test", user_id="user-1")
resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64, user=user)
assert resp is failed
runner.fail_abandoned.assert_awaited_once_with(db_response, user)
retrieve_harness.litellm_aretrieve.assert_not_called()
retrieve_harness.router_aretrieve.assert_not_called()
retrieve_harness.ensure_managed_files.assert_called_once()
@pytest.mark.asyncio
async def test_retrieve__executed_batch_with_a_fresh_heartbeat_is_left_running(retrieve_harness, executed_runner):
runner, _ = executed_runner
runner.fail_abandoned = AsyncMock()
db_response = make_batch(id="litellm-executed-batch", status="in_progress")
db_batch_object = MagicMock()
db_batch_object.updated_at = datetime.now(timezone.utc) - timedelta(seconds=30)
retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)
resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64)
assert resp is db_response
runner.fail_abandoned.assert_not_awaited()
@pytest.mark.asyncio
async def test_retrieve__executed_batch_without_db_row_404(retrieve_harness):
with pytest.raises(ProxyException) as exc:
await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64)
assert exc.value.code == "404"
retrieve_harness.litellm_aretrieve.assert_not_called()
retrieve_harness.router_aretrieve.assert_not_called()
# --------------------------------------------------------------------------- #
# Cross-cutting: enrichment route_type and failure-hook on provider error.
# --------------------------------------------------------------------------- #
@ -2299,6 +2534,35 @@ async def test_cancel__unified_no_router_500(cancel_harness):
assert exc.value.code == "500"
@pytest.mark.asyncio
async def test_cancel__executed_batch_routes_to_runner(cancel_harness, executed_runner):
runner, factory = executed_runner
caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-2")
resp = await call_cancel(cancel_harness, EXECUTED_BATCH_B64, user=caller)
runner.cancel.assert_awaited_once_with(EXECUTED_BATCH_B64, caller)
factory.assert_called_once_with(cancel_harness.router, cancel_harness.logging)
cancel_harness.router_acancel.assert_not_called()
cancel_harness.litellm_acancel.assert_not_called()
cancel_harness.creds_resolver.assert_not_called()
assert resp is runner.cancel.return_value
assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"
@pytest.mark.asyncio
async def test_cancel__executed_batch_no_router_500(cancel_harness, executed_runner):
runner, factory = executed_runner
with patch.object( # test-quality-ok: proxy_server module global is the endpoint's only injection point
proxy_server, "llm_router", None
):
with pytest.raises(ProxyException) as exc:
await call_cancel(cancel_harness, EXECUTED_BATCH_B64)
assert exc.value.code == "500"
factory.assert_not_called()
runner.cancel.assert_not_called()
# --------------------------------------------------------------------------- #
# SCENARIO 3 - fallback to custom_llm_provider. Rebuilds a CancelBatchRequest
# and forwards only {custom_llm_provider, batch_id}.
@ -2956,3 +3220,16 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h
assert exc_info.value.code == "403"
cancel_harness.router_acancel.assert_not_called()
@pytest.mark.asyncio
async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner):
runner, factory = executed_runner
with pytest.raises(ProxyException) as exc_info:
await call_cancel(cancel_harness, EXECUTED_BATCH_B64, user=_key_restricted_to("vertex-model"))
assert exc_info.value.code == "403"
factory.assert_not_called()
runner.cancel.assert_not_called()
cancel_harness.router_acancel.assert_not_called()

File diff suppressed because it is too large Load diff

View file

@ -359,3 +359,93 @@ def test_settings_store_refusal_stays_quiet_about_the_database_when_nothing_is_s
assert refused.value.shadows_db_value is False
assert "stored in the database" not in str(refused.value)
assert "config file" in str(refused.value)
def test_settings_store_keeps_a_resolved_runtime_value_when_a_db_row_repeats_it() -> None:
store: Final = SettingsStore("general_settings")
store.apply_db_row("general_settings", {"litellm_key_header_name": "os.environ/HDR"})
store.apply_runtime_values({"litellm_key_header_name": "X-Resolved-Header"})
store.apply_db_row("general_settings", {"litellm_key_header_name": "os.environ/HDR"})
assert store["litellm_key_header_name"] == "X-Resolved-Header"
def test_settings_store_drops_a_resolved_runtime_value_when_a_db_row_changes_it() -> None:
store: Final = SettingsStore("general_settings")
store.apply_db_row("general_settings", {"litellm_key_header_name": "os.environ/HDR"})
store.apply_runtime_values({"litellm_key_header_name": "X-Resolved-Header"})
store.apply_db_row("general_settings", {"litellm_key_header_name": "os.environ/OTHER"})
assert store["litellm_key_header_name"] == "os.environ/OTHER"
def test_settings_store_accepts_the_writes_it_does_not_report_as_rejected() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"litellm_key_header_name": "os.environ/HDR"})
store.apply_runtime_values({"litellm_key_header_name": "X-Resolved-Header"})
incoming: Final[dict[str, JsonValue]] = {"litellm_key_header_name": "X-Resolved-Header"}
assert store.rejected_writes(incoming) == ()
store["litellm_key_header_name"] = "X-Resolved-Header"
assert store["litellm_key_header_name"] == "X-Resolved-Header"
def test_settings_store_reports_a_rejected_write_the_store_itself_refuses() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"litellm_key_header_name": "os.environ/HDR"})
store.apply_runtime_values({"litellm_key_header_name": "X-Resolved-Header"})
assert store.rejected_writes({"litellm_key_header_name": "X-Other-Header"}) == ("litellm_key_header_name",)
with pytest.raises(ConfigOwnedKeyError):
store["litellm_key_header_name"] = "X-Other-Header"
def test_settings_store_reports_no_shadowing_when_the_database_repeats_the_config_template() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"litellm_key_header_name": "os.environ/HDR"})
store.apply_db_row("general_settings", {"litellm_key_header_name": "os.environ/HDR"})
store.apply_runtime_values({"litellm_key_header_name": "X-Resolved-Header"})
assert store.shadowed_db_keys() == ()
assert store.shadows_db_value("litellm_key_header_name") is False
def test_settings_store_still_reports_shadowing_when_the_database_holds_another_template() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({"litellm_key_header_name": "os.environ/HDR"})
store.apply_db_row("general_settings", {"litellm_key_header_name": "os.environ/OTHER"})
store.apply_runtime_values({"litellm_key_header_name": "X-Resolved-Header"})
assert store.shadowed_db_keys() == ("litellm_key_header_name",)
def test_settings_store_truthiness_stops_at_the_first_key() -> None:
store: Final = SettingsStore("general_settings")
store.load_yaml({f"key_{index}": index for index in range(25)})
resolutions: Final[list[str]] = []
original: Final = SettingsStore._resolution_for
def counted(self: SettingsStore, key: str):
resolutions.append(key)
return original(self, key)
with patch.object(SettingsStore, "_resolution_for", counted): # test-quality-ok: counting resolutions is the only way to observe that truthiness short-circuits
assert bool(store) is True
truthiness_resolutions: Final = len(resolutions)
resolutions.clear()
assert len(store) == 25
assert len(resolutions) == 25
assert truthiness_resolutions <= 1
def test_settings_store_truthiness_matches_emptiness() -> None:
store: Final = SettingsStore("general_settings")
assert bool(store) is False
store["max_parallel_requests"] = 3
assert bool(store) is True
del store["max_parallel_requests"]
assert bool(store) is False

View file

@ -616,3 +616,64 @@ async def test_connection_test_rejects_proxy_admin_viewer():
user_api_key_dict=UserAPIKeyAuth(api_key="hashed", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
)
assert exc_info.value.status_code == 403
def _real_proxy_config(file_general_settings: dict) -> "object":
from litellm.proxy.proxy_server import ProxyConfig
proxy_config = ProxyConfig()
proxy_config._load_yaml_settings_stores({"general_settings": file_general_settings})
proxy_config.get_config_state = MagicMock(
return_value={"general_settings": file_general_settings}
)
return proxy_config
@pytest.mark.asyncio
async def test_update_refuses_a_config_owned_coordination_redis_block(monkeypatch):
monkeypatch.setattr(litellm, "store_audit_logs", False)
mock_prisma = _prisma_with_general_settings({"master_key": "sk-1234"})
from_file = {"coordination_redis": {"host": "yaml-redis.example.com", "port": 6379}}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
patch("litellm.proxy.proxy_server.proxy_config", _real_proxy_config(from_file)), # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
):
with pytest.raises(HTTPException) as refused:
await update_coordination_redis_settings(
request=CoordinationRedisSettingsRequest(settings={"host": "db-redis.example.com", "port": 6380}),
user_api_key_dict=_admin_auth(),
litellm_changed_by=None,
)
assert refused.value.status_code == 400
assert refused.value.detail["keys"] == ["coordination_redis"]
mock_prisma.db.litellm_config.upsert.assert_not_called()
@pytest.mark.asyncio
async def test_update_still_persists_when_the_config_file_declares_no_block(monkeypatch):
monkeypatch.setattr(litellm, "store_audit_logs", False)
mock_prisma = _prisma_with_general_settings({"master_key": "sk-1234"})
async def _capture_invalidate(param_name: str) -> None:
return None
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
patch("litellm.proxy.proxy_server.proxy_config", _real_proxy_config({"master_key": "sk-1234"})), # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
patch( # test-quality-ok: the endpoint reads these proxy_server module globals at call time; there is no injection seam
"litellm.proxy.management_endpoints.coordination_redis_endpoints.invalidate_config_param",
new=_capture_invalidate,
),
):
await update_coordination_redis_settings(
request=CoordinationRedisSettingsRequest(settings={"host": "db-redis.example.com", "port": 6380}),
user_api_key_dict=_admin_auth(),
litellm_changed_by=None,
)
persisted = json.loads(mock_prisma.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"])
assert persisted["coordination_redis"] == {"host": "db-redis.example.com", "port": 6380}

View file

@ -4,10 +4,10 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.proxy.openai_files_endpoints.common_utils import (
apply_unified_file_ids,
get_credentials_for_model,
is_litellm_executed_batch,
map_raw_file_ids_to_unified,
)
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
@ -500,3 +500,17 @@ class TestCompletedBatchSafeToRetire:
def test_no_output_and_unknown_counts_is_not_safe(self):
assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False
@pytest.mark.parametrize(
"decoded_unified_batch_id, executed",
[
("litellm_proxy;model_id:my-vllm;llm_batch_id:litellm_batch_0123abcd", True),
("litellm_proxy;model_id:my-vllm;llm_batch_id:batch_0123abcd", False),
("litellm_proxy;model_id:my-vllm;generic_response_id:resp_0123abcd", False),
("litellm_proxy;model_id:my-vllm;llm_output_file_id:file-0123abcd", False),
("batch_0123abcd", False),
],
)
def test_is_litellm_executed_batch_reads_the_llm_batch_id_prefix(decoded_unified_batch_id: str, executed: bool):
assert is_litellm_executed_batch(decoded_unified_batch_id) is executed

View file

@ -609,6 +609,246 @@ def test_target_storage_with_target_models(
app.dependency_overrides.pop(ps.user_api_key_auth, None)
BATCH_JSONL_LINE = (
b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", '
b'"body": {"model": "my-vllm", "messages": [{"role": "user", "content": "hi"}]}}\n'
)
def _router_with_executed_batch_model() -> Router:
return Router(
model_list=[
{
"model_name": "my-vllm",
"litellm_params": {
"model": "hosted_vllm/qwen",
"api_key": "sk-vllm",
"api_base": "http://vllm.test/v1",
},
"model_info": {"id": "my-vllm-id"},
},
{
"model_name": "gemini-2.0-flash",
"litellm_params": {"model": "gemini/gemini-2.0-flash"},
"model_info": {"id": "gemini-2.0-flash-id"},
},
]
)
@pytest.fixture
def batch_upload_seams(mocker: MockerFixture, monkeypatch):
import litellm.proxy.proxy_server as ps
from litellm.proxy._types import LitellmUserRoles
llm_router = _router_with_executed_batch_model()
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
setup_proxy_logging_object(monkeypatch, llm_router)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user"
)
uploaded = OpenAIFileObject(
id="file-kept",
object="file",
purpose="batch",
created_at=0,
bytes=len(BATCH_JSONL_LINE),
filename="batch.jsonl",
status="uploaded",
)
stored = mocker.patch( # test-quality-ok: the route calls the storage service directly with no injection seam
"litellm.proxy.openai_files_endpoints.storage_backend_service.StorageBackendFileService.upload_file_to_storage_backend",
new=mocker.AsyncMock(return_value=uploaded),
)
provider_upload = mocker.patch( # test-quality-ok: the route calls litellm.acreate_file directly with no injection seam
"litellm.acreate_file", new=mocker.AsyncMock(return_value=uploaded)
)
try:
with respx.mock(assert_all_called=False) as upstream:
upstream_files_route = upstream.get("http://vllm.test/v1/files").mock(
return_value=httpx.Response(404, json={"detail": "Not Found"})
)
yield stored, provider_upload, upstream_files_route
finally:
app.dependency_overrides.pop(ps.user_api_key_auth, None)
def _upload_batch_file(headers: dict[str, str], form: dict[str, str]):
return client.post(
"/v1/files",
files={"file": ("batch.jsonl", BATCH_JSONL_LINE, "application/jsonl")},
data={"purpose": "batch", **form},
headers={"Authorization": "Bearer test-key", **headers},
)
@pytest.mark.parametrize(
"headers, form",
[({"x-litellm-model": "my-vllm"}, {}), ({}, {"target_model_names": "my-vllm"})],
ids=["x-litellm-model header", "target_model_names form field"],
)
def test_batch_upload_for_a_litellm_executed_model_is_kept_by_litellm(
batch_upload_seams, headers: dict[str, str], form: dict[str, str]
):
stored, provider_upload, _ = batch_upload_seams
response = _upload_batch_file(headers, form)
assert response.status_code == 200, response.text
provider_upload.assert_not_awaited()
stored.assert_awaited_once()
kwargs = stored.call_args.kwargs
assert kwargs["target_storage"] == "litellm_db"
assert tuple(kwargs["target_model_names"]) == ("my-vllm",)
assert kwargs["purpose"] == "batch"
@pytest.mark.parametrize(
"headers, form",
[({"x-litellm-model": "my-vllm"}, {}), ({}, {"target_model_names": "my-vllm"})],
ids=["x-litellm-model header", "target_model_names form field"],
)
def test_batch_upload_for_a_litellm_executed_model_the_key_cannot_call_is_refused_before_the_server_is_probed(
batch_upload_seams, headers: dict[str, str], form: dict[str, str]
):
import litellm.proxy.proxy_server as ps
stored, provider_upload, upstream_files_route = batch_upload_seams
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="restricted-user", models=["gemini-2.0-flash"]
)
response = _upload_batch_file(headers, form)
assert response.status_code == 403, response.text
assert "my-vllm" in response.text
assert upstream_files_route.call_count == 0
stored.assert_not_awaited()
provider_upload.assert_not_awaited()
def test_batch_upload_naming_an_executed_and_a_provider_model_is_rejected(batch_upload_seams):
stored, provider_upload, _ = batch_upload_seams
response = _upload_batch_file({}, {"target_model_names": "my-vllm,gemini-2.0-flash"})
assert response.status_code == 400, response.text
assert "my-vllm" in response.text
assert "target_model_names" in response.text
stored.assert_not_awaited()
provider_upload.assert_not_awaited()
@pytest.mark.parametrize("purpose", ["assistants", "user_data"])
def test_non_batch_upload_for_a_litellm_executed_model_is_rejected_with_the_purpose_to_use(
batch_upload_seams, purpose: str
):
stored, provider_upload, _ = batch_upload_seams
response = _upload_batch_file({"x-litellm-model": "my-vllm"}, {"purpose": purpose})
assert response.status_code == 400, response.text
error = response.json()["error"]
assert error["type"] == "invalid_request_error"
assert error["param"] == "purpose"
assert "purpose=batch" in error["message"]
assert f"purpose={purpose}" in error["message"]
stored.assert_not_awaited()
provider_upload.assert_not_awaited()
@pytest.mark.parametrize("purpose", ["batch", "assistants"])
@pytest.mark.parametrize(
"upstream_answer",
[httpx.Response(200, json={"object": "list", "data": []}), httpx.Response(405), httpx.ConnectError("refused")],
ids=["lists files", "files route without list", "unreachable"],
)
def test_upload_for_a_litellm_executed_model_goes_to_the_provider_unless_the_server_lacks_a_files_api(
batch_upload_seams, upstream_answer: httpx.Response | httpx.ConnectError, purpose: str
):
stored, provider_upload, upstream_files_route = batch_upload_seams
upstream_files_route.mock(side_effect=[upstream_answer])
response = _upload_batch_file({"x-litellm-model": "my-vllm"}, {"purpose": purpose})
assert response.status_code == 200, response.text
stored.assert_not_awaited()
provider_upload.assert_awaited_once()
assert provider_upload.call_args.kwargs["custom_llm_provider"] == "hosted_vllm"
assert provider_upload.call_args.kwargs["api_base"] == "http://vllm.test/v1"
@pytest.mark.parametrize(
"form",
[{}, {"target_model_names": "my-vllm"}, {"target_model_names": "gemini-2.0-flash"}],
ids=["no model", "litellm-executed model", "provider model"],
)
def test_upload_naming_litellm_db_as_target_storage_is_rejected(batch_upload_seams, form: dict[str, str]):
stored, provider_upload, upstream_files_route = batch_upload_seams
response = _upload_batch_file({}, {**form, "target_storage": "litellm_db"})
assert response.status_code == 400, response.text
error = response.json()["error"]
assert error["type"] == "invalid_request_error"
assert error["param"] == "target_storage"
assert "litellm_db" in error["message"]
assert upstream_files_route.call_count == 0
stored.assert_not_awaited()
provider_upload.assert_not_awaited()
@pytest.mark.parametrize("purpose", ["user_data", "batch"])
def test_upload_with_an_explicit_target_storage_goes_where_the_caller_said_without_probing_the_server(
batch_upload_seams, purpose: str
):
stored, provider_upload, upstream_files_route = batch_upload_seams
response = _upload_batch_file(
{}, {"purpose": purpose, "target_model_names": "my-vllm", "target_storage": "azure_storage"}
)
assert response.status_code == 200, response.text
assert upstream_files_route.call_count == 0
provider_upload.assert_not_awaited()
stored.assert_awaited_once()
kwargs = stored.call_args.kwargs
assert kwargs["target_storage"] == "azure_storage"
assert tuple(kwargs["target_model_names"]) == ("my-vllm",)
assert kwargs["purpose"] == purpose
def test_upload_with_an_explicit_target_storage_still_refuses_a_key_without_the_executed_model(batch_upload_seams):
import litellm.proxy.proxy_server as ps
stored, provider_upload, upstream_files_route = batch_upload_seams
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="restricted-user", models=["gemini-2.0-flash"]
)
response = _upload_batch_file({}, {"target_model_names": "my-vllm", "target_storage": "azure_storage"})
assert response.status_code == 403, response.text
assert "my-vllm" in response.text
assert upstream_files_route.call_count == 0
stored.assert_not_awaited()
provider_upload.assert_not_awaited()
def test_batch_upload_for_a_provider_model_still_goes_to_the_provider(batch_upload_seams):
stored, provider_upload, _ = batch_upload_seams
response = _upload_batch_file({"x-litellm-model": "gemini-2.0-flash"}, {})
assert response.status_code == 200, response.text
stored.assert_not_awaited()
provider_upload.assert_awaited_once()
assert provider_upload.call_args.kwargs["custom_llm_provider"] == "gemini"
@pytest.mark.skip(reason="mock respx fails on ci/cd - unclear why")
def test_create_file_and_call_chat_completion_e2e(
mocker: MockerFixture, monkeypatch, llm_router: Router

View file

@ -1,3 +1,5 @@
from unittest.mock import MagicMock
import pytest
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
@ -6,16 +8,24 @@ from litellm.proxy.openai_files_endpoints import storage_backend_service
from litellm.proxy.openai_files_endpoints.storage_backend_service import (
StorageBackendFileService,
)
from litellm.proxy.utils import PrismaClient
class _RecordingStorageBackend:
def __init__(self):
def __init__(self, delete_error: Exception | None = None):
self.upload_calls = []
self.delete_calls: list[str] = []
self.delete_error = delete_error
async def upload_file(self, **kwargs):
self.upload_calls.append(kwargs)
return "https://storage.example/blob-1"
async def delete_file(self, storage_url: str) -> None:
self.delete_calls.append(storage_url)
if self.delete_error is not None:
raise self.delete_error
class _FakeManagedFilesHook(BaseFileEndpoints):
def __init__(self):
@ -42,6 +52,11 @@ class _FakeManagedFilesHook(BaseFileEndpoints):
self.stored.append(kwargs)
class _FailingManagedFilesHook(_FakeManagedFilesHook):
async def store_unified_file_id(self, **kwargs):
raise RuntimeError("db down")
class _FakeProxyLogging:
def __init__(self, hook):
self._hook = hook
@ -57,7 +72,7 @@ def _file_data():
@pytest.mark.asyncio
async def test_upload_with_target_model_names_but_no_hook_raises_before_uploading(monkeypatch):
backend = _RecordingStorageBackend()
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend)
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend)
with pytest.raises(ProxyException) as exc_info:
await StorageBackendFileService.upload_file_to_storage_backend(
@ -80,7 +95,7 @@ async def test_upload_with_target_model_names_but_no_hook_raises_before_uploadin
@pytest.mark.asyncio
async def test_upload_without_target_model_names_skips_hook_requirement(monkeypatch):
backend = _RecordingStorageBackend()
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend)
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend)
file_object = await StorageBackendFileService.upload_file_to_storage_backend(
file_data=_file_data(),
@ -101,7 +116,7 @@ async def test_upload_without_target_model_names_skips_hook_requirement(monkeypa
@pytest.mark.asyncio
async def test_upload_with_target_model_names_and_hook_stores_unified_id(monkeypatch):
backend = _RecordingStorageBackend()
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend)
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend)
hook = _FakeManagedFilesHook()
file_object = await StorageBackendFileService.upload_file_to_storage_backend(
@ -125,3 +140,50 @@ async def test_upload_with_target_model_names_and_hook_stores_unified_id(monkeyp
"stored_id_matches_response": True,
"model_mappings": {"gpt-x": "https://storage.example/blob-1"},
}
@pytest.mark.asyncio
async def test_upload_hands_the_prisma_client_to_the_storage_backend_factory(monkeypatch: pytest.MonkeyPatch):
backend = _RecordingStorageBackend()
factory_calls: list[tuple[str, PrismaClient | None]] = []
def _factory(name: str, prisma_client: PrismaClient | None = None) -> _RecordingStorageBackend:
factory_calls.append((name, prisma_client))
return backend
monkeypatch.setattr(storage_backend_service, "get_storage_backend", _factory)
prisma_client = MagicMock()
await StorageBackendFileService.upload_file_to_storage_backend(
file_data=_file_data(),
target_storage="litellm_db",
target_model_names=[],
purpose="batch",
proxy_logging_obj=_FakeProxyLogging(hook=None),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
prisma_client=prisma_client,
)
assert factory_calls == [("litellm_db", prisma_client)]
@pytest.mark.asyncio
@pytest.mark.parametrize("delete_error", [None, OSError("blob locked")], ids=["delete succeeds", "delete fails"])
async def test_upload_deletes_the_uploaded_content_when_the_metadata_write_fails(
monkeypatch: pytest.MonkeyPatch, delete_error: Exception | None
):
backend = _RecordingStorageBackend(delete_error=delete_error)
monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name, prisma_client=None: backend)
with pytest.raises(RuntimeError, match="db down"):
await StorageBackendFileService.upload_file_to_storage_backend(
file_data=_file_data(),
target_storage="azure_storage",
target_model_names=["gpt-x"],
purpose="batch",
proxy_logging_obj=_FakeProxyLogging(hook=_FailingManagedFilesHook()),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
)
assert len(backend.upload_calls) == 1
assert backend.delete_calls == ["https://storage.example/blob-1"]

View file

@ -21,6 +21,15 @@ from uvicorn.importer import import_from_string
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
@pytest.fixture(autouse=True)
def fork_reservation():
"""Reserving is irreversible: it would forbid native routes in this pytest worker for good"""
with patch( # test-quality-ok: process-global native state, a real reservation would poison every later test in the worker
"litellm.rust_bridge.fork_guard.reserve_process_for_forking"
) as reserve:
yield reserve
@pytest.mark.xdist_group("proxy_cli")
class TestProxyInitializationHelpers:
@patch("importlib.metadata.version")
@ -1574,6 +1583,32 @@ class TestProxyInitializationHelpers:
assert captured["options"]["max_requests"] == 1000
assert captured["options"]["max_requests_jitter"] == 50
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation):
"""preload forks workers from the master, so native routes are forbidden there first"""
pytest.importorskip("gunicorn")
reserved_before_run: list = []
def capture_run(self):
reserved_before_run.append(fork_reservation.call_args)
with (
patch("gunicorn.app.base.BaseApplication.run", capture_run),
patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership
"litellm.proxy.proxy_cli.start_query_engine_reaper"
),
):
ProxyInitializationHelpers._run_gunicorn_server(
host="127.0.0.1",
port=4012,
app=MagicMock(),
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
)
assert [call.args for call in reserved_before_run] == [("the gunicorn master",)]
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_jitter_without_base_warns(self):
"""gunicorn path warns when jitter is set without --max_requests_before_restart"""

View file

@ -14737,3 +14737,99 @@ async def test_auth_cache_invalidation_subscriber_evicts_byok_credentials_cached
byok_credential_cache.flush_cache()
assert evicted, "the subscriber does not evict the BYOK credential cache on a peer worker's broadcast"
@pytest.mark.asyncio
async def test_delete_config_general_settings_refuses_a_key_the_config_file_owns(monkeypatch):
from litellm.proxy._types import ConfigFieldDelete
from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings
pc = ProxyConfig()
pc._load_yaml_settings_stores({"general_settings": {"max_request_size_mb": 42}})
monkeypatch.setattr(proxy_server_module, "proxy_config", pc)
monkeypatch.setattr(proxy_server_module, "prisma_client", _fake_prisma_with_config({"max_request_size_mb": 99}))
admin = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
with pytest.raises(HTTPException) as refused:
await delete_config_general_settings(
data=ConfigFieldDelete(field_name="max_request_size_mb", config_type="general_settings"),
user_api_key_dict=admin,
)
assert refused.value.status_code == 400
assert refused.value.detail["keys"] == ["max_request_size_mb"]
assert "config file" in refused.value.detail["error"]
assert pc.settings["max_request_size_mb"] == 42
@pytest.mark.asyncio
async def test_delete_config_general_settings_still_removes_a_key_the_database_owns(monkeypatch):
from litellm.proxy._types import ConfigFieldDelete
from litellm.proxy.proxy_server import ProxyConfig, delete_config_general_settings
pc = ProxyConfig()
pc._load_yaml_settings_stores({"general_settings": {}})
pc.settings.apply_db_row("general_settings", {"max_request_size_mb": 42})
monkeypatch.setattr(proxy_server_module, "proxy_config", pc)
monkeypatch.setattr(proxy_server_module, "prisma_client", _fake_prisma_with_config({"max_request_size_mb": 42}))
admin = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
await delete_config_general_settings(
data=ConfigFieldDelete(field_name="max_request_size_mb", config_type="general_settings"),
user_api_key_dict=admin,
)
assert "max_request_size_mb" not in pc.settings
@pytest.mark.asyncio
async def test_config_field_info_reports_the_declared_value_of_a_config_owned_secret(monkeypatch):
from litellm.proxy.proxy_server import ProxyConfig, get_config_general_settings
pc = ProxyConfig()
pc._load_yaml_settings_stores({"general_settings": {"master_key": "os.environ/PROXY_MASTER_KEY"}})
pc.settings.apply_runtime_values({"master_key": "sk-resolved-secret"})
monkeypatch.setattr(proxy_server_module, "proxy_config", pc)
monkeypatch.setattr(proxy_server_module, "prisma_client", _fake_prisma_with_config({}))
admin = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
info = await get_config_general_settings(field_name="master_key", user_api_key_dict=admin)
assert info.field_value == "os.environ/PROXY_MASTER_KEY"
assert info.source == "config"
assert info.editable is False
@pytest.mark.asyncio
async def test_config_field_info_still_reports_a_database_owned_value(monkeypatch):
from litellm.proxy.proxy_server import ProxyConfig, get_config_general_settings
pc = ProxyConfig()
pc._load_yaml_settings_stores({"general_settings": {}})
pc.settings.apply_db_row("general_settings", {"max_request_size_mb": 42})
monkeypatch.setattr(proxy_server_module, "proxy_config", pc)
monkeypatch.setattr(proxy_server_module, "prisma_client", _fake_prisma_with_config({"max_request_size_mb": 42}))
admin = UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN)
info = await get_config_general_settings(field_name="max_request_size_mb", user_api_key_dict=admin)
assert info.field_value == 42
assert info.source == "db"
@pytest.mark.asyncio
async def test_initialize_jwt_auth_leaves_the_declared_jwtauth_mapping_unresolved(monkeypatch):
from litellm.proxy.proxy_server import ProxyStartupEvent
declared = {"public_key_ttl": "600", "team_id_jwt_field": "os.environ/JWT_TEAM_FIELD"}
general_settings = {"litellm_jwtauth": declared}
monkeypatch.setattr(proxy_server_module, "get_secret", lambda value: "resolved-team-field")
ProxyStartupEvent._initialize_jwt_auth(
general_settings=general_settings,
prisma_client=None,
user_api_key_cache=DualCache(),
)
assert declared["team_id_jwt_field"] == "os.environ/JWT_TEAM_FIELD"
assert proxy_server_module.jwt_handler.litellm_jwtauth.team_id_jwt_field == "resolved-team-field"

View file

@ -2604,6 +2604,67 @@ def test_add_allowed_ip_writes_audit_log(mock_proxy_config, monkeypatch):
app.dependency_overrides.pop(user_api_key_auth, None)
def test_add_allowed_ip_hands_save_config_only_the_changed_general_setting(monkeypatch):
"""An allowed-IP write must not drag the config file's own general_settings into
the database row. This covers the route end of that contract: what /add/allowed_ip
hands save_config differs from the loaded config in allowed_ips and nothing else.
save_config's end -- that the row it writes holds only those changed keys -- is
covered by test_ProxyConfig_save_config_merges_changed_keys_without_copying_file_settings.
This lives here rather than in the e2e suite because /add/allowed_ip mutates the
live general_settings["allowed_ips"] that auth_utils._check_valid_ip reads, so on a
shared proxy the first call locks every later request out, cleanup included.
"""
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.config_resolvers.changed_section_keys import changed_section_keys
from litellm.proxy.config_resolvers.settings_store import SettingsStore
file_settings: Final = MappingProxyType({"max_parallel_requests": 100, "proxy_config_reload_interval_seconds": 7})
store: Final = SettingsStore("general_settings")
store.load_yaml(file_settings)
fake_prisma: Final = MagicMock()
fake_prisma.db.litellm_auditlog.create = AsyncMock()
save_config: Final = AsyncMock(side_effect=lambda new_config: new_config)
async def _get_config():
return {"general_settings": dict(file_settings)}
monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server_module, "store_model_in_db", True)
monkeypatch.setattr(proxy_server_module, "premium_user", True)
monkeypatch.setattr(proxy_server_module, "general_settings", store)
monkeypatch.setattr(proxy_server_module.proxy_config, "get_config", _get_config)
monkeypatch.setattr(proxy_server_module.proxy_config, "save_config", save_config)
async def _admin_auth():
return UserAPIKeyAuth(
user_id="config-admin",
api_key="hashed-admin-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
app.dependency_overrides[user_api_key_auth] = _admin_auth
try:
resp: Final = client.post("/add/allowed_ip", json={"ip": "203.0.113.77"})
assert resp.status_code == 200, resp.text
save_config.assert_awaited_once()
persisted: Final = save_config.await_args.kwargs["new_config"]["general_settings"]
changed, removed = changed_section_keys(file_settings, persisted)
assert dict(changed) == {"allowed_ips": ["203.0.113.77"]}
assert removed == frozenset()
assert store["allowed_ips"] == ["203.0.113.77"]
finally:
app.dependency_overrides.pop(user_api_key_auth, None)
def test_delete_allowed_ip_writes_deleted_audit_log(monkeypatch):
"""Removing an allowed IP must be audited as a deletion, symmetric with the
add path."""

View file

@ -12,7 +12,7 @@ from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator, AsyncIterator
from datetime import datetime
from typing import Any, Dict, List
from typing import Any, Dict, Final, List
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -179,6 +179,29 @@ async def _one_chunk() -> AsyncGenerator[object, None]:
yield "chunk"
class _AttributeStream:
_hidden_params = {"model_id": "m-1"}
model = "gpt-x"
def __init__(self) -> None:
self._chunks = ("chunk-1", "chunk-2")
self._index = 0
self.closed = False
def __aiter__(self) -> "_AttributeStream":
return self
async def __anext__(self) -> str:
if self._index >= len(self._chunks):
raise StopAsyncIteration
chunk = self._chunks[self._index]
self._index += 1
return chunk
async def aclose(self) -> None:
self.closed = True
@pytest.mark.asyncio
async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(proxy_logging):
async def gen():
@ -247,6 +270,68 @@ async def test_wrap_streaming_iterator_leaves_upstream_http_exception_unattribut
assert request_data == {}
@pytest.mark.asyncio
async def test_wrap_streaming_iterator_forwards_response_attributes_to_hook(proxy_logging):
async def prefix_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]:
async for chunk in response:
yield f"{response._hidden_params['model_id']}:{response.model}:{chunk}"
source = _AttributeStream()
wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(
callback=MagicMock(guardrail_name="g", event_hook="post_call"),
response=source,
hook=prefix_hook,
request_data={},
)
assert [chunk async for chunk in wrapped] == [
"m-1:gpt-x:chunk-1",
"m-1:gpt-x:chunk-2",
]
@pytest.mark.asyncio
async def test_wrap_streaming_iterator_forwards_aclose_to_upstream(proxy_logging):
async def close_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]:
first: Final = await response.__anext__()
yield first
await response.aclose()
source = _AttributeStream()
request_data: dict[str, object] = {}
wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(
callback=MagicMock(guardrail_name="g", event_hook="post_call"),
response=source,
hook=close_hook,
request_data=request_data,
)
assert [chunk async for chunk in wrapped] == ["chunk-1"]
assert source.closed is True
assert request_data == {}
@pytest.mark.asyncio
async def test_wrap_streaming_iterator_missing_attribute_still_raises(proxy_logging):
async def missing_attribute_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]:
_missing: Final = response.not_there
if False:
yield
request_data: dict[str, object] = {}
wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(
callback=MagicMock(guardrail_name="hook-bug", event_hook="post_call"),
response=_one_chunk(),
hook=missing_attribute_hook,
request_data=request_data,
)
with pytest.raises(AttributeError):
async for _ in wrapped:
pass
assert request_data["metadata"]["applied_guardrails"] == ["hook-bug"]
# ---------------------------------------------------------------------------
# async_post_call_streaming_hook
# ---------------------------------------------------------------------------

Some files were not shown because too many files have changed in this diff Show more