scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge remote-tracking branch 'origin/litellm_internal_staging' into devin/1787944648-registry-audit-rolling

Browse source
This commit is contained in:
Devin AI 2026-08-29 19:02:03 +00:00
commit 23703a5341
93 changed files with 4480 additions and 410 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 17271
"limit": 17270
},
"reportArgumentType": {
"limit": 2539

78
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import (
CallTypes,
@ -222,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object=file_object,
model_mappings=model_mappings,
flat_model_file_ids=list(model_mappings.values()),
created_by=user_api_key_dict.user_id,
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,
)
@ -238,7 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"unified_file_id": file_id,
"model_mappings": json.dumps(model_mappings),
"flat_model_file_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"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,
}
@ -342,7 +343,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_object": file_object.model_dump_json(),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": user_api_key_dict.user_id,
"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,
"status": file_object.status,
@ -473,19 +474,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
page_size: Final = min(limit or 20, 100)
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
batches = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=page_size + 1,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
matches: Final = await self._collect_listed_batches(
where_clause=where_clause,
after=after,
wanted=page_size + 1,
user_api_key_dict=user_api_key_dict,
)
return build_list_page(list(matches[:page_size]), has_more=len(matches) > page_size)
has_more = len(batches) > page_size
async def _collect_listed_batches(
self,
where_clause: Mapping[str, object],
after: Optional[str],
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
"""Read chunks newest-first until ``wanted`` batches survive parsing and
file-id resolution or the caller's rows run out, so a run of rows that will
not parse refills the page instead of emptying it. The first chunk is
``wanted`` rows, so a healthy page still costs one query; a scan that has to
continue widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE`` like ``afile_list``,
and every chunk advances the keyset cursor, so the walk ends once the
caller's rows are exhausted."""
matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks
cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row
chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk
while len(matches) < wanted:
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
)
matches = matches + await self._resolve_listed_rows(
rows=chunk, wanted=wanted - len(matches), user_api_key_dict=user_api_key_dict
)
if len(chunk) < chunk_size:
break
cursor_id = chunk[-1].unified_object_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return matches
async def _resolve_listed_rows(
self,
rows: "Sequence[PrismaManagedObjectRow]",
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
parsed_rows: Final = tuple(
(row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None
(row, batch_obj) for row in rows if (batch_obj := _parse_managed_batch_row(row)) is not None
)
unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified(
raw_file_ids=frozenset(
@ -496,19 +534,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
prisma_client=self.prisma_client,
)
resolved_batches: Final = [
await self._resolve_listed_batch(
resolved: Final[list[LiteLLMBatch]] = [] # mutable-ok: resolution stops as soon as the page is full
for row, batch_obj in parsed_rows:
if len(resolved) == wanted:
break
resolved_batch = await self._resolve_listed_batch(
row=row,
batch_obj=batch_obj,
unified_id_by_raw_id=unified_id_by_raw_id,
user_api_key_dict=user_api_key_dict,
)
for row, batch_obj in parsed_rows
]
return build_list_page(
[batch_obj for batch_obj in resolved_batches if batch_obj is not None],
has_more=has_more,
)
if resolved_batch is not None:
resolved.append(resolved_batch)
return tuple(resolved)
async def _resolve_listed_batch(
self,

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.61"
version = "0.1.62"
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.61"
version = "0.1.62"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.90"
version = "0.4.91"
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.90"
version = "0.4.91"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

340
litellm-rust/Cargo.lock generated
View file

@ -2,6 +2,36 @@
# It is not intended for manual editing.
version = 4
[[package]]
name = "aho-corasick"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba"
dependencies = [
"memchr",
]
[[package]]
name = "alloca"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4"
dependencies = [
"cc",
]
[[package]]
name = "anes"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
[[package]]
name = "anstyle"
version = "1.0.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
[[package]]
name = "arc-swap"
version = "1.9.2"
@ -506,6 +536,12 @@ dependencies = [
"either",
]
[[package]]
name = "cast"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
[[package]]
name = "cc"
version = "1.3.0"
@ -541,6 +577,58 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "ciborium"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
dependencies = [
"ciborium-io",
"ciborium-ll",
"serde",
]
[[package]]
name = "ciborium-io"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
[[package]]
name = "ciborium-ll"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
dependencies = [
"ciborium-io",
"half",
]
[[package]]
name = "clap"
version = "4.6.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca"
dependencies = [
"clap_builder",
]
[[package]]
name = "clap_builder"
version = "4.6.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889"
dependencies = [
"anstyle",
"clap_lex",
]
[[package]]
name = "clap_lex"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
[[package]]
name = "cmake"
version = "0.1.58"
@ -596,6 +684,72 @@ dependencies = [
"libc",
]
[[package]]
name = "criterion"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3"
dependencies = [
"alloca",
"anes",
"cast",
"ciborium",
"clap",
"criterion-plot",
"itertools",
"num-traits",
"oorandom",
"page_size",
"plotters",
"rayon",
"regex",
"serde",
"serde_json",
"tinytemplate",
"walkdir",
]
[[package]]
name = "criterion-plot"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea"
dependencies = [
"cast",
"itertools",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-epoch"
version = "0.9.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
[[package]]
name = "crunchy"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
[[package]]
name = "crypto-common"
version = "0.1.7"
@ -856,6 +1010,17 @@ dependencies = [
"tracing",
]
[[package]]
name = "half"
version = "2.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
dependencies = [
"cfg-if",
"crunchy",
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
@ -1179,6 +1344,15 @@ version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "itertools"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.18"
@ -1255,10 +1429,13 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"criterion",
"litellm-ai-gateway",
"litellm-core",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"serde",
"serde_json",
"tokio",
]
@ -1340,6 +1517,12 @@ version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
[[package]]
name = "oorandom"
version = "11.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
[[package]]
name = "openssl-probe"
version = "0.2.1"
@ -1352,6 +1535,16 @@ version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e"
[[package]]
name = "page_size"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@ -1376,6 +1569,34 @@ version = "0.3.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e"
[[package]]
name = "plotters"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747"
dependencies = [
"num-traits",
"plotters-backend",
"plotters-svg",
"wasm-bindgen",
"web-sys",
]
[[package]]
name = "plotters-backend"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a"
[[package]]
name = "plotters-svg"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670"
dependencies = [
"plotters-backend",
]
[[package]]
name = "portable-atomic"
version = "1.14.0"
@ -1486,6 +1707,16 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "pythonize"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89"
dependencies = [
"pyo3",
"serde",
]
[[package]]
name = "quinn"
version = "0.11.11"
@ -1613,12 +1844,61 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "rayon"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
dependencies = [
"either",
"rayon-core",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
dependencies = [
"crossbeam-deque",
"crossbeam-utils",
]
[[package]]
name = "regex"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
dependencies = [
"aho-corasick",
"memchr",
"regex-automata",
"regex-syntax",
]
[[package]]
name = "regex-automata"
version = "0.4.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2"
dependencies = [
"aho-corasick",
"memchr",
"regex-syntax",
]
[[package]]
name = "regex-lite"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
[[package]]
name = "regex-syntax"
version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "reqwest"
version = "0.12.28"
@ -1774,6 +2054,15 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "same-file"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
dependencies = [
"winapi-util",
]
[[package]]
name = "schannel"
version = "0.1.29"
@ -2099,6 +2388,16 @@ dependencies = [
"zerovec",
]
[[package]]
name = "tinytemplate"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
dependencies = [
"serde",
"serde_json",
]
[[package]]
name = "tinyvec"
version = "1.12.0"
@ -2363,6 +2662,16 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
[[package]]
name = "walkdir"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
dependencies = [
"same-file",
"winapi-util",
]
[[package]]
name = "want"
version = "0.3.1"
@ -2475,6 +2784,37 @@ dependencies = [
"rustls-pki-types",
]
[[package]]
name = "winapi"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
dependencies = [
"winapi-i686-pc-windows-gnu",
"winapi-x86_64-pc-windows-gnu",
]
[[package]]
name = "winapi-i686-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
[[package]]
name = "winapi-util"
version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "winapi-x86_64-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows-link"
version = "0.2.1"

1
litellm-rust/Cargo.toml
View file

@ -18,6 +18,7 @@ litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
axum = "0.7"
pyo3 = "0.29.0"
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", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }

15
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -9,10 +9,23 @@ repository.workspace = true
name = "_native"
crate-type = ["cdylib"]
[features]
default = ["extension-module"]
extension-module = ["pyo3/extension-module"]
[dependencies]
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
pyo3 = { workspace = true, features = ["extension-module"] }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
pythonize.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
[[bench]]
name = "serialization"
harness = false

103
litellm-rust/crates/python-bridge/benches/serialization.rs Normal file
View file

@ -0,0 +1,103 @@
use std::hint::black_box;
use std::time::Duration;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::{Value, json};
const PAYLOAD_SIZES: &[(&str, usize)] = &[
("1_KiB", 1024),
("64_KiB", 64 * 1024),
("1_MiB", 1024 * 1024),
("4_MiB", 4 * 1024 * 1024),
("16_MiB", 16 * 1024 * 1024),
];
fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Value {
let json = py.import("json").expect("Python json module should import");
let encoded: String = json
.call_method1("dumps", (value,))
.expect("payload should serialize")
.extract()
.expect("json.dumps should return a string");
serde_json::from_str(&encoded).expect("serialized JSON should parse")
}
fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value {
pythonize::depythonize(value).expect("payload should depythonize")
}
fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py<PyAny> {
let json = py.import("json").expect("Python json module should import");
let encoded = serde_json::to_string(value).expect("response should serialize");
json.call_method1("loads", (encoded,))
.expect("serialized response should parse in Python")
.unbind()
}
fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py<PyAny> {
pythonize::pythonize(py, value)
.expect("response should pythonize")
.unbind()
}
fn serialization(c: &mut Criterion) {
Python::initialize();
Python::attach(|py| {
for &(label, payload_bytes) in PAYLOAD_SIZES {
let data_uri = format!("data:image/png;base64,{}", "A".repeat(payload_bytes));
let document = PyDict::new(py);
document
.set_item("type", "image_url")
.expect("document type should be set");
document
.set_item("image_url", &data_uri)
.expect("document URL should be set");
let response = json!({
"pages": [{
"index": 0,
"markdown": "OCR text",
"images": [{"image_base64": data_uri}],
}],
"model": "mistral-ocr-latest",
"document_annotation": null,
"usage_info": {"pages_processed": 1},
"object": "ocr",
});
c.bench_with_input(
BenchmarkId::new("python_to_rust_json", label),
&document,
|b, document| {
b.iter(|| former_json_roundtrip_from_py(py, black_box(document.as_any())))
},
);
c.bench_with_input(
BenchmarkId::new("python_to_rust_pythonize", label),
&document,
|b, document| b.iter(|| pythonize_from_py(black_box(document.as_any()))),
);
c.bench_with_input(
BenchmarkId::new("rust_to_python_json", label),
&response,
|b, response| b.iter(|| former_json_roundtrip_to_py(py, black_box(response))),
);
c.bench_with_input(
BenchmarkId::new("rust_to_python_pythonize", label),
&response,
|b, response| b.iter(|| pythonize_to_py(py, black_box(response))),
);
}
});
}
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(20)
.warm_up_time(Duration::from_secs(1))
.measurement_time(Duration::from_secs(4));
targets = serialization
}
criterion_main!(benches);

48
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -19,6 +19,9 @@ use pyo3::types::{PyAny, PyDict};
use serde_json::{Map, Value};
mod gil;
mod marshal;
use marshal::{from_py, to_py};
pyo3::create_exception!(
_native,
@ -41,35 +44,18 @@ type MarshaledOcrInputs = (
Option<Duration>,
);
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
let json = py.import("json")?;
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string()))
}
fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
let json = py.import("json")?;
let encoded =
serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?;
Ok(json.call_method1("loads", (encoded,))?.unbind())
}
fn messages_response_to_py(
py: Python<'_>,
response: AnthropicMessagesResponse,
) -> PyResult<Py<PyAny>> {
let value =
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
json_to_py(py, value)
to_py(py, &response)
}
fn chat_completions_response_to_py(
py: Python<'_>,
response: ChatCompletionsResponse,
) -> PyResult<Py<PyAny>> {
let value =
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
json_to_py(py, value)
to_py(py, &response)
}
fn core_error_to_pyerr(err: CoreError) -> PyErr {
@ -116,7 +102,7 @@ fn optional_object_to_map(
value: Option<Py<PyAny>>,
) -> PyResult<Map<String, Value>> {
match value {
Some(value) => match py_to_json(py, value.bind(py))? {
Some(value) => match from_py(value.bind(py))? {
Value::Object(map) => Ok(map),
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
},
@ -139,7 +125,7 @@ fn marshal_headers(
headers: Option<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
let value = match headers {
Some(headers) => py_to_json(py, headers.bind(py))?,
Some(headers) => from_py(headers.bind(py))?,
None => Value::Object(Map::new()),
};
let Value::Object(headers) = value else {
@ -211,7 +197,7 @@ fn marshal_inputs(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledOcrInputs> {
let document = py_to_json(py, document.bind(py))?;
let document = from_py(document.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -262,7 +248,7 @@ fn ocr(
});
match result {
Ok(value) => json_to_py(py, value),
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
@ -307,7 +293,7 @@ fn aocr(
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
Python::attach(|py| to_py(py, &value))
})
}
@ -325,7 +311,7 @@ fn transcription(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let audio = from_py(audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -351,7 +337,7 @@ fn transcription(
))
});
match result {
Ok(value) => json_to_py(py, value),
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
@ -370,7 +356,7 @@ fn atranscription(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let audio = from_py(audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -394,7 +380,7 @@ fn atranscription(
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
Python::attach(|py| to_py(py, &value))
})
}
@ -406,7 +392,7 @@ fn marshal_messages_inputs(
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledMessagesInputs> {
let body = py_to_json(py, body.bind(py))?;
let body: Value = from_py(body.bind(py))?;
if !body.is_object() {
return Err(PyValueError::new_err("body must be a dict"));
}
@ -498,7 +484,7 @@ fn marshal_chat_completions_inputs(
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledChatCompletionsInputs> {
let messages = py_to_json(py, messages.bind(py))?;
let messages: Value = from_py(messages.bind(py))?;
if !messages.is_array() {
return Err(PyValueError::new_err("messages must be a list"));
}
@ -527,7 +513,7 @@ fn chat_completions_decline(
optional_params: Option<Py<PyAny>>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
let messages = py_to_json(py, messages.bind(py))?;
let messages = from_py(messages.bind(py))?;
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
Ok(chat_completions_decline_reason(
&model,

20
litellm-rust/crates/python-bridge/src/marshal.rs Normal file
View file

@ -0,0 +1,20 @@
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;
pub fn from_py<T>(value: &Bound<'_, PyAny>) -> PyResult<T>
where
T: DeserializeOwned,
{
pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string()))
}
pub fn to_py<T>(py: Python<'_>, value: &T) -> PyResult<Py<PyAny>>
where
T: Serialize + ?Sized,
{
pythonize::pythonize(py, value)
.map(Bound::unbind)
.map_err(|error| PyValueError::new_err(error.to_string()))
}

52
litellm-rust/crates/python-bridge/tests/marshal_boundary.rs Normal file
View file

@ -0,0 +1,52 @@
use std::fs;
use std::path::{Path, PathBuf};
const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[
"py.import(\"json\")",
"pythonize::",
"serde_json::to_string",
"serde_json::from_str",
];
fn source_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("src")
}
fn rust_sources(directory: &Path) -> Vec<PathBuf> {
fs::read_dir(directory)
.expect("bridge source directory should be readable")
.map(|entry| {
entry
.expect("bridge source entry should be readable")
.path()
})
.flat_map(|path| {
if path.is_dir() {
rust_sources(&path)
} else if path.extension().is_some_and(|extension| extension == "rs") {
vec![path]
} else {
Vec::new()
}
})
.collect()
}
#[test]
fn serialization_is_centralized_in_marshal_module() {
let root = source_root();
for path in rust_sources(&root) {
if path == root.join("marshal.rs") {
continue;
}
let source = fs::read_to_string(&path).expect("bridge source should be readable");
for disallowed in DISALLOWED_OUTSIDE_MARSHAL {
assert!(
!source.contains(disallowed),
"{} bypasses the typed marshal module with `{disallowed}`",
path.display()
);
}
}
}

85
litellm/litellm_core_utils/llm_request_utils.py
View file

@ -2,6 +2,7 @@ from collections.abc import Mapping
from typing import Final
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
def _form_field_value(value: object) -> str:
@ -13,18 +14,31 @@ def _form_field_value(value: object) -> str:
def _flatten_form_field(key: str, value: object) -> tuple[tuple[str, str], ...]:
if isinstance(value, Mapping):
return tuple(
item for subkey, subvalue in value.items() for item in _flatten_form_field(f"{key}[{subkey}]", subvalue)
)
if isinstance(value, (list, tuple)):
return tuple(item for entry in value for item in _flatten_form_field(f"{key}[]", entry))
if value is None:
return ()
serialized: Final = _form_field_value(value)
if not serialized:
return ()
return ((key, serialized),)
pending_fields: Final[ # mutable-ok: depth-capped stack walks nested JSON into multipart names
list[tuple[str, object, int]]
] = [ # mutable-ok: depth-capped stack walks nested JSON into multipart names
(key, value, 0)
]
flat_fields: Final[list[tuple[str, str]]] = [] # mutable-ok: local accumulator
while pending_fields:
current_key, current_value, depth = pending_fields.pop()
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError("form field nesting exceeds max depth")
if isinstance(current_value, Mapping):
pending_fields.extend(
(f"{current_key}[{subkey}]", subvalue, depth + 1)
for subkey, subvalue in reversed(tuple(current_value.items()))
)
continue
if isinstance(current_value, (list, tuple)):
pending_fields.extend((f"{current_key}[]", entry, depth + 1) for entry in reversed(tuple(current_value)))
continue
if current_value is None:
continue
serialized = _form_field_value(current_value)
if serialized:
flat_fields.append((current_key, serialized))
return tuple(flat_fields)
def _is_form_scalar(value: object) -> bool:
@ -32,23 +46,36 @@ def _is_form_scalar(value: object) -> bool:
def _flatten_form_data_field(key: str, value: object) -> tuple[tuple[str, str | tuple[str, ...]], ...]:
if isinstance(value, Mapping):
return tuple(
item
for subkey, subvalue in value.items()
for item in _flatten_form_data_field(f"{key}[{subkey}]", subvalue)
)
if isinstance(value, (list, tuple)):
if all(_is_form_scalar(entry) for entry in value):
serialized_fields: Final = tuple(field for entry in value if (field := _form_field_value(entry)))
return ((key, serialized_fields),) if serialized_fields else ()
return tuple(item for entry in value for item in _flatten_form_data_field(f"{key}[]", entry))
if value is None:
return ()
serialized: Final = _form_field_value(value)
if not serialized:
return ()
return ((key, serialized),)
pending_fields: Final[ # mutable-ok: depth-capped stack walks nested JSON into multipart names
list[tuple[str, object, int]]
] = [ # mutable-ok: depth-capped stack walks nested JSON into multipart names
(key, value, 0)
]
flat_fields: Final[list[tuple[str, str | tuple[str, ...]]]] = [] # mutable-ok: local accumulator
while pending_fields:
current_key, current_value, depth = pending_fields.pop()
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError("form field nesting exceeds max depth")
if isinstance(current_value, Mapping):
pending_fields.extend(
(f"{current_key}[{subkey}]", subvalue, depth + 1)
for subkey, subvalue in reversed(tuple(current_value.items()))
)
continue
if isinstance(current_value, (list, tuple)):
if all(_is_form_scalar(entry) for entry in current_value):
serialized_fields = tuple(field for entry in current_value if (field := _form_field_value(entry)))
if serialized_fields:
flat_fields.append((current_key, serialized_fields))
continue
pending_fields.extend((f"{current_key}[]", entry, depth + 1) for entry in reversed(tuple(current_value)))
continue
if current_value is None:
continue
serialized = _form_field_value(current_value)
if serialized:
flat_fields.append((current_key, serialized))
return tuple(flat_fields)
def flatten_form_field_values(*sources: Mapping[str, object] | None) -> tuple[tuple[str, str | tuple[str, ...]], ...]:

5
litellm/llms/base_llm/managed_resources/base_managed_resource.py
View file

@ -12,6 +12,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import SpecialEnums
@ -157,7 +158,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
"resource_object": resource_object,
"model_mappings": model_mappings,
"flat_model_resource_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"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,
}
@ -179,7 +180,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]):
"unified_resource_id": unified_resource_id,
"model_mappings": json.dumps(model_mappings),
"flat_model_resource_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"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,
}

49
litellm/llms/base_llm/managed_resources/isolation.py
View file

@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources.
Returns a Prisma filter and an ownership check that scope managed resources
to the caller's identity: proxy admins see everything, user-keyed callers
see records they created, and service-account keys (no user_id) fall back
to the resource's owning team. Callers with no admin role and no
identifying ids are denied so an empty user_id can never select an
unscoped query.
see records they created, service-account keys (no user_id) fall back to
the resource's owning team, and keys with neither a user_id nor a team_id
fall back to their own hashed token so they can still reach the resources
they created. Callers with no admin role and no identifying ids at all
are denied so an empty user_id can never select an unscoped query.
"""
from typing import Any, Final
@ -19,6 +20,32 @@ from litellm.proxy._types import (
)
def resolve_resource_owner_id(
user_api_key_dict: UserAPIKeyAuth,
) -> str | None:
"""Return the identity to stamp on (and match against) a managed
resource's ``created_by``.
A key with neither a user_id nor a team_id would otherwise stamp
``created_by=None`` and be locked out of its own resources, so it owns
them under its hashed token instead, using the ``key:`` scope prefix
already used by ``proxy/common_utils/resource_ownership.py``. ``None``
means the caller has no usable identity of its own and must fall back
to team scoping, or be denied.
"""
if user_api_key_dict.user_id is not None:
return user_api_key_dict.user_id
if user_api_key_dict.team_id is not None:
return None
token: Final = user_api_key_dict.token or user_api_key_dict.api_key
if token:
return f"key:{token}"
return None
def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]:
"""Build the OpenAI-style paginated list response shape used by managed
file/batch/vector-store listings. ``first_id`` and ``last_id`` are
@ -39,7 +66,8 @@ def build_owner_filter(
to records the caller is allowed to see.
- ``{}`` means no scoping (proxy admins).
- ``{"created_by": <user_id>}`` for user-keyed callers.
- ``{"created_by": <owner_id>}`` for user-keyed callers, and for keys
with no user_id and no team_id (owner id is their hashed token).
- ``{"team_id": <team_id>}`` for service-account callers
that have a team but no user_id.
- ``{"OR": [...]}`` when the caller has both — listing must include
@ -62,12 +90,13 @@ def build_owner_filter(
]
}
if user_id is not None:
return {"created_by": user_id}
if team_id is not None:
return {"team_id": team_id}
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
if owner_id is not None:
return {"created_by": owner_id}
return None
@ -86,8 +115,8 @@ def can_access_resource(
if _user_has_admin_view(user_api_key_dict):
return True
user_id: Final = user_api_key_dict.user_id
if user_id is not None and created_by is not None and created_by == user_id:
owner_id: Final = resolve_resource_owner_id(user_api_key_dict)
if owner_id is not None and created_by is not None and created_by == owner_id:
return True
team_id: Final = user_api_key_dict.team_id

8
litellm/llms/databricks/common_utils.py
View file

@ -13,6 +13,7 @@ Authentication priority:
import os
import re
from typing import Any, Final, Literal
from urllib.parse import urlsplit, urlunsplit
from litellm.llms.base_llm.chat.transformation import BaseLLMException
@ -224,11 +225,8 @@ class DatabricksBase:
"""
import requests
# Extract workspace URL from api_base
workspace_url = api_base.rstrip("/")
if "/serving-endpoints" in workspace_url:
workspace_url = workspace_url.replace("/serving-endpoints", "")
api_base_parts: Final = urlsplit(api_base)
workspace_url: Final = urlunsplit((api_base_parts.scheme, api_base_parts.netloc, "", "", ""))
token_url: Final = f"{workspace_url}/oidc/v1/token"
try:

9
litellm/llms/hosted_vllm/videos/__init__.py Normal file
View file

@ -0,0 +1,9 @@
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from .transformation import HostedVLLMVideoConfig
__all__ = ("HostedVLLMVideoConfig",)
def get_hosted_vllm_video_config(model: str | None) -> BaseVideoConfig:
return HostedVLLMVideoConfig()

206
litellm/llms/hosted_vllm/videos/transformation.py Normal file
View file

@ -0,0 +1,206 @@
"""Video generation for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/videos)."""
import json
from collections.abc import Mapping
from io import BufferedReader
from types import MappingProxyType
from typing import Final
from urllib.parse import urlparse
from httpx._types import FileTypes, RequestFiles
from litellm.images.utils import ImageEditRequestUtils
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
from litellm.types.videos.main import VideoCreateOptionalRequestParams
_EXCLUDED_FORM_KEYS: Final = frozenset(
{
"model",
"prompt",
"extra_headers",
"extra_query",
"extra_body",
"timeout",
"custom_llm_provider",
"input_reference",
"characters",
}
)
_VLLM_OMNI_VIDEO_PARAMS: Final = (
"image_reference",
"video_reference",
"audio_reference",
"width",
"height",
"num_frames",
"fps",
"num_inference_steps",
"guidance_scale",
"guidance_scale_2",
"boundary_ratio",
"flow_shift",
"true_cfg_scale",
"seed",
"generate_sound",
"sound_duration",
"negative_prompt",
"enable_frame_interpolation",
"frame_interpolation_exp",
"frame_interpolation_scale",
"frame_interpolation_model_path",
"lora",
"extra_params",
"aspect_ratio",
)
_REFERENCE_URL_KEYS: Final = MappingProxyType(
{
"image_reference": "image_url",
"video_reference": "video_url",
"audio_reference": "audio_url",
}
)
def _serialize_form_value(value: object) -> str:
if isinstance(value, str):
return value
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, (Mapping, list, tuple)):
return json.dumps(value)
return str(value)
def _maybe_json(value: object) -> object:
if not isinstance(value, str):
return value
stripped: Final = value.strip()
if not stripped or stripped[0] not in "{[":
return value
return json.loads(stripped)
def _reject_unsafe_media_url(url: str) -> None:
scheme: Final = urlparse(url).scheme.lower()
if scheme in ("", "data"):
return
if scheme not in ("http", "https"):
raise SSRFError(f"URL scheme '{scheme}' is not allowed")
validate_url(url)
def _reject_unsafe_urls_in_item(url_key: str, item: object) -> None:
if not isinstance(item, Mapping):
return
url: Final = item.get(url_key)
if isinstance(url, str):
_reject_unsafe_media_url(url)
def _reject_unsafe_media_urls(field_name: str, value: object) -> None:
url_key: Final = _REFERENCE_URL_KEYS.get(field_name)
if url_key is None:
return
parsed: Final = _maybe_json(value)
if isinstance(parsed, list):
for item in parsed:
_reject_unsafe_urls_in_item(url_key, item)
return
if isinstance(parsed, Mapping):
_reject_unsafe_urls_in_item(url_key, parsed)
def _form_value(key: str, value: object) -> str:
_reject_unsafe_media_urls(key, value)
return _serialize_form_value(value)
def _input_reference_file(reference: object) -> tuple[str, FileTypes]:
content_type: Final = ImageEditRequestUtils.get_image_content_type(reference)
if isinstance(reference, BufferedReader):
return ("input_reference", (reference.name, reference, content_type))
return ("input_reference", ("input_reference.png", reference, content_type))
class HostedVLLMVideoConfig(OpenAIVideoConfig):
"""
vLLM-Omni videos API is OpenAI-compatible but requires multipart/form-data.
https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/videos_api/
"""
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseVideoConfig contract
return [ # mutable-ok: BaseVideoConfig returns list
*super().get_supported_openai_params(model),
*_VLLM_OMNI_VIDEO_PARAMS,
]
def map_openai_params(
self,
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict: # mutable-ok: BaseVideoConfig contract; extra_body merge mutates this dict
return { # mutable-ok: VideoGenerationRequestUtils.update/pop extra_body onto this mapping
key: value for key, value in video_create_optional_params.items() if value is not None
}
def validate_environment(
self,
headers: dict, # mutable-ok: BaseVideoConfig contract
model: str,
api_key: str | None = None,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict: # mutable-ok: BaseVideoConfig contract
resolved_key: Final = (
(litellm_params.api_key if litellm_params is not None else None)
or api_key
or get_secret_str("HOSTED_VLLM_API_KEY")
or "fake-api-key"
)
return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict, # mutable-ok: BaseVideoConfig contract
) -> str:
resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE")
if resolved_api_base is None:
raise ValueError(
"api_base not set for Hosted VLLM videos API. "
"Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable"
)
trimmed: Final = resolved_api_base.rstrip("/")
if trimmed.endswith("/v1"):
return f"{trimmed}/videos"
return f"{trimmed}/v1/videos"
def transform_video_create_request(
self,
model: str,
prompt: str,
api_base: str,
video_create_optional_request_params: dict, # mutable-ok: BaseVideoConfig contract
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: BaseVideoConfig contract
) -> tuple[dict, RequestFiles, str]: # mutable-ok: BaseVideoConfig contract
data: Final = { # mutable-ok: BaseVideoConfig contract returns a data dict
"model": model,
"prompt": prompt,
**{ # mutable-ok: spread remaining Omni form fields into that data dict
key: _form_value(key, value)
for key, value in video_create_optional_request_params.items()
if key not in _EXCLUDED_FORM_KEYS and value is not None
},
}
input_reference: Final = video_create_optional_request_params.get("input_reference")
if input_reference is None:
return data, (), api_base
return data, (_input_reference_file(input_reference),), api_base

4
litellm/main.py
View file

@ -8612,9 +8612,9 @@ def _joined_streamed_citations(streamed_citations: "tuple[object, ...]") -> "lis
def _stream_builder_model_map_cost(response: ModelResponse) -> float | None:
model_name: Final = getattr(response, "model", None)
model_name: Final = response.model
usage: Final = getattr(response, "usage", None)
if not isinstance(model_name, str) or not model_name or not isinstance(usage, Usage):
if not model_name or not isinstance(usage, Usage):
return None
try:
prompt_cost, completion_tokens_cost = litellm.cost_per_token(model=model_name, usage_object=usage)

3
litellm/provider_endpoints_support_backup.json
View file

@ -1180,7 +1180,8 @@
"files": true,
"rerank": true,
"a2a": true,
"interactions": true
"interactions": true,
"video_generations": true
}
},
"huggingface": {

8
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -903,8 +903,9 @@ class MCPRequestHandler:
NotSessionBearer,
SessionBearerAdmitted,
SessionBearerInvalid,
SessionSigningConfigError,
active_session_signing_keys,
resolve_session_bearer,
session_keys_from_master_key,
)
from litellm.proxy.proxy_server import master_key
@ -913,7 +914,10 @@ class MCPRequestHandler:
await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route)
keys: Final = session_keys_from_master_key(master_key)
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp gateway session admission rejected: %s", keys.detail)
raise HTTPException(status_code=500, detail="Server misconfigured: mcp_session_token_signing is invalid")
result: Final = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc))
match result:
case SessionBearerAdmitted():

23
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -65,15 +65,16 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
SessionRefreshOpened,
SessionSigningConfigError,
active_session_signing_keys,
open_session_refresh_bearer,
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_REFRESH_TTL_SECONDS,
MintedSessionToken,
SessionAudience,
SessionKeys,
SessionPrincipal,
SessionSigningKeys,
mint_session_refresh_token,
mint_session_token,
)
@ -885,7 +886,7 @@ class _SingleUseGuard:
return "first" if count == 1 else "replayed"
def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response:
def _session_token_pair(principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime) -> Response:
access: Final = mint_session_token(principal, keys, now)
refresh: Final = mint_session_refresh_token(principal, keys, now)
if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken):
@ -912,7 +913,7 @@ class _ProxyCredentialTokenResponse(TypedDict):
def _proxy_credential_response(
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionKeys, now: datetime
minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime
) -> Response:
"""The proxy-API token response: the access token is the very credential ``lite
login`` stores (accepted on every proxy route with user and team attribution), and
@ -998,7 +999,10 @@ async def aggregate_token(
if master_key is None:
verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured")
return _oauth_error(500, "server_error", "the gateway has no master key configured")
keys: Final = session_keys_from_master_key(master_key)
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp_gateway_dcr token grant rejected: %s", keys.detail)
return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid")
now: Final = datetime.now(timezone.utc)
issue: Final = _GrantIssuer(
request=request,
@ -1043,7 +1047,7 @@ class _GrantIssuer:
self,
request: Request,
resource: str | None,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
reload_user: ReloadUser,
mint_proxy_credential: MintProxyCredential,
@ -1146,7 +1150,7 @@ async def _refresh_token_grant(
refresh_token: str | None,
client_id: str,
resource: str | None,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
issue: _GrantIssuer,
) -> Response:
@ -1182,7 +1186,10 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non
if master_key is None:
verbose_logger.error("mcp_gateway_dcr revoke rejected: no master_key configured")
return _oauth_error(500, "server_error", "the gateway has no master key configured")
keys: Final = session_keys_from_master_key(master_key)
keys: Final = active_session_signing_keys(master_key)
if isinstance(keys, SessionSigningConfigError):
verbose_logger.error("mcp_gateway_dcr revoke rejected: %s", keys.detail)
return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid")
now: Final = datetime.now(timezone.utc)
opened: Final = open_session_refresh_bearer(token, keys, now, expected_client_id=client_id)
if isinstance(opened, SessionRefreshOpened):

102
litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py
View file

@ -20,13 +20,16 @@ from datetime import datetime
from functools import lru_cache
from typing import Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, SecretStr
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
AsymmetricSessionKeys,
OpenedSessionToken,
SessionExpired,
SessionKeys,
SessionPrincipal,
SessionRotatedPublicKey,
SessionSigningKeys,
is_session_refresh_token,
is_session_token,
open_session_refresh_token,
@ -68,6 +71,99 @@ def session_keys_from_master_key(master_key: str) -> SessionKeys:
return SessionKeys(signing_key=SecretStr(signing))
class SessionSigningPreviousKey(BaseModel):
"""One retired key in ``mcp_session_token_signing.previous_public_keys``: its ``kid``
and the PEM public half (inline or an ``os.environ/`` reference)."""
model_config = ConfigDict(frozen=True, extra="forbid")
kid: str = Field(min_length=1)
public_key: str = Field(min_length=1)
class MCPSessionTokenSigningSettings(BaseModel):
"""The ``general_settings.mcp_session_token_signing`` block: opt-in asymmetric signing
for the gateway session tokens. Absent, the gateway keeps the backward-compatible
HS256 key derived from ``master_key``. ``private_key`` and each ``public_key`` accept
a PEM string inline or an ``os.environ/<NAME>`` (or secret manager) reference."""
model_config = ConfigDict(frozen=True, extra="forbid")
algorithm: Literal["RS256"]
kid: str = Field(min_length=1)
private_key: str = Field(min_length=1)
previous_public_keys: tuple[SessionSigningPreviousKey, ...] = ()
class SessionSigningConfigError(BaseModel):
"""``mcp_session_token_signing`` is present but unusable (bad shape, unresolvable
secret reference, or a key that is not a loadable RSA PEM); the caller fails closed
with a server error instead of silently falling back to HS256."""
model_config = ConfigDict(frozen=True)
tag: Literal["session_signing_config_error"] = "session_signing_config_error"
detail: str
def _resolve_key_material(value: str) -> str | None:
if not value.startswith("os.environ/"):
return value
from litellm.secret_managers.main import get_secret_str # noqa: PLC0415 # heavy import kept off the pure path
return get_secret_str(value)
def resolve_session_signing_keys(
master_key: str,
raw_settings: object | None,
) -> SessionSigningKeys | SessionSigningConfigError:
"""Turn the operator's ``mcp_session_token_signing`` setting into signing key material.
``None`` (the setting absent) keeps the backward-compatible HS256 key derived from
``master_key``. A present setting must fully validate into RS256 material; any defect
is a ``SessionSigningConfigError`` value so token issuance and admission fail closed
rather than minting under a key the operator did not intend.
"""
if raw_settings is None:
return session_keys_from_master_key(master_key)
try:
settings: Final = MCPSessionTokenSigningSettings.model_validate(raw_settings)
except ValidationError as exc:
return SessionSigningConfigError(detail=f"mcp_session_token_signing is malformed: {exc}")
private_pem: Final = _resolve_key_material(settings.private_key)
if private_pem is None:
return SessionSigningConfigError(detail="mcp_session_token_signing.private_key reference did not resolve")
resolved_previous: Final = tuple(
(previous.kid, _resolve_key_material(previous.public_key)) for previous in settings.previous_public_keys
)
unresolved: Final = tuple(kid for kid, pem in resolved_previous if pem is None)
if unresolved:
return SessionSigningConfigError(
detail=f"mcp_session_token_signing.previous_public_keys reference did not resolve for kid(s): {', '.join(unresolved)}"
)
try:
return AsymmetricSessionKeys(
private_key_pem=SecretStr(private_pem),
kid=settings.kid,
previous_public_keys=tuple(
SessionRotatedPublicKey(kid=kid, public_key_pem=pem)
for kid, pem in resolved_previous
if pem is not None
),
)
except ValidationError as exc:
return SessionSigningConfigError(
detail=f"mcp_session_token_signing keys are not usable RSA PEM material: {exc}"
)
def active_session_signing_keys(master_key: str) -> SessionSigningKeys | SessionSigningConfigError:
"""Wiring helper for the token endpoint and the admission edge: resolve the signing
keys from the live ``general_settings.mcp_session_token_signing`` block, or derive the
default HS256 key from ``master_key`` when the block is absent."""
from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load
return resolve_session_signing_keys(master_key, general_settings.get("mcp_session_token_signing"))
class NotSessionBearer(BaseModel):
"""The bearer is not session-shaped; admission continues on its normal path."""
@ -116,7 +212,7 @@ def is_session_bearer_shaped(authorization_value: str) -> bool:
def resolve_session_bearer(
authorization_value: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> SessionBearerResult:
"""Classify an ``Authorization`` value presented at the aggregate MCP edge.
@ -166,7 +262,7 @@ SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid
def open_session_refresh_bearer(
refresh_value: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
expected_client_id: str,
) -> SessionRefreshResult:

181
litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py
View file

@ -8,8 +8,11 @@ is therefore a stable REFERENCE, not an authorization: admission reloads the liv
record and policy on every request, so deactivating the user (or their team) kills
outstanding sessions immediately without a revocation store.
Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT,
the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp``
Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + a JWT signed with
the injected key material: HS256 under the default master-key-derived secret (the same
signing approach as :mod:`.envelope`), or RS256 under an operator-provided RSA private
key (:class:`AsymmetricSessionKeys`) so downstream validators hold only the public half.
Claims are ``iss``/``iat``/``exp``
plus ``jti`` (per-mint uniqueness, so two tokens minted in the same second never
collide and a future revocation list has a stable handle), ``kind``, ``user_id``, and
``client_id``; ``client_id`` binds the refresh token
@ -31,11 +34,16 @@ injected ``now``); the strict pydantic claims model is the sole, total type gate
from __future__ import annotations
import secrets
from collections import Counter
from datetime import datetime, timedelta
from functools import lru_cache
from typing import Final, Literal, TypeAlias
import jwt
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
from cryptography.exceptions import UnsupportedAlgorithm
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError, field_validator, model_validator
SESSION_TOKEN_PREFIX: Final = "llm_session_"
"""Marker prefix on every serialized session ACCESS token so the admission edge can cheaply
@ -71,6 +79,11 @@ limits while bounding hostile input before JWT parsing."""
_SESSION_JWT_ALGORITHM: Final = "HS256"
_SESSION_RSA_ALGORITHM: Final = "RS256"
_MIN_RSA_KEY_BITS: Final = 2048
"""RFC 7518 section 3.3: RS256 requires a key of at least 2048 bits."""
SessionTokenKind = Literal["session", "session_refresh"]
"""Which credential a session token is. Stamped into the signed claims and required to match
on open, so a signature-valid token of one kind cannot be replayed as the other even if its
@ -120,6 +133,85 @@ class SessionKeys(BaseModel):
signing_key: SecretStr = Field(min_length=32)
class SessionRotatedPublicKey(BaseModel):
"""The public half of a retired signing key, kept verifiable under its ``kid`` during a
rotation window so tokens minted before the rotation stay valid until they expire."""
model_config = ConfigDict(frozen=True)
kid: str = Field(min_length=1)
public_key_pem: str = Field(min_length=1)
@field_validator("public_key_pem")
@classmethod
def _pem_is_an_rsa_public_key(cls, value: str) -> str:
try:
loaded: Final = serialization.load_pem_public_key(value.encode())
except (ValueError, TypeError, UnsupportedAlgorithm) as exc:
raise ValueError(f"public_key_pem is not a loadable PEM public key: {exc}") from exc
if not isinstance(loaded, rsa.RSAPublicKey):
raise ValueError("public_key_pem must be an RSA public key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError
if loaded.key_size < _MIN_RSA_KEY_BITS:
raise ValueError(f"public_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits")
return value
class AsymmetricSessionKeys(BaseModel):
"""Injected RS256 key material: the issuer-held RSA private key and the stable ``kid``
stamped into every minted token's JOSE header, plus the public halves of previously
rotated keys that verification still accepts while their tokens age out. Downstream
validators never need the private key: :func:`session_public_key_pem` yields the
public half to distribute."""
model_config = ConfigDict(frozen=True)
private_key_pem: SecretStr
kid: str = Field(min_length=1)
previous_public_keys: tuple[SessionRotatedPublicKey, ...] = ()
@field_validator("private_key_pem")
@classmethod
def _pem_is_a_strong_rsa_private_key(cls, value: SecretStr) -> SecretStr:
try:
loaded: Final = serialization.load_pem_private_key(value.get_secret_value().encode(), password=None)
except (ValueError, TypeError, UnsupportedAlgorithm) as exc:
raise ValueError(f"private_key_pem is not a loadable unencrypted PEM private key: {exc}") from exc
if not isinstance(loaded, rsa.RSAPrivateKey):
raise ValueError("private_key_pem must be an unencrypted RSA private key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError
if loaded.key_size < _MIN_RSA_KEY_BITS:
raise ValueError(f"private_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits")
return value
@model_validator(mode="after")
def _kids_are_unique(self) -> AsymmetricSessionKeys:
kids: Final = (self.kid, *(previous.kid for previous in self.previous_public_keys))
duplicates: Final = tuple(kid for kid, count in Counter(kids).items() if count > 1)
if duplicates:
raise ValueError(
f"every kid must be unique across the current and previous keys; duplicated: {', '.join(duplicates)}"
)
return self
SessionSigningKeys: TypeAlias = SessionKeys | AsymmetricSessionKeys
"""Every key material shape the mints and openers accept: the default master-key-derived
HS256 secret, or operator-configured RS256 RSA keys."""
@lru_cache(maxsize=8)
def _public_key_pem_from_private(private_key_pem: str) -> str:
loaded: Final = serialization.load_pem_private_key(private_key_pem.encode(), password=None)
return (
loaded.public_key()
.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
.decode()
)
def session_public_key_pem(keys: AsymmetricSessionKeys) -> str:
"""The PEM public half of the current RS256 signing key: the only material a downstream
validator (an external gateway verifying ``kid``-matched tokens) ever needs."""
return _public_key_pem_from_private(keys.private_key_pem.get_secret_value())
class MintedSessionToken(BaseModel):
"""A minted session token: the client-held bearer value and when it expires."""
@ -221,7 +313,7 @@ def is_session_refresh_token(candidate: str) -> bool:
def mint_session_token(
principal: SessionPrincipal,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> MintedSessionToken | SessionTokenMintError:
"""Mint the short-lived session ACCESS token for ``principal``.
@ -241,7 +333,7 @@ def mint_session_token(
def mint_session_refresh_token(
principal: SessionPrincipal,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> MintedSessionToken | SessionTokenMintError:
"""Mint the long-lived session REFRESH token for ``principal``.
@ -262,7 +354,7 @@ def mint_session_refresh_token(
def open_session_token(
candidate: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> OpenedSessionToken | SessionTokenOpenError:
"""Validate a session ACCESS ``candidate`` and recover the principal.
@ -275,7 +367,7 @@ def open_session_token(
def open_session_refresh_token(
candidate: str,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> OpenedSessionToken | SessionTokenOpenError:
"""Validate a session REFRESH ``candidate`` and recover the principal.
@ -292,7 +384,7 @@ def _mint(
prefix: str,
principal: SessionPrincipal,
expires_at: datetime,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> MintedSessionToken | SessionTokenTooLarge:
"""Sign the claims for either token kind and enforce the size cap. Shared by both mints
@ -309,20 +401,33 @@ def _mint(
audience=principal.audience,
team_id=principal.team_id,
)
token: Final = prefix + jwt.encode(
claims.model_dump(exclude_none=True), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM
)
token: Final = prefix + _sign_claims(claims, keys)
size_bytes: Final = len(token.encode("utf-8"))
if size_bytes > MAX_SESSION_TOKEN_BYTES:
return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES)
return MintedSessionToken(token=SecretStr(token), expires_at=expires_at)
def _sign_claims(claims: _SessionClaims, keys: SessionSigningKeys) -> str:
"""Sign the claim set under whichever key material was injected: RS256 with the ``kid``
in the JOSE header (so a validator can pick the right public key), or the default
HS256 secret with no header extras (byte-compatible with every pre-RS256 token)."""
payload: Final = claims.model_dump(exclude_none=True)
if isinstance(keys, AsymmetricSessionKeys):
return jwt.encode(
payload,
keys.private_key_pem.get_secret_value(),
algorithm=_SESSION_RSA_ALGORITHM,
headers={"kid": keys.kid},
)
return jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM)
def _open(
candidate: str,
prefix: str,
expected_kind: SessionTokenKind,
keys: SessionKeys,
keys: SessionSigningKeys,
now: datetime,
) -> OpenedSessionToken | SessionTokenOpenError:
"""Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an
@ -337,7 +442,7 @@ def _open(
return SessionMalformed()
if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES:
return SessionMalformed()
claims: Final = _decode_claims(candidate.removeprefix(prefix), keys.signing_key)
claims: Final = _decode_claims(candidate.removeprefix(prefix), keys)
if not isinstance(claims, _SessionClaims):
return claims
if claims.kind != expected_kind:
@ -356,14 +461,51 @@ def _open(
)
class _VerificationMaterial(BaseModel):
model_config = ConfigDict(frozen=True)
key: SecretStr
algorithm: Literal["HS256", "RS256"]
def _verification_material(
compact: str,
keys: SessionSigningKeys,
) -> _VerificationMaterial | SessionBadSignature | SessionMalformed:
"""Pick the single key and algorithm the candidate is allowed to verify under.
HS256 mode has exactly one secret. RS256 mode routes by the JOSE header ``kid``: the
current key's derived public half, or a retired key's stored public half during a
rotation window. An unknown or missing ``kid`` is ``SessionBadSignature`` (a foreign
key), and an undecodable header is ``SessionMalformed``. The algorithm is pinned per
key shape, never read from the header, so an HS256 token can never be verified
against a public key or vice versa.
"""
if isinstance(keys, SessionKeys):
return _VerificationMaterial(key=keys.signing_key, algorithm=_SESSION_JWT_ALGORITHM)
try:
header: Final = jwt.get_unverified_header(compact)
except jwt.InvalidTokenError:
return SessionMalformed()
kid: Final = header.get("kid")
if kid == keys.kid:
return _VerificationMaterial(key=SecretStr(session_public_key_pem(keys)), algorithm=_SESSION_RSA_ALGORITHM)
for previous in keys.previous_public_keys:
if previous.kid == kid:
return _VerificationMaterial(key=SecretStr(previous.public_key_pem), algorithm=_SESSION_RSA_ALGORITHM)
return SessionBadSignature()
def _decode_claims(
compact: str,
signing_key: SecretStr,
keys: SessionSigningKeys,
) -> _SessionClaims | SessionBadSignature | SessionMalformed:
"""Verify the HS256 signature and shape of an attacker-controlled compact JWT.
"""Verify the signature and shape of an attacker-controlled compact JWT.
``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_BYTES`` by the caller.
PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim
The accepted algorithm is pinned by :func:`_verification_material` from the injected
key shape, so ``alg`` confusion (``none``, or HS256 signed with a public key as the
secret) fails before or at signature verification. PyJWT's ``iat``/``nbf``/``exp``
validators are disabled: they raise on hostile claim
types and, for ``iat``/``nbf``, compare against the wall clock rather than the injected
``now`` (``exp`` is checked by the caller against ``now``). Apart from a signature
mismatch, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces
@ -371,11 +513,14 @@ def _decode_claims(
``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid
token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate.
"""
material: Final = _verification_material(compact, keys)
if not isinstance(material, _VerificationMaterial):
return material
try:
payload: Final = jwt.decode(
compact,
signing_key.get_secret_value(),
algorithms=[_SESSION_JWT_ALGORITHM],
material.key.get_secret_value(),
algorithms=[material.algorithm],
issuer=SESSION_ISSUER,
options={
"verify_exp": False,

10
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -301,12 +301,12 @@ class LakeraAIGuardrail(CustomGuardrail):
explicit sync below a hot reload that changes mode would pass validation but
keep dispatching on the stale event_hook.
"""
new_event_hook: Final = getattr(litellm_params, "mode", None) or self.event_hook
prospective_payload: Final = getattr(litellm_params, "payload", None)
prospective_breakdown: Final = getattr(litellm_params, "breakdown", None)
new_event_hook: Final = litellm_params.mode or self.event_hook
prospective_payload: Final = litellm_params.payload
prospective_breakdown: Final = litellm_params.breakdown
self._validate_advisory_config(
on_flagged=getattr(litellm_params, "on_flagged", None) or self.on_flagged,
advisory_system_message=getattr(litellm_params, "advisory_system_message", None),
on_flagged=litellm_params.on_flagged or self.on_flagged,
advisory_system_message=litellm_params.advisory_system_message,
payload=self.payload if prospective_payload is None else prospective_payload,
breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown,
)

2
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -121,7 +121,7 @@ class QualifireGuardrail(CustomGuardrail):
the live instance untouched instead of raising after it's already been
corrupted. Mirrors LakeraAIGuardrail's own override of this same method.
"""
prospective_on_flagged: Final = getattr(litellm_params, "on_flagged", None) or self.on_flagged
prospective_on_flagged: Final = litellm_params.on_flagged or self.on_flagged
self._validate_on_flagged(prospective_on_flagged)
super().update_in_memory_litellm_params(litellm_params=litellm_params)

17
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -413,14 +413,15 @@ class GuardrailRegistry:
raise Exception(f"Error getting guardrail from DB: {e}")
def _apply_configured_bool_override(instance: CustomGuardrail, litellm_params: LitellmParams, param_name: str) -> None:
"""Override ``instance.<param_name>`` only when ``litellm_params`` explicitly
sets it, preserving whatever default the guardrail's own constructor chose
def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: LitellmParams) -> None:
"""Override the parallel/raw-scan flags only when ``litellm_params`` explicitly
sets them, preserving whatever default the guardrail's own constructor chose
otherwise (its constructor default may be True, so blindly copying an
absent/None config value would silently clobber it back to False)."""
configured: Final = getattr(litellm_params, param_name, None)
if configured is not None:
setattr(instance, param_name, bool(configured))
if litellm_params.run_in_parallel is not None:
instance.run_in_parallel = bool(litellm_params.run_in_parallel)
if litellm_params.scan_raw_request is not None:
instance.scan_raw_request = bool(litellm_params.scan_raw_request)
class InMemoryGuardrailHandler:
@ -544,8 +545,7 @@ class InMemoryGuardrailHandler:
"skip_tool_message_in_guardrail are enabled together, which excludes every message from "
"scanning, so no request content would ever be scanned. Remove one of the two."
)
for override_param in ("run_in_parallel", "scan_raw_request"):
_apply_configured_bool_override(custom_guardrail_callback, litellm_params, override_param)
_apply_configured_bool_overrides(custom_guardrail_callback, litellm_params)
parsed_guardrail: Final = Guardrail(
guardrail_id=guardrail.get("guardrail_id"),
@ -803,7 +803,6 @@ class InMemoryGuardrailHandler:
previous_guardrail: Final = self.IN_MEMORY_GUARDRAILS.get(guardrail_id)
previous_source: Final = self._sources.get(guardrail_id, source)
# Remove from memory if exists (also removes from callbacks)
if guardrail_id in self.IN_MEMORY_GUARDRAILS:
self.delete_in_memory_guardrail(guardrail_id)

46
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -1372,8 +1372,25 @@ class DBHealthCache(TypedDict):
db_health_cache: DBHealthCache = {"status": "unknown", "last_updated": datetime.now()}
# Bounds each DB round-trip on the probe path so a hung connection during a
# failover cannot make the probe fail by timeout (k8s default timeoutSeconds: 5).
DB_READINESS_CHECK_TIMEOUT_SECONDS: Final = 2.0
# One deadline for the whole probe-path DB check (initial check + reconnect +
# re-check, including reconnect lock waits), kept under timeoutSeconds: 5.
DB_READINESS_PROBE_DEADLINE_SECONDS: Final = 4.0
async def _db_health_readiness_check():
async def _db_health_readiness_check() -> DBHealthCache:
try:
return await asyncio.wait_for(
_db_health_readiness_check_unbounded(),
timeout=DB_READINESS_PROBE_DEADLINE_SECONDS,
)
except asyncio.TimeoutError:
return {"status": "disconnected", "last_updated": db_health_cache["last_updated"]}
async def _db_health_readiness_check_unbounded() -> DBHealthCache:
from litellm.proxy.proxy_server import prisma_client
global db_health_cache
@ -1387,7 +1404,7 @@ async def _db_health_readiness_check():
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
return db_health_cache
await prisma_client.health_check()
await asyncio.wait_for(prisma_client.health_check(), timeout=DB_READINESS_CHECK_TIMEOUT_SECONDS)
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
return db_health_cache
except Exception as e:
@ -1395,8 +1412,15 @@ async def _db_health_readiness_check():
if PrismaDBExceptionHandler.is_database_transport_error(e):
try:
verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect")
await prisma_client.attempt_db_reconnect(reason="health_readiness_check")
await prisma_client.health_check()
await prisma_client.attempt_db_reconnect(
reason="health_readiness_check",
timeout_seconds=DB_READINESS_CHECK_TIMEOUT_SECONDS,
lock_timeout_seconds=DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
await asyncio.wait_for(
prisma_client.health_check(),
timeout=DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
verbose_proxy_logger.info("_db_health_readiness_check: reconnect succeeded")
db_health_cache = {
"status": "connected",
@ -1580,7 +1604,14 @@ async def _get_health_readiness_details(
# serve requests that depend on persisted state (keys, budgets,
# spend logs). Return 503 so orchestrators take this pod out of
# rotation; "Not connected" (no DB configured at all) stays 200.
if response is not None and db_health_status["status"] != "connected":
# With allow_requests_on_db_unavailable the proxy keeps serving
# during a DB outage, so the pod must stay in rotation (200) and
# report the DB state through the body instead.
if (
response is not None
and db_health_status["status"] != "connected"
and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
):
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
return {
"status": "healthy",
@ -1671,7 +1702,10 @@ async def _resolve_public_readiness_db(response: Response) -> str:
return "Not connected"
db_health_status: Final = await _db_health_readiness_check()
if db_health_status["status"] != "connected":
if (
db_health_status["status"] != "connected"
and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
):
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
return db_health_status["status"]

7
litellm/proxy/litellm_pre_call_utils.py
View file

@ -1948,6 +1948,13 @@ async def add_litellm_data_to_request(
for key, value in data["litellm_metadata"].items():
if key not in data[_metadata_variable_name]:
data[_metadata_variable_name][key] = value
if _metadata_variable_name == "metadata":
data["metadata"]["tags"] = LiteLLMProxyRequestSetup._merge_tags( # pyright: ignore[reportPrivateUsage] # same-module helper, budget blocks the unsuppressed idiom sibling call sites use
request_tags=data["metadata"].get("tags"),
tags_to_add=data["litellm_metadata"].get("tags"),
)
if _metadata_variable_name == "metadata":
data.pop("litellm_metadata", None)
data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data=data,

14
litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py
View file

@ -9,6 +9,7 @@ from typing import Final
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL
from litellm.proxy._types import ProxyErrorTypes, ProxyException
SUGGEST_TOOL: Final = {
"type": "function",
@ -60,6 +61,18 @@ class AiPolicySuggester:
system_prompt: Final = self._build_system_prompt(templates)
user_prompt: Final = self._build_user_prompt(attack_examples, description)
model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL
custom_llm_provider: Final = model.split("/", 1)[0] if "/" in model else None
supported_params: Final = litellm.get_supported_openai_params(
model=model,
custom_llm_provider=custom_llm_provider,
)
if supported_params is not None and "tools" not in supported_params:
raise ProxyException(
message=(f"AI policy suggestion requires tool calling; model '{model}' does not support it"),
type=ProxyErrorTypes.validation_error.value,
param="model",
code=400,
)
try:
response: Final = await litellm.acompletion(
@ -74,6 +87,7 @@ class AiPolicySuggester:
"function": {"name": "select_policy_templates"},
},
temperature=0.2,
drop_params=True,
)
tool_calls: Final = response.choices[0].message.tool_calls

3
litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
View file

@ -49,6 +49,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.managed_resources.isolation import (
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
@ -686,7 +687,7 @@ async def _mint_or_reuse_object(
"file_object": json.dumps(body_snapshot),
"model_object_id": namespaced_model_object_id,
"file_purpose": file_purpose,
"created_by": user_api_key_dict.user_id,
"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,
},

2
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -185,7 +185,7 @@ class PipelineExecutor:
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = getattr(callback, "scan_raw_request", False)
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None

22
litellm/proxy/proxy_server.py
View file

@ -16,7 +16,7 @@ import threading
import time
import traceback
import warnings
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, MutableMapping, Sequence
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType, UnionType
from typing import (
@ -12792,6 +12792,7 @@ async def _fetch_db_models_for_search(
size: int,
sort_by: str | None,
is_byok_outside_caller_teams: Callable[[dict[str, JsonValue]], bool],
model_name: str | None = None,
) -> tuple[list[dict[str, Any]], int]:
"""
Run the bounded DB query that backs `/v2/model/info?search=`. Returns
@ -12808,7 +12809,9 @@ async def _fetch_db_models_for_search(
filter for `team_public_model_name` instead and keep the DB cost
bounded by `search`.
"""
db_where_condition: Final[dict[str, Any]] = {"model_name": {"contains": search_lower, "mode": "insensitive"}}
db_where_condition: Final[dict[str, Any]] = {
"model_name": {"contains": search_lower, "mode": "insensitive"} if model_name is None else model_name
}
if db_model_ids_in_router:
db_where_condition["model_id"] = {"not": {"in": list(db_model_ids_in_router)}}
@ -12855,6 +12858,7 @@ async def _apply_search_filter_to_models(
page: int = 1,
size: int = 50,
sort_by: str | None = None,
model_name: str | None = None,
) -> tuple[list[dict[str, Any]], int | None]:
"""
Apply search filter to models, querying database for additional matching models.
@ -12875,6 +12879,11 @@ async def _apply_search_filter_to_models(
sort_by: Sort field. When set, results must be sorted across the
full match set, so the DB fetch is capped at
``_SORTED_SEARCH_DB_FETCH_CAP`` instead of one page.
model_name: Exact ``model_name`` the caller already narrowed
``all_models`` to (``?model=``). The DB query matches it
exactly instead of the substring, and is skipped when the
substring cannot occur in it, otherwise rows from other model
groups leak into the result and the count.
Returns:
Tuple of (filtered_models, total_count). total_count is None if not searching.
@ -12932,7 +12941,8 @@ async def _apply_search_filter_to_models(
# Query database for additional models with search term
db_models: list[dict[str, Any]] = []
if prisma_client is not None:
exact_name_can_match: Final = model_name is None or search_lower in model_name.lower()
if prisma_client is not None and exact_name_can_match:
try:
db_models, db_models_total_count = await _fetch_db_models_for_search(
prisma_client=prisma_client,
@ -12944,6 +12954,7 @@ async def _apply_search_filter_to_models(
size=size,
sort_by=sort_by,
is_byok_outside_caller_teams=_is_byok_outside_caller_teams,
model_name=model_name,
)
search_total_count = router_models_count + db_models_total_count
except Exception as e:
@ -13497,7 +13508,7 @@ async def model_info_v2(
all_models += [user_model]
if model is not None:
all_models = [m for m in all_models if m["model_name"] == model]
all_models = [m for m in all_models if _deployment_matches_allowed_model_names(m, frozenset((model,)))]
# Apply search filter if provided
all_models, search_total_count = await _apply_search_filter_to_models(
@ -13509,6 +13520,7 @@ async def model_info_v2(
page=page,
size=size,
sort_by=sortBy,
model_name=model,
)
if user_models_only:
@ -14023,7 +14035,7 @@ async def model_metrics_exceptions(
return {"data": response, "exception_types": list(exception_types)}
def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: set[str]) -> bool:
def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: Collection[str]) -> bool:
"""Match a router deployment against allowed public model names.
Team-scoped rows store an internal routing key in ``model_name``; callers

34
litellm/proxy/utils.py
View file

@ -1416,7 +1416,7 @@ class ProxyLogging:
mutation is discarded and a warning is logged so the misconfiguration
is visible instead of silently forwarding unredacted content.
"""
scans_raw_request: Final = getattr(callback, "scan_raw_request", False)
scans_raw_request: Final = callback.scan_raw_request
should_use_raw_snapshot: Final = scans_raw_request and raw_request_snapshot is not None
input_data: Final = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot) if should_use_raw_snapshot else data
@ -1453,7 +1453,7 @@ class ProxyLogging:
"scan_raw_request is for block-only guardrails and this mutation is being "
"discarded. Remove scan_raw_request from this guardrail's config if it needs "
"to mask/rewrite content.",
getattr(callback, "guardrail_name", None) or callback.__class__.__name__,
callback.guardrail_name or callback.__class__.__name__,
)
if scans_raw_request:
if result is not None:
@ -1778,7 +1778,7 @@ class ProxyLogging:
# guarantee must hold even under litellm.safe_memory_mode, which
# otherwise makes deep copies return the original object.
needs_raw_request_snapshot: Final = any(
isinstance(cb, CustomGuardrail) and getattr(cb, "scan_raw_request", False)
isinstance(cb, CustomGuardrail) and cb.scan_raw_request
for cb in ProxyLogging._callback_capabilities().resolved_callbacks
)
raw_request_snapshot: Final[dict | None] = ( # mutable-ok: same request-payload shape as data
@ -1938,7 +1938,7 @@ class ProxyLogging:
"""
def _input_for(callback: CustomGuardrail) -> dict: # mutable-ok: same request-payload shape as data
if not getattr(callback, "scan_raw_request", False) or raw_request_snapshot is None:
if not callback.scan_raw_request or raw_request_snapshot is None:
return data
return independent_snapshot(raw_request_snapshot)
@ -1962,11 +1962,7 @@ class ProxyLogging:
# deployment-level guardrail sharing this name would see no marker
# via _pre_call_hook_already_ran and re-run it a second time on
# live kwargs.
if (
getattr(callback, "scan_raw_request", False)
and not isinstance(result, BaseException)
and result is not None
):
if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None:
callback.mark_pre_call_hook_ran(data)
raised: Final = tuple(result for result in results if isinstance(result, BaseException))
blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None)
@ -5580,12 +5576,8 @@ class PrismaClient:
return True
acquire_task: Final = asyncio.create_task(_acquire_reconnect_lock())
done, _pending = await asyncio.wait(
{acquire_task},
timeout=lock_timeout_seconds,
return_when=asyncio.FIRST_COMPLETED,
)
if acquire_task not in done:
async def _abandon_acquire_task() -> None:
acquire_task.cancel()
try:
await acquire_task
@ -5600,6 +5592,18 @@ class PrismaClient:
self._db_reconnect_lock.release()
except RuntimeError:
pass
try:
done, _pending = await asyncio.wait(
{acquire_task},
timeout=lock_timeout_seconds,
return_when=asyncio.FIRST_COMPLETED,
)
except asyncio.CancelledError:
await asyncio.shield(_abandon_acquire_task())
raise
if acquire_task not in done:
await _abandon_acquire_task()
verbose_proxy_logger.debug(
"Skipping DB reconnect attempt due to lock acquisition timeout. reason=%s timeout=%ss",
reason,

4
litellm/utils.py
View file

@ -9122,6 +9122,10 @@ class ProviderConfigManager:
from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig
return RunwayMLVideoConfig()
elif LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config
return get_hosted_vllm_video_config(model)
return None
@staticmethod

3
provider_endpoints_support.json
View file

@ -1277,7 +1277,8 @@
"files": true,
"rerank": true,
"a2a": true,
"interactions": true
"interactions": true,
"video_generations": true
}
},
"huggingface": {

4
pyproject.toml
View file

@ -67,8 +67,8 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.90",
"litellm-enterprise==0.1.61",
"litellm-proxy-extras==0.4.91",
"litellm-enterprise==0.1.62",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",

1
tests/e2e/coverage_registry/guardrail.yaml
View file

@ -32,3 +32,4 @@
- {id: guardrail.mcp_security.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/mcp_security", rationale: "MCP protocol security"}
- {id: guardrail.llm_as_a_judge.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/llm_as_a_judge", rationale: "LLM-based judgment guardrail"}
- {id: guardrail.litellm_content_filter.pre_mcp_call.blocks, module: guardrail, tier: P1, hook_point: pre_mcp_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/litellm_content_filter/content_filter.py:_scan_mcp_tool_call_arguments", rationale: "A general content-filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments before it reaches the upstream MCP server; a clean argument passes"}
- {id: guardrail.dispatch.pre_call.rejects_unknown_name, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "proxy guardrail dispatch (per-request `guardrails` selector)", rationale: "A request naming a guardrail this proxy does not serve must fail closed with a 4xx; today it is silently served unguarded, so a typo'd name drops the protection the caller asked for"}

1
tests/e2e/coverage_registry/logging.yaml
View file

@ -24,3 +24,4 @@
- {id: logging.focus.success.writes_object, module: logging, tier: P1, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/focus/focus_logger.py", rationale: "Cost mgmt multi-destination export"}
- {id: logging.niche_integrations.success.logs_spend, module: logging, tier: P2, event: success, assertions: [logs_spend], exercised_on: [chat_completions], source: grammar, rationale: "SMOKE cohort: athina/galileo/deepeval/langtrace/weave/lunary/humanloop/traceloop/helicone/argilla/newrelic/sqs/supabase/dynamodb/agentops/lago/etc"}
- {id: logging.niche_integrations.failure.logs_spend, module: logging, tier: P2, event: failure, assertions: [logs_spend], exercised_on: [chat_completions], source: grammar, rationale: "SMOKE niche failure path"}
- {id: logging.langfuse.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions], source: "integrations/langfuse/langfuse_otel.py", rationale: "Team-scoped Langfuse delivery via /team/callback; LangChain-ecosystem evals spend"}

64
tests/e2e/guardrails/guardrails_client.py
View file

@ -68,11 +68,27 @@ class BlockCodeExecutionParamsBody(GuardrailParamsBase):
guardrail: Literal["block_code_execution"] = "block_code_execution"
class PresidioParamsBody(GuardrailParamsBase):
"""Presidio PII guardrail params. `presidio_filter_scope="input"` keeps the
registration to a single callback on the configured mode; the default
("both") also registers a second post_call output-masking callback, which a
pre_call- or logging_only-scoped test must not drag in. `output_parse_pii`
stays unset/False: True would unmask the response back to the caller."""
guardrail: Literal["presidio"] = "presidio"
presidio_analyzer_api_base: str
presidio_anonymizer_api_base: str
presidio_filter_scope: Literal["input", "output", "both"] | None = None
presidio_language: str | None = None
output_parse_pii: bool | None = None
GuardrailParamsBody = (
ContentFilterParamsBody
| BedrockGuardrailParamsBody
| OpenAIModerationParamsBody
| BlockCodeExecutionParamsBody
| PresidioParamsBody
)
@ -253,6 +269,30 @@ class GuardrailsClient:
),
)
def chat_stream_raw(
self,
key: str,
model: str,
text: str,
*,
guardrails: list[str] | None = None,
max_tokens: int = 64,
) -> StreamingResponse:
"""Drive /chat/completions with stream=true, returning the raw HTTP
outcome (status, headers, SSE events) via the shared ProxyClient stream
sender - a streamed guardrail block is judged on status and stream
shape, not a typed body."""
return self.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=text)],
max_tokens=max_tokens,
stream=True,
guardrails=guardrails,
),
)
def messages(
self,
key: str,
@ -318,7 +358,7 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient:
return GuardrailsClient(proxy=proxy)
def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatResponse]:
def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]:
"""Retry a call that a guardrail should reject until it is, returning the last result.
Registering a guardrail is a control-plane write; the data-plane worker that
@ -337,3 +377,25 @@ def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatR
time.sleep(POLL_INTERVAL)
last = call()
return last
#: Statuses a stream poll keeps retrying through instead of returning as "the
#: block": network failures (-1), key propagation (401), rate limits (429) -
#: transient rig noise, not a guardrail verdict.
_TRANSIENT_STREAM_STATUSES = frozenset({-1, 401, 429})
def poll_until_blocked_stream(call: Callable[[], StreamingResponse]) -> StreamingResponse:
"""poll_until_blocked for raw/streamed sends, which return a StreamingResponse
instead of a Result: retry while the call still succeeds (the data-plane worker
has not picked the new guardrail up yet) or fails with a transient status,
returning the first guardrail-shaped non-2xx outcome or the last result at
the deadline."""
deadline = time.monotonic() + POLL_TIMEOUT
last = call()
while time.monotonic() < deadline:
if not last.ok and last.status_code not in _TRANSIENT_STREAM_STATUSES:
return last
time.sleep(POLL_INTERVAL)
last = call()
return last

92
tests/e2e/guardrails/test_bedrock_guardrail_e2e.py
View file

@ -1,9 +1,12 @@
"""Live e2e: Bedrock ApplyGuardrail pre_call blocks denied input on chat.
"""Live e2e: Bedrock ApplyGuardrail blocks on chat, pre_call and post_call.
Registers a default-on bedrock guardrail via POST /guardrails with identifier/
pre_call registers a bedrock guardrail via POST /guardrails with identifier/
version from env, then sends a prompt the guardrail's configured policy denies.
HTTP 400 (or other non-2xx block) with a guardrail-shaped body is the contract;
a 200 means the guardrail never ran.
a 200 means the guardrail never ran. post_call scans the MODEL OUTPUT only, so
its test makes the model echo the word the guardrail's word policy denies
(BEDROCK_GUARDRAIL_BLOCKED_WORD, default FORBIDDENWORD) and the block must
arrive without leaking the model's text.
No AWS keys are passed: the gateway signs ApplyGuardrail with its own
pod-identity role, since the static AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY
@ -12,17 +15,37 @@ env vars are deliberately absent from the gateway (they hijack RDS IAM auth).
from __future__ import annotations
import json
import os
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import UnknownApiError
from guardrails_client import GuardrailsClient, poll_until_blocked
from guardrails_client import (
BedrockGuardrailParamsBody,
GuardrailsClient,
poll_until_blocked,
)
from lifecycle import ResourceManager
from pydantic import JsonValue, TypeAdapter
pytestmark = pytest.mark.e2e
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
def _without_assessments(value: JsonValue) -> JsonValue:
"""The assessments echo guardrail CONFIG, not content: the stage guardrail's
topic policy is itself named after the denied word, so its label lands in
every assessment listing and would trip a leak check aimed at model output."""
if isinstance(value, dict):
return {key: _without_assessments(child) for key, child in value.items() if key != "assessments"}
if isinstance(value, list):
return [_without_assessments(item) for item in value]
return value
MODEL = "gemini-2.5-flash"
# Matches the word/topic policy the guardrail this suite points at actually denies.
# Content filters are not assumed: the guardrail resource carries no contentPolicy,
@ -42,23 +65,17 @@ class TestBedrockGuardrail:
version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
name = f"e2e-bedrock-guard-{unique_marker()}"
guardrail_id = client.create_bedrock_guardrail(
name, identifier=identifier, version=version
)
guardrail_id = client.create_bedrock_guardrail(name, identifier=identifier, version=version)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
# Selected per request rather than registered default_on, so an upstream
# ApplyGuardrail failure surfaces here instead of 403ing every other suite
# running against this proxy.
result = poll_until_blocked(
lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name])
)
result = poll_until_blocked(lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name]))
match result:
case UnknownApiError(status_code=status, body=body):
assert status in {400, 403}, (
f"expected a guardrail block status, got {status}: {body[:400]}"
)
assert status in {400, 403}, f"expected a guardrail block status, got {status}: {body[:400]}"
body_lower = body.lower()
assert any(
token in body_lower
@ -72,6 +89,49 @@ class TestBedrockGuardrail:
)
), f"block body should name the guardrail reason; got: {body[:400]}"
case _:
pytest.fail(
f"bedrock default-on guardrail did not block harmful prompt; got {result}"
pytest.fail(f"bedrock default-on guardrail did not block harmful prompt; got {result}")
@pytest.mark.covers(
"guardrail.bedrock.post_call.blocks",
exercised_on=["chat_completions"],
)
def test_bedrock_post_call_blocks_denied_model_output(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
blocked_word = os.environ.get("BEDROCK_GUARDRAIL_BLOCKED_WORD", "FORBIDDENWORD")
name = f"e2e-bedrock-post-{unique_marker()}"
guardrail_id = client.register(
name,
BedrockGuardrailParamsBody(
mode="post_call",
default_on=False,
guardrailIdentifier=identifier,
guardrailVersion=version,
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
# post_call scans OUTPUT only, so the denied word has to come out of the
# model: ask it to echo the word verbatim. The word in the prompt itself
# is not scanned in this mode.
prompt = f"Reply with exactly this one word and nothing else: {blocked_word}"
result = poll_until_blocked(lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128))
match result:
case UnknownApiError(status_code=status, body=body):
# A policy block is a 400 naming the verdict; a failed
# ApplyGuardrail call surfaces as 403 "guardrail request
# failed", which must not count as a block.
assert status == 400, f"expected the guardrail block status 400, got {status}: {body[:400]}"
body_lower = body.lower()
assert any(token in body_lower for token in ("violated", "blocked", "intervened")), (
f"block body should name the guardrail verdict; got: {body[:400]}"
)
assert blocked_word not in json.dumps(_without_assessments(_JSON.validate_json(body))), (
f"the blocked model output must not leak into the error body; got: {body[:400]}"
)
case _:
pytest.fail(f"bedrock post_call guardrail did not block denied model output; got {result}")

41
tests/e2e/guardrails/test_guardrail_dispatch_e2e.py Normal file
View file

@ -0,0 +1,41 @@
"""Live e2e: the per-request `guardrails` selector must fail closed.
A request that names a guardrail is a caller asking for protection. When the
proxy does not serve that name (a typo, a deleted guardrail, or a worker that
never loaded it), answering 200 silently drops the protection the caller asked
for; the contract this test pins is a 4xx naming the unknown guardrail.
"""
from __future__ import annotations
import pytest
from e2e_config import unique_marker
from e2e_http import UnknownApiError, ValidationError
from guardrails_client import GuardrailsClient
pytestmark = pytest.mark.e2e
MODEL = "gemini-2.5-flash"
@pytest.mark.skip(
reason=(
"stage red: product gap, a request naming a guardrail the proxy does not "
"serve is silently served unguarded (200) instead of failing closed"
)
)
@pytest.mark.covers(
"guardrail.dispatch.pre_call.rejects_unknown_name",
exercised_on=["chat_completions"],
)
def test_request_naming_an_unknown_guardrail_fails_closed(client: GuardrailsClient, scoped_key: str) -> None:
result = client.chat(scoped_key, MODEL, "say hi", guardrails=[f"e2e-no-such-guardrail-{unique_marker()}"])
match result:
case UnknownApiError(status_code=status, body=body):
assert status == 400, f"expected a 400 for an unknown guardrail name, got {status}: {body[:400]}"
assert "guardrail" in body.lower(), f"the rejection should name the guardrail; got: {body[:400]}"
case ValidationError(message=message):
assert "guardrail" in message.lower(), f"the rejection should name the guardrail; got: {message[:400]}"
case _:
pytest.fail(f"a request naming an unknown guardrail must fail closed with a 4xx; got {result}")

47
tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py
View file

@ -7,7 +7,9 @@ before the upstream model runs; a prompt that trips the policy must be rejected
with HTTP 400 naming the moderation policy, and the same guardrail must let a
benign prompt through. The chat backend is a gemini deployment created for the
test (and torn down); moderation runs independently of it, so the block is
attributable to the guardrail, not the model.
attributable to the guardrail, not the model. The same pre_call contract is
also exercised through /v1/messages (Anthropic format): a flagged prompt is
rejected with a 400 naming moderation and a benign one passes.
"""
from __future__ import annotations
@ -69,3 +71,46 @@ class TestOpenAIModerationGuardrail:
"the same moderation guardrail must let a benign prompt through, but the "
f"call returned no choices: {allowed}"
)
@pytest.mark.covers(
"guardrail.openai_moderations.pre_call.blocks",
exercised_on=["messages"],
)
def test_moderation_blocks_flagged_input_on_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
model = client.create_backend_model(resources, prefix="e2e-moderation-msg-backend")
name = f"e2e-openai-moderation-msg-{unique_marker()}"
guardrail_id = client.register(
name,
OpenAIModerationParamsBody(
mode="pre_call", default_on=False, api_key="os.environ/OPENAI_API_KEY"
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
blocked = poll_until_blocked(
lambda: client.messages(scoped_key, model, FLAGGED_PROMPT, guardrails=[name])
)
match blocked:
case UnknownApiError(status_code=400, body=body):
assert "moderation" in body.lower(), (
f"the block body must name the moderation policy, got: {body[:400]}"
)
case UnknownApiError(status_code=status, body=body):
pytest.fail(
f"expected a 400 moderation block on /v1/messages, got {status}: {body[:400]}"
)
case _:
pytest.fail(
f"openai moderation did not block a flagged /v1/messages prompt; got {blocked}"
)
allowed = unwrap(
client.messages(scoped_key, model, BENIGN_PROMPT, guardrails=[name], max_tokens=64)
)
assert allowed.content or allowed.choices, (
"the same moderation guardrail must let a benign /v1/messages prompt through, but "
f"the response carried neither content nor choices: {allowed}"
)

184
tests/e2e/guardrails/test_presidio_masking_e2e.py Normal file
View file

@ -0,0 +1,184 @@
"""Live e2e: the Presidio PII guardrail masks, per its configured hook point.
pre_call: the guardrail calls the Presidio analyzer/anonymizer on the request
messages BEFORE the model runs, so the model only ever sees placeholders like
<EMAIL_ADDRESS>. A prompt asking the model to repeat a fake email + phone back
must come back with the placeholders echoed and the raw PII absent, on
/chat/completions and on /v1/messages (Anthropic format).
The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE /
PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
Each guardrail registers with presidio_filter_scope="input" so only the
configured hook's callback exists (the default "both" adds a second post_call
output masker), and is deleted on teardown.
"""
from __future__ import annotations
import os
import time
from collections.abc import Callable
import pytest
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_http import Result, Success
from guardrails_client import GuardrailsClient, PresidioParamsBody
from lifecycle import ResourceManager
from models import AnthropicMessagesResponse, ChatResponse
pytestmark = pytest.mark.e2e
MODEL = "gemini-2.5-flash"
# A guardrail created via POST /guardrails reaches the worker that served the
# create immediately, but every other worker only picks it up on its next
# periodic DB sync (~30s), so the first requests can be served unguarded.
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
# Presidio's anonymizer replaces a detected entity with its unnumbered type
# placeholder, e.g. <EMAIL_ADDRESS>. The pre_call assertions match on the bare
# token because the model is echoing the masked prompt and may not preserve the
# angle brackets; the logged payload keeps the placeholder verbatim.
MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS"
MASKED_PHONE_TOKEN = "PHONE_NUMBER"
# Fictional NANP 555 number; a standard format Presidio's phone recognizer detects.
FAKE_PHONE = "+1 415-555-0134"
def _presidio_bases() -> tuple[str, str]:
analyzer = os.environ.get("PRESIDIO_ANALYZER_API_BASE", "").strip()
anonymizer = os.environ.get("PRESIDIO_ANONYMIZER_API_BASE", "").strip()
if not analyzer or not anonymizer:
pytest.fail(
"Presidio e2e requires PRESIDIO_ANALYZER_API_BASE and PRESIDIO_ANONYMIZER_API_BASE "
"(the running Presidio analyzer/anonymizer services); missing env is a hard failure, not a skip"
)
return analyzer, anonymizer
def _register_presidio(
client: GuardrailsClient,
resources: ResourceManager,
*,
name: str,
) -> None:
analyzer, anonymizer = _presidio_bases()
guardrail_id = client.register(
name,
PresidioParamsBody(
mode="pre_call",
default_on=False,
presidio_analyzer_api_base=analyzer,
presidio_anonymizer_api_base=anonymizer,
presidio_filter_scope="input",
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
def _fake_email() -> str:
return f"jane.doe.{unique_marker()}@example.com"
def _pii_prompt(marker: str, email: str) -> str:
return (
f"{marker} Repeat this sentence back to me exactly, word for word: "
f"My email address is {email} and my phone number is {FAKE_PHONE}."
)
def _first_content(response: ChatResponse) -> str:
if not response.choices:
return ""
message = response.choices[0].message
return (message.content if message else None) or ""
def _messages_text(response: AnthropicMessagesResponse) -> str:
"""The text of a /v1/messages answer, whichever shape the proxy produced
(Anthropic-native content blocks or OpenAI-normalized choices)."""
parts: list[str] = []
for block in response.content or []:
if block.text:
parts.append(block.text)
for choice in response.choices or []:
if choice.message and choice.message.content:
parts.append(choice.message.content)
return "\n".join(parts)
def _assert_eventually_masked[R: BaseModel](
fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str
) -> None:
"""Retry the call until the response comes back masked, to the propagation
deadline. An unmasked early response is in-flight guardrail propagation, not
a failure, and neither is a transient non-Success (a replica that has not
reloaded the guardrail answers 404, the live model can rate-limit) - only a
response that still carries the raw PII at the deadline is."""
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
result = fetch()
match result:
case Success(data=data):
content = extract(data)
last = content
masked = MASKED_EMAIL_TOKEN in content and MASKED_PHONE_TOKEN in content and email not in content
if masked:
assert FAKE_PHONE not in content, (
f"the raw phone number must be masked before the model sees it, but the "
f"response echoed it: {content[:300]!r}"
)
return
case _:
last = f"<non-Success result: {result}>"
if time.monotonic() >= deadline:
pytest.fail(
f"presidio pre_call guardrail never masked the PII within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
class TestPresidioPreCallMasking:
@pytest.mark.covers(
"guardrail.presidio.pre_call.masks",
exercised_on=["chat_completions"],
)
def test_pre_call_masks_pii_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-pre-chat-{unique_marker()}"
_register_presidio(client, resources, name=name)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
_assert_eventually_masked(
lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
_first_content,
email=email,
)
@pytest.mark.covers(
"guardrail.presidio.pre_call.masks",
exercised_on=["messages"],
)
def test_pre_call_masks_pii_on_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-pre-msg-{unique_marker()}"
_register_presidio(client, resources, name=name)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
_assert_eventually_masked(
lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
_messages_text,
email=email,
)

87
tests/e2e/guardrails/test_streaming_guardrail_e2e.py Normal file
View file

@ -0,0 +1,87 @@
"""Live e2e: a Bedrock guardrail in during_call mode blocks a streamed chat.
during_call runs the Bedrock ApplyGuardrail INPUT scan in an asyncio.gather
alongside the LLM call (common_request_processing.py); when the scan flags the
prompt, the raised block cancels the LLM task before the stream ever starts, so
the client sees a non-2xx JSON error - not an SSE stream, not an in-stream
error frame - and zero content chunks are delivered.
The prompt deliberately contains the exact word the guardrail's word policy
denies (BEDROCK_GUARDRAIL_BLOCKED_WORD, default FORBIDDENWORD), so the INPUT
scan intervenes deterministically. Identifier/version come from
BEDROCK_GUARDRAIL_IDENTIFIER / BEDROCK_GUARDRAIL_VERSION like the rest of the
bedrock suite; no AWS keys are passed (the gateway signs with pod identity).
The guardrail registers default_on=False and is selected per request, so an
upstream ApplyGuardrail failure surfaces here instead of 403ing other suites.
"""
from __future__ import annotations
import os
import pytest
from e2e_config import unique_marker
from guardrails_client import (
BedrockGuardrailParamsBody,
GuardrailsClient,
poll_until_blocked_stream,
)
from lifecycle import ResourceManager
pytestmark = pytest.mark.e2e
MODEL = "gemini-2.5-flash"
class TestBedrockDuringCallStreaming:
@pytest.mark.covers(
"guardrail.bedrock.during.blocks",
exercised_on=["chat_completions"],
)
def test_during_call_blocks_stream_before_first_chunk(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"]
version = os.environ["BEDROCK_GUARDRAIL_VERSION"]
blocked_word = os.environ.get("BEDROCK_GUARDRAIL_BLOCKED_WORD", "FORBIDDENWORD")
name = f"e2e-bedrock-during-{unique_marker()}"
guardrail_id = client.register(
name,
BedrockGuardrailParamsBody(
mode="during_call",
default_on=False,
guardrailIdentifier=identifier,
guardrailVersion=version,
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
# The denied word sits in the INPUT: during_call scans the request
# messages while the model call runs, and the flag must win the race
# by cancelling the stream outright.
prompt = f"Please use the word {blocked_word} in a sentence."
result = poll_until_blocked_stream(
lambda: client.chat_stream_raw(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=64)
)
assert not result.ok, (
f"the during_call guardrail never blocked the streamed request; got a "
f"{result.status_code} with {result.chunks} chunks"
)
assert result.status_code == 400, (
f"a during_call block surfaces as HTTP 400 before the stream starts, got "
f"{result.status_code}: {result.body[:400]}"
)
assert result.chunks == 0 and not result.stream_events, (
f"no content chunk may be delivered on a during_call block, but "
f"{result.chunks} chunks arrived: {result.stream_events[:3]}"
)
assert "text/event-stream" not in (result.content_type or ""), (
f"the block must be a JSON error response, not an SSE stream; got content-type {result.content_type!r}"
)
body_lower = result.body.lower()
assert any(token in body_lower for token in ("guardrail", "violated", "blocked", "bedrock", "intervened")), (
f"block body should name the guardrail reason; got: {result.body[:400]}"
)

4
tests/e2e/logging/conftest.py
View file

@ -47,6 +47,4 @@ def dd_logs() -> DdLogsReader:
def datadog_creds() -> None:
"""Require Datadog shipping credentials. Hard-fail when absent; never skip."""
if not (os.getenv("DD_API_KEY") and os.getenv("DD_SITE")):
pytest.fail(
"Datadog e2e requires DD_API_KEY and DD_SITE; missing credentials is a hard failure, not a skip"
)
pytest.fail("Datadog e2e requires DD_API_KEY and DD_SITE; missing credentials is a hard failure, not a skip")

33
tests/e2e/logging/datadog_reader.py
View file

@ -97,15 +97,22 @@ class DdLogsReader:
indexed ``message`` empty, so a plain full-text query matches nothing;
``*:`` extends the scan to every attribute (the marker sits in the
prompt, e.g. ``messages.content``, wherever the route's payload puts
it). More than one hit for one call IS the duplicate-delivery bug, so
this never collapses to a single event. A 429 backs off and retries -
the search budget is org-wide, so another consumer can empty it under
us - while any other failure stays a hard fail."""
it)."""
return self.events_for_query(f"*:*{marker}*")
def events_for_query(self, query: str) -> list[DdLogEvent]:
"""Every ingested event the search query matches (failure payloads
carry no prompt to mark, so failure scenarios query indexed attributes
like ``@model_group:...`` instead of a body marker). More than one hit
for one call IS the duplicate-delivery bug, so this never collapses to
a single event. A 429 backs off and retries - the search budget is
org-wide, so another consumer can empty it under us - while any other
failure stays a hard fail."""
for _ in range(_RATE_LIMIT_RETRIES):
result = post(
URL(f"https://api.{self.site}/api/v2/logs/events/search"),
headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key),
json=_SearchRequest(filter=_SearchFilter(query=f"*:*{marker}*")),
json=_SearchRequest(filter=_SearchFilter(query=query)),
response_type=_SearchResponse,
timeout=30.0,
)
@ -123,6 +130,10 @@ class DdLogsReader:
)
def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]:
"""``poll_events_for_query`` over the every-attribute marker scan."""
return self.poll_events_for_query(f"*:*{marker}*")
def poll_events_for_query(self, query: str) -> list[DdLogEvent]:
"""Poll until at least one matching event is searchable (the callback
flushes in periodic batches and DataDog ingestion adds seconds of lag),
then keep re-reading for DD_SETTLE_SECONDS so a late duplicate cannot
@ -132,15 +143,13 @@ class DdLogsReader:
request budget. At the deadline the last result is returned as-is."""
deadline = time.monotonic() + POLL_TIMEOUT
while time.monotonic() < deadline:
events = self.events_for_marker(marker)
events = self.events_for_query(query)
if events:
return self._settled_events_for_marker(marker, events)
return self._settled_events_for_query(query, events)
time.sleep(DD_SEARCH_INTERVAL)
return self.events_for_marker(marker)
return self.events_for_query(query)
def _settled_events_for_marker(
self, marker: str, events: list[DdLogEvent]
) -> list[DdLogEvent]:
def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]:
"""Re-read at every search interval until the settle window closes; a
duplicate ends the watch early because more waiting cannot clear it.
@ -151,7 +160,7 @@ class DdLogsReader:
last_nonempty = events
while time.monotonic() < settle_deadline:
time.sleep(DD_SEARCH_INTERVAL)
latest = self.events_for_marker(marker)
latest = self.events_for_query(query)
if not latest:
continue
if len(latest) > 1:

220
tests/e2e/logging/gcs_reader.py Normal file
View file

@ -0,0 +1,220 @@
"""Read-back for the gcs_bucket logging test against the real GCS bucket.
The proxy ships StandardLoggingPayload objects with its own service account
(litellm_settings.callbacks: ["gcs_bucket"] + GCS_BUCKET_NAME), and the test
reads them back through the GCS JSON API. Auth is a self-signed service-account
JWT (RS256 via PyJWT + cryptography, both litellm proxy dependencies the
runner installs) minted per request and sent directly as the Bearer token -
Google accepts that for storage.googleapis.com with no token exchange, which
keeps every HTTP read inside ``e2e_http``.
The default gcs_bucket mode batches payloads into ``{date}/batch-{id}.ndjson``
objects; unbatched mode writes ``{date}/{response_id}`` per call. The reader
handles both: it polls the day's listing, downloads the direct object when
present, and otherwise scans batch objects fresh enough to hold the call.
Missing configuration is a hard failure, never a skip.
"""
from __future__ import annotations
import os
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from urllib.parse import quote
import jwt
import pytest
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, Headers, probe
_GCS_API = "https://storage.googleapis.com"
#: Tolerance for clock skew between this host and GCS object timestamps.
_SKEW = timedelta(seconds=120)
#: How long to keep re-reading after the first match before trusting the
#: exactly-one assertion: past one full gcs_bucket flush interval (~20s), so
#: a duplicate shipped by a later flush is seen, plus listing-latency margin.
GCS_SETTLE_SECONDS = 45.0
class _ServiceAccount(BaseModel):
model_config = ConfigDict(extra="ignore")
client_email: str
private_key: str
class _GcsAuthHeaders(Headers):
authorization: str = Field(serialization_alias="Authorization")
class _GcsObject(BaseModel):
model_config = ConfigDict(extra="ignore")
name: str
updated: datetime | None = None
class _GcsListResponse(BaseModel):
model_config = ConfigDict(extra="ignore")
items: list[_GcsObject] = []
next_page_token: str | None = Field(default=None, validation_alias="nextPageToken")
class _GcsListParams(BaseModel):
prefix: str
max_results: int = Field(default=1000, serialization_alias="maxResults")
page_token: str | None = Field(default=None, serialization_alias="pageToken")
class _GcsMediaParams(BaseModel):
alt: str = "media"
class GcsLogRecord(BaseModel):
"""The StandardLoggingPayload fields the gcs scenario pins."""
model_config = ConfigDict(extra="ignore")
id: str
status: str
model_group: str | None = None
response_cost: float | None = None
total_tokens: int | None = None
error_str: str | None = None
def _mint_bearer(account: _ServiceAccount) -> str:
"""Self-signed service-account JWT: for Google APIs a token whose ``aud``
is the service endpoint authorizes directly, no oauth2 token exchange.
Minted per request so a long session never outlives one token's expiry."""
now = int(time.time())
claims: dict[str, str | int] = {
"iss": account.client_email,
"sub": account.client_email,
"aud": f"{_GCS_API}/",
"iat": now,
"exp": now + 3600,
}
return jwt.encode(claims, account.private_key, algorithm="RS256")
@dataclass(frozen=True, slots=True)
class GcsLogReader:
bucket: str
account: _ServiceAccount
def _headers(self) -> _GcsAuthHeaders:
return _GcsAuthHeaders(authorization=f"Bearer {_mint_bearer(self.account)}")
def _list(self, prefix: str) -> list[_GcsObject]:
"""Every object under ``prefix``, following ``nextPageToken`` - the
shared day prefix accumulates all of the proxy's traffic, and a fresh
record past the 1000-object page cap must still be seen."""
items: list[_GcsObject] = []
page_token: str | None = None
while True:
result = probe(
URL(f"{_GCS_API}/storage/v1/b/{self.bucket}/o"),
headers=self._headers(),
params=_GcsListParams(prefix=prefix, page_token=page_token),
)
if result.status_code != 200:
pytest.fail(
f"GCS object listing for gs://{self.bucket}/{prefix} failed "
f"({result.status_code}): {result.body[:300]}"
)
page = _GcsListResponse.model_validate_json(result.body)
items.extend(page.items)
page_token = page.next_page_token
if not page_token:
return items
def _download(self, name: str) -> str:
result = probe(
URL(f"{_GCS_API}/storage/v1/b/{self.bucket}/o/{quote(name, safe='')}"),
headers=self._headers(),
params=_GcsMediaParams(),
)
if result.status_code != 200:
pytest.fail(
f"GCS object download gs://{self.bucket}/{name} failed ({result.status_code}): {result.body[:300]}"
)
return result.body
def records_for_response_id(self, response_id: str, *, since: datetime) -> list[GcsLogRecord]:
"""Every payload written for ``response_id``: the direct
``{date}/{response_id}`` object plus any hit inside batch NDJSON
objects updated after ``since``. More than one hit is the
duplicate-delivery bug, so this never collapses to a single record."""
records: list[GcsLogRecord] = []
window_start = since - _SKEW
for day_offset in (-1, 0, 1):
day = (since + timedelta(days=day_offset)).strftime("%Y-%m-%d")
for obj in self._list(f"{day}/"):
if obj.name == f"{day}/{response_id}":
records.append(GcsLogRecord.model_validate_json(self._download(obj.name)))
continue
is_fresh_batch = f"{day}/batch-" in obj.name and obj.updated is not None and obj.updated >= window_start
if is_fresh_batch:
records.extend(
GcsLogRecord.model_validate_json(line)
for line in self._download(obj.name).splitlines()
if response_id in line
)
return records
def poll_records_for_response_id(self, response_id: str, *, since: datetime) -> list[GcsLogRecord]:
"""Poll until the payload is readable (the gcs_bucket callback flushes
on a ~20s timer), then keep re-reading for GCS_SETTLE_SECONDS - past a
full flush interval - so a duplicate shipped by a later flush cannot
hide from the exactly-one assertion. A duplicate ends the settle early
because more waiting cannot clear it."""
deadline = time.monotonic() + POLL_TIMEOUT
while time.monotonic() < deadline:
records = self.records_for_response_id(response_id, since=since)
if records:
return self._settled_records(response_id, since=since, first=records)
time.sleep(POLL_INTERVAL)
return []
def _settled_records(self, response_id: str, *, since: datetime, first: list[GcsLogRecord]) -> list[GcsLogRecord]:
"""Re-read at every poll interval until the settle window closes; a
transiently empty re-read never downgrades what was already seen."""
settle_deadline = time.monotonic() + GCS_SETTLE_SECONDS
latest = first
while time.monotonic() < settle_deadline and len(latest) <= 1:
time.sleep(POLL_INTERVAL)
latest = self.records_for_response_id(response_id, since=since) or latest
return latest
def utc_now() -> datetime:
return datetime.now(timezone.utc)
def build_gcs_reader() -> GcsLogReader:
bucket = os.environ.get("GCS_BUCKET_NAME", "")
if not bucket:
pytest.fail(
"GCS_BUCKET_NAME must be set: the gcs test reads the proxy's gcs_bucket "
"delivery back from the real bucket (the cluster secret manager injects "
"it; locally set it in tests/e2e/.env)"
)
raw = ""
credentials_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "")
if credentials_path and Path(credentials_path).is_file():
raw = Path(credentials_path).read_text()
else:
raw = os.environ.get("VERTEXAI_CREDENTIALS", "")
if not raw:
pytest.fail(
"GCS read-back needs a service-account key: set "
"GOOGLE_APPLICATION_CREDENTIALS (path) or VERTEXAI_CREDENTIALS (JSON), "
"as the cluster secret manager does"
)
return GcsLogReader(bucket=bucket, account=_ServiceAccount.model_validate_json(raw))

40
tests/e2e/logging/logging_client.py
View file

@ -480,12 +480,8 @@ class LoggingClient:
stream=True if stream else None,
)
if stream:
return self.proxy.transport.stream(
"/v1/messages", headers=self.proxy.transport.bearer(key), json=body
)
return self.proxy.transport.send(
"/v1/messages", headers=self.proxy.transport.bearer(key), json=body
)
return self.proxy.transport.stream("/v1/messages", headers=self.proxy.transport.bearer(key), json=body)
return self.proxy.transport.send("/v1/messages", headers=self.proxy.transport.bearer(key), json=body)
def responses_raw(
self, key: str, model: str, text: str, *, max_output_tokens: int = 64, stream: bool = False
@ -499,12 +495,8 @@ class LoggingClient:
model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None
)
if stream:
return self.proxy.transport.stream(
"/v1/responses", headers=self.proxy.transport.bearer(key), json=body
)
return self.proxy.transport.send(
"/v1/responses", headers=self.proxy.transport.bearer(key), json=body
)
return self.proxy.transport.stream("/v1/responses", headers=self.proxy.transport.bearer(key), json=body)
return self.proxy.transport.send("/v1/responses", headers=self.proxy.transport.bearer(key), json=body)
def scrape_metrics(self) -> str:
return self.proxy.probe("/metrics", params=NoBody()).body
@ -530,9 +522,7 @@ class LoggingClient:
return False
return True
rows = self.proxy.poll_logs_for_key(
key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs)
)
rows = self.proxy.poll_logs_for_key(key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs))
for row in rows:
if _matches(row):
return row
@ -593,9 +583,7 @@ class LoggingClient:
deadline = time.monotonic() + POLL_TIMEOUT
last: LangfuseObservation | None = None
while time.monotonic() < deadline:
last = self.find_langfuse_observation(
creds, key_alias=key_alias, prompt_marker=prompt_marker
)
last = self.find_langfuse_observation(creds, key_alias=key_alias, prompt_marker=prompt_marker)
if last is not None:
cost = observation_spend(last)
if not require_positive_cost or (cost is not None and cost > 0):
@ -611,9 +599,7 @@ class LoggingClient:
prompt_marker: str,
) -> list[LangfuseObservation]:
"""Generation plus any sibling/child observations (guardrail spans, etc.)."""
gen = self.poll_langfuse_observation(
creds, key_alias=key_alias, prompt_marker=prompt_marker
)
gen = self.poll_langfuse_observation(creds, key_alias=key_alias, prompt_marker=prompt_marker)
if gen is None or not gen.trace_id:
return [] if gen is None else [gen]
return self.list_langfuse_observations(creds, trace_id=gen.trace_id) or [gen]
@ -636,3 +622,15 @@ def first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> St
def build_logging_client(proxy: ProxyClient) -> LoggingClient:
return LoggingClient(proxy=proxy)
def readiness_details_body(client: LoggingClient) -> str:
"""/health/readiness/details, tolerating the 503 it serves while the
ephemeral stack's DB leg blips: the recorded state the logging suites check
here is the callback list, which the body carries either way."""
result = client.proxy.probe("/health/readiness/details", params=NoBody())
db_blip = result.status_code == 503 and '"db":"disconnected"' in result.body
assert result.status_code == 200 or db_blip, (
f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
)
return result.body

115
tests/e2e/logging/s3_reader.py Normal file
View file

@ -0,0 +1,115 @@
"""Read-back for the s3 logging tests against the real S3 bucket the proxy
ships StandardLoggingPayload objects to (litellm_settings.callbacks: ["s3_v2"]).
Delivery is judged on what actually landed in the bucket: the proxy writes
with its own credentials exactly as in production, and the tests list and
download the objects back with boto3 (already a litellm proxy dependency, so
the e2e runner image carries it; it is an AWS SDK, not a raw HTTP client, so
the e2e_http-only transport rule is untouched). The bucket comes from
S3_LOGS_BUCKET_NAME - on the cluster the secret manager injects it, locally
tests/e2e/.env provides it. Missing configuration is a hard failure, never a
skip.
"""
from __future__ import annotations
import os
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING
import boto3
import pytest
from pydantic import BaseModel, ConfigDict
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
if TYPE_CHECKING:
from types_boto3_s3.client import S3Client
#: How long to keep re-reading after the first match before trusting the
#: exactly-one assertion: past one full s3_v2 flush interval (~10s), so a
#: duplicate shipped by a LATER flush is seen, plus listing-latency margin.
#: The DataDog reader settles the same way (DD_SETTLE_SECONDS).
S3_SETTLE_SECONDS = 25.0
class S3LogRecord(BaseModel):
"""The StandardLoggingPayload fields the s3 scenarios pin."""
model_config = ConfigDict(extra="ignore")
id: str
status: str
model_group: str | None = None
response_cost: float | None = None
total_tokens: int | None = None
error_str: str | None = None
@dataclass(frozen=True, slots=True)
class S3LogReader:
bucket: str
client: S3Client
def list_keys(self, prefix: str) -> list[str]:
response = self.client.list_objects_v2(Bucket=self.bucket, Prefix=prefix)
return [obj["Key"] for obj in response.get("Contents", []) if "Key" in obj]
def read_record(self, key: str) -> S3LogRecord:
body = self.client.get_object(Bucket=self.bucket, Key=key)["Body"].read()
return S3LogRecord.model_validate_json(body)
def records_matching(self, *, prefix: str, predicate: Callable[[S3LogRecord], bool]) -> list[S3LogRecord]:
return [record for record in map(self.read_record, self.list_keys(prefix)) if predicate(record)]
def poll_records(self, *, prefix: str, predicate: Callable[[S3LogRecord], bool]) -> list[S3LogRecord]:
"""Poll until at least one matching object is listed (the s3_v2
callback flushes on a ~10s timer), then keep re-reading for
S3_SETTLE_SECONDS - past a full flush interval - so a duplicate
shipped by a later flush cannot hide from the exactly-one assertion.
One blind spot is inherent: a duplicate write that reuses the exact
same object key overwrites the first object and no listing can see
it; distinct-key duplicates are what this catches. At the deadline an
empty list is returned and the caller's assertion carries the failure
message."""
deadline = time.monotonic() + POLL_TIMEOUT
while time.monotonic() < deadline:
records = self.records_matching(prefix=prefix, predicate=predicate)
if records:
return self._settled_records(prefix=prefix, predicate=predicate, first=records)
time.sleep(POLL_INTERVAL)
return []
def _settled_records(
self, *, prefix: str, predicate: Callable[[S3LogRecord], bool], first: list[S3LogRecord]
) -> list[S3LogRecord]:
"""Re-read at every poll interval until the settle window closes; a
duplicate ends the watch early because more waiting cannot clear it.
A transiently empty re-read never downgrades what was already seen."""
settle_deadline = time.monotonic() + S3_SETTLE_SECONDS
latest = first
while time.monotonic() < settle_deadline and len(latest) <= 1:
time.sleep(POLL_INTERVAL)
latest = self.records_matching(prefix=prefix, predicate=predicate) or latest
return latest
def build_s3_reader() -> S3LogReader:
bucket = os.environ.get("S3_LOGS_BUCKET_NAME", "")
if not bucket:
pytest.fail(
"S3_LOGS_BUCKET_NAME must be set: the s3 tests read the proxy's s3_v2 "
"delivery back from the real bucket (the cluster secret manager injects "
"it; locally set it in tests/e2e/.env to the same bucket "
"s3_callback_params.s3_bucket_name names)"
)
region = os.environ.get("AWS_REGION_NAME") or os.environ.get("AWS_REGION") or "us-east-1"
return S3LogReader(
bucket=bucket,
# boto3.client's overload set covers every AWS service; the ones without
# installed stubs type as Unknown, so the member is "partially unknown"
# even though the s3 overload itself resolves to S3Client.
client=boto3.client("s3", region_name=region), # pyright: ignore[reportUnknownMemberType]
)

128
tests/e2e/logging/test_datadog_log_e2e.py
View file

@ -19,15 +19,16 @@ received).
from __future__ import annotations
import math
import time
import pytest
from pydantic import BaseModel, ConfigDict
from datadog_reader import DdLogEvent, DdLogsReader
from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker
from e2e_http import NoBody
from lifecycle import ResourceManager
from logging_client import LoggingClient, first_ok
from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body
from models import LiteLLMParamsBody
pytestmark = pytest.mark.e2e
@ -46,19 +47,17 @@ class _DdMessagePayload(BaseModel):
status: str
call_type: str
stream: bool | None = None
error_str: str | None = None
def _assert_datadog_configured(client: LoggingClient) -> None:
"""Recorded state: the proxy reports the DataDog callback among its active
callbacks, so a missing destination config fails here, before any
delivery-based assertion can time out confusingly."""
result = client.proxy.probe("/health/readiness/details", params=NoBody())
assert result.status_code == 200, (
f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
)
assert DD_LOGGER_NAME in result.body, (
body = readiness_details_body(client)
assert DD_LOGGER_NAME in body, (
f"the proxy must report the {DD_LOGGER_NAME} callback active "
f"(callbacks + DD_* env in the compose config); got: {result.body[:400]}"
f"(callbacks + DD_* env in the compose config); got: {body[:400]}"
)
@ -89,18 +88,14 @@ def _assert_exactly_one_event(
# indexed event status from the parsed payload's status attribute
# ("success") and normalizes it to its OK severity - so "ok" is what a
# successfully ingested success event looks like on the search API.
assert event.status == "ok", (
f"success events must index at DataDog's ok severity, got {event.status!r}"
)
assert event.status == "ok", f"success events must index at DataDog's ok severity, got {event.status!r}"
payload = _DdMessagePayload.model_validate(event.attributes)
assert payload.status == "success", f"payload status must be success, got {payload.status!r}"
assert payload.model_group == model_group, (
f"payload model_group must be {model_group!r}, got {payload.model_group!r}"
)
assert payload.call_type == call_type, (
f"payload call_type must be {call_type!r}, got {payload.call_type!r}"
)
assert payload.call_type == call_type, f"payload call_type must be {call_type!r}, got {payload.call_type!r}"
assert payload.total_tokens > 0, f"payload must count real tokens, got {payload.total_tokens}"
# Relative tolerance, not bit-equality: the cost round-trips through
# DataDog's attribute indexing, whose float serialization may drift in the
@ -109,9 +104,7 @@ def _assert_exactly_one_event(
f"payload response_cost {payload.response_cost} must equal the anchor cost {cost_anchor}"
)
if expect_stream:
assert payload.stream is True, (
f"a streamed call's payload must record stream=true, got {payload.stream!r}"
)
assert payload.stream is True, f"a streamed call's payload must record stream=true, got {payload.stream!r}"
return payload
@ -211,7 +204,9 @@ class TestDataDogLogDelivery:
marker = unique_marker()
outcome = first_ok(
client,
lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", stream=True, max_tokens=16),
lambda: client.chat_raw(
key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", stream=True, max_tokens=16
),
)
assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}"
assert outcome.chunks > 0, "the stream must deliver at least one event"
@ -231,9 +226,7 @@ class TestDataDogLogDelivery:
cost_anchor=spend_row.spend,
expect_stream=True,
)
assert spend_row.total_tokens is not None, (
"the spend row must record total_tokens for the token cross-check"
)
assert spend_row.total_tokens is not None, "the spend row must record total_tokens for the token cross-check"
assert spend_row.total_tokens == payload.total_tokens, (
f"the spend row and the DataDog event must agree on tokens: "
f"{spend_row.total_tokens} vs {payload.total_tokens}"
@ -255,7 +248,9 @@ class TestDataDogLogDelivery:
marker = unique_marker()
outcome = first_ok(
client,
lambda: client.messages_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16, stream=True),
lambda: client.messages_raw(
key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16, stream=True
),
)
assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}"
assert outcome.chunks > 0, "the stream must deliver at least one event"
@ -275,9 +270,7 @@ class TestDataDogLogDelivery:
cost_anchor=spend_row.spend,
expect_stream=True,
)
assert spend_row.total_tokens is not None, (
"the spend row must record total_tokens for the token cross-check"
)
assert spend_row.total_tokens is not None, "the spend row must record total_tokens for the token cross-check"
assert spend_row.total_tokens == payload.total_tokens, (
f"the spend row and the DataDog event must agree on tokens: "
f"{spend_row.total_tokens} vs {payload.total_tokens}"
@ -319,10 +312,89 @@ class TestDataDogLogDelivery:
cost_anchor=spend_row.spend,
expect_stream=True,
)
assert spend_row.total_tokens is not None, (
"the spend row must record total_tokens for the token cross-check"
)
assert spend_row.total_tokens is not None, "the spend row must record total_tokens for the token cross-check"
assert spend_row.total_tokens == payload.total_tokens, (
f"the spend row and the DataDog event must agree on tokens: "
f"{spend_row.total_tokens} vs {payload.total_tokens}"
)
def _assert_exactly_one_failure_event(events: list[DdLogEvent], *, model_group: str) -> _DdMessagePayload:
"""The enforced behavior for a failed call: the intake holds exactly one
event for the deployment, sourced from litellm, indexed at an error-grade
severity (DataDog derives it from the payload's status="failure"; observed
as its "emergency" bucket), whose payload carries the provider error and
no cost."""
assert events, "no DataDog log event for the failed call reached the intake within the deadline"
assert len(events) == 1, (
f"expected exactly ONE DataDog log event for the failed call, got {len(events)} - "
"more than one event for one call is the duplicate-delivery bug"
)
event = events[0]
assert "source:litellm" in event.tags, (
f"the ingested event must carry the litellm source (shipped as ddsource), got tags {event.tags!r}"
)
assert event.status in ("error", "emergency"), (
f"failure events must index at an error-grade severity, got {event.status!r}"
)
payload = _DdMessagePayload.model_validate(event.attributes)
assert payload.status == "failure", f"payload status must be failure, got {payload.status!r}"
assert payload.model_group == model_group, (
f"payload model_group must be {model_group!r}, got {payload.model_group!r}"
)
assert not payload.response_cost, f"a failed call must not be billed, got response_cost={payload.response_cost!r}"
return payload
class TestDataDogFailureDelivery:
@pytest.mark.covers("logging.datadog.failure.exports_metric", exercised_on=["chat_completions"])
def test_failed_chat_completions_emits_one_error_event(
self, client: LoggingClient, dd_logs: DdLogsReader, resources: ResourceManager
) -> None:
"""A /chat/completions call that fails at the provider must reach the
DataDog logs intake as exactly one error-grade event carrying the
provider error - failure metrics drive alerting and SLOs, so a dropped
failure event is an invisible outage.
A deployment with an invalid upstream key lets the request pass proxy
auth and fail at the provider (the same lever as the OTEL error test).
Failure payloads carry no prompt to mark, so the read-back queries the
indexed @model_group attribute of the per-run unique deployment name;
proxy-side 401s during key propagation never reach the provider and
ship no payload, so exactly one provider failure exists for it."""
_assert_datadog_configured(client)
model_name = f"dd-err-{unique_marker()}"
model_id = client.create_model(
model_name,
LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY),
)
resources.defer(lambda: client.delete_model(model_id))
key = client.key_with_alias(f"dd-err-key-{unique_marker()}", models=[model_name])
resources.defer(lambda: client.delete_key(key))
deadline = time.monotonic() + client.proxy.poll_timeout
while True:
outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16)
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
assert outcome.status_code != -1, (
"network failure between the test and the proxy while provoking the provider "
"failure; retrying now could double-log the failure payload and falsely trip "
f"the exactly-one assertion - fix the rig connectivity first: {outcome.body[:200]}"
)
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
break
time.sleep(client.proxy.poll_interval)
assert "AnthropicException" in outcome.body, (
"never saw the upstream provider failure before the deadline; the key may still be "
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
)
assert outcome.status_code == 401, (
f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}"
)
events = dd_logs.poll_events_for_query(f"@model_group:{model_name}")
payload = _assert_exactly_one_failure_event(events, model_group=model_name)
assert payload.error_str is not None and "AnthropicException" in payload.error_str, (
f"the event must carry the provider error, got error_str={payload.error_str!r}"
)

97
tests/e2e/logging/test_gcs_log_e2e.py Normal file
View file

@ -0,0 +1,97 @@
"""Live e2e: gcs_bucket log delivery for successful calls.
Covers logging.gcs_bucket.success.writes_object: one successful
/chat/completions call must land in the real GCS bucket as exactly one
StandardLoggingPayload record (GCS is the audit-trail parallel to S3 for GCP
deployments). Delivery is judged on what is actually readable in the bucket:
the proxy writes with its production service account, and the test reads the
record back through the GCS JSON API - covering both the batched NDJSON layout
(the default) and the per-request object layout.
Both halves of the contract are asserted: the recorded state (the proxy
reports the GCSBucketLogger callback active via /health/readiness/details -
note gcs_bucket is enterprise-gated, so this also requires a license) and the
enforced behavior (the record in the bucket, cost cross-checked against the
x-litellm-response-cost header of the very response the caller received).
"""
from __future__ import annotations
import math
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from gcs_reader import GcsLogReader, build_gcs_reader, utc_now
from lifecycle import ResourceManager
from logging_client import LoggingClient, completion_response_id, first_ok, readiness_details_body
pytestmark = pytest.mark.e2e
#: The active gcs_bucket callback's name in /health/readiness/details success_callbacks.
GCS_LOGGER_NAME = "GCSBucketLogger"
@pytest.fixture(scope="session")
def gcs_logs() -> GcsLogReader:
return build_gcs_reader()
def _assert_gcs_configured(client: LoggingClient) -> None:
"""Recorded state: the proxy reports the gcs_bucket callback among its
active callbacks, so a missing destination config (or a missing enterprise
license - gcs_bucket refuses to initialize without one) fails here, before
any delivery-based assertion can time out confusingly."""
body = readiness_details_body(client)
assert GCS_LOGGER_NAME in body, (
f"the proxy must report the {GCS_LOGGER_NAME} callback active "
f"(litellm_settings.callbacks: ['gcs_bucket'] + GCS_BUCKET_NAME env + enterprise license); "
f"got: {body[:400]}"
)
class TestGcsLogDelivery:
@pytest.mark.covers("logging.gcs_bucket.success.writes_object", exercised_on=["chat_completions"])
def test_chat_completions_writes_one_success_record(
self, client: LoggingClient, gcs_logs: GcsLogReader, resources: ResourceManager
) -> None:
"""One successful non-streaming /chat/completions call must be
readable back from the bucket as exactly one payload record carrying
the model group, the token counts, and the same cost the caller's
response header reported."""
_assert_gcs_configured(client)
alias = f"gcs-chat-{unique_marker()}"
key = client.key_with_alias(alias, models=[CHEAP_ANTHROPIC_MODEL])
resources.defer(lambda: client.delete_key(key))
since = utc_now()
marker = unique_marker()
outcome = first_ok(
client,
lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16),
)
assert outcome.response_cost is not None and outcome.response_cost > 0, (
f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}"
)
body_id = completion_response_id(outcome.body)
assert body_id is not None, "the completion body must carry an id (it names the gcs record)"
records = gcs_logs.poll_records_for_response_id(body_id, since=since)
assert records, f"no gcs record for response {body_id} was readable from the bucket within the deadline"
assert len(records) == 1, (
f"expected exactly ONE gcs record for the call, got {len(records)} - "
"more than one record for one call is the duplicate-delivery bug"
)
record = records[0]
assert record.id == body_id, f"record id must be the response id, got {record.id!r}"
assert record.status == "success", f"payload status must be success, got {record.status!r}"
assert record.model_group == CHEAP_ANTHROPIC_MODEL, (
f"payload model_group must be {CHEAP_ANTHROPIC_MODEL!r}, got {record.model_group!r}"
)
assert record.total_tokens is not None and record.total_tokens > 0, (
f"payload must count real tokens, got {record.total_tokens!r}"
)
assert record.response_cost is not None and math.isclose(
record.response_cost, outcome.response_cost, rel_tol=1e-9
), f"payload response_cost {record.response_cost!r} must equal the header cost {outcome.response_cost}"

35
tests/e2e/logging/test_otel_trace_e2e.py
View file

@ -23,9 +23,8 @@ import pytest
from pydantic import BaseModel, ConfigDict, ValidationError
from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker
from e2e_http import NoBody
from lifecycle import ResourceManager
from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok
from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body
from models import LiteLLMParamsBody
from otel_client import JaegerSpan, JaegerTrace, OtelReader
@ -48,11 +47,7 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None:
"""Recorded state: the proxy reports the OTEL v2 logger among its active
callbacks, so a missing/failed destination config fails here, before any
traffic-based assertion can time out confusingly."""
result = client.proxy.probe("/health/readiness/details", params=NoBody())
assert result.status_code == 200, (
f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}"
)
details = _ReadinessDetails.model_validate_json(result.body)
details = _ReadinessDetails.model_validate_json(readiness_details_body(client))
assert OTEL_V2_LOGGER_NAME in details.success_callbacks, (
f"the proxy must report the {OTEL_V2_LOGGER_NAME} callback active "
f"(LITELLM_OTEL_V2 + arize_phoenix preset in the compose config); got: {details.success_callbacks}"
@ -164,17 +159,14 @@ def served_genai_spans(trace: JaegerTrace, genai_span: str) -> list[JaegerSpan]:
these tests fail whenever the upstream 429s, 529s, or hands back a stale
credential on the first try."""
return [
span
for span in trace.spans
if span.operation_name == genai_span and _tag(span, ERROR_STATUS_TAG) != "ERROR"
span for span in trace.spans if span.operation_name == genai_span and _tag(span, ERROR_STATUS_TAG) != "ERROR"
]
def one_served_genai_span(trace: JaegerTrace, genai_span: str) -> JaegerSpan:
served = served_genai_spans(trace, genai_span)
assert len(served) == 1, (
f"a streamed call must produce exactly ONE served gen-AI span, got {len(served)}; "
f"spans: {trace.span_names()}"
f"a streamed call must produce exactly ONE served gen-AI span, got {len(served)}; spans: {trace.span_names()}"
)
return served[0]
@ -190,8 +182,7 @@ def _assert_real_ttft(hits: list[JaegerTrace], *, genai_span: str) -> None:
"(nothing tagged with its call id was found)"
)
assert len(hits) == 1, (
f"expected exactly ONE trace for the call, got {len(hits)}: "
f"{[(t.trace_id, t.span_names()) for t in hits]}"
f"expected exactly ONE trace for the call, got {len(hits)}: {[(t.trace_id, t.span_names()) for t in hits]}"
)
trace = hits[0]
span = one_served_genai_span(trace, genai_span)
@ -280,9 +271,7 @@ def _assert_error_span_contract(span: JaegerSpan) -> None:
"the span status description must carry the same untruncated message as error.message"
)
stack = _tag(span, "litellm.provider.error.stack_trace")
assert isinstance(stack, str) and stack, (
"the error span must carry a non-empty litellm.provider.error.stack_trace"
)
assert isinstance(stack, str) and stack, "the error span must carry a non-empty litellm.provider.error.stack_trace"
class TestOtelTraceCompleteness:
@ -313,9 +302,7 @@ class TestOtelTraceCompleteness:
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
outcome = first_ok(
client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16)
)
outcome = first_ok(client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16))
assert outcome.call_id is not None, "success response must carry x-litellm-call-id"
hits = otel_reader.poll_traces_for_call(
@ -520,9 +507,7 @@ class TestOtelTraceCompleteness:
route = "/v1/responses"
_assert_otel_destination_configured(client)
key = client.key_with_alias(
f"otel-stream-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]
)
key = client.key_with_alias(f"otel-stream-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL])
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
@ -660,9 +645,7 @@ class TestOtelTraceCompleteness:
route = "/v1/responses"
_assert_otel_destination_configured(client)
key = client.key_with_alias(
f"otel-ttft-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]
)
key = client.key_with_alias(f"otel-ttft-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL])
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()

170
tests/e2e/logging/test_s3_log_e2e.py Normal file
View file

@ -0,0 +1,170 @@
"""Live e2e: s3_v2 log delivery for successful and failed calls.
Covers logging.s3.success.writes_object and logging.s3.failure.writes_object:
one /chat/completions call must land in the real S3 bucket as exactly one
StandardLoggingPayload object (the primary audit trail; the batch flush must
neither drop nor duplicate it), and a failed call must be persisted the same
way for compliance. Delivery is judged on what is actually in the bucket: the
proxy writes with its production credentials and the test lists and reads the
objects back.
Both halves of the contract are asserted: the recorded state (the proxy
reports the S3Logger callback active via /health/readiness/details) and the
enforced behavior (the object in the bucket, with the cost cross-checked
against the x-litellm-response-cost header of the very response the caller
received).
The suite requires ``s3_callback_params.s3_use_key_prefix: true`` on the proxy,
which keys objects as ``{key_alias}/{date}/time-..._{id}.json`` - a unique key
alias per test turns the poll into a cheap prefix listing.
"""
from __future__ import annotations
import math
import time
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from lifecycle import ResourceManager
from logging_client import (
INVALID_UPSTREAM_API_KEY,
LoggingClient,
completion_response_id,
first_ok,
readiness_details_body,
)
from models import LiteLLMParamsBody
from s3_reader import S3LogReader, build_s3_reader
pytestmark = pytest.mark.e2e
#: The active s3_v2 callback's name in /health/readiness/details success_callbacks.
S3_LOGGER_NAME = "S3Logger"
@pytest.fixture(scope="session")
def s3_logs() -> S3LogReader:
return build_s3_reader()
def _assert_s3_configured(client: LoggingClient) -> None:
"""Recorded state: the proxy reports the s3_v2 callback among its active
callbacks, so a missing destination config fails here, before any
delivery-based assertion can time out confusingly."""
body = readiness_details_body(client)
assert S3_LOGGER_NAME in body, (
f"the proxy must report the {S3_LOGGER_NAME} callback active "
f"(litellm_settings.callbacks: ['s3_v2'] + s3_callback_params in the proxy config); "
f"got: {body[:400]}"
)
class TestS3LogDelivery:
@pytest.mark.covers("logging.s3.success.writes_object", exercised_on=["chat_completions"])
def test_chat_completions_writes_one_success_object(
self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager
) -> None:
"""One successful non-streaming /chat/completions call must land in
the bucket as exactly one payload object carrying the model group, the
token counts, and the same cost the caller's response header reported."""
_assert_s3_configured(client)
alias = f"s3-chat-{unique_marker()}"
key = client.key_with_alias(alias, models=[CHEAP_ANTHROPIC_MODEL])
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
outcome = first_ok(
client,
lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16),
)
assert outcome.response_cost is not None and outcome.response_cost > 0, (
f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}"
)
body_id = completion_response_id(outcome.body)
assert body_id is not None, "the completion body must carry an id (it names the s3 object)"
records = s3_logs.poll_records(prefix=f"{alias}/", predicate=lambda r: r.id == body_id)
assert records, (
f"no s3 object for response {body_id} under prefix {alias}/ reached the bucket within the deadline"
)
assert len(records) == 1, (
f"expected exactly ONE s3 object for the call, got {len(records)} - "
"more than one object for one call is the duplicate-delivery bug"
)
record = records[0]
assert record.status == "success", f"payload status must be success, got {record.status!r}"
assert record.model_group == CHEAP_ANTHROPIC_MODEL, (
f"payload model_group must be {CHEAP_ANTHROPIC_MODEL!r}, got {record.model_group!r}"
)
assert record.total_tokens is not None and record.total_tokens > 0, (
f"payload must count real tokens, got {record.total_tokens!r}"
)
assert record.response_cost is not None and math.isclose(
record.response_cost, outcome.response_cost, rel_tol=1e-9
), f"payload response_cost {record.response_cost!r} must equal the header cost {outcome.response_cost}"
@pytest.mark.covers("logging.s3.failure.writes_object", exercised_on=["chat_completions"])
def test_chat_completions_failure_writes_one_object(
self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager
) -> None:
"""A call that fails at the provider must be persisted to the bucket as
exactly one failure payload carrying the provider error - failed calls
are part of the audit trail, not an exemption from it.
A deployment with an invalid upstream key lets the request pass proxy
auth and fail at the provider (the same lever as the OTEL error test).
Proxy-side rejections during key/model propagation can also ship
failure payloads under this alias, but without a model_group and
without the provider error, so the read-back keys on both: only
provider-reaching calls carry them, and with this key every one of
those is the AnthropicException that ends the send loop."""
_assert_s3_configured(client)
model_name = f"s3-err-{unique_marker()}"
model_id = client.create_model(
model_name,
LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY),
)
resources.defer(lambda: client.delete_model(model_id))
alias = f"s3-err-key-{unique_marker()}"
key = client.key_with_alias(alias, models=[model_name])
resources.defer(lambda: client.delete_key(key))
deadline = time.monotonic() + client.proxy.poll_timeout
while True:
outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16)
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
assert outcome.status_code != -1, (
"network failure between the test and the proxy while provoking the provider "
"failure; retrying now could double-log the failure payload and falsely trip "
f"the exactly-one assertion - fix the rig connectivity first: {outcome.body[:200]}"
)
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
break
time.sleep(client.proxy.poll_interval)
assert "AnthropicException" in outcome.body, (
"never saw the upstream provider failure before the deadline; the key may still be "
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
)
assert outcome.status_code == 401, (
f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}"
)
records = s3_logs.poll_records(
prefix=f"{alias}/",
predicate=lambda r: (
r.status == "failure" and r.model_group == model_name and "AnthropicException" in (r.error_str or "")
),
)
assert records, (
f"no failure object for {model_name} under prefix {alias}/ reached the bucket within the deadline"
)
assert len(records) == 1, f"expected exactly ONE failure object for the call, got {len(records)}"
record = records[0]
assert record.error_str is not None and "AnthropicException" in record.error_str, (
f"the persisted failure must carry the provider error, got error_str={record.error_str!r}"
)
assert not record.response_cost, f"a failed call must not be billed, got response_cost={record.response_cost!r}"

123
tests/e2e/logging/test_team_langfuse_callback_e2e.py Normal file
View file

@ -0,0 +1,123 @@
"""Live e2e: team-scoped Langfuse callback delivery and isolation.
Covers logging.langfuse.success.logs_spend: a team configured with a Langfuse
callback via POST /team/{id}/callback must deliver its members' calls to the
real Langfuse project (generation readable back through Langfuse's own API,
with the cost agreeing with the x-litellm-response-cost header), while traffic
from keys outside the team must NOT reach that project - the isolation is the
point of team-scoped callbacks.
Both halves of the contract are asserted: the recorded state (the /team/callback
registration itself answers success) and the enforced behavior (the generation
at the destination for the team key, and its absence for the non-team key).
"""
from __future__ import annotations
import time
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from lifecycle import ResourceManager
from logging_client import (
LangfuseCreds,
LoggingClient,
costs_agree,
first_ok,
load_langfuse_creds,
observation_spend,
)
pytestmark = pytest.mark.e2e
#: How long to keep re-checking that the non-team call never surfaces in
#: Langfuse after the team call's generation has already been ingested; the
#: positive observation bounds the pipeline's latency, so a wrong delivery
#: would be visible within the same order of magnitude.
ISOLATION_SETTLE_SECONDS = 30.0
ISOLATION_CHECK_INTERVAL_SECONDS = 5.0
@pytest.fixture(scope="session")
def langfuse_creds() -> LangfuseCreds:
return load_langfuse_creds()
class TestTeamLangfuseCallback:
@pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["chat_completions"])
def test_team_callback_delivers_and_isolates(
self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager
) -> None:
team_id = client.create_team(f"lf-team-{unique_marker()}", models=[CHEAP_ANTHROPIC_MODEL])
resources.defer(lambda: client.delete_team(team_id))
# Recorded state: the registration endpoint itself must answer success
# (add_team_langfuse_callback asserts it).
client.add_team_langfuse_callback(team_id, langfuse_creds)
team_alias = f"lf-team-key-{unique_marker()}"
team_key = client.key_with_alias(team_alias, models=[CHEAP_ANTHROPIC_MODEL], team_id=team_id)
resources.defer(lambda: client.delete_key(team_key))
solo_alias = f"lf-solo-key-{unique_marker()}"
solo_key = client.key_with_alias(solo_alias, models=[CHEAP_ANTHROPIC_MODEL])
resources.defer(lambda: client.delete_key(solo_key))
# Enforced behavior, positive half, with one propagation retry: a
# worker still holding the pre-callback team object can serve the
# first call without shipping it, and by the time the first Langfuse
# poll has timed out the team cache TTL has lapsed, so a second call
# must deliver.
team_marker = ""
team_outcome = None
observation = None
for _attempt in range(2):
team_marker = unique_marker()
team_outcome = first_ok(
client,
lambda marker=team_marker: client.chat_raw(
team_key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16
),
)
assert team_outcome.response_cost is not None and team_outcome.response_cost > 0, (
f"the response must report x-litellm-response-cost, got {team_outcome.response_cost!r}"
)
observation = client.poll_langfuse_observation(
langfuse_creds,
key_alias=team_alias,
prompt_marker=team_marker,
require_positive_cost=True,
)
if observation is not None:
break
solo_marker = unique_marker()
_ = first_ok(
client,
lambda: client.chat_raw(
solo_key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {solo_marker}", max_tokens=16
),
)
assert observation is not None, (
f"the team key's call (marker {team_marker}) never reached Langfuse within the deadline, "
"even after a fresh call past the team-object cache TTL"
)
assert team_outcome is not None and team_outcome.response_cost is not None
cost = observation_spend(observation)
assert cost is not None and costs_agree(team_outcome.response_cost, cost), (
f"Langfuse calculatedTotalCost {cost!r} must agree with the header cost {team_outcome.response_cost}"
)
# Enforced behavior, negative half: the non-team call must never show
# up in this project. The positive generation above has already been
# ingested, which bounds the pipeline latency, so keep re-checking for
# a settle window rather than trusting a single instant.
settle_deadline = time.monotonic() + ISOLATION_SETTLE_SECONDS
while True:
leaked = client.find_langfuse_observation(langfuse_creds, key_alias=solo_alias, prompt_marker=solo_marker)
assert leaked is None, (
f"a non-team key's call (marker {solo_marker}) reached the team's Langfuse "
f"project: {leaked.id} - team callbacks must not apply outside the team"
)
if time.monotonic() >= settle_deadline:
break
time.sleep(ISOLATION_CHECK_INTERVAL_SECONDS)

87
tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py
View file

@ -2635,6 +2635,93 @@ async def test_list_batches_unparseable_row_does_not_truncate_pagination():
assert len(seen) == len(set(seen))
@pytest.mark.asyncio
async def test_list_batches_fills_a_page_past_a_full_page_of_unparseable_rows():
"""A page whose rows all fail to parse must still let the caller advance.
``has_more`` came from the raw fetch while ``last_id`` came from the parsed
survivors, so a full page of corrupt rows answered ``data: []``,
``last_id: None``, ``has_more: True``, and a client following ``last_id``
could not move past them.
"""
from litellm.proxy._types import UserAPIKeyAuth
rows = [_managed_batch_row(i) for i in range(5)]
for corrupt_row in rows[2:4]:
corrupt_row.file_object = "{ not valid json"
prisma_client = _fake_managed_object_table(rows)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
pages = await _walk_batch_pages(
proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=1
)
assert [[batch.id for batch in page["data"]] for page in pages] == [
[rows[4].unified_object_id],
[rows[1].unified_object_id],
[rows[0].unified_object_id],
]
assert [page["has_more"] for page in pages] == [True, True, False]
_DEEP_BATCH_SCAN_ROW_COUNT = 2000
_DEEP_BATCH_SCAN_QUERY_BUDGET = 10
@pytest.mark.asyncio
async def test_list_batches_bounds_the_queries_a_deep_unparseable_run_costs():
"""A tiny limit behind thousands of corrupt rows must not turn one request into thousands of queries."""
from litellm.proxy._types import UserAPIKeyAuth
rows = [_managed_batch_row(0)] + [
_managed_batch_row(index, file_object="{ not valid json")
for index in range(1, _DEEP_BATCH_SCAN_ROW_COUNT + 1)
]
prisma_client = _fake_managed_object_table(rows)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
page = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=1
)
assert [batch.id for batch in page["data"]] == [rows[0].unified_object_id]
assert page["has_more"] is False
assert (
prisma_client.db.litellm_managedobjecttable.find_many.call_count
<= _DEEP_BATCH_SCAN_QUERY_BUDGET
)
@pytest.mark.asyncio
async def test_list_batches_reads_one_chunk_when_the_first_one_fills_the_page():
"""The widened chunk must stay off the common path, where the newest rows already fill the page."""
from litellm.proxy._types import UserAPIKeyAuth
rows = [_managed_batch_row(index) for index in range(_DEEP_BATCH_SCAN_ROW_COUNT)]
prisma_client = _fake_managed_object_table(rows)
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=prisma_client
)
page = await proxy_managed_files.list_user_batches(
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2
)
assert [batch.id for batch in page["data"]] == [
rows[-1].unified_object_id,
rows[-2].unified_object_id,
]
assert page["has_more"] is True
assert prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1
@pytest.mark.asyncio
async def test_return_unified_file_id_includes_expires_at():
from litellm.types.llms.openai import OpenAIFileObject

107
tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py
View file

@ -10,11 +10,14 @@ with deployment credentials, bypassing the managed files access-control hooks.
import base64
import pytest
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.utils import LiteLLMBatch
def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth:
@ -161,6 +164,108 @@ async def test_service_account_blocked_from_other_team_file():
assert exc_info.value.status_code == 403
# --- Keyless key must not be locked out of the batch it created ---
def _make_unified_batch_id() -> str:
raw = "litellm_proxy;model_id:my-model-id;llm_batch_id:batch_raw_123"
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
def _make_managed_files_instance_with_object_store():
"""Managed-files hook backed by an in-memory stand-in for the managed
object table, so create and retrieve exercise the same stored row."""
from litellm_enterprise.proxy.hooks.managed_files import (
_PROXY_LiteLLMManagedFiles,
)
store = {}
async def upsert(where, data):
store[where["unified_object_id"]] = SimpleNamespace(**data["create"])
async def find_first(where):
return store.get(where["unified_object_id"])
mock_prisma = MagicMock()
mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock(side_effect=upsert)
mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(
side_effect=find_first
)
return (
_PROXY_LiteLLMManagedFiles(
internal_usage_cache=DualCache(),
prisma_client=mock_prisma,
),
store,
)
async def _store_batch(managed_files, unified_batch_id: str, creator: UserAPIKeyAuth):
await managed_files.store_unified_object_id(
unified_object_id=unified_batch_id,
file_object=LiteLLMBatch(
id="batch_raw_123",
completion_window="24h",
created_at=0,
endpoint="/v1/chat/completions",
input_file_id="file-1",
object="batch",
status="validating",
),
litellm_parent_otel_span=None,
model_object_id="batch_raw_123",
file_purpose="batch",
user_api_key_dict=creator,
)
@pytest.mark.asyncio
async def test_keyless_key_can_retrieve_the_batch_it_created():
"""Regression: a key with no user_id and no team_id (what `/key/generate`
by a proxy admin and service-account keys produce) stamped
`created_by=None` and was then denied its own managed batch with
"User None does not have access"."""
unified_batch_id = _make_unified_batch_id()
managed_files, store = _make_managed_files_instance_with_object_store()
keyless = UserAPIKeyAuth(api_key="sk-keyless", parent_otel_span=None)
await _store_batch(managed_files, unified_batch_id, keyless)
assert store[unified_batch_id].created_by == f"key:{keyless.token}"
data = {"batch_id": unified_batch_id}
await managed_files.async_pre_call_hook(
user_api_key_dict=keyless,
cache=DualCache(),
data=data,
call_type=CallTypes.aretrieve_batch.value,
)
assert data["batch_id"] == "batch_raw_123"
assert data["model"] == "my-model-id"
@pytest.mark.asyncio
async def test_other_keyless_key_still_denied_the_batch():
unified_batch_id = _make_unified_batch_id()
managed_files, _ = _make_managed_files_instance_with_object_store()
await _store_batch(
managed_files,
unified_batch_id,
UserAPIKeyAuth(api_key="sk-creator", parent_otel_span=None),
)
with pytest.raises(HTTPException) as exc_info:
await managed_files.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-other", parent_otel_span=None),
cache=DualCache(),
data={"batch_id": unified_batch_id},
call_type=CallTypes.aretrieve_batch.value,
)
assert exc_info.value.status_code == 403
# --- Option C fix test: check_batch_cost bypasses managed files hook ---

24
tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
View file

@ -527,13 +527,33 @@ async def test_afile_list_orders_newest_first_and_breaks_ties_on_the_cursor_colu
@pytest.mark.asyncio
async def test_afile_list_denies_a_caller_without_a_user_or_team():
async def test_afile_list_scopes_a_keyless_key_to_its_own_hashed_token():
caller = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None)
managed_files, table = _make_managed_files_over_rows(
[
_make_managed_file_row("unified-mine", created_by=f"key:{caller.token}"),
_make_managed_file_row("unified-theirs", created_by="other-user"),
]
)
response = await managed_files.afile_list(
purpose=None,
litellm_parent_otel_span=None,
user_api_key_dict=caller,
)
assert [file.id for file in response.data] == ["unified-mine"]
assert table.find_many_calls[0]["where"] == {"created_by": f"key:{caller.token}"}
@pytest.mark.asyncio
async def test_afile_list_denies_a_caller_with_no_identity_at_all():
managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")])
response = await managed_files.afile_list(
purpose=None,
litellm_parent_otel_span=None,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None),
user_api_key_dict=UserAPIKeyAuth(parent_otel_span=None),
)
assert response.data == []

14
tests/test_litellm/litellm_core_utils/test_llm_request_utils.py
View file

@ -1,4 +1,5 @@
import httpx
import pytest
from litellm.litellm_core_utils.llm_request_utils import (
flatten_form_field_values,
@ -80,9 +81,7 @@ def test_flatten_form_field_values_later_source_wins_on_collision():
def test_flatten_form_field_values_keeps_scalar_lists_as_repeated_fields():
assert flatten_form_field_values(
{"loras": ["a", "b", "c"], "generation_config": {"tags": [1, 2]}, "seed": 42}
) == (
assert flatten_form_field_values({"loras": ["a", "b", "c"], "generation_config": {"tags": [1, 2]}, "seed": 42}) == (
("loras", ("a", "b", "c")),
("generation_config[tags]", ("1", "2")),
("seed", "42"),
@ -97,3 +96,12 @@ def test_flatten_form_field_values_scalar_list_survives_update_into_multipart():
assert names.count("loras") == 2
assert names.count("model") == 1
def test_flatten_form_field_values_rejects_over_deep_nesting():
nested: object = "leaf"
for _ in range(102):
nested = {"k": nested}
assert isinstance(nested, dict)
with pytest.raises(ValueError, match="max depth"):
flatten_form_field_values(nested)

44
tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py
View file

@ -7,6 +7,7 @@ import pytest
from litellm.llms.base_llm.managed_resources.isolation import (
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
@ -154,3 +155,46 @@ def test_access_identity_less_caller_always_denied(created_by, resource_team_id)
)
is False
)
# ---------------------------------------------------------------------------
# keyless keys (no user_id, no team_id) own their resources by hashed token
# ---------------------------------------------------------------------------
def test_owner_id_prefers_user_id_then_falls_back_to_token():
assert resolve_resource_owner_id(UserAPIKeyAuth(user_id="alice")) == "alice"
assert resolve_resource_owner_id(UserAPIKeyAuth(team_id="team-eng")) is None
assert resolve_resource_owner_id(UserAPIKeyAuth()) is None
keyless = UserAPIKeyAuth(api_key="sk-keyless")
assert resolve_resource_owner_id(keyless) == f"key:{keyless.token}"
def test_keyless_key_can_access_its_own_resource():
"""Regression for the self-lockout: a key generated by a proxy admin (or a
service-account key) has no user_id and no team_id, so it used to stamp
`created_by=None` and then be denied its own batches and files."""
keyless = UserAPIKeyAuth(api_key="sk-keyless")
owner_id = resolve_resource_owner_id(keyless)
assert build_owner_filter(keyless) == {"created_by": owner_id}
assert (
can_access_resource(keyless, created_by=owner_id, resource_team_id=None) is True
)
def test_keyless_key_denied_another_keyless_keys_resource():
"""The #27004 isolation invariant: two distinct keyless keys must not see
each other's resources."""
creator = UserAPIKeyAuth(api_key="sk-creator")
other = UserAPIKeyAuth(api_key="sk-other")
assert (
can_access_resource(
other,
created_by=resolve_resource_owner_id(creator),
resource_team_id=None,
)
is False
)

18
tests/test_litellm/llms/databricks/test_databricks_partner_integration.py
View file

@ -245,6 +245,24 @@ class TestOAuthM2M:
assert "/serving-endpoints" not in call_url
assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
def test_oauth_m2m_strips_ai_gateway_path(self):
"""OAuth M2M derives the token URL from the workspace origin."""
databricks_base = DatabricksBase()
mock_response = Mock()
mock_response.status_code = 200
mock_response.json.return_value = {"access_token": "token"}
with patch("requests.post", return_value=mock_response) as mock_post:
databricks_base._get_oauth_m2m_token(
api_base="https://adb-123.azuredatabricks.net/ai-gateway/mlflow/v1",
client_id="id",
client_secret="secret",
)
call_url = mock_post.call_args[0][0]
assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token"
class TestValidateEnvironmentWithOAuth:
"""Test OAuth M2M is used when credentials are available."""

312
tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py Normal file
View file

@ -0,0 +1,312 @@
"""Tests for hosted_vllm video generation (vLLM-Omni /v1/videos)."""
import json
from io import BytesIO
import httpx
import pytest
import litellm
from litellm.litellm_core_utils.url_utils import SSRFError
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config
from litellm.llms.hosted_vllm.videos.transformation import (
HostedVLLMVideoConfig,
_serialize_form_value,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.types.videos.main import VideoObject
from litellm.utils import ProviderConfigManager
def test_provider_config_registration():
config = ProviderConfigManager.get_provider_video_config(
model="hosted_vllm/MiniMax-H3",
provider=LlmProviders.HOSTED_VLLM,
)
assert config is not None
assert isinstance(config, HostedVLLMVideoConfig)
assert isinstance(get_hosted_vllm_video_config("MiniMax-H3"), HostedVLLMVideoConfig)
def test_get_complete_url_appends_videos():
config = HostedVLLMVideoConfig()
assert (
config.get_complete_url(model="MiniMax-H3", api_base="http://localhost:8091", litellm_params={})
== "http://localhost:8091/v1/videos"
)
assert (
config.get_complete_url(model="MiniMax-H3", api_base="http://localhost:8091/v1", litellm_params={})
== "http://localhost:8091/v1/videos"
)
assert (
config.get_complete_url(model="MiniMax-H3", api_base="http://localhost:8091/v1/", litellm_params={})
== "http://localhost:8091/v1/videos"
)
def test_get_complete_url_requires_api_base():
config = HostedVLLMVideoConfig()
with pytest.raises(ValueError, match="api_base not set"):
config.get_complete_url(model="MiniMax-H3", api_base=None, litellm_params={})
def test_validate_environment_defaults_to_fake_api_key():
config = HostedVLLMVideoConfig()
headers = config.validate_environment(
headers={},
model="MiniMax-H3",
litellm_params=GenericLiteLLMParams(),
)
assert headers.get("Authorization") == "Bearer fake-api-key"
def test_validate_environment_uses_provided_api_key():
config = HostedVLLMVideoConfig()
headers = config.validate_environment(
headers={"X-Test": "1"},
model="MiniMax-H3",
litellm_params=GenericLiteLLMParams(api_key="my-custom-key"),
)
assert headers.get("Authorization") == "Bearer my-custom-key"
assert headers.get("X-Test") == "1"
def test_transform_video_create_request_uses_multipart_form_fields():
"""vLLM-Omni rejects JSON create bodies. Extra Omni fields must be form parts."""
config = HostedVLLMVideoConfig()
extra_params = {"task": "t2va", "duration": 10.0, "audio_flow_shift": 3.0}
data, files, url = config.transform_video_create_request(
model="MiniMax-H3",
prompt="three cats march into a bedroom playing tiny brass instruments",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={
"width": 1280,
"height": 720,
"fps": 24,
"num_inference_steps": 20,
"flow_shift": 12,
"seed": 1101,
"aspect_ratio": "16:9",
"extra_params": extra_params,
"extra_headers": {"X-Ignored": "yes"},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert url == "http://localhost:8091/v1/videos"
assert files == ()
assert data["model"] == "MiniMax-H3"
assert data["prompt"] == "three cats march into a bedroom playing tiny brass instruments"
assert data["width"] == "1280"
assert data["height"] == "720"
assert data["fps"] == "24"
assert data["num_inference_steps"] == "20"
assert data["flow_shift"] == "12"
assert data["seed"] == "1101"
assert data["aspect_ratio"] == "16:9"
assert json.loads(data["extra_params"]) == extra_params
assert "extra_headers" not in data
def test_transform_video_create_request_keeps_openai_size_and_seconds():
config = HostedVLLMVideoConfig()
data, files, _ = config.transform_video_create_request(
model="Wan2.2",
prompt="a mountain lake at sunrise",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={"seconds": "8", "size": "1280x720"},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert files == ()
assert data["seconds"] == "8"
assert data["size"] == "1280x720"
def test_transform_video_create_request_attaches_input_reference_file():
config = HostedVLLMVideoConfig()
reference = BytesIO(b"fake-png")
reference.name = "input.png"
data, files, _ = config.transform_video_create_request(
model="Wan2.2",
prompt="animate this image",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={"input_reference": reference, "width": 832},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert data["width"] == "832"
assert "input_reference" not in data
reference_parts = [value for name, value in files if name == "input_reference"]
assert len(reference_parts) == 1
filename, content, content_type = reference_parts[0]
assert filename == "input_reference.png"
assert content is reference
assert content_type == "image/png"
def test_serialize_form_value_does_not_quote_plain_strings():
assert _serialize_form_value("16:9") == "16:9"
assert _serialize_form_value(True) == "true"
assert _serialize_form_value({"task": "t2va"}) == json.dumps({"task": "t2va"})
def test_map_openai_params_passes_through_omni_fields():
config = HostedVLLMVideoConfig()
mapped = config.map_openai_params(
video_create_optional_params={
"width": 1280,
"extra_params": {"task": "t2va"},
"aspect_ratio": "16:9",
"extra_body": None,
},
model="MiniMax-H3",
drop_params=False,
)
assert mapped["width"] == 1280
assert mapped["extra_params"] == {"task": "t2va"}
assert mapped["aspect_ratio"] == "16:9"
assert "extra_body" not in mapped
def test_get_supported_openai_params_includes_omni_extensions():
config = HostedVLLMVideoConfig()
supported = config.get_supported_openai_params("MiniMax-H3")
assert "prompt" in supported
assert "input_reference" in supported
assert "width" in supported
assert "extra_params" in supported
assert "aspect_ratio" in supported
assert "image_reference" in supported
assert "audio_reference" in supported
def _http_handler_for(handler) -> HTTPHandler:
return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))
def test_video_generation_posts_multipart_not_json():
captured: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
captured.append(request)
return httpx.Response(
200,
json={
"id": "video-123",
"object": "video",
"status": "queued",
"created_at": 1701234567,
},
)
response = litellm.video_generation(
model="hosted_vllm/MiniMax-H3",
prompt="three cats march into a bedroom playing tiny brass instruments",
api_base="http://localhost:8091",
api_key="test-key",
client=_http_handler_for(handler),
extra_body={
"width": 1280,
"height": 720,
"fps": 24,
"extra_params": {"task": "t2va", "duration": 10.0},
},
)
assert isinstance(response, VideoObject)
assert response.status == "queued"
assert len(captured) == 1
request = captured[0]
assert str(request.url) == "http://localhost:8091/v1/videos"
assert request.headers["authorization"] == "Bearer test-key"
body = request.content
assert b'name="prompt"' in body
assert b"three cats march into a bedroom playing tiny brass instruments" in body
assert b'name="width"' in body
assert b"1280" in body
assert b'name="extra_params"' in body
assert b"t2va" in body
assert request.headers.get("content-type", "").startswith("multipart/form-data")
def test_http_image_reference_is_forwarded_not_downloaded():
config = HostedVLLMVideoConfig()
data, files, _ = config.transform_video_create_request(
model="MiniMax-H3",
prompt="a person singing",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={
"image_reference": {"image_url": "http://1.1.1.1/face.png"},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert files == ()
payload = json.loads(data["image_reference"])
assert payload["image_url"] == "http://1.1.1.1/face.png"
def test_data_url_image_reference_is_forwarded():
data_url = "data:image/png;base64,AAAA"
config = HostedVLLMVideoConfig()
data, files, _ = config.transform_video_create_request(
model="MiniMax-H3",
prompt="a person singing",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={"image_reference": {"image_url": data_url}},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert files == ()
assert json.loads(data["image_reference"])["image_url"] == data_url
def test_metadata_url_in_image_reference_is_rejected():
config = HostedVLLMVideoConfig()
with pytest.raises(SSRFError, match="blocked address"):
config.transform_video_create_request(
model="MiniMax-H3",
prompt="a person singing",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={
"image_reference": {"image_url": "http://169.254.169.254/latest/meta-data/"},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)
def test_file_scheme_media_reference_is_rejected():
config = HostedVLLMVideoConfig()
with pytest.raises(SSRFError, match="scheme"):
config.transform_video_create_request(
model="MiniMax-H3",
prompt="a person singing",
api_base="http://localhost:8091/v1/videos",
video_create_optional_request_params={
"video_reference": {"video_url": "file:///etc/passwd"},
},
litellm_params=GenericLiteLLMParams(),
headers={},
)

91
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py
View file

@ -3,6 +3,9 @@
from datetime import datetime, timedelta, timezone
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import SecretStr
from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import (
envelope_keys_from_master_key,
@ -13,17 +16,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
SessionBearerInvalid,
SessionRefreshInvalid,
SessionRefreshOpened,
SessionSigningConfigError,
is_session_bearer_shaped,
open_session_refresh_bearer,
resolve_session_bearer,
resolve_session_signing_keys,
session_keys_from_master_key,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
SESSION_TTL_SECONDS,
AsymmetricSessionKeys,
MintedSessionToken,
SessionKeys,
SessionPrincipal,
mint_session_refresh_token,
mint_session_token,
session_public_key_pem,
)
NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
@ -133,3 +141,86 @@ def test_refresh_grant_rejects_a_different_client():
def test_refresh_grant_rejects_access_token_presented_as_refresh():
result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc")
assert isinstance(result, SessionRefreshInvalid)
def _rsa_private_pem() -> str:
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
return key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode()
def test_absent_signing_setting_keeps_the_master_key_hs256_default():
resolved = resolve_session_signing_keys(MASTER_KEY, None)
assert isinstance(resolved, SessionKeys)
assert resolved.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value()
def test_rs256_signing_setting_resolves_inline_pem_material():
pem = _rsa_private_pem()
resolved = resolve_session_signing_keys(
MASTER_KEY,
{"algorithm": "RS256", "kid": "2026-01", "private_key": pem},
)
assert isinstance(resolved, AsymmetricSessionKeys)
assert resolved.kid == "2026-01"
minted = mint_session_token(PRINCIPAL, resolved, NOW)
assert isinstance(minted, MintedSessionToken)
admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW)
assert isinstance(admitted, SessionBearerAdmitted)
def test_rs256_signing_setting_resolves_env_reference(monkeypatch):
monkeypatch.setenv("MCP_SESSION_PRIVATE_KEY", _rsa_private_pem())
resolved = resolve_session_signing_keys(
MASTER_KEY,
{"algorithm": "RS256", "kid": "2026-01", "private_key": "os.environ/MCP_SESSION_PRIVATE_KEY"},
)
assert isinstance(resolved, AsymmetricSessionKeys)
def test_rs256_signing_setting_resolves_previous_public_keys():
old_pem = _rsa_private_pem()
old_keys = AsymmetricSessionKeys(private_key_pem=SecretStr(old_pem), kid="2025-06")
resolved = resolve_session_signing_keys(
MASTER_KEY,
{
"algorithm": "RS256",
"kid": "2026-01",
"private_key": _rsa_private_pem(),
"previous_public_keys": [{"kid": "2025-06", "public_key": session_public_key_pem(old_keys)}],
},
)
assert isinstance(resolved, AsymmetricSessionKeys)
minted = mint_session_token(PRINCIPAL, old_keys, NOW)
assert isinstance(minted, MintedSessionToken)
admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW)
assert isinstance(admitted, SessionBearerAdmitted)
@pytest.mark.parametrize(
"raw",
[
{"algorithm": "HS512", "kid": "k", "private_key": "irrelevant"},
{"algorithm": "RS256", "kid": "k"},
{"algorithm": "RS256", "kid": "k", "private_key": "not a pem"},
{"algorithm": "RS256", "kid": "k", "private_key": "os.environ/UNSET_MCP_SESSION_KEY_VAR"},
{"algorithm": "RS256", "kid": "k", "private_key": "x", "unexpected": True},
"not-a-mapping",
],
)
def test_defective_signing_setting_fails_closed_never_falls_back_to_hs256(raw):
resolved = resolve_session_signing_keys(MASTER_KEY, raw)
assert isinstance(resolved, SessionSigningConfigError)
def test_signing_config_error_detail_never_leaks_key_material():
pem = _rsa_private_pem()
resolved = resolve_session_signing_keys(
MASTER_KEY,
{"algorithm": "RS256", "kid": "k", "private_key": pem, "unexpected": True},
)
assert isinstance(resolved, SessionSigningConfigError)
assert pem.splitlines()[1] not in resolved.detail

187
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py
View file

@ -4,6 +4,8 @@ from datetime import datetime, timedelta, timezone
import jwt
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from pydantic import SecretStr, ValidationError
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
@ -13,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
SESSION_REFRESH_TTL_SECONDS,
SESSION_TOKEN_PREFIX,
SESSION_TTL_SECONDS,
AsymmetricSessionKeys,
MintedSessionToken,
NotASessionToken,
OpenedSessionToken,
@ -21,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
SessionKeys,
SessionMalformed,
SessionPrincipal,
SessionRotatedPublicKey,
SessionTokenTooLarge,
is_session_refresh_token,
is_session_token,
@ -28,8 +32,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i
mint_session_token,
open_session_refresh_token,
open_session_token,
session_public_key_pem,
)
def _rsa_private_pem(bits: int = 2048) -> str:
key = rsa.generate_private_key(public_exponent=65537, key_size=bits)
return key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode()
_RSA_PEM_A = _rsa_private_pem()
_RSA_PEM_B = _rsa_private_pem()
NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
KEYS = SessionKeys(signing_key=SecretStr("k" * 32))
OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32))
@ -264,3 +282,172 @@ def test_signed_claims_with_a_non_string_team_are_rejected():
def test_principal_rejects_an_unknown_audience_at_construction():
with pytest.raises(ValidationError):
SessionPrincipal(user_id="user-123", client_id="llm_client_abc", audience="mcp")
RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01")
OTHER_RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_B), kid="2025-06")
def test_rs256_access_round_trip_with_kid_and_alg_pinned_in_header():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
header = jwt.get_unverified_header(token.removeprefix(SESSION_TOKEN_PREFIX))
assert header["alg"] == "RS256"
assert header["kid"] == "2026-01"
opened = open_session_token(token, RSA_KEYS, NOW)
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == PRINCIPAL
def test_rs256_refresh_round_trip():
minted = mint_session_refresh_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
opened = open_session_refresh_token(token, RSA_KEYS, NOW)
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == PRINCIPAL
def test_rs256_token_verifies_with_public_key_only():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
public_pem = session_public_key_pem(RSA_KEYS)
assert "PUBLIC KEY" in public_pem
assert "PRIVATE" not in public_pem
claims = jwt.decode(
minted.token.get_secret_value().removeprefix(SESSION_TOKEN_PREFIX),
public_pem,
algorithms=["RS256"],
issuer=SESSION_ISSUER,
options={"verify_exp": False},
)
assert claims["user_id"] == "user-123"
def test_rs256_tampered_signature_is_bad_signature():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
assert isinstance(open_session_token(tampered, RSA_KEYS, NOW), SessionBadSignature)
def test_rs256_expired_token_is_expired():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, after), SessionExpired)
def test_hs256_token_is_rejected_in_rs256_mode():
assert isinstance(open_session_token(_mint_access(), RSA_KEYS, NOW), SessionBadSignature)
def test_hs256_token_claiming_the_current_kid_is_rejected_by_alg_pinning():
token = SESSION_TOKEN_PREFIX + jwt.encode(
_valid_claims(),
KEYS.signing_key.get_secret_value(),
algorithm="HS256",
headers={"kid": RSA_KEYS.kid},
)
assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionMalformed)
def test_rs256_token_is_rejected_in_hs256_mode():
minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
assert isinstance(open_session_token(minted.token.get_secret_value(), KEYS, NOW), SessionMalformed)
def test_rs256_token_from_an_unknown_kid_is_bad_signature():
minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, NOW), SessionBadSignature)
def test_rs256_token_signed_by_a_foreign_key_claiming_the_current_kid_is_bad_signature():
token = SESSION_TOKEN_PREFIX + jwt.encode(
_valid_claims(),
_RSA_PEM_B,
algorithm="RS256",
headers={"kid": RSA_KEYS.kid},
)
assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature)
def test_alg_none_token_with_the_current_kid_is_rejected_in_rs256_mode():
unsigned = jwt.api_jws.encode(
b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none", headers={"kid": RSA_KEYS.kid}
)
assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, RSA_KEYS, NOW), SessionMalformed)
def test_rotation_previous_public_key_still_verifies_until_removed():
minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
rotated = AsymmetricSessionKeys(
private_key_pem=SecretStr(_RSA_PEM_A),
kid="2026-01",
previous_public_keys=(
SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)),
),
)
opened = open_session_token(token, rotated, NOW)
assert isinstance(opened, OpenedSessionToken)
assert opened.principal == PRINCIPAL
assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature)
def test_rotation_window_still_enforces_expiry_and_tamper_on_the_previous_key():
minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW)
assert isinstance(minted, MintedSessionToken)
token = minted.token.get_secret_value()
rotated = AsymmetricSessionKeys(
private_key_pem=SecretStr(_RSA_PEM_A),
kid="2026-01",
previous_public_keys=(
SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)),
),
)
after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1)
assert isinstance(open_session_token(token, rotated, after), SessionExpired)
tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb")
assert isinstance(open_session_token(tampered, rotated, NOW), SessionBadSignature)
def test_weak_or_garbage_private_key_pem_rejected_at_construction():
with pytest.raises(ValidationError):
AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem(bits=1024)), kid="weak")
with pytest.raises(ValidationError):
AsymmetricSessionKeys(private_key_pem=SecretStr("not a pem"), kid="junk")
with pytest.raises(ValidationError):
SessionRotatedPublicKey(kid="junk", public_key_pem="not a pem")
with pytest.raises(ValidationError):
SessionRotatedPublicKey(kid="private-half", public_key_pem=_RSA_PEM_A)
def test_weak_rotated_public_key_rejected_at_construction():
weak_public = (
serialization.load_pem_private_key(_rsa_private_pem(bits=1024).encode(), password=None)
.public_key()
.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)
.decode()
)
with pytest.raises(ValidationError):
SessionRotatedPublicKey(kid="2024-01", public_key_pem=weak_public)
def test_duplicate_kids_rejected_at_construction():
previous = SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS))
with pytest.raises(ValidationError):
AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2025-06", previous_public_keys=(previous,))
with pytest.raises(ValidationError):
AsymmetricSessionKeys(
private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01", previous_public_keys=(previous, previous)
)
def test_asymmetric_keys_repr_never_leaks_the_private_key():
assert _RSA_PEM_A not in repr(RSA_KEYS)

