diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 962d1266fd7..ef88ae574fb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 17271 + "limit": 17270 }, "reportArgumentType": { "limit": 2539 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 39f8de0b0cc..570b306d6df 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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, diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 3653aba67ef..cac98b69793 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -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==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 0ef5cd1e856..d5741d479bf 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -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==", diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ce28f737334..4388e561026 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 6d63be05d00..481ea3f8f66 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 20a9ba789ce..0c4a753f762 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs new file mode 100644 index 00000000000..8a90cf667d0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -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 { + 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 { + 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); diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index c6f81cf6916..f9e75f45f75 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -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, ); -fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { - 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> { - 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> { - 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> { - 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>, ) -> PyResult> { 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>, ) -> PyResult> { 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>, timeout_seconds: Option, ) -> PyResult { - 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>, timeout_seconds: Option, ) -> PyResult> { - 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>, timeout_seconds: Option, ) -> PyResult> { - 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>, timeout_seconds: Option, ) -> PyResult { - 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>, timeout_seconds: Option, ) -> PyResult { - 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>, custom_llm_provider: Option, ) -> PyResult> { - 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, diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..c3d0638427c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,20 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde::Serialize; +use serde::de::DeserializeOwned; + +pub fn from_py(value: &Bound<'_, PyAny>) -> PyResult +where + T: DeserializeOwned, +{ + pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) +} + +pub fn to_py(py: Python<'_>, value: &T) -> PyResult> +where + T: Serialize + ?Sized, +{ + pythonize::pythonize(py, value) + .map(Bound::unbind) + .map_err(|error| PyValueError::new_err(error.to_string())) +} diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs new file mode 100644 index 00000000000..6a6ede22e85 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -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 { + 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() + ); + } + } +} diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index c833d57b6a9..04824a5bf39 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -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, ...]], ...]: diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2a59eddf88a..cced330d873 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -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, } diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index e1b204214d7..6a71e8e9223 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -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": }`` for user-keyed callers. + - ``{"created_by": }`` for user-keyed callers, and for keys + with no user_id and no team_id (owner id is their hashed token). - ``{"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 diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b8dc98f2582..7695b1cb35e 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -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: diff --git a/litellm/llms/hosted_vllm/videos/__init__.py b/litellm/llms/hosted_vllm/videos/__init__.py new file mode 100644 index 00000000000..89aa5ef2e8b --- /dev/null +++ b/litellm/llms/hosted_vllm/videos/__init__.py @@ -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() diff --git a/litellm/llms/hosted_vllm/videos/transformation.py b/litellm/llms/hosted_vllm/videos/transformation.py new file mode 100644 index 00000000000..96cbfc3cf70 --- /dev/null +++ b/litellm/llms/hosted_vllm/videos/transformation.py @@ -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 diff --git a/litellm/main.py b/litellm/main.py index cafa1e4718f..c341db08155 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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) diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 86c14fb4cd8..ead26ab65c5 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1180,7 +1180,8 @@ "files": true, "rerank": true, "a2a": true, - "interactions": true + "interactions": true, + "video_generations": true } }, "huggingface": { diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index d5461f01ed8..425f82794e6 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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(): diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 314c80adbc4..853a07972c1 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -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): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index 70a04ac290a..df2bbdba345 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -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/`` (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: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 2c7b970ca0e..6824f96f927 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7791adeb41e..bcaffa8e91c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -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, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index daeb91eb2bd..f834426d619 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -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) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 90d5f6f4970..dc13c09dd38 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -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.`` 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) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 88aa55fd4a9..8b57bdca2fe 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -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"] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 194ef5d756c..b7688790e54 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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, diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 4bc53678c23..1096954536a 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 23cfef6576c..567d8375737 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -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, }, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index a5619821197..50cb813c6fa 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 597acdf661f..3a70750528f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ca03b3b8042..eab56c31c39 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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, diff --git a/litellm/utils.py b/litellm/utils.py index 5cd9bfc5f32..fa2226dbf2c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1d8d374c2c4..7c7d508856f 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1277,7 +1277,8 @@ "files": true, "rerank": true, "a2a": true, - "interactions": true + "interactions": true, + "video_generations": true } }, "huggingface": { diff --git a/pyproject.toml b/pyproject.toml index a0db4d49467..34c1fec1c11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index f66a73e7daf..c64fd6150af 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -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"} diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 856636c3dbc..1f2f1d64711 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -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"} diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index c158fc89c81..f03e70df84a 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -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 diff --git a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py index dd61e630d7d..449803f3c80 100644 --- a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py @@ -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}") diff --git a/tests/e2e/guardrails/test_guardrail_dispatch_e2e.py b/tests/e2e/guardrails/test_guardrail_dispatch_e2e.py new file mode 100644 index 00000000000..793974ccdb1 --- /dev/null +++ b/tests/e2e/guardrails/test_guardrail_dispatch_e2e.py @@ -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}") diff --git a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py index d117832221d..43deb279bc8 100644 --- a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py @@ -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}" + ) diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py new file mode 100644 index 00000000000..6d927292975 --- /dev/null +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -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 +. 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. . 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 = "" + 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"" + 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, + ) diff --git a/tests/e2e/guardrails/test_streaming_guardrail_e2e.py b/tests/e2e/guardrails/test_streaming_guardrail_e2e.py new file mode 100644 index 00000000000..911ddf9304b --- /dev/null +++ b/tests/e2e/guardrails/test_streaming_guardrail_e2e.py @@ -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]}" + ) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 60536ea01d4..621595e6b46 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -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") diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index 7d882a7fa81..d0f478185c2 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -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: diff --git a/tests/e2e/logging/gcs_reader.py b/tests/e2e/logging/gcs_reader.py new file mode 100644 index 00000000000..60622c121ac --- /dev/null +++ b/tests/e2e/logging/gcs_reader.py @@ -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)) diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index d76f7b356b2..f0f7ad7eaa4 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -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 diff --git a/tests/e2e/logging/s3_reader.py b/tests/e2e/logging/s3_reader.py new file mode 100644 index 00000000000..d605dec6096 --- /dev/null +++ b/tests/e2e/logging/s3_reader.py @@ -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] + ) diff --git a/tests/e2e/logging/test_datadog_log_e2e.py b/tests/e2e/logging/test_datadog_log_e2e.py index 94811c6217e..a4821ed058b 100644 --- a/tests/e2e/logging/test_datadog_log_e2e.py +++ b/tests/e2e/logging/test_datadog_log_e2e.py @@ -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}" + ) diff --git a/tests/e2e/logging/test_gcs_log_e2e.py b/tests/e2e/logging/test_gcs_log_e2e.py new file mode 100644 index 00000000000..17ad1507049 --- /dev/null +++ b/tests/e2e/logging/test_gcs_log_e2e.py @@ -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}" diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py index 52cb691e2b7..9f08fa6c4e7 100644 --- a/tests/e2e/logging/test_otel_trace_e2e.py +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -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() diff --git a/tests/e2e/logging/test_s3_log_e2e.py b/tests/e2e/logging/test_s3_log_e2e.py new file mode 100644 index 00000000000..7a1ee1e6536 --- /dev/null +++ b/tests/e2e/logging/test_s3_log_e2e.py @@ -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}" diff --git a/tests/e2e/logging/test_team_langfuse_callback_e2e.py b/tests/e2e/logging/test_team_langfuse_callback_e2e.py new file mode 100644 index 00000000000..89cd45c9f16 --- /dev/null +++ b/tests/e2e/logging/test_team_langfuse_callback_e2e.py @@ -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) diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 2d845a445b5..e7d7fdaef81 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -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 diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index c75c8099ea1..ad46798b788 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -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 --- diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index eddfc4fbd34..f3ad8a8592e 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -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 == [] diff --git a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py index 3a09702de45..765c47547ce 100644 --- a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py @@ -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) diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py index b5fcd9d8219..1746926c689 100644 --- a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py +++ b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py @@ -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 + ) diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index 86fdd89acf6..39198bb20f3 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -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.""" diff --git a/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py b/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py new file mode 100644 index 00000000000..eb90430f303 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py @@ -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={}, + ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 8fa7c15d2d3..00ff06ea082 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index 36280530eac..2a59e6c1baa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -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) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index dcf122745d2..e3f71692c78 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -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( diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index e3893a66094..93dc429168f 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -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"}, + } diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 2f4018b55ab..038d061350f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -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=` 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 diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index c51c1c3f73b..9fb31d2a6db 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -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" diff --git a/tests/test_litellm/proxy/test_pricing_field_strip.py b/tests/test_litellm/proxy/test_pricing_field_strip.py index 1707f5bbc05..a84c6ba2b8a 100644 --- a/tests/test_litellm/proxy/test_pricing_field_strip.py +++ b/tests/test_litellm/proxy/test_pricing_field_strip.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8e470cc663b..949088ea3ba 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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=&search=`: 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(): """ diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 719d7cc73f5..41b0eb3cf95 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -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 diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 0971ce09d79..9d2a27ce9d3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -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 diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 72bb6756d24..b33bd912be9 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -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" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index 411e8402e11..7231c126a63 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -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, ); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index a5fbc433ea3..a9f7c54698a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -38,6 +38,7 @@ export const useModelsInfo = ( sortBy?: string, sortOrder?: string, excludeAutoRouters: boolean = false, + modelName?: string, ) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ @@ -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), }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 4d0b1c466a4..65faa85e29e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -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(); + + 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(); + + expect(lastModelsInfoCall().modelName).toBeUndefined(); + }); + + it("keeps the exact model group alongside a typed search", async () => { + render(); + + 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(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 1a9d33a50bc..be2cf22d71a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts index 292b27618bd..c4ea206022c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts @@ -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)); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts index 2cfad341d25..e83a81a53cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts @@ -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 }; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx index 9d40ea32185..552a4f57b24 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx @@ -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(null); + const { modelGroup, setModelGroup } = useModelGroupFilterRouting(); const { availableModelGroups, availableModelAccessGroups } = useModelDashboardData(); const { openModel, openTeam } = useModelDetailRouting(); return ( setModelGroup(group === ALL_MODEL_GROUPS_VALUE ? null : group)} availableModelGroups={availableModelGroups} availableModelAccessGroups={availableModelAccessGroups} setSelectedModelId={openModel} diff --git a/ui/litellm-dashboard/src/components/CodeBlock.tsx b/ui/litellm-dashboard/src/components/CodeBlock.tsx index 88ef9acf08e..a39d7417302 100644 --- a/ui/litellm-dashboard/src/components/CodeBlock.tsx +++ b/ui/litellm-dashboard/src/components/CodeBlock.tsx @@ -20,10 +20,10 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => { }; return ( -
+
@@ -1738,9 +1748,9 @@ const TeamInfoView: React.FC = ({

Default Member Models

{info.default_team_member_models.map((model, index) => ( - + {model} - + ))}
diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 0d41d199b22..7817d4e7cea 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -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( + {}} 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( + {}} 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( {currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} - + )) ) : (

No models specified

@@ -996,9 +997,9 @@ export default function KeyInfoView({
{currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} - + )) ) : (

No models specified

diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx index 055e400167b..a0ccb7f5323 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx @@ -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( + + + , + ); + describe("JsonViewer", () => { it("should render a placeholder and no tree when the log entry carries no payload", () => { - render(); + 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(); + 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(); + 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)); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx index 86b897f526f..980573a65cc 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx @@ -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 No data; return ( @@ -24,8 +27,8 @@ export function JsonViewer({ data }: JsonViewerProps) { borderRadius: 4, }} > -
- +
+
); diff --git a/ui/litellm-dashboard/src/utils/entityLinks.test.ts b/ui/litellm-dashboard/src/utils/entityLinks.test.ts new file mode 100644 index 00000000000..47161a903ed --- /dev/null +++ b/ui/litellm-dashboard/src/utils/entityLinks.test.ts @@ -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(); + }, + ); +}); diff --git a/ui/litellm-dashboard/src/utils/entityLinks.ts b/ui/litellm-dashboard/src/utils/entityLinks.ts index 675ac8d0554..ad257ec7969 100644 --- a/ui/litellm-dashboard/src/utils/entityLinks.ts +++ b/ui/litellm-dashboard/src/utils/entityLinks.ts @@ -1,5 +1,11 @@ import { migratedHref } from "@/utils/migratedPages"; +const MODEL_GRANT_SENTINELS: ReadonlySet = 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)}`; +} diff --git a/uv.lock b/uv.lock index 019c5a70f9e..8ef72116466 100644 --- a/uv.lock +++ b/uv.lock @@ -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]]