176
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -1,10 +1,10 @@
import asyncio
import json
import time
from datetime import datetime, timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
@ -15,7 +15,6 @@ from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, Prisma
import litellm
import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module
from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64
from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.health_endpoints._health_endpoints import (
@ -145,7 +144,11 @@ async def test_db_health_transport_error_never_raises(transport_error):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(reason="health_readiness_check")
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check",
timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
lock_timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
@pytest.mark.asyncio
@ -175,7 +178,11 @@ async def test_db_health_transport_error_reconnect_succeeds(transport_error):
result = await _db_health_readiness_check()
assert result["status"] == "connected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(reason="health_readiness_check")
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check",
timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
lock_timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS,
)
assert mock_prisma.health_check.call_count == 2
@ -2276,6 +2283,159 @@ async def test_health_readiness_returns_503_when_db_disconnected():
assert result == {"status": "healthy", "db": "disconnected"}
@pytest.mark.asyncio
async def test_health_readiness_returns_200_when_db_down_and_allow_requests_on_db_unavailable():
"""
Regression test for https://github.com/BerriAI/litellm/issues/34934.
allow_requests_on_db_unavailable keeps the proxy serving through a DB
outage, so the readiness probe must keep the pod in rotation (200) and
report the DB state through the body, not the status code. Otherwise
K8s pulls every replica before the request-layer fail-open can run.
"""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import health_readiness
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope"))
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
response = Response()
with (
patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
result = await health_readiness(response=response)
assert response.status_code == 200
assert result == {"status": "healthy", "db": "disconnected"}
@pytest.mark.asyncio
async def test_health_readiness_details_returns_200_when_db_down_and_allow_requests_on_db_unavailable():
"""
The detailed readiness payload (public via
allow_public_health_readiness_details, or /health/readiness/details)
must honor the same flag so probes pointed at it also stay 200.
"""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import (
_get_health_readiness_details,
)
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope"))
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
response = Response()
with (
patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
),
patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
result = await _get_health_readiness_details(response=response)
assert response.status_code == 200
assert result["db"] == "disconnected"
@pytest.mark.asyncio
async def test_db_health_readiness_check_bounds_hung_health_check():
"""
A connection that hangs mid-failover must not stall the probe past the
kubelet's timeoutSeconds; the DB round-trip is bounded and reported as
disconnected instead.
"""
from litellm.proxy.health_endpoints._health_endpoints import (
_db_health_readiness_check,
)
async def hang():
await asyncio.sleep(60)
mock_prisma = MagicMock()
mock_prisma.health_check = hang
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still down"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
with patch( # test-quality-ok: lowers the module-level probe timeout so the hung-call test finishes fast
"litellm.proxy.health_endpoints._health_endpoints.DB_READINESS_CHECK_TIMEOUT_SECONDS",
0.05,
):
start = time.monotonic()
with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
):
result = await _db_health_readiness_check()
elapsed = time.monotonic() - start
assert result["status"] == "disconnected"
assert elapsed < 5
@pytest.mark.asyncio
async def test_db_health_readiness_check_overall_deadline_bounds_hung_reconnect():
"""
The whole probe-path DB check (initial check + reconnect + re-check,
including reconnect lock waits) runs under one deadline, so a reconnect
that hangs on the lock still returns disconnected within the deadline.
"""
from litellm.proxy.health_endpoints._health_endpoints import (
_db_health_readiness_check,
)
async def hang(**kwargs):
await asyncio.sleep(60)
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=httpx.ConnectError("down"))
mock_prisma.attempt_db_reconnect = hang
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
with patch( # test-quality-ok: lowers the module-level probe timeout so the hung-call test finishes fast
"litellm.proxy.health_endpoints._health_endpoints.DB_READINESS_PROBE_DEADLINE_SECONDS",
0.05,
):
start = time.monotonic()
with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", mock_prisma
):
result = await _db_health_readiness_check()
elapsed = time.monotonic() - start
assert result["status"] == "disconnected"
assert elapsed < 5
@pytest.mark.asyncio
async def test_health_readiness_returns_200_when_db_connected():
"""Happy path: connected DB keeps the legacy 200."""
@ -2746,13 +2906,13 @@ def test_test_model_connection_accepts_image_edit_mode(monkeypatch):
app = FastAPI()
app.include_router(_health_endpoints_module.router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
client = TestClient(app)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam
patch( # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
respx.mock(assert_all_called=True) as respx_mock,
):
respx_mock.post(host="api.openai.com", path="/v1/images/edits").respond(

77
tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py
View file

@ -7,6 +7,9 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import (
SUGGEST_TOOL,
AiPolicySuggester,
@ -234,6 +237,7 @@ class TestAiPolicySuggester:
call_kwargs = mock_acompletion.call_args.kwargs
assert call_kwargs["model"] == "gpt-4o-mini"
assert call_kwargs["temperature"] == 0.2
assert call_kwargs["drop_params"] is True
assert len(call_kwargs["tools"]) == 1
assert call_kwargs["tools"][0]["function"]["name"] == "select_policy_templates"
assert (
@ -242,3 +246,76 @@ class TestAiPolicySuggester:
assert len(call_kwargs["messages"]) == 2
assert call_kwargs["messages"][0]["role"] == "system"
assert call_kwargs["messages"][1]["role"] == "user"
class TestSuggesterRejectsModelsWithoutToolCalling:
@pytest.mark.asyncio
async def test_a_tools_less_model_is_rejected(self, local_model_cost_map):
with pytest.raises(ProxyException) as exc:
await AiPolicySuggester().suggest(
templates=SAMPLE_TEMPLATES,
attack_examples=["Ignore all previous instructions"],
description="Block prompt injection attempts",
model="perplexity/sonar",
)
assert int(exc.value.code) == 400
assert exc.value.param == "model"
assert "tool calling" in exc.value.message
def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map):
supported_params = litellm.get_supported_openai_params(
model="amazon.nova-pro-v1:0",
custom_llm_provider="bedrock",
)
assert supported_params is not None
assert "tools" in supported_params
assert "tool_choice" not in supported_params
class TestSuggesterToleratesAModelThatRefusesItsSamplingParams:
"""The model is operator-supplied, so it can be a reasoning model whose only accepted
temperature is 1. This call pins temperature=0.2 for tool-selection determinism, which such
a model rejects outright: without drop_params litellm raises UnsupportedParamsError and the
whole suggestion fails rather than degrading. Every other internal LLM call in the proxy
already opts in through judge_acompletion; this one was the exception.
"""
@pytest.mark.asyncio
async def test_a_reasoning_model_gets_past_param_mapping(self, monkeypatch, local_model_cost_map):
"""Drives the real entry point with no patching and no network. Which exception escapes is
the discriminator: param mapping runs before any credential check, so UnsupportedParamsError
means the call died on the pinned temperature, while AuthenticationError means it survived
that and got as far as needing a key. Asserting the latter is what the caller observes.
"""
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
with pytest.raises(litellm.AuthenticationError):
await AiPolicySuggester().suggest(
templates=SAMPLE_TEMPLATES,
attack_examples=["My SSN is 123-45-6789"],
description="",
model="gpt-5.6-terra",
)
def test_the_pinned_temperature_is_what_such_a_model_refuses(self, local_model_cost_map):
"""The other half of the discriminator above: the same temperature this call pins is
exactly what the model rejects, and drop_params is what removes it."""
from litellm.utils import get_optional_params
optional_params = get_optional_params(
model="gpt-5.6-terra",
custom_llm_provider="openai",
temperature=0.2,
tools=[SUGGEST_TOOL],
tool_choice={"type": "function", "function": {"name": "select_policy_templates"}},
drop_params=True,
)
assert "temperature" not in optional_params
assert optional_params["tools"] == [SUGGEST_TOOL]
assert optional_params["tool_choice"] == {
"type": "function",
"function": {"name": "select_policy_templates"},
}

53
tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
View file

@ -154,6 +154,59 @@ async def test_model_info_v2_translates_team_model_name(monkeypatch):
assert "model_name_team-abc-123_4a6b8" not in names
@pytest.mark.asyncio
async def test_model_info_v2_exact_model_filter_matches_team_public_name(monkeypatch):
"""`/v2/model/info?model=<public name>` must keep the team-scoped row whose
`model_name` is the internal routing key: the dashboard links team model
chips with the public name, and the exact filter ran before translation."""
global_row = {
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": "normal-id-1", "db_model": False},
}
router = MagicMock()
router.model_list = [_team_row(), global_row]
monkeypatch.setattr(ps, "llm_router", router)
monkeypatch.setattr(ps, "user_model", None)
monkeypatch.setattr(ps, "prisma_client", MagicMock())
monkeypatch.setattr(ps.proxy_config, "get_config", AsyncMock(return_value={}))
monkeypatch.setattr(
ps,
"_apply_search_filter_to_models",
AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))),
)
monkeypatch.setattr(
ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model
)
import litellm.proxy.agent_endpoints.model_list_helpers as mlh
monkeypatch.setattr(
mlh,
"append_agents_to_model_info",
AsyncMock(side_effect=lambda models, **kw: models),
)
admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN)
resp = await ps.model_info_v2(
user_api_key_dict=admin,
model="team-claude-sonnet",
user_models_only=False,
include_team_models=False,
debug=False,
page=1,
size=50,
search=None,
modelId=None,
teamId=None,
sortBy=None,
sortOrder="asc",
)
assert [m["model_name"] for m in resp["data"]] == ["team-claude-sonnet"]
assert resp["total_count"] == 1
@pytest.mark.asyncio
async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch):
"""/v1/model/info list path (no litellm_model_id) must include team-scoped

94
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -1009,11 +1009,10 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
}
for metadata_key in ("metadata", "litellm_metadata"):
cleaned_metadata = updated.get(metadata_key) or {}
for stripped_key in stripped_keys:
assert stripped_key not in cleaned_metadata
assert cleaned_metadata.get("safe_user_metadata") == "kept"
assert "litellm_metadata" not in updated
for stripped_key in stripped_keys:
assert stripped_key not in updated["metadata"]
assert updated["metadata"]["safe_user_metadata"] == "kept"
requester_metadata = updated["metadata"]["requester_metadata"]
for stripped_key in stripped_keys:
@ -1576,10 +1575,7 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o
header.lower()
for header in updated["proxy_server_request"]["body"]["metadata"]["headers"]
}
assert "litellm-disable-message-redaction" in {
header.lower()
for header in (updated.get("litellm_metadata") or {}).get("headers", {})
}
assert "litellm_metadata" not in updated
@pytest.mark.asyncio
@ -6658,9 +6654,9 @@ async def test_add_litellm_data_to_request_strips_caller_supplied_callback_crede
assert "gcs_bucket_name" not in updated
assert updated["dd_api_key"] == "team-dd-key"
assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "team-dd-key"}
for metadata_key in ("metadata", "litellm_metadata"):
assert "dd_site" not in updated[metadata_key]
assert "dd_agent_host" not in updated[metadata_key]
assert "litellm_metadata" not in updated
assert "dd_site" not in updated["metadata"]
assert "dd_agent_host" not in updated["metadata"]
assert "dd_site" not in updated["litellm_params"]["metadata"]
assert updated["metadata"]["safe_user_metadata"] == "kept"
@ -7510,10 +7506,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
version="test-version",
)
for bucket in ("metadata", "litellm_metadata"):
assert "attempted_fallbacks" not in updated[bucket]
assert "original_model_group" not in updated[bucket]
assert updated[bucket]["client_key"] == "client_value"
assert "litellm_metadata" not in updated
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"
@pytest.mark.asyncio
@ -7535,10 +7531,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js
version="test-version",
)
assert isinstance(updated["litellm_metadata"], dict)
assert "attempted_fallbacks" not in updated["litellm_metadata"]
assert "original_model_group" not in updated["litellm_metadata"]
assert updated["litellm_metadata"]["client_key"] == "client_value"
assert "litellm_metadata" not in updated
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
assert updated["metadata"]["client_key"] == "client_value"
@pytest.mark.asyncio
@ -7562,9 +7558,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite
version="test-version",
)
assert updated["litellm_metadata"]["model_info"] == {"input_cost_per_token": 0.0}
assert "attempted_fallbacks" not in updated["litellm_metadata"]
assert "original_model_group" not in updated["litellm_metadata"]
assert "litellm_metadata" not in updated
assert updated["metadata"]["model_info"] == {"input_cost_per_token": 0.0}
assert "attempted_fallbacks" not in updated["metadata"]
assert "original_model_group" not in updated["metadata"]
@pytest.mark.asyncio
@ -7601,7 +7598,8 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
litellm_metadata made the router hand downstream a scrubbed copy, so the proxy's
post_call write-backs (guardrail telemetry, applied guardrails) landed in a dict the
spend row never read. After the boundary strip plus the in-place scrub, the object the
router forwards is the proxy's own request_data bucket."""
router forwards is the proxy's own request_data bucket; on chat routes that bucket is
``metadata``, since the boundary folds client ``litellm_metadata`` into it."""
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
data = {
@ -7617,7 +7615,9 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
general_settings={},
version="test-version",
)
proxy_bucket = request_data["litellm_metadata"]
proxy_bucket = request_data["metadata"]
assert "attempted_fallbacks" not in proxy_bucket
assert "original_model_group" not in proxy_bucket
router = litellm.Router(
model_list=[
{
@ -7630,7 +7630,7 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
original_acompletion = router._acompletion
async def _spy(*args, **spy_kwargs):
forwarded_buckets.append(spy_kwargs["litellm_metadata"])
forwarded_buckets.append(spy_kwargs["metadata"])
return await original_acompletion(*args, **spy_kwargs)
router._acompletion = _spy
@ -7639,7 +7639,45 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
assert forwarded_buckets == [proxy_bucket]
assert forwarded_buckets[0] is proxy_bucket
assert "attempted_fallbacks" not in proxy_bucket
assert "original_model_group" not in proxy_bucket
assert proxy_bucket["attempted_fallbacks"] == 0
assert proxy_bucket.get("original_model_group") != "spoofed-group"
proxy_bucket["standard_logging_guardrail_information"] = [{"guardrail_name": "postcall-guard"}]
assert forwarded_buckets[0]["standard_logging_guardrail_information"] == [{"guardrail_name": "postcall-guard"}]
@pytest.mark.asyncio
async def test_add_litellm_data_to_request_folds_litellm_metadata_into_metadata_on_chat_routes():
data = {
"model": "gpt-3.5-turbo",
"metadata": {"tags": ["from-metadata"]},
"litellm_metadata": {"trace_id": "abc", "tags": ["from-litellm-metadata"]},
}
updated = await add_litellm_data_to_request(
data=data,
request=_make_chat_request_mock(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert "litellm_metadata" not in updated
assert updated["metadata"]["trace_id"] == "abc"
assert updated["metadata"]["tags"] == ["from-metadata", "from-litellm-metadata"]
@pytest.mark.asyncio
async def test_add_litellm_data_to_request_keeps_litellm_metadata_on_litellm_metadata_routes():
data = {"model": "claude-sonnet-5", "litellm_metadata": {"trace_id": "abc"}}
updated = await add_litellm_data_to_request(
data=data,
request=_make_request_mock("/v1/messages", {"Content-Type": "application/json"}),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
assert updated["litellm_metadata"]["trace_id"] == "abc"

9
tests/test_litellm/proxy/test_pricing_field_strip.py
View file

@ -285,8 +285,8 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in():
async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata():
"""``litellm_metadata`` may arrive as a JSON-encoded string (multipart/
form-data or ``extra_body``). The strip has to run after the proxy parses
it into a dict; otherwise the ``isinstance(dict)`` guard skips the field
and ``model_info`` survives the strip via the string path.
it into a dict but before the chat-route fold into ``metadata``; otherwise
``model_info`` survives via the string path and lands in the folded bucket.
"""
import json
@ -305,9 +305,8 @@ async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata()
version="test-version",
)
parsed_metadata = updated.get("litellm_metadata")
assert isinstance(parsed_metadata, dict)
assert "model_info" not in parsed_metadata
assert "litellm_metadata" not in updated
assert "model_info" not in updated["metadata"]
@pytest.mark.asyncio

47
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -2126,6 +2126,53 @@ async def test_apply_search_filter_bounds_db_fetch_by_page_and_cap():
assert take < 10_000, "sorted search must cap below the full match set"
@pytest.mark.asyncio
async def test_apply_search_filter_honours_exact_model_name_in_db_query():
"""
`/v2/model/info?model=<group>&search=<term>`: the router list is already
narrowed to the exact group, so the DB count and fetch must be too, or
other groups' rows leak into the page and inflate total_count.
"""
from litellm.proxy.proxy_server import _apply_search_filter_to_models
prisma_client = MagicMock()
prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=0)
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
proxy_config = MagicMock()
proxy_config.decrypt_model_list_from_db = lambda rows: []
await _apply_search_filter_to_models(
all_models=[],
search="sonnet",
prisma_client=prisma_client,
proxy_config=proxy_config,
model_name="anthropic-sonnet-5",
)
where = prisma_client.db.litellm_proxymodeltable.count.call_args.kwargs["where"]
assert where["model_name"] == "anthropic-sonnet-5"
assert prisma_client.db.litellm_proxymodeltable.find_many.call_args.kwargs["where"] == where
prisma_client.db.litellm_proxymodeltable.count.reset_mock()
_, total_count = await _apply_search_filter_to_models(
all_models=[],
search="opus",
prisma_client=prisma_client,
proxy_config=proxy_config,
model_name="anthropic-sonnet-5",
)
prisma_client.db.litellm_proxymodeltable.count.assert_not_called()
assert total_count == 0
await _apply_search_filter_to_models(
all_models=[],
search="sonnet",
prisma_client=prisma_client,
proxy_config=proxy_config,
)
where = prisma_client.db.litellm_proxymodeltable.count.call_args.kwargs["where"]
assert where["model_name"] == {"contains": "sonnet", "mode": "insensitive"}
@pytest.mark.asyncio
async def test_filter_models_by_team_id_excludes_viewer_direct_access():
"""

63
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py
View file

@ -89,9 +89,7 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails(
prisma_client._cleanup_engine_watcher = MagicMock()
writer = MagicMock()
writer.query_raw = AsyncMock(
side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]
)
writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]])
monkeypatch.setattr(
PrismaClient,
"writer_db",
@ -171,9 +169,7 @@ async def test_run_reconnect_cycle_passes_writer_generation_to_recreate(
writer = MagicMock()
writer._engine_generation = 7
writer.query_raw = AsyncMock(
side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]
)
writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]])
monkeypatch.setattr(
PrismaClient,
"writer_db",
@ -229,9 +225,7 @@ async def test_attempt_reconnect_inside_lock_runs_cycle_and_resets_counter(
prisma_client._consecutive_reconnect_failures = 2
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=1
)
ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="test", timeout_seconds=1)
pinned = {
"returned": ok,
"cycle_called": prisma_client._run_reconnect_cycle.await_count,
@ -254,9 +248,7 @@ async def test_attempt_reconnect_inside_lock_skips_when_in_cooldown(
prisma_client._db_last_reconnect_attempt_ts = time.time()
prisma_client._run_reconnect_cycle = AsyncMock()
ok = await prisma_client._attempt_reconnect_inside_lock(
force=False, reason="test", timeout_seconds=1
)
ok = await prisma_client._attempt_reconnect_inside_lock(force=False, reason="test", timeout_seconds=1)
assert ok is False
assert prisma_client._run_reconnect_cycle.await_count == 0
@ -269,9 +261,7 @@ async def test_attempt_reconnect_inside_lock_increments_failure_counter_on_error
prisma_client._consecutive_reconnect_failures = 0
prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("boom"))
ok = await prisma_client._attempt_reconnect_inside_lock(
force=True, reason="failing_test", timeout_seconds=1
)
ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="failing_test", timeout_seconds=1)
assert ok is False
assert prisma_client._consecutive_reconnect_failures == 1
@ -316,9 +306,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false(
by replacing ``asyncio.wait`` with a callable that returns the loser
task as still-pending after it's already been completed elsewhere.
"""
completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task(
_no_op_returning_true()
)
completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task(_no_op_returning_true())
# Ensure the inner task has finished before attempt_db_reconnect sees it.
await completed_task
@ -329,7 +317,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false(
monkeypatch.setattr(
asyncio,
"create_task",
lambda coro, *a, **kw: (coro.close() or completed_task),
lambda coro, *a, **kw: coro.close() or completed_task,
)
prisma_client._db_last_reconnect_attempt_ts = 0.0
@ -465,9 +453,7 @@ async def test_db_health_watchdog_loop_triggers_reconnect_on_timeout(
await prisma_client._db_health_watchdog_loop()
pinned = {
"reconnect_called": prisma_client.attempt_db_reconnect.await_count,
"reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs[
"reason"
],
"reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs["reason"],
"wait_for_calls": call_count["n"],
"loop_exited_clean": True,
}
@ -522,10 +508,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once(
from litellm.proxy.db.prisma_client import PrismaWrapper
def token_db_url(created: datetime) -> str:
token = (
f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}"
f"&X-Amz-Expires=900&X-Amz-Signature=abc"
)
token = f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}&X-Amz-Expires=900&X-Amz-Signature=abc"
return f"postgresql://user:{urllib.parse.quote(token, safe='')}@host:5432/db"
# Old engine (PID 111) carries an expired token; in-flight queries on it
@ -577,9 +560,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once(
# In-flight transport-error path fires while the refresh holds the
# wrapper's reconnection lock mid-recreate.
reconnect_task = asyncio.create_task(
prisma_client.attempt_db_reconnect(
reason="in_flight_transport_error", force=True
)
prisma_client.attempt_db_reconnect(reason="in_flight_transport_error", force=True)
)
await asyncio.sleep(0.05)
release_connect.set()
@ -1096,3 +1077,27 @@ async def test_unrelated_reconnect_failure_does_not_erase_the_burst_record(
"cycles_after": prisma_client._run_reconnect_cycle.await_count,
}
assert pinned == {"cycles_before": 2, "cycles_after": 2}
@pytest.mark.asyncio
async def test_attempt_db_reconnect_cancelled_while_waiting_does_not_strand_lock(
prisma_client: PrismaClient,
) -> None:
"""A reconnect cancelled while waiting on the lock (e.g. the readiness
probe deadline firing) must abandon its lock-acquisition task instead of
leaving it to grab the lock later with no owner to release it."""
prisma_client._db_last_reconnect_attempt_ts = 0.0
prisma_client._attempt_reconnect_inside_lock = AsyncMock(return_value=True)
await prisma_client._db_reconnect_lock.acquire()
waiting_reconnect: Final = asyncio.create_task(
prisma_client.attempt_db_reconnect(reason="probe_deadline", lock_timeout_seconds=30.0)
)
await asyncio.sleep(0.05)
waiting_reconnect.cancel()
with pytest.raises(asyncio.CancelledError):
await waiting_reconnect
prisma_client._db_reconnect_lock.release()
await asyncio.sleep(0.05)
assert prisma_client._db_reconnect_lock.locked() is False

35
tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -9,11 +9,14 @@ import pytest
from fastapi import HTTPException
import litellm
from litellm.caching.caching import DualCache
from litellm.exceptions import RejectedRequestError
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypesLiteral
def _load(module: str, name: str):
@ -473,7 +476,13 @@ class _RedactingGuardrail(CustomGuardrail):
kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call)
super().__init__(guardrail_name="redactor", **kwargs)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override]
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> dict | None:
for msg in data.get("messages", []):
if "SECRET" in msg.get("content", ""):
msg["content"] = msg["content"].replace("SECRET", "[REDACTED]")
@ -488,7 +497,13 @@ class _BlockOnSecretGuardrail(CustomGuardrail):
kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call)
super().__init__(guardrail_name="blocker", **kwargs)
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override]
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> dict | None:
if any("SECRET" in msg.get("content", "") for msg in data.get("messages", [])):
raise HTTPException(status_code=400, detail="blocked: SECRET detected")
return None
@ -560,7 +575,13 @@ async def test_scan_raw_request_guardrail_does_not_undo_later_masking(
separate marker (PII_TOKEN) that only the redactor reacts to."""
class _PiiRedactor(_RedactingGuardrail):
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override]
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> dict | None:
for msg in data.get("messages", []):
if "PII_TOKEN" in msg.get("content", ""):
msg["content"] = msg["content"].replace("PII_TOKEN", "[REDACTED]")
@ -692,7 +713,13 @@ async def test_scan_raw_request_warns_when_guardrail_mutation_discarded(
super().__init__(**kwargs)
self.scan_raw_request = True
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override]
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> dict | None:
for msg in data.get("messages", []):
msg["content"] = msg["content"].replace("SECRET", "[REDACTED]")
return data

75
tests/test_litellm/router_strategy/test_router_tag_routing.py
View file

@ -3072,3 +3072,78 @@ async def test_non_router_tags_still_pick_the_matching_tier_deployment():
)
assert response._hidden_params["model_id"] == "tier-gemini-flash-us"
def _chat_completions_request_mock():
from unittest.mock import MagicMock
from fastapi import Request
request_mock = MagicMock(spec=Request)
request_mock.url = MagicMock()
request_mock.url.path = "/v1/chat/completions"
request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions"
request_mock.method = "POST"
request_mock.query_params = {}
request_mock.headers = {"Content-Type": "application/json"}
request_mock.client = MagicMock()
request_mock.client.host = "127.0.0.1"
return request_mock
def _team_a_and_default_router():
return litellm.Router(
model_list=[
{
"model_name": "gpt-5.4-mini",
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock", "tags": ["team-a"]},
"model_info": {"id": "team-a-deployment"},
},
{
"model_name": "gpt-5.4-mini",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "mock", "tags": ["default"]},
"model_info": {"id": "default-deployment"},
},
],
enable_tag_filtering=True,
)
@pytest.mark.asyncio()
@pytest.mark.parametrize(
"team_metadata,body_extra",
[
({"tags": ["team-a"]}, {}),
({}, {"tags": ["team-a"]}),
],
ids=["team-tags", "body-tags"],
)
async def test_chat_request_carrying_litellm_metadata_still_routes_on_proxy_merged_tags(team_metadata, body_extra):
from unittest.mock import MagicMock
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
router = _team_a_and_default_router()
data = {
"model": "gpt-5.4-mini",
"messages": [{"role": "user", "content": "hi"}],
"litellm_metadata": {"trace_id": "abc"},
**body_extra,
}
request_kwargs = await add_litellm_data_to_request(
data=data,
request=_chat_completions_request_mock(),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata=team_metadata),
proxy_config=MagicMock(),
general_settings={},
version="test-version",
)
deployment = await router.async_get_available_deployment(
model="gpt-5.4-mini",
request_kwargs=request_kwargs,
messages=request_kwargs["messages"],
)
assert deployment["model_info"]["id"] == "team-a-deployment"

2
ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts
View file

@ -118,6 +118,7 @@ describe("useModelsInfo", () => {
// exclude_auto_routers defaults off: only the Models + Endpoints table opts in, so
// every other consumer of this hook keeps seeing auto-routers.
false,
undefined,
);
expect(modelInfoCall).toHaveBeenCalledTimes(1);
});
@ -145,6 +146,7 @@ describe("useModelsInfo", () => {
// exclude_auto_routers defaults off: only the Models + Endpoints table opts in, so
// every other consumer of this hook keeps seeing auto-routers.
false,
undefined,
);
});

3
ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts
View file

@ -38,6 +38,7 @@ export const useModelsInfo = (
sortBy?: string,
sortOrder?: string,
excludeAutoRouters: boolean = false,
modelName?: string,
) => {
const { accessToken, userId, userRole } = useAuthorized();
return useQuery<PaginatedModelInfoResponse>({
@ -48,6 +49,7 @@ export const useModelsInfo = (
page,
size,
...(search && { search }),
...(modelName && { modelName }),
...(modelId && { modelId }),
...(teamId && { teamId }),
...(sortBy && { sortBy }),
@ -70,6 +72,7 @@ export const useModelsInfo = (
sortBy,
sortOrder,
excludeAutoRouters,
modelName,
),
enabled: Boolean(accessToken && userId && userRole),
});

29
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx
View file

@ -33,6 +33,7 @@ interface ModelsInfoArgs {
teamId?: string;
sortBy?: string;
sortOrder?: string;
modelName?: string;
}
const modelsInfoCalls: ModelsInfoArgs[] = [];
@ -47,12 +48,14 @@ type UseModelsInfoArgs = [
teamId?: string,
sortBy?: string,
sortOrder?: string,
excludeAutoRouters?: boolean,
modelName?: string,
];
vi.mock("../../hooks/models/useModels", () => ({
useModelsInfo: (...args: UseModelsInfoArgs) => {
const [page, size, search, , teamId, sortBy, sortOrder] = args;
const call: ModelsInfoArgs = { page, size, search, teamId, sortBy, sortOrder };
const [page, size, search, , teamId, sortBy, sortOrder, , modelName] = args;
const call: ModelsInfoArgs = { page, size, search, teamId, sortBy, sortOrder, modelName };
modelsInfoCalls.push(call);
return { ...modelsInfoResult, refetch: mockRefetch };
},
@ -260,6 +263,28 @@ describe("AllModelsTab", () => {
expect(within(table).queryByText("gpt-4")).not.toBeInTheDocument();
});
it("asks the server for the exact selected model group so deployments beyond the first page are found", () => {
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
expect(lastModelsInfoCall().modelName).toBe("claude-opus");
expect(lastModelsInfoCall().search).toBeUndefined();
});
it.each(["all", "wildcard"])("sends no exact model name for the %s pseudo group", (group) => {
render(<AllModelsTab {...defaultProps} selectedModelGroup={group} />);
expect(lastModelsInfoCall().modelName).toBeUndefined();
});
it("keeps the exact model group alongside a typed search", async () => {
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
fireEvent.change(screen.getByPlaceholderText("Search model names…"), { target: { value: "opus" } });
await waitFor(() => expect(lastModelsInfoCall().search).toBe("opus"));
expect(lastModelsInfoCall().modelName).toBe("claude-opus");
});
it("resets search, filters, team and sorting from the drawer reset button", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} selectedModelGroup="gpt-4" />);

6
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
View file

@ -81,6 +81,11 @@ const AllModelsTab = ({
}, [modelNameSearch, debouncedUpdateSearch]);
const teamIdForQuery = selectedTeamValue === PERSONAL_TEAM_VALUE ? undefined : selectedTeamValue;
const isConcreteModelGroup =
Boolean(selectedModelGroup) &&
selectedModelGroup !== ALL_MODEL_GROUPS_VALUE &&
selectedModelGroup !== WILDCARD_MODEL_GROUP_VALUE;
const modelNameForQuery = isConcreteModelGroup ? selectedModelGroup ?? undefined : undefined;
const sortBy = useMemo(() => {
if (sorting.length === 0) return undefined;
@ -108,6 +113,7 @@ const AllModelsTab = ({
// Auto-routers are routing constructs, not deployments; the sibling Auto-Routers tab
// lists and manages them. Excluded server-side so total_count stays honest.
true,
modelNameForQuery,
);
const isLoading = isLoadingModelsInfo || isLoadingModelCostMap;

28
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts
View file

@ -1,7 +1,7 @@
import { act, renderHook, waitFor } from "@testing-library/react";
import { withNuqsTestingAdapter, type UrlUpdateEvent } from "nuqs/adapters/testing";
import { describe, expect, it, vi } from "vitest";
import { useModelDetailRouting } from "./detailNavigation";
import { useModelDetailRouting, useModelGroupFilterRouting } from "./detailNavigation";
describe("useModelDetailRouting", () => {
it("openModel sets ?model= with a history push", async () => {
@ -54,3 +54,29 @@ describe("useModelDetailRouting", () => {
expect(result.current.teamId).toBeNull();
});
});
describe("useModelGroupFilterRouting", () => {
it("reads the selected group from ?model_group=", () => {
const { result } = renderHook(() => useModelGroupFilterRouting(), {
wrapper: withNuqsTestingAdapter({ searchParams: "?model_group=gpt-4.1" }),
});
expect(result.current.modelGroup).toBe("gpt-4.1");
});
it("writes the selected group to ?model_group= and clears it on null", async () => {
const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>();
const { result } = renderHook(() => useModelGroupFilterRouting(), {
wrapper: withNuqsTestingAdapter({ onUrlUpdate }),
});
await act(async () => {
result.current.setModelGroup("claude-sonnet-5");
});
await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled());
expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("model_group")).toBe("claude-sonnet-5");
await act(async () => {
result.current.setModelGroup(null);
});
await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.has("model_group")).toBe(false));
});
});

20
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts
View file

@ -1,4 +1,4 @@
import { parseAsString, useQueryStates } from "nuqs";
import { parseAsString, useQueryState, useQueryStates } from "nuqs";
import { useCallback } from "react";
export interface ModelDetailRouting {
@ -41,3 +41,21 @@ export function useModelDetailRouting(): ModelDetailRouting {
close,
};
}
export interface ModelGroupFilterRouting {
modelGroup: string | null;
setModelGroup: (modelGroup: string | null) => void;
}
export function useModelGroupFilterRouting(): ModelGroupFilterRouting {
const [modelGroup, setParam] = useQueryState("model_group", parseAsString);
const setModelGroup = useCallback(
(next: string | null) => {
void setParam(next);
},
[setParam],
);
return { modelGroup, setModelGroup };
}

13
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx
View file

@ -1,19 +1,22 @@
"use client";
import { useState } from "react";
import AllModelsTab from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTab";
import { ALL_MODEL_GROUPS_VALUE } from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTable";
import { useModelDashboardData } from "@/app/(dashboard)/models-and-endpoints/useModelDashboardData";
import { useModelDetailRouting } from "@/app/(dashboard)/models-and-endpoints/detailNavigation";
import {
useModelDetailRouting,
useModelGroupFilterRouting,
} from "@/app/(dashboard)/models-and-endpoints/detailNavigation";
export default function AllModelsPanel() {
const [selectedModelGroup, setSelectedModelGroup] = useState<string | null>(null);
const { modelGroup, setModelGroup } = useModelGroupFilterRouting();
const { availableModelGroups, availableModelAccessGroups } = useModelDashboardData();
const { openModel, openTeam } = useModelDetailRouting();
return (
<AllModelsTab
selectedModelGroup={selectedModelGroup}
setSelectedModelGroup={setSelectedModelGroup}
selectedModelGroup={modelGroup}
setSelectedModelGroup={(group) => setModelGroup(group === ALL_MODEL_GROUPS_VALUE ? null : group)}
availableModelGroups={availableModelGroups}
availableModelAccessGroups={availableModelAccessGroups}
setSelectedModelId={openModel}

7
ui/litellm-dashboard/src/components/CodeBlock.tsx
View file

@ -20,10 +20,10 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => {
};
return (
<div className="relative rounded-lg border border-border overflow-hidden">
<div className="relative rounded-lg border border-border bg-muted overflow-hidden">
<button
onClick={copyToClipboard}
className="absolute top-3 right-3 p-2 rounded-md bg-muted hover:bg-accent text-muted-foreground z-raised"
className="absolute top-3 right-3 p-2 rounded-md border border-border bg-background text-muted-foreground hover:bg-accent hover:text-foreground z-raised"
aria-label="Copy code"
>
{copied ? <CheckIcon size={16} /> : <ClipboardIcon size={16} />}
@ -36,8 +36,9 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => {
padding: "1.5rem",
borderRadius: "0.5rem",
fontSize: "0.9rem",
backgroundColor: "#fafafa",
background: "transparent",
}}
codeTagProps={{ style: { background: "transparent" } }}
showLineNumbers
>
{code}

39
ui/litellm-dashboard/src/components/networking.test.ts
View file

@ -104,6 +104,45 @@ describe("loginCall - storeLoginToken integration", () => {
});
});
describe("modelInfoCall", () => {
let currentFetch: typeof global.fetch;
beforeEach(() => {
currentFetch = global.fetch;
});
afterEach(() => {
global.fetch = currentFetch;
});
it("sends the exact model name as the model query param and leaves search alone", async () => {
const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue({ data: [] }) } as any);
global.fetch = mockFetch as any;
await Networking.modelInfoCall(
"token",
"user",
"Admin",
2,
25,
undefined,
undefined,
undefined,
undefined,
undefined,
true,
"gpt-4",
);
const parsed = new URL(mockFetch.mock.calls[0][0] as string, "http://example.com");
expect(parsed.pathname).toBe("/v2/model/info");
expect(parsed.searchParams.get("model")).toBe("gpt-4");
expect(parsed.searchParams.has("search")).toBe(false);
expect(parsed.searchParams.get("page")).toBe("2");
expect(parsed.searchParams.get("exclude_auto_routers")).toBe("true");
});
});
describe("daily activity helpers", () => {
const startTime = new Date("2025-02-12T00:00:00.000Z");
const endTime = new Date("2025-02-19T00:00:00.000Z");

4
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -1677,6 +1677,7 @@ export const modelInfoCall = async (
sortBy?: string,
sortOrder?: string,
excludeAutoRouters?: boolean,
modelName?: string,
) => {
/**
* Get all models on proxy
@ -1690,6 +1691,9 @@ export const modelInfoCall = async (
if (search && search.trim()) {
params.append("search", search.trim());
}
if (modelName && modelName.trim()) {
params.append("model", modelName.trim());
}
if (modelId && modelId.trim()) {
params.append("modelId", modelId.trim());
}

20
ui/litellm-dashboard/src/components/shared/table_cells/status_badge.test.tsx
View file

@ -1,9 +1,12 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, expect, it } from "vitest";
import { describe, expect, it, vi } from "vitest";
import { StatusBadge, type StatusTone } from "./status_badge";
const push = vi.fn();
vi.mock("next/navigation", () => ({ useRouter: () => ({ push }) }));
describe("StatusBadge", () => {
const toneClasses: Record<StatusTone, string[]> = {
success: ["border-success/20", "bg-success/10", "text-success"],
@ -39,4 +42,19 @@ describe("StatusBadge", () => {
await user.hover(screen.getByText("Blocked"));
expect(await screen.findByText("This key was blocked by SCIM")).toBeInTheDocument();
});
it("renders a tinted anchor that navigates client-side when href is given", async () => {
const user = userEvent.setup();
render(<StatusBadge tone="info" label="gpt-4.1" href="/models-and-endpoints?model_group=gpt-4.1" />);
const link = screen.getByRole("link", { name: "gpt-4.1" });
expect(link).toHaveAttribute("href", "/models-and-endpoints?model_group=gpt-4.1");
expect(link).toHaveClass("text-info");
await user.click(link);
expect(push).toHaveBeenCalledWith("/models-and-endpoints?model_group=gpt-4.1");
});
it("renders no anchor without an href", () => {
render(<StatusBadge tone="info" label="gpt-4.1" />);
expect(screen.queryByRole("link")).not.toBeInTheDocument();
});
});

39
ui/litellm-dashboard/src/components/shared/table_cells/status_badge.tsx
View file

@ -2,6 +2,7 @@
import * as React from "react";
import { useEntityLinkClick } from "@/components/shared/EntityLink";
import { Badge } from "@/components/ui/badge";
import { cn } from "@/lib/cva.config";
@ -23,15 +24,17 @@ interface StatusBadgeProps {
tooltip?: React.ReactNode;
dataTestId?: string;
className?: string;
href?: string;
}
export function StatusBadge({ tone, label, tooltip, dataTestId, className }: StatusBadgeProps) {
const badge = (
<Badge
variant="outline"
data-testid={dataTestId}
className={cn("whitespace-nowrap font-normal", TONE_CLASS[tone], className)}
>
export function StatusBadge({ tone, label, tooltip, dataTestId, className, href }: StatusBadgeProps) {
const badgeClassName = cn("whitespace-nowrap font-normal", TONE_CLASS[tone], className);
const badge = href ? (
<LinkedStatusBadge href={href} dataTestId={dataTestId} className={badgeClassName}>
{label}
</LinkedStatusBadge>
) : (
<Badge variant="outline" data-testid={dataTestId} className={badgeClassName}>
{label}
</Badge>
);
@ -41,3 +44,25 @@ export function StatusBadge({ tone, label, tooltip, dataTestId, className }: Sta
}
return <CellTooltip content={tooltip} trigger={badge} />;
}
interface LinkedStatusBadgeProps {
href: string;
dataTestId?: string;
className: string;
children: string;
}
function LinkedStatusBadge({ href, dataTestId, className, children }: LinkedStatusBadgeProps) {
const handleClick = useEntityLinkClick(href);
return (
<Badge
variant="outline"
data-testid={dataTestId}
className={cn("cursor-pointer hover:underline", className)}
render={<a href={href} onClick={handleClick} />}
>
{children}
</Badge>
);
}

33
ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
View file

@ -21,7 +21,10 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
}),
}));
vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) }));
vi.mock("@/components/networking", () => ({
serverRootPath: "",
teamInfoCall: vi.fn(),
teamMemberDeleteCall: vi.fn(),
teamMemberAddCall: vi.fn(),
@ -278,6 +281,36 @@ describe("TeamInfoView", () => {
});
});
it("links direct and access-group model badges to the models page filtered to that group", async () => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
models: ["gpt-4.1"],
access_group_models: ["claude-sonnet-5"],
access_group_details: [{ access_group_id: "ag-1", access_group_name: "prod", models: ["claude-sonnet-5"] }],
}),
);
renderWithProviders(<TeamInfoView {...defaultProps} />);
expect(await screen.findByRole("link", { name: "gpt-4.1" })).toHaveAttribute(
"href",
expect.stringContaining("/models-and-endpoints?model_group=gpt-4.1"),
);
expect(screen.getByRole("link", { name: "claude-sonnet-5" })).toHaveAttribute(
"href",
expect.stringContaining("/models-and-endpoints?model_group=claude-sonnet-5"),
);
});
it("keeps the all-proxy-models badge non-clickable", async () => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["all-proxy-models"] }));
renderWithProviders(<TeamInfoView {...defaultProps} />);
expect(await screen.findByText("All proxy models")).toBeInTheDocument();
expect(screen.queryByRole("link", { name: "All proxy models" })).not.toBeInTheDocument();
});
it("should display loading state while fetching team data", () => {
vi.mocked(networking.teamInfoCall).mockImplementation(() => new Promise(() => {}));

20
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
View file

@ -22,7 +22,9 @@ import type { ObjectPermission } from "@/components/object_permission_types";
import { isProxyAdminRole } from "@/utils/roles";
import { ArrowLeftIcon } from "@heroicons/react/outline";
import { StatusBadge, type StatusTone } from "@/components/shared/table_cells/status_badge";
import { BadgeLink } from "@/components/shared/BadgeLink";
import { Badge } from "@/components/ui/badge";
import { modelGroupHref } from "@/utils/entityLinks";
import { Card } from "@/components/ui/card";
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
import { Input as UIInput } from "@/components/ui/input";
@ -53,6 +55,7 @@ import {
computeTeamModelBadges,
normalizeTeamModelSelection,
TeamAccessGroupModelGrant,
TeamModelBadge,
TeamModelBadgeKind,
} from "./teamModelAccess";
import MetadataKeyValueFields, {
@ -111,6 +114,9 @@ const TEAM_MODEL_BADGE_TONES: Record<TeamModelBadgeKind, StatusTone> = {
"access-group": "success",
};
const teamModelBadgeHref = (badge: TeamModelBadge): string | undefined =>
badge.kind === "direct" || badge.kind === "access-group" ? modelGroupHref(badge.label) : undefined;
export interface TeamMembership {
user_id: string;
team_id: string;
@ -1006,7 +1012,11 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
(badge, index) => (
<SimpleTooltip key={`${badge.kind}-${badge.label}-${index}`} content={badge.tooltip}>
<span>
<StatusBadge tone={TEAM_MODEL_BADGE_TONES[badge.kind]} label={badge.label} />
<StatusBadge
tone={TEAM_MODEL_BADGE_TONES[badge.kind]}
label={badge.label}
href={teamModelBadgeHref(badge)}
/>
</span>
</SimpleTooltip>
),
@ -1727,9 +1737,9 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<p className="font-medium">Models</p>
<div className="flex flex-wrap gap-2 mt-1">
{info.models.map((model, index) => (
<Badge key={index} variant="secondary">
<BadgeLink key={index} href={modelGroupHref(model)}>
{model}
</Badge>
</BadgeLink>
))}
</div>
</div>
@ -1738,9 +1748,9 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<p className="font-medium">Default Member Models</p>
<div className="flex flex-wrap gap-2 mt-1">
{info.default_team_member_models.map((model, index) => (
<Badge key={index} variant="secondary">
<BadgeLink key={index} href={modelGroupHref(model)}>
{model}
</Badge>
</BadgeLink>
))}
</div>
</div>

26
ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx
View file

@ -561,6 +561,32 @@ describe("KeyInfoView", () => {
);
});
it("links each model chip to the models page filtered to that model group", async () => {
const keyData = { ...MOCK_KEY_DATA, models: ["gpt-4.1", "anthropic/*"] };
renderWithProviders(
<KeyInfoView keyData={keyData} onClose={() => {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />,
);
expect(await screen.findByRole("link", { name: "gpt-4.1" })).toHaveAttribute(
"href",
expect.stringContaining("/models-and-endpoints?model_group=gpt-4.1"),
);
expect(screen.getByRole("link", { name: "anthropic/*" })).toHaveAttribute(
"href",
expect.stringContaining("/models-and-endpoints?model_group=anthropic%2F*"),
);
});
it("keeps the all-proxy-models grant chip non-clickable", async () => {
const keyData = { ...MOCK_KEY_DATA, models: ["all-proxy-models"] };
renderWithProviders(
<KeyInfoView keyData={keyData} onClose={() => {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />,
);
expect((await screen.findAllByText("all-proxy-models")).length).toBeGreaterThan(0);
expect(screen.queryByRole("link", { name: "all-proxy-models" })).not.toBeInTheDocument();
});
it("renders no team link when the key has no team", async () => {
renderWithProviders(
<KeyInfoView

11
ui/litellm-dashboard/src/components/templates/key_info_view.tsx
View file

@ -12,7 +12,8 @@ import { Card } from "@/components/ui/card";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import { EntityLink } from "@/components/shared/EntityLink";
import { teamDetailHref } from "@/utils/entityLinks";
import { modelGroupHref, teamDetailHref } from "@/utils/entityLinks";
import { BadgeLink } from "@/components/shared/BadgeLink";
import { KeyInfoHeader } from "./KeyInfoHeader";
import KeySavingsTab from "./KeySavingsTab";
import { useEffect, useState } from "react";
@ -660,9 +661,9 @@ export default function KeyInfoView({
<div className="mt-2 flex flex-wrap gap-2">
{currentKeyData.models && currentKeyData.models.length > 0 ? (
currentKeyData.models.map((model, index) => (
<Badge key={index} variant="secondary" className="min-w-0 break-words">
<BadgeLink key={index} href={modelGroupHref(model)} className="min-w-0 break-words">
{model}
</Badge>
</BadgeLink>
))
) : (
<p className="text-sm">No models specified</p>
@ -996,9 +997,9 @@ export default function KeyInfoView({
<div className="flex flex-wrap gap-2 mt-1">
{currentKeyData.models && currentKeyData.models.length > 0 ? (
currentKeyData.models.map((model, index) => (
<span key={index} className="px-2 py-1 bg-info/15 rounded-sm text-xs">
<BadgeLink key={index} href={modelGroupHref(model)} className="min-w-0 break-words">
{model}
</span>
</BadgeLink>
))
) : (
<p className="text-sm">No models specified</p>

32
ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx
View file

@ -1,17 +1,26 @@
import { render, screen } from "@testing-library/react";
import { ThemeProvider } from "next-themes";
import { darkStyles, defaultStyles } from "react-json-view-lite";
import { describe, expect, it } from "vitest";
import { JsonViewer } from "./JsonViewer";
const renderWithTheme = (theme: "light" | "dark", data: unknown) =>
render(
<ThemeProvider attribute="class" defaultTheme={theme} enableSystem={false}>
<JsonViewer data={data} mode="formatted" />
</ThemeProvider>,
);
describe("JsonViewer", () => {
it("should render a placeholder and no tree when the log entry carries no payload", () => {
render(<JsonViewer data={null} mode="formatted" />);
renderWithTheme("light", null);
expect(screen.getByText("No data")).toBeInTheDocument();
expect(screen.queryByRole("tree")).not.toBeInTheDocument();
});
it("should render the payload as a tree exposing its keys", () => {
render(<JsonViewer data={{ model: "claude-opus-4-5", stream: true }} mode="formatted" />);
renderWithTheme("light", { model: "claude-opus-4-5", stream: true });
expect(screen.getByRole("tree")).toBeInTheDocument();
expect(screen.getByText(/model/)).toBeInTheDocument();
@ -20,9 +29,26 @@ describe("JsonViewer", () => {
});
it("should treat an empty payload as data rather than showing the placeholder", () => {
render(<JsonViewer data={{}} mode="formatted" />);
renderWithTheme("light", {});
expect(screen.getByRole("tree")).toBeInTheDocument();
expect(screen.queryByText("No data")).not.toBeInTheDocument();
});
it("should style the tree with the light palette when the dashboard theme is light", () => {
renderWithTheme("light", { model: "claude-opus-4-5" });
expect(screen.getByRole("tree")).toHaveClass(...defaultStyles.container.split(" "));
});
it("should style the tree with the dark palette when the dashboard theme is dark", () => {
renderWithTheme("dark", { model: "claude-opus-4-5" });
const tree = screen.getByRole("tree");
expect(tree).toHaveClass(...darkStyles.container.split(" "));
defaultStyles.container
.split(" ")
.filter((className) => !darkStyles.container.split(" ").includes(className))
.forEach((lightOnlyClassName) => expect(tree).not.toHaveClass(lightOnlyClassName));
});
});

9
ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx
View file

@ -1,4 +1,5 @@
import { JsonView, defaultStyles } from "react-json-view-lite";
import { useTheme } from "next-themes";
import { JsonView, darkStyles, defaultStyles } from "react-json-view-lite";
import "react-json-view-lite/dist/index.css";
import { JSON_MAX_HEIGHT, SPACING_LARGE } from "./constants";
@ -12,6 +13,8 @@ interface JsonViewerProps {
* Uses an interactive tree component for easy navigation.
*/
export function JsonViewer({ data }: JsonViewerProps) {
const { resolvedTheme } = useTheme();
if (!data) return <span className="text-muted-foreground">No data</span>;
return (
@ -24,8 +27,8 @@ export function JsonViewer({ data }: JsonViewerProps) {
borderRadius: 4,
}}
>
<div className="**:[[role='tree']]:bg-background! **:[[role='tree']]:text-foreground">
<JsonView data={data} style={defaultStyles} clickToExpandNode={true} />
<div className="**:[[role='tree']]:bg-transparent!">
<JsonView data={data} style={resolvedTheme === "dark" ? darkStyles : defaultStyles} clickToExpandNode={true} />
</div>
</div>
);

19
ui/litellm-dashboard/src/utils/entityLinks.test.ts Normal file
View file

@ -0,0 +1,19 @@
import { describe, expect, it, vi } from "vitest";
vi.mock("@/components/networking", () => ({ serverRootPath: "" }));
import { modelGroupHref } from "./entityLinks";
describe("modelGroupHref", () => {
it("targets the models page filtered to the encoded model group", () => {
expect(modelGroupHref("gpt-4.1")).toMatch(/\/models-and-endpoints\?model_group=gpt-4\.1$/);
expect(modelGroupHref("openai/*")).toMatch(/\?model_group=openai%2F\*$/);
});
it.each(["all-proxy-models", "all-team-models", "no-default-models"])(
"returns no href for the %s grant sentinel",
(sentinel) => {
expect(modelGroupHref(sentinel)).toBeUndefined();
},
);
});

11
ui/litellm-dashboard/src/utils/entityLinks.ts
View file

@ -1,5 +1,11 @@
import { migratedHref } from "@/utils/migratedPages";
const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
"all-proxy-models",
"all-team-models",
"no-default-models",
]);
export function teamDetailHref(teamId: string): string {
return `${migratedHref("teams")}?team=${encodeURIComponent(teamId)}`;
}
@ -15,3 +21,8 @@ export function userDetailHref(userId: string): string {
export function orgDetailHref(orgId: string): string {
return `${migratedHref("organizations")}?org=${encodeURIComponent(orgId)}`;
}
export function modelGroupHref(modelGroup: string): string | undefined {
if (MODEL_GRANT_SENTINELS.has(modelGroup)) return undefined;
return `${migratedHref("models-and-endpoints")}?model_group=${encodeURIComponent(modelGroup)}`;
}

6
uv.lock generated
View file

@ -10,7 +10,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-08-25T23:16:47.126855Z"
exclude-newer = "2026-08-26T18:33:25.773031Z"
exclude-newer-span = "P3D"
[manifest]
@ -4665,12 +4665,12 @@ proxy-dev = [
[[package]]
name = "litellm-enterprise"
version = "0.1.61"
version = "0.1.62"
source = { editable = "enterprise" }
[[package]]
name = "litellm-proxy-extras"
version = "0.4.90"
version = "0.4.91"
source = { editable = "litellm-proxy-extras" }
[[package]]