From 88de192dcf55e26f0f2cabb4d769a825dd8dbc7f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:29:04 -0700 Subject: [PATCH 01/30] test: bind management E2E callers and isolate JWT actors --- .github/e2e-stack/assert_tests_ran.py | 15 + .github/e2e-stack/oidc-profile.sh | 5 + .github/e2e-stack/select_tests.py | 2 + .github/workflows/test-e2e-changed.yml | 2 +- .../test_e2e_changed_gate.py | 23 ++ tests/e2e/CONTRIBUTING.md | 6 +- .../e2e/coverage_registry/management_cases.py | 151 +++++++++ tests/e2e/coverage_registry/mgmt.yaml | 8 + tests/e2e/e2e_http.py | 31 +- tests/e2e/idp.py | 254 ++++++++++++++- tests/e2e/junit_properties.py | 3 +- tests/e2e/management/conftest.py | 19 +- tests/e2e/management/jwt_actors.py | 175 ++++++++++ tests/e2e/management/management_client.py | 122 ++++--- .../e2e/management/test_jwt_management_e2e.py | 233 +++++++++++-- tests/e2e/models.py | 28 +- tests/e2e/proxy_client.py | 116 ++++--- tests/e2e/test_e2e_http.py | 10 + tests/e2e/test_idp.py | 122 ++++++- tests/e2e/test_proxy_client.py | 307 +++++++++++++++++- tests/e2e/transport.py | 17 +- tests/e2e/ui/oidcSetup.ts | 30 ++ tests/e2e/ui/playwright.oidc.config.ts | 22 ++ 23 files changed, 1541 insertions(+), 160 deletions(-) create mode 100755 .github/e2e-stack/oidc-profile.sh create mode 100644 tests/e2e/coverage_registry/management_cases.py create mode 100644 tests/e2e/management/jwt_actors.py create mode 100644 tests/e2e/ui/oidcSetup.ts create mode 100644 tests/e2e/ui/playwright.oidc.config.ts diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index c4348c20873..2303c42f4fb 100644 --- a/.github/e2e-stack/assert_tests_ran.py +++ b/.github/e2e-stack/assert_tests_ran.py @@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET from pathlib import Path from typing import Final +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e")) +from coverage_registry.management_cases import MANAGEMENT_CASES + def main() -> int: selected: Final = tuple(sys.argv[2:]) @@ -16,6 +19,17 @@ def main() -> int: case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) ) missing: Final = tuple(path for path in selected if path not in passed) + required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected) + passed_nodes: Final = frozenset( + prop.get("value") + for case in cases + if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) + for prop in case.findall("./properties/property") + if prop.get("name") == "management_node" + ) + missing_nodes: Final = required_nodes - passed_nodes + for node in sorted(missing_nodes): + _ = sys.stdout.write(f"::error::required management case did not pass: {node}\n") for path in selected: collected: Final = sum(case.get("file") == path for case in cases) skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) @@ -27,6 +41,7 @@ def main() -> int: if ( selected and not missing + and not missing_nodes and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error")) ): return 0 diff --git a/.github/e2e-stack/oidc-profile.sh b/.github/e2e-stack/oidc-profile.sh new file mode 100755 index 00000000000..84eaaaf8051 --- /dev/null +++ b/.github/e2e-stack/oidc-profile.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +set -euo pipefail +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "${REPO_ROOT}" +exec uv run --no-sync python tests/e2e/idp.py "$@" diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index 238818a0d36..982e93cf642 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile( HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" r"|^tests/e2e/idp_realm\.json$" + r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$" + r"|^tests/e2e/coverage_registry/management_cases\.py$" r"|^tests/e2e/gateway/" r"|^\.github/e2e-stack/" r"|^\.github/workflows/test-e2e-changed\.yml$" diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 1db597ff673..c9f08deb36e 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -183,7 +183,7 @@ jobs: log="${RUNNER_TEMP}/e2e-pass-${pass}.log" echo "::group::pass ${pass} of 3" set +e - uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \ + uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \ -o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1 status=$? uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}" diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 588402e3996..101816c7f11 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -81,6 +81,25 @@ def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None assert result.returncode == 1 +@pytest.mark.parametrize("omitted_role", ("proxy_admin", "team_member", "internal_user_viewer")) +def test_one_passing_management_case_cannot_hide_a_missing_actor(tmp_path: Path, omitted_role: str) -> None: + suite: Final = ET.Element("testsuite") + path: Final = "tests/e2e/management/test_jwt_management_e2e.py" + case: Final = ET.SubElement(suite, "testcase", file=path) + properties: Final = ET.SubElement(case, "properties") + _ = ET.SubElement( + properties, + "property", + name="management_node", + value=f"{path}::TestJwtManagement::test_actor_subject_and_database_role[proxy_admin_viewer]", + ) + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + result: Final = subprocess.run([sys.executable, str(GATE), str(report), path], capture_output=True, text=True) + assert result.returncode == 1 + assert f"test_actor_subject_and_database_role[{omitted_role}]" in result.stdout + + def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_path: Path) -> None: env_path: Final = tmp_path / ".env" @@ -141,6 +160,10 @@ def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them( ( "tests/e2e/proxy_client.py", "tests/e2e/conftest.py", + "tests/e2e/management/management_client.py", + "tests/e2e/management/jwt_actors.py", + "tests/e2e/management/conftest.py", + "tests/e2e/coverage_registry/management_cases.py", "tests/e2e/pytest.ini", "tests/e2e/gateway/stage_mirror_ci_config.yml", ".github/e2e-stack/up.sh", diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 44564a51e26..78c05ea4b30 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -60,7 +60,11 @@ The suites run against a live proxy, so bring one up first by running the litell Keycloak's password grant is a test-only provisioning shortcut, not a production login recommendation. The `litellm-e2e-admin` client adds the proxy's admin scope; the normal client does not. Never reuse this permissive realm outside an isolated test stack. - Management tests can use the shared `idp` and `jwt_identity` fixtures. Each test gets a unique Keycloak group/user and a matching proxy user/team. Setup and fallback cleanup use the master key; the operations and read-backs being tested must explicitly use `caller_key=idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)` (or a member token). See `management/test_jwt_management_e2e.py` for create/read/update/clear/delete and tenant-denial examples. A group claim alone is not database team membership: permission tests explicitly add the member and prove an allowed read before asserting the denied write. + Management tests can bind a credential once with `client.with_caller(Caller(...))`; direct calls, delegated helpers and replica read-backs then retain that caller. Explicit `caller_key` arguments override the binding. Keep the original master-backed client for bootstrap and cleanup. `actor_factory` lazily provisions database roles and tenant memberships, with `database_role` tokens carrying no groups and `group_scoped` actors retaining the existing team route gate. Token minting is explicit through `actor.mint_caller(idp)`. The factory runs requests without backend retries and reports cleanup failures. `coverage_registry/management_cases.py` records exact canary nodes and non-secret actor labels; the CI execution assertion rejects a missing or skipped actor row + + For the opt-in browser profile, start the existing IdP first, then run `.github/e2e-stack/oidc-profile.sh "$PROXY_BASE_URL" `. The wrapper creates a confidential client with an exact `/sso/callback` redirect and S256 PKCE, passes the client secret only through the child process environment, and removes the client on exit. It uses the existing generic OIDC handler with `GENERIC_USER_ID_ATTRIBUTE=sub`. Preserve the IdP's PostgreSQL data across restarts + + `tests/e2e/ui/playwright.oidc.config.ts` uses an already running OIDC stack and separate storage/output files. Supply `E2E_OIDC_UI_URL`, `JWT_ISSUER`, `E2E_OIDC_USERNAME` and `E2E_OIDC_PASSWORD` for a seeded actor. Its setup follows the real login and callback path. The current Python canary qualifies browser-client configuration and token/userinfo identity mapping; browser journey specs under `ui/oidc/` are a separate coverage step Every successful IdP create immediately registers cleanup, including partial setup failures. Cleanup failures emit warnings. Tokens are minted on demand, and the expiration test waits relative to the token's actual `exp` with a bounded clock-drift check. To check first-attempt behavior locally, run both files with `--reruns 0`: diff --git a/tests/e2e/coverage_registry/management_cases.py b/tests/e2e/coverage_registry/management_cases.py new file mode 100644 index 00000000000..15dc7d333c5 --- /dev/null +++ b/tests/e2e/coverage_registry/management_cases.py @@ -0,0 +1,151 @@ +from dataclasses import dataclass +from typing import Final, Literal + +CredentialKind = Literal["master", "idp_admin", "direct_jwt", "virtual_key", "dashboard_session"] +DependencyProfile = Literal["management_only", "real_oidc_browser", "external_provider_required"] + + +@dataclass(frozen=True, slots=True) +class ManagementCase: + node: str + credential_kind: CredentialKind + actor: str + profile: str + method: Literal["GET", "POST"] + path: str + operation_family: str + dependency_profile: DependencyProfile = "management_only" + + +JWT_FILE: Final = "tests/e2e/management/test_jwt_management_e2e.py" +JWT_CLASS: Final = f"{JWT_FILE}::TestJwtManagement" +ACTORS: Final = ( + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", +) +MANAGEMENT_CASES: Final = tuple( + ManagementCase( + node=f"{JWT_CLASS}::test_actor_subject_and_database_role[{role}]", + credential_kind="direct_jwt", + actor=role, + profile="database_role", + method="GET", + path="/user/info", + operation_family="identity", + ) + for role in ACTORS +) + ( + ManagementCase( + node=f"{JWT_CLASS}::test_admin_viewer_reads_but_cannot_update", + credential_kind="direct_jwt", + actor="proxy_admin_viewer", + profile="database_role", + method="POST", + path="/key/update", + operation_family="denial", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[direct_jwt]", + credential_kind="direct_jwt", + actor="proxy_admin", + profile="group_scoped", + method="POST", + path="/key/generate", + operation_family="key_lifecycle", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]", + credential_kind="virtual_key", + actor="proxy_admin", + profile="database_role", + method="POST", + path="/key/generate", + operation_family="key_lifecycle", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_two_actor_sets_keep_tenants_and_keys_isolated", + credential_kind="direct_jwt", + actor="team_member", + profile="group_scoped", + method="GET", + path="/key/info", + operation_family="tenant_isolation", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_member_cannot_write_and_another_team_cannot_read_the_key", + credential_kind="direct_jwt", + actor="team_member", + profile="group_scoped", + method="POST", + path="/key/update", + operation_family="tenant_isolation", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_multi_group_actor_keeps_exact_memberships", + credential_kind="master", + actor="bootstrap", + profile="group_scoped", + method="GET", + path="/team/info", + operation_family="memberships", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_successful_actor_cleanup_removes_owned_state", + credential_kind="master", + actor="bootstrap", + profile="failure_cleanup", + method="GET", + path="/team/info", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[group]", + credential_kind="idp_admin", + actor="idp_admin", + profile="failure_cleanup", + method="POST", + path="/groups", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_partial_setup_removes_previously_created_identities[user]", + credential_kind="idp_admin", + actor="idp_admin", + profile="failure_cleanup", + method="POST", + path="/users", + operation_family="cleanup", + ), + ManagementCase( + node=f"{JWT_CLASS}::test_oidc_browser_profile_identity_mapping", + credential_kind="direct_jwt", + actor="internal_user", + profile="oidc_configuration", + method="GET", + path="/protocol/openid-connect/userinfo", + operation_family="oidc_identity", + ), +) + + +def canonical_node(node: str) -> str: + return node if node.startswith("tests/e2e/") else f"tests/e2e/{node}" + + +def case_properties(node: str) -> tuple[tuple[str, str], ...]: + case: Final = next((case for case in MANAGEMENT_CASES if case.node == canonical_node(node)), None) + if case is None: + return () + return ( + ("management_node", case.node), + ("credential_kind", case.credential_kind), + ("actor", case.actor), + ("auth_profile", case.profile), + ("dependency_profile", case.dependency_profile), + ) diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index d1227fe7c0c..31ad61ba3e2 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -90,3 +90,11 @@ - {id: mgmt.mcp_toolset.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3098", rationale: "Narrowing the tools to one entry reads back exactly that entry"} - {id: mgmt.mcp_toolset.update.clear_persists, module: mgmt, tier: P0, surface: api, assertions: [clear_persists], source: "mcp_management_endpoints.py:3098", fail_before_fix: proven, rationale: "An explicit null clears the stored description; the update used to drop null and keep the old value"} - {id: mgmt.mcp_toolset.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "mcp_management_endpoints.py:3149", rationale: "A deleted toolset is gone by id and from the list on every replica"} + +- {id: mgmt.user.jwt.database_roles, module: mgmt, tier: P0, surface: api, assertions: [database_roles], source: "auth/handle_jwt.py", rationale: "User-only JWT subjects retain their seeded database roles and memberships"} +- {id: mgmt.key.jwt.viewer_denied, module: mgmt, tier: P0, surface: api, assertions: [viewer_denied], source: "auth/route_checks.py", rationale: "An admin viewer can read a key but cannot update it or change stored state"} +- {id: mgmt.user.oidc.identity_mapping, module: mgmt, tier: P0, surface: api, assertions: [identity_mapping], source: "tests/e2e/idp.py", rationale: "IdP configuration canary only: confidential-client token and userinfo subjects match the seeded user; application SSO is separate"} +- {id: mgmt.team.jwt.tenant_isolation, module: mgmt, tier: P0, surface: api, assertions: [tenant_isolation], source: "auth/handle_jwt.py", rationale: "Isolated team actors read their own key and receive 403 for the other tenant key"} +- {id: mgmt.team.jwt.multiple_memberships, module: mgmt, tier: P0, surface: api, assertions: [multiple_memberships], source: "auth/handle_jwt.py", rationale: "A multi-group actor has exactly the configured memberships without admin scope"} +- {id: mgmt.user.jwt.cleanup, module: mgmt, tier: P0, surface: api, assertions: [cleanup], source: "management_endpoints/internal_user_endpoints.py", rationale: "Owned users teams organizations keys and IdP objects disappear after successful cleanup"} +- {id: mgmt.user.jwt.partial_cleanup, module: mgmt, tier: P0, surface: api, assertions: [partial_cleanup], source: "auth/handle_jwt.py", rationale: "Partial identity setup removes the group and user created before failure"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index ce069720c6e..67370c98274 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -16,9 +16,11 @@ requests itself imports. from __future__ import annotations import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Generator, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from dataclasses import dataclass -from typing import Final, Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast +from typing import Final, Generic, Literal, NewType, Protocol, TypeVar, cast import pytest import requests @@ -36,8 +38,8 @@ class Headers(BaseModel): class AuthHeaders(Headers): # litellm accepts either; set whichever the call needs, leave the other None. - authorization: str | None = None - x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key") + authorization: str | None = Field(default=None, repr=False) + x_litellm_api_key: str | None = Field(default=None, alias="x-litellm-api-key", repr=False) class AnthropicHeaders(AuthHeaders): @@ -292,6 +294,22 @@ def _params(params: BaseModel | None) -> dict[str, str]: TRANSIENT_STATUSES: frozenset[int] = frozenset({529}) RETRY_ATTEMPTS: int = 3 +_QUALIFICATION: Final[ContextVar[bool]] = ContextVar("e2e_qualification", default=False) + + +def retry_attempts(default: int) -> int: + return 1 if _QUALIFICATION.get() else default + + +@contextmanager +def without_retries() -> Generator[None]: + token: Final = _QUALIFICATION.set(True) + try: + yield + finally: + _QUALIFICATION.reset(token) + + RETRY_BACKOFF_SECONDS: float = 0.5 @@ -319,7 +337,7 @@ def request_with_retry[T: RetryableResponse]( hang should surface as a hang instead of doubling the wall clock. Every retry prints, so flakiness stays visible in the run log instead of vanishing into green.""" - for attempt in range(1, RETRY_ATTEMPTS): + for attempt in range(1, retry_attempts(RETRY_ATTEMPTS)): resp = issue() if resp.status_code not in TRANSIENT_STATUSES: return resp @@ -414,6 +432,7 @@ def get_external[R: BaseModel]( url: str, *, response_type: type[R], + headers: BaseModel | None = None, timeout: float = 30.0, ) -> Result[R]: """GET an absolute URL outside the proxy (e.g. a public /.well-known document). @@ -422,7 +441,7 @@ def get_external[R: BaseModel]( try: resp = requests.get( url, - headers={"Accept": "application/json"}, + headers={"Accept": "application/json", **(_headers(headers) if headers is not None else {})}, timeout=timeout, ) except requests.RequestException as exc: diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 6d2fc84eb27..12db91bbd88 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -2,11 +2,17 @@ from __future__ import annotations +import base64 import os import secrets +import signal +import subprocess +import sys import warnings from collections.abc import Callable -from dataclasses import dataclass, field +from contextlib import ExitStack +from dataclasses import dataclass, field, replace +from types import FrameType from typing import Final, Literal import pytest @@ -14,11 +20,15 @@ from e2e_http import ( AuthHeaders, ExternalWrite, NetworkError, + NoBody, Result, Success, + UnknownApiError, delete_external, + get_external, post_form_external, post_json_external, + unwrap, ) from pydantic import BaseModel, Field @@ -46,7 +56,9 @@ class TokenGrantForm(BaseModel): grant_type: Literal["password"] = "password" client_id: str username: str - password: str + password: str = Field(repr=False) + client_secret: str | None = Field(default=None, repr=False) + scope: str | None = None class TokenResponse(BaseModel): @@ -63,7 +75,7 @@ class GroupCreateBody(BaseModel): class PasswordCredential(BaseModel): type: Literal["password"] = "password" - value: str + value: str = Field(repr=False) temporary: bool = False @@ -101,8 +113,20 @@ class Identity: user_id: str username: str password: str = field(repr=False) - group: str - group_id: str + groups: tuple[str, ...] + group_ids: tuple[str, ...] + + @property + def group(self) -> str: + if len(self.groups) != 1: + raise ValueError("A single-group identity is required") + return self.groups[0] + + @property + def group_id(self) -> str: + if len(self.group_ids) != 1: + raise ValueError("A single-group identity is required") + return self.group_ids[0] @dataclass(frozen=True, slots=True) @@ -111,6 +135,10 @@ class Keycloak: realm: str admin_username: str admin_password: str = field(repr=False) + strict_cleanup: bool = False + + def with_strict_cleanup(self) -> Keycloak: + return replace(self, strict_cleanup=True) @property def issuer(self) -> str: @@ -150,7 +178,9 @@ class Keycloak: f"group {name}", ) - def create_user(self, *, username: str, email: str, password: str, group: str) -> str: + def create_user( + self, *, username: str, email: str, password: str, group: str | None = None, groups: tuple[str, ...] = () + ) -> str: return created_id( post_json_external( self._admin_url("/users"), @@ -158,7 +188,7 @@ class Keycloak: json=UserCreateBody( username=username, email=email, - groups=(group,), + groups=(group,) if group is not None else groups, credentials=(PasswordCredential(value=password),), ), ), @@ -171,14 +201,28 @@ class Keycloak: def delete_group(self, group_id: str) -> None: self._delete(f"/groups/{group_id}") + def assert_absent(self, kind: Literal["users", "groups", "clients"], resource_id: str) -> None: + result: Final = get_external( + self._admin_url(f"/{kind}/{resource_id}"), + headers=self._admin_headers(), + response_type=NoBody, + ) + assert isinstance(result, UnknownApiError) and result.status_code == 404, ( + f"Owned IdP {kind} still exists: {result}" + ) + def _delete(self, path: str) -> None: try: headers: Final = self._admin_headers() except pytest.fail.Exception as exc: + if self.strict_cleanup: + raise RuntimeError(f"Keycloak cleanup could not authenticate for {path}") from exc warnings.warn(f"Keycloak cleanup could not authenticate for {path}: {exc}", RuntimeWarning, stacklevel=2) return result: Final = delete_external(self._admin_url(path), headers=headers) if result.status_code not in (204, 404): + if self.strict_cleanup: + raise RuntimeError(f"Keycloak cleanup failed for {path}: HTTP {result.status_code}") warnings.warn( f"Keycloak cleanup failed for {path}: HTTP {result.status_code} {result.body[:300]}", RuntimeWarning, @@ -188,15 +232,34 @@ class Keycloak: def provision(self, *, marker: str, group: str, defer: Callable[[Callable[[], object]], None]) -> Identity: """Create `group` and a user in it, credentialed with a password generated for this test alone, and hand back the identity a token can be minted for.""" - group_id: Final = self.create_group(group) - defer(lambda: self.delete_group(group_id)) + return self.provision_groups(marker=marker, groups=(group,), defer=defer) + + def provision_groups( + self, *, marker: str, groups: tuple[str, ...], defer: Callable[[Callable[[], object]], None] + ) -> Identity: + def provision_group(name: str) -> str: + created: Final = self.create_group(name) + defer(lambda: self.delete_group(created)) + return created + + group_ids: Final = tuple(provision_group(group) for group in groups) + return self.provision_user(marker=marker, groups=groups, group_ids=group_ids, defer=defer) + + def provision_user( + self, + *, + marker: str, + groups: tuple[str, ...], + group_ids: tuple[str, ...], + defer: Callable[[Callable[[], object]], None], + ) -> Identity: username: Final = f"e2e-jwt-user-{marker}" password: Final = secrets.token_urlsafe(24) user_id: Final = self.create_user( - username=username, email=f"{username}@example.com", password=password, group=group + username=username, email=f"{username}@example.com", password=password, groups=groups ) defer(lambda: self.delete_user(user_id)) - return Identity(user_id=user_id, username=username, password=password, group=group, group_id=group_id) + return Identity(user_id=user_id, username=username, password=password, groups=groups, group_ids=group_ids) def access_token( self, identity: Identity, *, client_id: str = TESTS_CLIENT_ID, issuer_host: str | None = None @@ -211,6 +274,65 @@ class Keycloak: ) return self._token(result, f"a token for {identity.username}") + def discovery(self) -> Discovery: + return unwrap(get_external(f"{self.issuer}/.well-known/openid-configuration", response_type=Discovery)) + + def browser_client(self, *, callback_url: str, defer: Callable[[Callable[[], object]], None]) -> BrowserClient: + client: Final = BrowserClient( + client_id=f"e2e-browser-{secrets.token_hex(8)}", + secret=secrets.token_urlsafe(32), + callback_url=callback_url, + ) + resource_id: Final = created_id( + post_json_external( + self._admin_url("/clients"), + headers=self._admin_headers(), + json=BrowserClientBody( + clientId=client.client_id, + secret=client.secret, + redirectUris=(callback_url,), + ), + ), + "browser client", + ) + defer(lambda: self._delete(f"/clients/{resource_id}")) + configured: Final = unwrap( + get_external( + self._admin_url(f"/clients/{resource_id}"), + headers=self._admin_headers(), + response_type=BrowserClientBody, + ) + ) + assert configured.redirect_uris == (callback_url,) + assert configured.standard_flow_enabled and not configured.public_client + assert configured.attributes.pkce == "S256" + return client + + def browser_token(self, identity: Identity, client: BrowserClient) -> str: + return self._token( + post_form_external( + self.token_url(self.realm), + form=TokenGrantForm( + client_id=client.client_id, + client_secret=client.secret, + username=identity.username, + password=identity.password, + scope="openid email", + ), + response_type=TokenResponse, + ), + "browser-profile identity mapping", + ) + + def userinfo(self, token: str) -> UserInfo: + return unwrap( + get_external( + f"{self.issuer}/protocol/openid-connect/userinfo", + headers=AuthHeaders(authorization=f"Bearer {token}"), + response_type=UserInfo, + ) + ) + def keycloak_from_env() -> Keycloak: admin_username: Final = os.environ.get(KEYCLOAK_ADMIN_USER_ENV, "").strip() @@ -226,3 +348,113 @@ def keycloak_from_env() -> Keycloak: admin_username=admin_username, admin_password=admin_password, ) + + +class TokenClaims(BaseModel): + sub: str + iss: str + aud: str | tuple[str, ...] + exp: int + scope: str = "" + groups: tuple[str, ...] = () + + +class Discovery(BaseModel): + issuer: str + authorization_endpoint: str + token_endpoint: str + userinfo_endpoint: str + jwks_uri: str + + +class UserInfo(BaseModel): + sub: str + email: str + + +class BrowserAttributes(BaseModel): + pkce: str = Field(default="S256", alias="pkce.code.challenge.method") + + +class AudienceConfig(BaseModel): + audience: str = Field(default="litellm-e2e", alias="included.custom.audience") + access_token: str = Field(default="true", alias="access.token.claim") + id_token: str = Field(default="false", alias="id.token.claim") + + +class AudienceMapper(BaseModel): + name: str = "litellm-audience" + protocol: str = "openid-connect" + mapper: str = Field(default="oidc-audience-mapper", alias="protocolMapper") + config: AudienceConfig = Field(default_factory=AudienceConfig) + + +class BrowserClientBody(BaseModel): + client_id: str = Field(alias="clientId") + secret: str = Field(repr=False) + redirect_uris: tuple[str, ...] = Field(alias="redirectUris") + enabled: bool = True + public_client: bool = Field(default=False, alias="publicClient") + standard_flow_enabled: bool = Field(default=True, alias="standardFlowEnabled") + direct_access_grants_enabled: bool = Field(default=True, alias="directAccessGrantsEnabled") + default_client_scopes: tuple[str, ...] = Field(default=("email", "basic"), alias="defaultClientScopes") + attributes: BrowserAttributes = Field(default_factory=BrowserAttributes) + protocol_mappers: tuple[AudienceMapper, ...] = Field(default=(AudienceMapper(),), alias="protocolMappers") + + +@dataclass(frozen=True, slots=True) +class BrowserClient: + client_id: str + secret: str = field(repr=False) + callback_url: str + + def environment(self, discovery: Discovery) -> dict[str, str]: + return { + "GENERIC_CLIENT_ID": self.client_id, + "GENERIC_CLIENT_SECRET": self.secret, + "GENERIC_USER_ID_ATTRIBUTE": "sub", + "GENERIC_AUTHORIZATION_ENDPOINT": discovery.authorization_endpoint, + "GENERIC_TOKEN_ENDPOINT": discovery.token_endpoint, + "GENERIC_USERINFO_ENDPOINT": discovery.userinfo_endpoint, + "GENERIC_CLIENT_USE_PKCE": "true", + "GENERIC_SCOPE": "openid email", + } + + +def token_claims(token: str) -> TokenClaims: + payload: Final = token.split(".")[1] + return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) + + +def run_oidc_profile(proxy_url: str, command: list[str]) -> int: + idp: Final = keycloak_from_env().with_strict_cleanup() + with ExitStack() as cleanup: + + def terminate(signum: int, frame: FrameType | None) -> None: + raise SystemExit(128 + signum) + + previous: Final = signal.signal(signal.SIGTERM, terminate) + cleanup.callback(signal.signal, signal.SIGTERM, previous) + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer) + environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url} + with subprocess.Popen(command, env=environment) as child: + try: + return child.wait() + finally: + if child.poll() is None: + child.terminate() + try: + child.wait(timeout=5) + except subprocess.TimeoutExpired: + child.kill() + child.wait() + + +if __name__ == "__main__": + if len(sys.argv) < 3: + raise SystemExit("Usage: idp.py PROXY_URL COMMAND [ARG ...]; requires a running test IdP") + raise SystemExit(run_oidc_profile(sys.argv[1], sys.argv[2:])) diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index c5971c5362c..b9f5da871ae 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -19,6 +19,7 @@ from __future__ import annotations from collections.abc import Iterable import pytest +from coverage_registry.management_cases import case_properties # Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing # at runtime names this suite's place in the repo. test_junit_properties.py @@ -94,7 +95,7 @@ def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: ("package", package_from_nodeid(item.nodeid)), ("covers", ",".join(covers_from_item(item))), ("source", source_from_item(item)), - ) + ) + case_properties(item.nodeid) def attach_result_properties(item: pytest.Item) -> None: diff --git a/tests/e2e/management/conftest.py b/tests/e2e/management/conftest.py index bd69c8c0ff3..5a11b634085 100644 --- a/tests/e2e/management/conftest.py +++ b/tests/e2e/management/conftest.py @@ -5,8 +5,14 @@ holds the shared ProxyClient so `resources` / `scoped_key` clean up keys, teams, users, and orgs this suite creates. """ -import pytest +from collections.abc import Generator +from typing import Final +import pytest +from e2e_http import without_retries +from idp import Keycloak +from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory from management_client import ManagementClient, build_client from proxy_client import ProxyClient @@ -21,3 +27,14 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> ManagementClient: return build_client(proxy) + + +@pytest.fixture +def actor_factory(proxy: ProxyClient, idp: Keycloak) -> Generator[ActorFactory]: + bootstrap: Final = build_client(proxy) + resources: Final = ResourceManager(client=proxy, strict_cleanup=True) + with without_retries(): + try: + yield ActorFactory(bootstrap=bootstrap, idp=idp, resources=resources) + finally: + resources.teardown() diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py new file mode 100644 index 00000000000..909d1652ada --- /dev/null +++ b/tests/e2e/management/jwt_actors.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Final, Literal + +from e2e_config import unique_marker +from e2e_http import NoBody, unwrap +from idp import ADMIN_CLIENT_ID, TESTS_CLIENT_ID, Identity, Keycloak +from lifecycle import ResourceManager +from management.management_client import ManagementClient +from models import ( + KeyGenerateBody, + KeyGenerateResponse, + OrgDeleteBody, + OrgDeleteResponse, + OrgMemberAddBody, + OrgMemberEntry, + OrgNewBody, + TeamDeleteBody, + TeamMemberAddBody, + TeamMemberEntry, + TeamNewBody, + UserNewBody, + UserRole, +) +from proxy_client import Caller + +ActorRole = Literal[ + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", +] +ActorProfile = Literal["database_role", "group_scoped"] + + +@dataclass(frozen=True, slots=True) +class Tenant: + organization_id: str + team_id: str + group_id: str + + +@dataclass(frozen=True, slots=True) +class Actor: + identity: Identity + role: ActorRole + global_role: UserRole + profile: ActorProfile + tenants: tuple[Tenant, ...] + + def mint_caller(self, idp: Keycloak) -> Caller: + return Caller( + credential=idp.access_token( + self.identity, client_id=ADMIN_CLIENT_ID if self.role == "proxy_admin" else TESTS_CLIENT_ID + ), + kind="direct_jwt", + role=self.role, + tenant=self.tenants[0].organization_id if self.tenants else None, + ) + + +@dataclass(frozen=True, slots=True) +class ActorFactory: + bootstrap: ManagementClient + idp: Keycloak + resources: ResourceManager + + def __post_init__(self) -> None: + if self.bootstrap.proxy.caller is not None: + raise ValueError("Actor bootstrap requires a separately held master client") + + def key(self, tenant: Tenant | None = None, *, user_id: str | None = None) -> KeyGenerateResponse: + created: Final = unwrap( + self.bootstrap.generate_key( + KeyGenerateBody( + team_id=tenant.team_id if tenant is not None else None, + user_id=user_id, + key_alias=f"e2e-actor-key-{unique_marker()}", + ) + ) + ) + self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key)) + return created + + def tenant(self) -> Tenant: + marker: Final = unique_marker() + organization_id: Final = self.bootstrap.create_org(OrgNewBody(organization_alias=f"e2e-organization-{marker}")) + self.resources.defer( + lambda: unwrap( + self.bootstrap.proxy.transport.delete( + "/organization/delete", + headers=self.bootstrap.proxy.management_headers(), + json=OrgDeleteBody(organization_ids=[organization_id]), + response_type=OrgDeleteResponse, + ) + ) + ) + team_id: Final = self.bootstrap.proxy.create_team( + TeamNewBody(team_alias=f"e2e-team-{marker}", organization_id=organization_id) + ) + self.resources.defer( + lambda: unwrap( + self.bootstrap.proxy.transport.post( + "/team/delete", + headers=self.bootstrap.proxy.management_headers(), + json=TeamDeleteBody(team_ids=[team_id]), + response_type=NoBody, + ) + ) + ) + self.bootstrap.delete_team_member(team_id, self.bootstrap.user_info().user_id) + group_id: Final = self.idp.create_group(team_id) + self.resources.defer(lambda: self.idp.with_strict_cleanup().delete_group(group_id)) + return Tenant(organization_id=organization_id, team_id=team_id, group_id=group_id) + + def create( + self, role: ActorRole, *, tenants: tuple[Tenant, ...] = (), profile: ActorProfile = "database_role" + ) -> Actor: + if role in ("team_admin", "team_member", "organization_admin") and not tenants: + raise ValueError("A membership actor requires a tenant") + identity: Final = self.idp.with_strict_cleanup().provision_user( + marker=unique_marker(), + groups=tuple(tenant.team_id for tenant in tenants) if profile == "group_scoped" else (), + group_ids=tuple(tenant.group_id for tenant in tenants) if profile == "group_scoped" else (), + defer=self.resources.defer, + ) + global_role: Final[UserRole] = ( + role + if role in ("proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer") + else "internal_user" + ) + self.bootstrap.create_user( + UserNewBody( + user_id=identity.user_id, + user_email=f"{identity.username}@example.com", + user_role=global_role, + auto_create_key=False, + ) + ) + self.resources.defer(lambda: self.bootstrap.delete_user_strict(identity.user_id)) + for tenant in tenants: + unwrap( + self.bootstrap.proxy.transport.post( + "/organization/member_add", + headers=self.bootstrap.proxy.management_headers(), + json=OrgMemberAddBody( + organization_id=tenant.organization_id, + member=OrgMemberEntry( + user_id=identity.user_id, + role="org_admin" if role == "organization_admin" else "internal_user", + ), + ), + response_type=NoBody, + ) + ) + unwrap( + self.bootstrap.proxy.transport.post( + "/team/member_add", + headers=self.bootstrap.proxy.management_headers(), + json=TeamMemberAddBody( + team_id=tenant.team_id, + member=TeamMemberEntry( + user_id=identity.user_id, + role="admin" if role == "team_admin" else "user", + ), + ), + response_type=NoBody, + ) + ) + return Actor(identity=identity, role=role, global_role=global_role, profile=profile, tenants=tenants) diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index e17b92a13ed..a0243e868e2 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -7,7 +7,8 @@ llm-only key hitting a management route). from __future__ import annotations import time -from dataclasses import dataclass +import warnings +from dataclasses import dataclass, field, replace import jwt from e2e_config import MASTER_KEY @@ -20,6 +21,7 @@ from e2e_http import ( StreamingResponse, Success, UnknownApiError, + retry_attempts, unwrap, ) from models import ( @@ -81,7 +83,7 @@ from models import ( UserNewResponse, UserUpdateBody, ) -from proxy_client import ProxyClient +from proxy_client import Caller, ProxyClient MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied" ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route" @@ -98,7 +100,7 @@ class DashboardSession: its bearer on every subsequent call, the claims it renders the signed-in user from, and where it lands the browser.""" - session_key: str + session_key: str = field(repr=False) claims: UiSessionClaims redirect_url: str @@ -106,7 +108,10 @@ class DashboardSession: @dataclass(frozen=True, slots=True) class ManagementClient: proxy: ProxyClient - master_key: str + master_key: str = field(repr=False) + + def with_caller(self, caller: Caller) -> ManagementClient: + return replace(self, proxy=self.proxy.with_caller(caller)) def llm_only_key(self) -> str: return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])) @@ -117,7 +122,7 @@ class ManagementClient: dashboard creates it under the session key their sign-in minted). Returns the outcome rather than unwrapping it, so a caller can poll a route that is only transiently refusing.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) return self.proxy.transport.post( "/key/generate", headers=headers, @@ -131,9 +136,9 @@ class ManagementClient: sign-in minted, never the master key). Returns the outcome rather than unwrapping it, so a caller can poll a route that is only transiently refusing; `update_key_models` is the unwrapping shorthand.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) last: Result[NoBody] = NetworkError(message="/key/update was never attempted") - for attempt in range(_KEY_WRITE_ATTEMPTS): + for attempt in range(retry_attempts(_KEY_WRITE_ATTEMPTS)): last = self.proxy.transport.post( "/key/update", headers=headers, @@ -144,6 +149,7 @@ class ManagementClient: case UnknownApiError(body=error_body) if any( marker in error_body.lower() for marker in _TRANSIENT_BACKEND_MARKERS ): + warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2) time.sleep(0.5 * (attempt + 1)) continue case _: @@ -153,10 +159,10 @@ class ManagementClient: def update_key_models(self, key: str, models: list[str]) -> None: _ = unwrap(self.update_key(KeyUpdateBody(key=key, models=models))) - def key_info_as(self, key: str, *, caller_key: str) -> Result[KeyInfoResponse]: + def key_info_as(self, key: str, *, caller_key: str | None = None) -> Result[KeyInfoResponse]: return self.proxy.transport.get( "/key/info", - headers=self.proxy.transport.bearer(caller_key), + headers=self.proxy.management_headers(caller_key), params=KeyInfoParams(key=key), response_type=KeyInfoResponse, ) @@ -167,7 +173,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/key/delete", - headers=self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key), + headers=self.proxy.management_headers(caller_key), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -179,7 +185,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/model/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -190,7 +196,7 @@ class ManagementClient: Connection button, probing the live provider with the supplied params.""" return self.proxy.transport.post( "/health/test_connection", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=ConnectionTestResponse, timeout=120.0, @@ -200,7 +206,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/key/block", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyBlockBody(key=key), response_type=NoBody, ) @@ -209,7 +215,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/key/regenerate", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyRegenerateBody(key=key, grace_period=grace_period), response_type=KeyGenerateResponse, ) @@ -219,7 +225,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( f"/key/{key}/reset_spend", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=KeyResetSpendBody(reset_to=reset_to), response_type=KeyResetSpendResponse, ) @@ -228,7 +234,7 @@ class ManagementClient: def key_list(self, key_alias: str, *, caller_key: str | None = None) -> Result[KeyListResponse]: """GET /key/list, the Virtual Keys page's own inventory call. `caller_key` is who is asking: the master key by default, or a virtual key.""" - headers = self.proxy.transport.master if caller_key is None else self.proxy.transport.bearer(caller_key) + headers = self.proxy.management_headers(caller_key) return self.proxy.transport.get( "/key/list", headers=headers, @@ -266,7 +272,7 @@ class ManagementClient: team_id = unwrap( self.proxy.transport.post( "/team/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=TeamNewResponse, ) @@ -276,10 +282,10 @@ class ManagementClient: def update_team(self, body: TeamUpdateBody) -> None: last: Result[NoBody] | None = None - for attempt in range(5): + for attempt in range(retry_attempts(5)): last = self.proxy.transport.post( "/team/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -289,6 +295,7 @@ class ManagementClient: case UnknownApiError(body=body_text) if ( "connecting to redis" in body_text.lower() or "name resolution" in body_text.lower() ): + warnings.warn(f"Transient backend response on attempt {attempt + 1}", RuntimeWarning, stacklevel=2) time.sleep(0.5 * (attempt + 1)) continue case _: @@ -299,7 +306,7 @@ class ManagementClient: def delete_team(self, team_id: str) -> None: _ = self.proxy.transport.post( "/team/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -308,7 +315,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/team/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -320,7 +327,7 @@ class ManagementClient: for entry in unwrap( self.proxy.transport.get( "/team/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=NoBody(), response_type=TeamListResponse, ) @@ -328,14 +335,16 @@ class ManagementClient: ) def team_info_status(self, team_id: str) -> ProbeResult: - return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id)) + return self.proxy.transport.probe( + "/team/info", params=TeamInfoParams(team_id=team_id), headers=self.proxy.management_headers() + ) def _wait_for_team(self, team_id: str) -> None: last: Result[TeamInfoResponse] | None = None - for _ in range(_TEAM_READY_ATTEMPTS): + for _ in range(retry_attempts(_TEAM_READY_ATTEMPTS)): last = self.proxy.transport.get( "/team/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=TeamInfoParams(team_id=team_id), response_type=TeamInfoResponse, ) @@ -343,25 +352,29 @@ class ManagementClient: case Success(): return case _: + warnings.warn("Repeating team read while the team becomes available", RuntimeWarning, stacklevel=2) time.sleep(_TEAM_READY_SLEEP_SECONDS) assert last is not None raise AssertionError(last) def add_team_member(self, team_id: str, user_id: str) -> None: last: Result[NoBody] | None = None - for attempt in range(_TEAM_READY_ATTEMPTS): + for attempt in range(retry_attempts(_TEAM_READY_ATTEMPTS)): last = self.proxy.transport.post( "/team/member_add", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)), response_type=NoBody, ) match last: case Success(): return - case UnknownApiError(body=body) if ( - "doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS + case UnknownApiError(body=body) if "doesn't exist" in body and attempt + 1 < retry_attempts( + _TEAM_READY_ATTEMPTS ): + warnings.warn( + "Retrying team membership while the team becomes available", RuntimeWarning, stacklevel=2 + ) time.sleep(_TEAM_READY_SLEEP_SECONDS) continue case _: @@ -373,7 +386,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/team/member_delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id), response_type=NoBody, ) @@ -383,7 +396,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/user/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=UserNewResponse, ) @@ -393,7 +406,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/customer/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=CustomerNewBody(user_id=user_id), response_type=CustomerResponse, ) @@ -404,7 +417,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/customer/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=CustomerInfoParams(end_user_id=end_user_id), response_type=CustomerResponse, ) @@ -413,7 +426,7 @@ class ManagementClient: def delete_customer(self, user_id: str) -> None: _ = self.proxy.transport.post( "/customer/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=CustomerDeleteBody(user_ids=[user_id]), response_type=NoBody, ) @@ -422,7 +435,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/user/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -431,7 +444,7 @@ class ManagementClient: def delete_user(self, user_id: str) -> None: _ = self.proxy.transport.post( "/user/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=NoBody, ) @@ -442,17 +455,17 @@ class ManagementClient: _ = unwrap( self.proxy.transport.post( "/user/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=UserDeleteResponse, ) ) - def user_info(self, user_id: str) -> UserInfoResponse: + def user_info(self, user_id: str | None = None) -> UserInfoResponse: return unwrap( self.proxy.transport.get( "/user/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserInfoParams(user_id=user_id), response_type=UserInfoResponse, ) @@ -462,7 +475,7 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/user/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -472,7 +485,7 @@ class ManagementClient: listing = unwrap( self.proxy.transport.get( "/user/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=UserListParams(user_ids=user_id), response_type=UserListResponse, ) @@ -483,7 +496,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/organization/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=OrgNewResponse, ) @@ -493,7 +506,7 @@ class ManagementClient: _ = unwrap( self.proxy.transport.patch( "/organization/update", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -502,7 +515,7 @@ class ManagementClient: def delete_org(self, organization_id: str) -> None: _ = self.proxy.transport.delete( "/organization/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=OrgDeleteBody(organization_ids=[organization_id]), response_type=NoBody, ) @@ -511,19 +524,24 @@ class ManagementClient: return unwrap( self.proxy.transport.get( "/organization/info", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=OrgInfoParams(organization_id=organization_id), response_type=OrgInfoResponse, ) ) def org_info_status(self, organization_id: str) -> ProbeResult: - return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id)) + return self.proxy.transport.probe( + "/organization/info", + params=OrgInfoParams(organization_id=organization_id), + headers=self.proxy.management_headers(), + ) + def create_tag(self, body: TagNewBody) -> None: _ = unwrap( self.proxy.transport.post( "/tag/new", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=NoBody, ) @@ -532,7 +550,7 @@ class ManagementClient: def delete_tag(self, name: str) -> None: _ = self.proxy.transport.post( "/tag/delete", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=TagDeleteBody(name=name), response_type=NoBody, ) @@ -542,7 +560,7 @@ class ManagementClient: unwrap( self.proxy.transport.get( "/tag/list", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), params=NoBody(), response_type=TagListResponse, ) @@ -553,7 +571,7 @@ class ManagementClient: return unwrap( self.proxy.transport.post( "/v1/mcp/server", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=McpServerRow, ) @@ -565,7 +583,7 @@ class ManagementClient: return unwrap( self.proxy.transport.put( "/v1/mcp/server", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=body, response_type=McpServerRow, ) @@ -576,7 +594,7 @@ class ManagementClient: unwrap it while a deferred teardown can ignore an already-deleted server.""" return self.proxy.transport.delete( f"/v1/mcp/server/{server_id}", - headers=self.proxy.transport.master, + headers=self.proxy.management_headers(), json=NoBody(), response_type=NoBody, ) diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py index 22306da8eb8..0d23f954158 100644 --- a/tests/e2e/management/test_jwt_management_e2e.py +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -2,60 +2,249 @@ from __future__ import annotations -from typing import Final +from typing import Final, Literal import pytest -from e2e_config import CHEAP_OPENAI_MODEL, unique_marker +from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker from e2e_http import UnauthorizedError, UnknownApiError, unwrap -from idp import ADMIN_CLIENT_ID, Identity, Keycloak +from idp import ADMIN_CLIENT_ID, Identity, Keycloak, token_claims from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory, ActorRole from management_client import ManagementClient -from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserNewBody +from models import KeyGenerateBody, KeyUpdateBody, TeamNewBody, UserInfoParams, UserInfoResponse, UserNewBody +from proxy_client import Caller pytestmark = pytest.mark.e2e class TestJwtManagement: + @pytest.mark.parametrize( + "role", + ( + "proxy_admin", + "proxy_admin_viewer", + "organization_admin", + "team_admin", + "team_member", + "internal_user", + "internal_user_viewer", + "unrelated_user", + ), + ) + @pytest.mark.covers("mgmt.user.jwt.database_roles") + def test_actor_subject_and_database_role(self, actor_factory: ActorFactory, role: ActorRole) -> None: + tenants: Final = ( + (actor_factory.tenant(),) if role in ("organization_admin", "team_admin", "team_member") else () + ) + actor: Final = actor_factory.create(role, tenants=tenants) + caller: Final = actor.mint_caller(actor_factory.idp) + claims: Final = token_claims(caller.credential) + assert claims.sub == actor.identity.user_id + assert claims.iss == actor_factory.idp.issuer + assert claims.aud == "litellm-e2e" or "litellm-e2e" in claims.aud + assert actor.identity.groups == () + assert ("litellm_proxy_admin" in claims.scope.split()) == (role == "proxy_admin") + stored: Final = actor_factory.bootstrap.user_info(actor.identity.user_id) + assert stored.user_id == actor.identity.user_id + assert stored.user_info.user_role == actor.global_role + bound: Final = actor_factory.bootstrap.with_caller(caller) + own: Final = unwrap( + bound.proxy.transport.get( + "/user/info", + headers=bound.proxy.management_headers(), + params=UserInfoParams(), + response_type=UserInfoResponse, + ) + ) + assert own.user_id == actor.identity.user_id + assert own.user_info.user_role == actor.global_role + for tenant in tenants: + info = actor_factory.bootstrap.team_info(tenant.team_id) + assert info.organization_id == tenant.organization_id + assert {(member.user_id, member.role) for member in info.members_with_roles} == { + (actor.identity.user_id, "admin" if role == "team_admin" else "user") + } + assert { + (member.user_id, member.user_role) + for member in actor_factory.bootstrap.org_info(tenant.organization_id).members + } == {(actor.identity.user_id, "org_admin" if role == "organization_admin" else "internal_user")} + + @pytest.mark.covers("mgmt.key.jwt.viewer_denied") + def test_admin_viewer_reads_but_cannot_update(self, actor_factory: ActorFactory) -> None: + actor: Final = actor_factory.create("proxy_admin_viewer") + viewer: Final = actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) + alias: Final = f"e2e-viewer-{unique_marker()}" + key: Final = actor_factory.key().key + unwrap(actor_factory.bootstrap.update_key(KeyUpdateBody(key=key, key_alias=alias))) + assert viewer.proxy.key_info(key).key_alias == alias + denied: Final = viewer.update_key(KeyUpdateBody(key=key, key_alias="forbidden")) + assert isinstance(denied, UnknownApiError) and denied.status_code == 403, f"viewer write was accepted: {denied}" + assert "proxy_admin_viewer" in denied.body and "/key/update" in denied.body + assert actor_factory.bootstrap.proxy.key_info(key).key_alias == alias + + @pytest.mark.covers("mgmt.user.oidc.identity_mapping") + def test_oidc_browser_profile_identity_mapping(self, actor_factory: ActorFactory) -> None: + actor: Final = actor_factory.create("internal_user") + idp: Final = actor_factory.idp.with_strict_cleanup() + discovery: Final = idp.discovery() + assert discovery.issuer == idp.issuer + assert discovery.jwks_uri == idp.jwks_url + callback: Final = f"{PROXY_BASE_URL}/sso/callback" + browser: Final = idp.browser_client(callback_url=callback, defer=actor_factory.resources.defer) + token: Final = idp.browser_token(actor.identity, browser) + assert token_claims(token).sub == actor.identity.user_id + userinfo: Final = idp.userinfo(token) + assert userinfo.sub == actor.identity.user_id + assert userinfo.email == f"{actor.identity.username}@example.com" + assert browser.environment(discovery)["GENERIC_USER_ID_ATTRIBUTE"] == "sub" + @pytest.mark.covers("mgmt.key.jwt.lifecycle") + @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key")) def test_admin_creates_reads_updates_clears_and_deletes_a_key( - self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager + self, + client: ManagementClient, + idp: Keycloak, + jwt_identity: Identity, + resources: ResourceManager, + actor_factory: ActorFactory, + credential_kind: Literal["direct_jwt", "virtual_key"], ) -> None: - admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + actor: Final = actor_factory.create("proxy_admin") + virtual_key: Final = ( + actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None + ) + admin: Final = ( + virtual_key if virtual_key is not None else idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + ) + bound: Final = client.with_caller(Caller(credential=admin, kind=credential_kind, role="proxy_admin")) alias: Final = f"e2e-jwt-key-{unique_marker()}" created: Final = unwrap( - client.generate_key( + bound.generate_key( KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), - caller_key=admin, ) ) resources.defer(lambda: client.proxy.delete_key(created.key)) - original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + original: Final = unwrap(bound.key_info_as(created.key)).info assert original.key_alias == alias and original.team_id == jwt_identity.group assert original.models == [CHEAP_OPENAI_MODEL] updated_alias: Final = f"{alias}-updated" - unwrap( - client.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120), caller_key=admin) - ) - updated: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + unwrap(bound.update_key(KeyUpdateBody(key=created.key, key_alias=updated_alias, rpm_limit=120))) + updated: Final = unwrap(bound.key_info_as(created.key)).info assert updated.key_alias == updated_alias and updated.rpm_limit == 120 assert updated.models == [CHEAP_OPENAI_MODEL], "omitted models must preserve the restriction" - unwrap(client.update_key(KeyUpdateBody(key=created.key, models=[]), caller_key=admin)) - cleared: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info + unwrap(bound.update_key(KeyUpdateBody(key=created.key, models=[]))) + cleared: Final = unwrap(bound.key_info_as(created.key)).info assert cleared.models == [] and cleared.rpm_limit == 120 - assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 1 - client.delete_key_strict(created.key, caller_key=admin) - assert unwrap(client.key_list(updated_alias, caller_key=admin)).total_count == 0 + assert unwrap(bound.key_list(updated_alias)).total_count == 1 + bound.delete_key_strict(created.key) + assert unwrap(bound.key_list(updated_alias)).total_count == 0 + + @pytest.mark.covers("mgmt.team.jwt.tenant_isolation") + def test_two_actor_sets_keep_tenants_and_keys_isolated(self, actor_factory: ActorFactory) -> None: + first: Final = actor_factory.tenant() + second: Final = actor_factory.tenant() + assert first.organization_id != second.organization_id and first.team_id != second.team_id + actors: Final = tuple( + actor_factory.create("team_member", tenants=(tenant,), profile="group_scoped") for tenant in (first, second) + ) + assert actors[0].identity.user_id != actors[1].identity.user_id + callers: Final = tuple( + actor_factory.bootstrap.with_caller(actor.mint_caller(actor_factory.idp)) for actor in actors + ) + keys: Final = tuple(actor_factory.key(tenant) for tenant in (first, second)) + assert keys[0].key != keys[1].key + assert callers[0].proxy.key_info(keys[0].key).team_id == first.team_id + assert callers[1].proxy.key_info(keys[1].key).team_id == second.team_id + for caller, other_key in ((callers[0], keys[1].key), (callers[1], keys[0].key)): + hidden = caller.key_info_as(other_key) + assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403 + assert tuple(actor.identity.groups for actor in actors) == ((first.team_id,), (second.team_id,)) + + @pytest.mark.covers("mgmt.team.jwt.multiple_memberships") + def test_multi_group_actor_keeps_exact_memberships(self, actor_factory: ActorFactory) -> None: + tenants: Final = (actor_factory.tenant(), actor_factory.tenant()) + actor: Final = actor_factory.create("team_member", tenants=tenants, profile="group_scoped") + claims: Final = token_claims(actor.mint_caller(actor_factory.idp).credential) + assert set(claims.groups) == {tenant.team_id for tenant in tenants} + assert "litellm_proxy_admin" not in claims.scope.split() + assert actor.identity.groups == tuple(tenant.team_id for tenant in tenants) + for tenant in tenants: + assert { + (entry.user_id, entry.role) + for entry in actor_factory.bootstrap.team_info(tenant.team_id).members_with_roles + } == {(actor.identity.user_id, "user")} + + @pytest.mark.covers("mgmt.user.jwt.cleanup") + def test_successful_actor_cleanup_removes_owned_state(self, actor_factory: ActorFactory) -> None: + resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True) + factory: Final = ActorFactory(bootstrap=actor_factory.bootstrap, idp=actor_factory.idp, resources=resources) + try: + tenant: Final = factory.tenant() + actor: Final = factory.create("team_member", tenants=(tenant,), profile="group_scoped") + key: Final = factory.key(tenant) + alias: Final = factory.bootstrap.proxy.key_info(key.key).key_alias + assert alias is not None + finally: + resources.teardown() + assert factory.bootstrap.user_count(actor.identity.user_id) == 0 + assert factory.bootstrap.key_alias_count(alias) == 0 + assert factory.bootstrap.team_info_status(tenant.team_id).status_code == 404 + assert factory.bootstrap.org_info_status(tenant.organization_id).status_code == 404 + factory.idp.assert_absent("users", actor.identity.user_id) + factory.idp.assert_absent("groups", tenant.group_id) + + @pytest.mark.parametrize("stage", ("group", "user")) + @pytest.mark.covers("mgmt.user.jwt.partial_cleanup") + def test_partial_setup_removes_previously_created_identities( + self, + actor_factory: ActorFactory, + stage: Literal["group", "user"], + ) -> None: + idp: Final = actor_factory.idp.with_strict_cleanup() + resources: Final = ResourceManager(client=actor_factory.bootstrap.proxy, strict_cleanup=True) + marker: Final = unique_marker() + group_id: Final = idp.create_group(f"e2e-partial-{marker}") + resources.defer(lambda: idp.delete_group(group_id)) + try: + identity: Final = ( + idp.provision_user( + marker=marker, + groups=(f"e2e-partial-{marker}",), + group_ids=(group_id,), + defer=resources.defer, + ) + if stage == "user" + else None + ) + if identity is None: + with pytest.raises(pytest.fail.Exception, match="HTTP 409"): + idp.create_group(f"e2e-partial-{marker}") + else: + with pytest.raises(pytest.fail.Exception, match="HTTP 409"): + idp.create_user( + username=identity.username, + email=f"{identity.username}@example.com", + password=identity.password, + groups=identity.groups, + ) + finally: + resources.teardown() + idp.assert_absent("groups", group_id) + if identity is not None: + idp.assert_absent("users", identity.user_id) @pytest.mark.covers("mgmt.key.jwt.member_denied", "mgmt.key.jwt.other_team_denied") def test_member_cannot_write_and_another_team_cannot_read_the_key( self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager ) -> None: admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + bound: Final = client.with_caller(Caller(credential=admin, kind="direct_jwt", role="proxy_admin")) member: Final = idp.access_token(jwt_identity) + member_client: Final = client.with_caller(Caller(credential=member, kind="direct_jwt", role="team_member")) alias: Final = f"e2e-jwt-owned-{unique_marker()}" created: Final = unwrap( client.generate_key(KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group), caller_key=admin) @@ -63,14 +252,14 @@ class TestJwtManagement: resources.defer(lambda: client.proxy.delete_key(created.key)) client.add_team_member(jwt_identity.group, jwt_identity.user_id) - assert unwrap(client.key_info_as(created.key, caller_key=member)).info.key_alias == alias + assert unwrap(member_client.key_info_as(created.key)).info.key_alias == alias - refused: Final = client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden"), caller_key=member) + refused: Final = member_client.update_key(KeyUpdateBody(key=created.key, key_alias="forbidden")) assert isinstance(refused, UnauthorizedError), f"member write was accepted: {refused}" assert "does not have permissions for endpoint" in refused.body.lower(), ( f"expected a permission denial: {refused}" ) - assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.key_alias == alias + assert unwrap(bound.key_info_as(created.key)).info.key_alias == alias marker: Final = unique_marker() outsider: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer) @@ -88,4 +277,4 @@ class TestJwtManagement: assert isinstance(hidden, UnknownApiError) and hidden.status_code == 403, ( f"another team must not read this key: {hidden}" ) - assert unwrap(client.key_info_as(created.key, caller_key=admin)).info.team_id == jwt_identity.group + assert unwrap(bound.key_info_as(created.key)).info.team_id == jwt_identity.group diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 3cab0334dea..6fca1268ebc 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1091,13 +1091,13 @@ class UiLoginBody(BaseModel): class UiLoginResponse(BaseModel): - token: str + token: str = Field(repr=False) redirect_url: str class UiSessionClaims(BaseModel): user_id: str - key: str + key: str = Field(repr=False) user_role: str login_method: Literal["sso", "username_password"] exp: int @@ -1135,6 +1135,7 @@ class TeamInfoParams(BaseModel): class TeamData(BaseModel): + organization_id: str | None = None team_alias: str | None = None models: list[str] = [] members_with_roles: list[TeamMemberEntry] = [] @@ -1175,6 +1176,7 @@ class UserNewBody(BaseModel): user_email: str user_role: UserRole user_id: str | None = None + auto_create_key: bool | None = None class UserNewResponse(BaseModel): @@ -1187,7 +1189,7 @@ class UserUpdateBody(BaseModel): class UserInfoParams(BaseModel): - user_id: str + user_id: str | None = None class UserData(BaseModel): @@ -1240,16 +1242,36 @@ class OrgInfoParams(BaseModel): organization_id: str +class OrgMembership(BaseModel): + user_id: str + user_role: str + + class OrgInfoResponse(BaseModel): organization_id: str organization_alias: str | None = None models: list[str] = [] + members: tuple[OrgMembership, ...] = () + + +class OrgMemberEntry(BaseModel): + user_id: str + role: Literal["org_admin", "internal_user"] + + +class OrgMemberAddBody(BaseModel): + organization_id: str + member: OrgMemberEntry class OrgDeleteBody(BaseModel): organization_ids: list[str] +class OrgDeleteResponse(RootModel[tuple[OrgInfoResponse, ...]]): + pass + + # ---------- tags (management) ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 1fe2ec905ef..48a6110dc0b 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -11,14 +11,23 @@ from __future__ import annotations import time import warnings from collections.abc import Callable, Mapping -from dataclasses import dataclass -from functools import reduce +from dataclasses import dataclass, field, replace from datetime import datetime +from functools import reduce from types import MappingProxyType -from typing import Final - -from pydantic import BaseModel +from typing import Final, Literal +from e2e_config import ( + CONTROL_PLANE_BASE_URL, + MASTER_KEY, + POLL_INTERVAL, + POLL_TIMEOUT, + PROXY_BASE_URL, + PROXY_REPLICA_URLS, + REQUEST_TIMEOUT, + SLOW_PROVIDER_TIMEOUT_SECONDS, + settle_propagation, +) from e2e_http import ( AnthropicHeaders, AuthHeaders, @@ -55,6 +64,7 @@ from models import ( KeyInfoParams, KeyInfoResponse, LiteLLMParamsBody, + MemorySummaryResponse, ModelDeleteBody, ModelInfoBody, ModelInfoEntry, @@ -63,7 +73,6 @@ from models import ( ModelNewBody, ModelNewResponse, ModelsListParams, - MemorySummaryResponse, ModelsListResponse, ModelUpdateBody, OcrBody, @@ -76,23 +85,13 @@ from models import ( TeamDeleteBody, TeamNewBody, TeamNewResponse, - UserDeleteBody, - UserDeleteResponse, ToolsetCreateBody, ToolsetRow, ToolsetUpdateBody, + UserDeleteBody, + UserDeleteResponse, ) -from e2e_config import ( - CONTROL_PLANE_BASE_URL, - MASTER_KEY, - POLL_INTERVAL, - POLL_TIMEOUT, - PROXY_BASE_URL, - PROXY_REPLICA_URLS, - REQUEST_TIMEOUT, - SLOW_PROVIDER_TIMEOUT_SECONDS, - settle_propagation, -) +from pydantic import BaseModel from transport import HttpTransport, SplitTransport, Transport, is_control_plane_path RowsPredicate = Callable[[list[SpendLogRow]], bool] @@ -421,11 +420,23 @@ def converge_timeout_message(*, what: str, replica: str, timeout: float, last_re ) +CredentialKind = Literal["master", "direct_jwt", "virtual_key", "dashboard_session"] + + +@dataclass(frozen=True, slots=True) +class Caller: + credential: str = field(repr=False) + kind: CredentialKind + role: str + tenant: str | None = None + + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport replicas: Mapping[str, Transport] control_replicas: Mapping[str, Transport] + caller: Caller | None = None poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -433,13 +444,24 @@ class ProxyClient: model_servable_interval: float = MODEL_SERVABLE_INTERVAL model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT + def with_caller(self, caller: Caller) -> ProxyClient: + return replace(self, caller=caller) + + def management_headers(self, caller_key: str | None = None, *, transport: Transport | None = None) -> AuthHeaders: + selected: Final = self.transport if transport is None else transport + if caller_key is not None: + return selected.bearer(caller_key) + if self.caller is not None: + return selected.bearer(self.caller.credential) + return selected.master + # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- def generate_key(self, body: KeyGenerateBody) -> str: return unwrap( self.transport.post( "/key/generate", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=KeyGenerateResponse, ) @@ -448,7 +470,7 @@ class ProxyClient: def delete_key(self, key: str) -> None: _ = self.transport.post( "/key/delete", - headers=self.transport.master, + headers=self.management_headers(), json=KeyDeleteBody(keys=[key]), response_type=NoBody, ) @@ -458,7 +480,7 @@ class ProxyClient: return _ = self.transport.post( "/customer/delete", - headers=self.transport.master, + headers=self.management_headers(), json=CustomerDeleteBody(user_ids=user_ids), response_type=NoBody, ) @@ -467,7 +489,7 @@ class ProxyClient: return unwrap( self.transport.get( "/key/info", - headers=self.transport.master, + headers=self.management_headers(), params=KeyInfoParams(key=key), response_type=KeyInfoResponse, ) @@ -477,7 +499,7 @@ class ProxyClient: return { url: transport.get( "/debug/memory/summary", - headers=transport.master, + headers=self.management_headers(transport=transport), params=NoBody(), response_type=MemorySummaryResponse, ) @@ -524,11 +546,12 @@ class ProxyClient: {replica: outcome.result for replica, outcome in outcomes.items() if isinstance(outcome, Converged)} ) - @staticmethod def _body_poller[R: BaseModel]( - transport: Transport, path: str, params: BaseModel, response_type: type[R] + self, transport: Transport, path: str, params: BaseModel, response_type: type[R] ) -> Poller[Result[R]]: - return lambda: transport.get(path, headers=transport.master, params=params, response_type=response_type) + return lambda: transport.get( + path, headers=self.management_headers(transport=transport), params=params, response_type=response_type + ) def model_info(self) -> list[ModelInfoEntry]: """Every configured deployment with the price the proxy resolved for it @@ -536,7 +559,7 @@ class ProxyClient: return unwrap( self.transport.get( "/model/info", - headers=self.transport.master, + headers=self.management_headers(), params=NoBody(), response_type=ModelInfoResponse, ) @@ -546,7 +569,7 @@ class ProxyClient: return unwrap( self.transport.get( "/public/litellm_model_cost_map", - headers=self.transport.master, + headers=self.management_headers(), params=NoBody(), response_type=CostMap, ) @@ -607,7 +630,7 @@ class ProxyClient: model_id = unwrap( self.transport.post( "/model/new", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ModelNewResponse, ) @@ -623,7 +646,7 @@ class ProxyClient: def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" - headers: Final = self.transport.master if listed_for is None else self.transport.bearer(listed_for) + headers: Final = self.management_headers(listed_for) outcome: Final = await_servable_everywhere( {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, @@ -666,7 +689,7 @@ class ProxyClient: unwrap( self.transport.post( "/model/update", - headers=self.transport.master, + headers=self.management_headers(), json=ModelUpdateBody( litellm_params=litellm_params, model_info=ModelInfoBody(id=model_id), @@ -678,7 +701,7 @@ class ProxyClient: def delete_model(self, model_id: str) -> None: result = self.transport.post( "/model/delete", - headers=self.transport.master, + headers=self.management_headers(), json=ModelDeleteBody(id=model_id), response_type=NoBody, ) @@ -747,11 +770,10 @@ class ProxyClient: f"GET {path} on {replica} still answers {self.poll_timeout}s after the delete; last read: {last}" ) - @staticmethod - def _reader[R: BaseModel](transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: + def _reader[R: BaseModel](self, transport: Transport, path: str, response_type: type[R]) -> ReplicaRead[Result[R]]: return lambda request_timeout: transport.get( path, - headers=transport.master, + headers=self.management_headers(transport=transport), params=NoBody(), response_type=response_type, timeout=request_timeout, @@ -763,7 +785,7 @@ class ProxyClient: return unwrap( self.transport.post( "/v1/mcp/toolset", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ToolsetRow, ) @@ -775,7 +797,7 @@ class ProxyClient: return unwrap( self.transport.put( "/v1/mcp/toolset", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=ToolsetRow, ) @@ -786,7 +808,7 @@ class ProxyClient: can unwrap it while a deferred teardown can ignore an already-deleted row.""" return self.transport.delete( f"/v1/mcp/toolset/{toolset_id}", - headers=self.transport.master, + headers=self.management_headers(), json=NoBody(), response_type=NoBody, ) @@ -795,7 +817,7 @@ class ProxyClient: unwrap( self.transport.post( "/credentials", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=CredentialCreateResponse, ) @@ -804,7 +826,7 @@ class ProxyClient: def delete_credential(self, credential_name: str) -> None: result = self.transport.delete( f"/credentials/{credential_name}", - headers=self.transport.master, + headers=self.management_headers(), json=NoBody(), response_type=NoBody, ) @@ -815,7 +837,7 @@ class ProxyClient: return unwrap( self.transport.post( "/team/new", - headers=self.transport.master, + headers=self.management_headers(), json=body, response_type=TeamNewResponse, ) @@ -824,7 +846,7 @@ class ProxyClient: def delete_team(self, team_id: str) -> None: result = self.transport.post( "/team/delete", - headers=self.transport.master, + headers=self.management_headers(), json=TeamDeleteBody(team_ids=[team_id]), response_type=NoBody, ) @@ -836,7 +858,7 @@ class ProxyClient: a user the proxy only upserts after a successful auth.""" result = self.transport.post( "/user/delete", - headers=self.transport.master, + headers=self.management_headers(), json=UserDeleteBody(user_ids=[user_id]), response_type=UserDeleteResponse, ) @@ -909,7 +931,7 @@ class ProxyClient: def spend_logs(self, params: SpendLogsParams) -> list[SpendLogRow]: result = self.transport.get( "/spend/logs", - headers=self.transport.master, + headers=self.management_headers(), params=params, response_type=SpendLogs, ) @@ -924,7 +946,7 @@ class ProxyClient: return unwrap( self.transport.get( "/spend/logs/v2", - headers=self.transport.master, + headers=self.management_headers(), params=SpendLogsPageParams( start_date=start.strftime("%Y-%m-%d %H:%M:%S"), end_date=end.strftime("%Y-%m-%d %H:%M:%S"), @@ -977,7 +999,7 @@ class ProxyClient: # ---- route probe ---------------------------------------------------- def probe(self, path: str, *, params: NoBody) -> ProbeResult: - return self.transport.probe(path, params=params) + return self.transport.probe(path, params=params, headers=self.management_headers()) def build_proxy_client( diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 81cd6c8d3d1..7201da84924 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -29,6 +29,7 @@ from e2e_http import ( request_with_retry, streaming_outcome, wire_body, + without_retries, ) from pydantic import BaseModel, TypeAdapter @@ -56,6 +57,15 @@ def _issue_from(responses: Sequence[FakeResponse]) -> Callable[[], FakeResponse] class TestTransientRetryPolicy: + def test_qualification_disables_retries_and_restores_the_default(self) -> None: + responses: Final = (FakeResponse(529), FakeResponse(200)) + sleep: Final = SleepRecorder() + with without_retries(): + assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[0] + assert sleep.delays == () + assert request_with_retry(_issue_from(responses), sleep=sleep) is responses[1] + assert sleep.delays == (0.5,) + def test_transient_set_is_only_statuses_the_proxy_cannot_emit(self) -> None: assert TRANSIENT_STATUSES == frozenset({529}) assert 429 not in TRANSIENT_STATUSES diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py index 33a09a0f13a..8a92dfa527c 100644 --- a/tests/e2e/test_idp.py +++ b/tests/e2e/test_idp.py @@ -4,12 +4,19 @@ these carry no `e2e` marker and run everywhere.""" from __future__ import annotations +import os +import signal +import subprocess +import sys +import time +from builtins import ExceptionGroup from collections.abc import Callable, Generator from contextlib import ExitStack, contextmanager from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path from queue import SimpleQueue from threading import Thread -from typing import Final +from typing import Final, Literal import pytest from e2e_http import ExternalWrite @@ -18,6 +25,8 @@ from idp import ( KEYCLOAK_ADMIN_USER_ENV, KEYCLOAK_REALM_ENV, KEYCLOAK_URL_ENV, + BrowserClientBody, + Discovery, Keycloak, PasswordCredential, UserCreateBody, @@ -60,24 +69,48 @@ def _idp_server( ) -> Generator[tuple[Keycloak, SimpleQueue[str]]]: """Exercise provisioning failures through the same HTTP transport as live tests.""" deletions: SimpleQueue[str] = SimpleQueue() + clients: SimpleQueue[BrowserClientBody] = SimpleQueue() class Handler(BaseHTTPRequestHandler): def log_message(self, format: str, *args: object) -> None: pass def do_POST(self) -> None: - self.rfile.read(int(self.headers.get("Content-Length", "0"))) + body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) if self.path.endswith("/token"): self.send_response(admin_status) self.end_headers() self.wfile.write(b'{"access_token":"synthetic-harness-token"}') else: + if self.path.endswith("/clients"): + clients.put(BrowserClientBody.model_validate_json(body)) self.send_response(user_status if self.path.endswith("/users") else 201) self.send_header("Location", f"{self.path}/resource-1") self.end_headers() if user_status != 201 and self.path.endswith("/users"): self.wfile.write(b"injected create failure") + def do_GET(self) -> None: + self.send_response(200) + self.end_headers() + if "/clients/" in self.path: + client: Final = clients.get_nowait() + clients.put(client) + self.wfile.write(client.model_dump_json(by_alias=True).encode()) + else: + issuer: Final = f"http://127.0.0.1:{server.server_port}/realms/test" + self.wfile.write( + Discovery( + issuer=issuer, + authorization_endpoint=f"{issuer}/auth", + token_endpoint=f"{issuer}/token", + userinfo_endpoint=f"{issuer}/userinfo", + jwks_uri=f"{issuer}/certs", + ) + .model_dump_json() + .encode() + ) + def do_DELETE(self) -> None: deletions.put(self.path) self.send_response(delete_status) @@ -115,6 +148,54 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No assert deletions.empty() +@pytest.mark.parametrize("exit_mode", ("normal", "parent", "group")) +def test_oidc_launcher_removes_client_on_exit_and_termination( + tmp_path: Path, exit_mode: Literal["normal", "parent", "group"] +) -> None: + ready: Final = tmp_path / "ready" + child_command: Final = ( + "import os,time; from pathlib import Path; " + 'assert os.environ["GENERIC_CLIENT_SECRET"]; ' + 'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; ' + f"Path({str(ready)!r}).touch(); " + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") + ) + with _idp_server() as (idp, deletions): + with subprocess.Popen( + [ + sys.executable, + str(Path(__file__).with_name("idp.py")), + "http://127.0.0.1:9999", + sys.executable, + "-c", + child_command, + ], + env={ + **os.environ, + KEYCLOAK_URL_ENV: idp.base_url, + KEYCLOAK_REALM_ENV: idp.realm, + KEYCLOAK_ADMIN_USER_ENV: idp.admin_username, + KEYCLOAK_ADMIN_PASSWORD_ENV: idp.admin_password, + }, + start_new_session=True, + ) as process: + try: + deadline: Final = time.monotonic() + 15 + while not ready.exists() and time.monotonic() < deadline and process.poll() is None: + time.sleep(0.05) + assert ready.exists(), "OIDC child did not start" + if exit_mode == "parent": + process.terminate() + elif exit_mode == "group": + os.killpg(process.pid, signal.SIGTERM) + assert process.wait(timeout=10) == (7 if exit_mode == "normal" else 143) + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=5) + assert deletions.get(timeout=5) == "/admin/realms/test/clients/resource-1" + assert deletions.empty() + + def test_successful_provisioning_cleans_up_user_before_group() -> None: with _idp_server() as (idp, deletions): with ExitStack() as cleanup: @@ -134,6 +215,43 @@ def test_cleanup_failure_is_visible() -> None: idp.delete_group("group") +def test_strict_cleanup_reports_each_failure_and_continues() -> None: + from lifecycle import ResourceManager + from proxy_client import build_proxy_client + + with _idp_server(delete_status=500) as (idp, deletions): + resources: Final = ResourceManager(client=build_proxy_client(), strict_cleanup=True) + strict: Final = idp.with_strict_cleanup() + resources.defer(lambda: strict.delete_group("group")) + resources.defer(lambda: strict.delete_user("user")) + with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as error: + resources.teardown() + assert len(error.value.exceptions) == 2 + assert deletions.get_nowait() == "/admin/realms/test/users/user" + assert deletions.get_nowait() == "/admin/realms/test/groups/group" + + +@pytest.mark.parametrize("groups", ((), ("one",), ("one", "two"))) +def test_provisioning_records_zero_one_or_multiple_groups(groups: tuple[str, ...]) -> None: + with _idp_server() as (idp, deletions): + with ExitStack() as cleanup: + + def defer(callback: Callable[[], object]) -> None: + cleanup.callback(callback) + + identity: Final = idp.provision_groups( + marker="memberships", + groups=groups, + defer=defer, + ) + assert identity.groups == groups + assert len(identity.group_ids) == len(groups) + assert deletions.get_nowait() == "/admin/realms/test/users/resource-1" + for _ in groups: + assert deletions.get_nowait() == "/admin/realms/test/groups/resource-1" + assert deletions.empty() + + def test_expired_admin_credentials_do_not_abort_remaining_cleanups() -> None: with _idp_server(admin_status=401) as (idp, _): cleanup: Final = ExitStack() diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 1b0133f12cb..879bd88980c 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -11,19 +11,53 @@ injected clock, so nothing here monkeypatches anything. from __future__ import annotations -from collections.abc import Iterable, Mapping +import json +from builtins import ExceptionGroup +from collections.abc import Callable, Generator, Iterable, Mapping +from contextlib import contextmanager from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import chain, repeat +from queue import SimpleQueue +from threading import Thread from types import MappingProxyType from typing import Final, cast import pytest from e2e_config import parse_replica_urls -from e2e_http import Result, Success -from models import KeyInfo, KeyInfoResponse, ModelListEntry, ModelsListResponse +from e2e_http import NoBody, Result, Success, without_retries +from idp import Keycloak +from lifecycle import ResourceManager +from management.jwt_actors import ActorFactory +from management.management_client import ManagementClient +from models import ( + ConnectionTestBody, + CredentialCreateBody, + KeyGenerateBody, + KeyInfo, + KeyInfoResponse, + KeyUpdateBody, + LiteLLMParamsBody, + McpServerCreateBody, + McpServerUpdateBody, + ModelListEntry, + ModelsListResponse, + OrgNewBody, + OrgUpdateBody, + SpendLogsParams, + TagNewBody, + TeamNewBody, + TeamUpdateBody, + ToolsetCreateBody, + ToolsetUpdateBody, + UserNewBody, + UserUpdateBody, +) from proxy_client import ( - ConvergeOutcome, + Caller, Converged, + ConvergeOutcome, + CredentialKind, EverywhereConverged, ModelsPoller, NeverConvergedOn, @@ -42,6 +76,108 @@ from proxy_client import ( ) from transport import Transport + +@contextmanager +def caller_boundary( + status: int = 200, bodies: SimpleQueue[bytes] | None = None, *, delete_status: int | None = None +) -> Generator[tuple[ManagementClient, SimpleQueue[str]]]: + received: Final[SimpleQueue[str]] = SimpleQueue() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: object) -> None: + pass + + def do_GET(self) -> None: + received.put(self.headers.get("Authorization", "")) + self.send_response(delete_status if self.path == "/key/delete" and delete_status is not None else status) + self.end_headers() + self.wfile.write( + b'{"key":"owned","info":{"key_alias":"owned"},"data":[{"id":"owned"}],"team_id":"owned","team_info":{},"model_id":"owned"}' + ) + + def do_POST(self) -> None: + body: Final = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + if bodies is not None: + bodies.put(body) + self.do_GET() + + do_PATCH = do_POST + do_PUT = do_POST + do_DELETE = do_POST + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + url: Final = f"http://127.0.0.1:{server.server_port}" + proxy: Final = build_proxy_client( + base_url=url, control_plane_base_url=url, replica_urls=(url,), master_key="bootstrap" + ) + try: + yield ManagementClient(proxy=proxy, master_key="bootstrap"), received + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +class TestBoundManagementCaller: + def test_actor_key_cleanup_reports_failure_and_continues(self) -> None: + with caller_boundary(delete_status=500) as (bootstrap, received), without_retries(): + resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True) + remaining: SimpleQueue[str] = SimpleQueue() + resources.defer(lambda: remaining.put("cleaned")) + factory: Final = ActorFactory( + bootstrap=bootstrap, + idp=Keycloak(base_url="http://unused.test", realm="test", admin_username="test", admin_password="test"), + resources=resources, + ) + assert factory.key().key == "owned" + with pytest.raises(ExceptionGroup, match="Resource cleanup failed") as failure: + resources.teardown() + assert len(failure.value.exceptions) == 1 + assert remaining.get_nowait() == "cleaned" + assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap") + + @pytest.mark.parametrize("kind", ("direct_jwt", "virtual_key", "dashboard_session")) + def test_direct_delegated_and_replica_reads_keep_the_bound_caller(self, kind: CredentialKind) -> None: + with caller_boundary() as (bootstrap, received): + caller: Final = Caller(credential="synthetic-caller", kind=kind, role="internal_user", tenant="tenant-a") + bound: Final = bootstrap.with_caller(caller) + bound.update_key(KeyUpdateBody(key="owned", key_alias="updated")) + bound.proxy.key_info("owned") + bound.proxy.read_back_everywhere( + "/key/info", + params=KeyUpdateBody(key="owned"), + response_type=KeyInfoResponse, + converged=lambda result: isinstance(result, Success), + ) + bound.proxy.read_body_back_everywhere( + "/key/info", KeyInfoResponse, settled=lambda result: result.info.key_alias == "owned" + ) + assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer synthetic-caller",) * 4 + assert received.empty() + bootstrap.proxy.key_info("owned") + assert received.get_nowait() == "Bearer bootstrap" + + def test_explicit_override_wins_without_rebinding_or_changing_master(self) -> None: + with caller_boundary() as (bootstrap, received): + bound: Final = bootstrap.with_caller(Caller(credential="bound", kind="direct_jwt", role="internal_user")) + bound.update_key(KeyUpdateBody(key="owned"), caller_key="override") + bound.proxy.key_info("owned") + assert received.get_nowait() == "Bearer override" + assert received.get_nowait() == "Bearer bound" + assert bound.master_key == "bootstrap" + + def test_credentials_are_absent_from_binding_and_header_diagnostics(self) -> None: + with caller_boundary() as (bootstrap, _): + caller: Final = Caller(credential="private-value", kind="direct_jwt", role="internal_user") + bound: Final = bootstrap.with_caller(caller) + assert "private-value" not in repr(caller) + assert "private-value" not in repr(bound) + assert "private-value" not in repr(bound.proxy.management_headers()) + assert "bootstrap" not in repr(bound) + + MODEL: Final = "gpt-under-test" _NO_TRANSPORTS: Final = cast(Transport, None) TIMEOUT: Final = 10.0 @@ -275,3 +411,166 @@ class TestReplicasFor: client: Final = ProxyClient(transport=_NO_TRANSPORTS, replicas={}, control_replicas={}) with pytest.raises(AssertionError, match="no replica is configured"): _ = client.replicas_for("/v1/models") + + +MANAGEMENT_OPERATIONS: Final[tuple[tuple[str, Callable[[ManagementClient], object]], ...]] = ( + ("generate_key", lambda c: c.generate_key(KeyGenerateBody())), + ("llm_only_key", lambda c: c.llm_only_key()), + ("update_key", lambda c: c.update_key(KeyUpdateBody(key="owned"))), + ("update_key_models", lambda c: c.update_key_models("owned", [])), + ("key_info", lambda c: c.key_info_as("owned")), + ("delete_key_strict", lambda c: c.delete_key_strict("owned")), + ("delete_model_strict", lambda c: c.delete_model_strict("owned")), + ( + "connection_test", + lambda c: c.connection_test( + ConnectionTestBody(litellm_params=LiteLLMParamsBody(model="synthetic"), mode="chat") + ), + ), + ("block_key", lambda c: c.block_key("owned")), + ("regenerate_key", lambda c: c.regenerate_key("owned")), + ("reset_key_spend", lambda c: c.reset_key_spend("owned", 0)), + ("key_list", lambda c: c.key_list("owned")), + ("key_alias_count", lambda c: c.key_alias_count("owned")), + ("create_team", lambda c: c.create_team(TeamNewBody(team_alias="owned"))), + ("update_team", lambda c: c.update_team(TeamUpdateBody(team_id="owned", team_alias="updated"))), + ("delete_team", lambda c: c.delete_team("owned")), + ("team_info", lambda c: c.team_info("owned")), + ("team_list_ids", lambda c: c.team_list_ids()), + ("team_info_status", lambda c: c.team_info_status("owned")), + ("add_team_member", lambda c: c.add_team_member("owned", "user")), + ("delete_team_member", lambda c: c.delete_team_member("owned", "user")), + ("create_user", lambda c: c.create_user(UserNewBody(user_email="actor@example.com", user_role="internal_user"))), + ("create_customer", lambda c: c.create_customer("owned")), + ("customer_info", lambda c: c.customer_info("owned")), + ("delete_customer", lambda c: c.delete_customer("owned")), + ("update_user", lambda c: c.update_user(UserUpdateBody(user_id="owned", user_role="internal_user"))), + ("delete_user", lambda c: c.delete_user("owned")), + ("delete_user_strict", lambda c: c.delete_user_strict("owned")), + ("user_info", lambda c: c.user_info("owned")), + ("user_count", lambda c: c.user_count("owned")), + ("user_list_ids", lambda c: c.user_list_ids("owned")), + ("create_org", lambda c: c.create_org(OrgNewBody(organization_alias="owned"))), + ("update_org", lambda c: c.update_org(OrgUpdateBody(organization_id="owned", organization_alias="updated"))), + ("delete_org", lambda c: c.delete_org("owned")), + ("org_info", lambda c: c.org_info("owned")), + ("org_info_status", lambda c: c.org_info_status("owned")), + ("create_tag", lambda c: c.create_tag(TagNewBody(name="owned"))), + ("delete_tag", lambda c: c.delete_tag("owned")), + ("tag_list", lambda c: c.tag_list()), + ("create_mcp_server", lambda c: c.create_mcp_server(McpServerCreateBody(alias="owned", url="http://example.test"))), + ("update_mcp_server", lambda c: c.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None))), + ("delete_mcp_server", lambda c: c.delete_mcp_server("owned")), + ("proxy.generate_key", lambda c: c.proxy.generate_key(KeyGenerateBody())), + ("proxy.delete_key", lambda c: c.proxy.delete_key("owned")), + ("proxy.delete_customers", lambda c: c.proxy.delete_customers(["owned"])), + ("proxy.key_info", lambda c: c.proxy.key_info("owned")), + ("proxy.memory_summary", lambda c: c.proxy.memory_summary_everywhere()), + ("proxy.model_info", lambda c: c.proxy.model_info()), + ("proxy.model_cost_map", lambda c: c.proxy.model_cost_map()), + ("proxy.create_model", lambda c: c.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic"))), + ("proxy.update_model", lambda c: c.proxy.update_model("owned", LiteLLMParamsBody(model="synthetic"))), + ("proxy.delete_model", lambda c: c.proxy.delete_model("owned")), + ("proxy.create_toolset", lambda c: c.proxy.create_toolset(ToolsetCreateBody(toolset_name="owned", tools=[]))), + ("proxy.update_toolset", lambda c: c.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None))), + ("proxy.delete_toolset", lambda c: c.proxy.delete_toolset("owned")), + ( + "proxy.create_credential", + lambda c: c.proxy.create_credential(CredentialCreateBody(credential_name="owned", credential_values={})), + ), + ("proxy.delete_credential", lambda c: c.proxy.delete_credential("owned")), + ("proxy.create_team", lambda c: c.proxy.create_team(TeamNewBody(team_alias="owned"))), + ("proxy.delete_team", lambda c: c.proxy.delete_team("owned")), + ("proxy.delete_user", lambda c: c.proxy.delete_user("owned")), + ("proxy.spend_logs", lambda c: c.proxy.spend_logs(SpendLogsParams(api_key="owned"))), + ("proxy.probe", lambda c: c.proxy.probe("/user/info", params=NoBody())), +) + + +@pytest.mark.parametrize( + ("name", "operation"), MANAGEMENT_OPERATIONS, ids=tuple(name for name, _ in MANAGEMENT_OPERATIONS) +) +@pytest.mark.parametrize("kind", ("master", "direct_jwt", "virtual_key", "dashboard_session")) +def test_management_operations_send_the_selected_credential( + name: str, + operation: Callable[[ManagementClient], object], + kind: CredentialKind, +) -> None: + with caller_boundary(status=401) as (bootstrap, received), without_retries(): + client: Final = ( + bootstrap + if kind == "master" + else bootstrap.with_caller(Caller(credential=f"synthetic-{kind}", kind=kind, role="internal_user")) + ) + try: + operation(client) + except AssertionError: + pass + expected: Final = "Bearer bootstrap" if kind == "master" else f"Bearer synthetic-{kind}" + assert received.get_nowait() == expected, name + assert received.empty(), "an unauthorized request must not be retried" + + +class TestSplitCallerPropagation: + def test_control_and_data_replica_readers_keep_the_caller(self) -> None: + with caller_boundary() as (data, data_headers), caller_boundary() as (control, control_headers): + data_url: Final = next(iter(data.proxy.replicas)) + control_url: Final = next(iter(control.proxy.replicas)) + proxy: Final = build_proxy_client( + base_url=data_url, + control_plane_base_url=control_url, + replica_urls=(data_url,), + master_key="bootstrap", + ).with_caller(Caller(credential="tenant-token", kind="direct_jwt", role="team_member")) + proxy.key_info("owned") + proxy.read_body_back_everywhere( + "/key/info", KeyInfoResponse, settled=lambda info: info.info.key_alias == "owned" + ) + proxy.read_back_everywhere( + "/key/info", + params=NoBody(), + response_type=KeyInfoResponse, + converged=lambda result: isinstance(result, Success), + ) + assert control_headers.get_nowait() == "Bearer tenant-token" + assert control_headers.get_nowait() == "Bearer tenant-token" + assert data_headers.get_nowait() == "Bearer tenant-token" + assert control_headers.empty() and data_headers.empty() + + def test_successful_team_and_model_polling_uses_the_bound_caller(self) -> None: + with caller_boundary() as (bootstrap, received): + bound: Final = bootstrap.with_caller(Caller(credential="caller", kind="direct_jwt", role="proxy_admin")) + bound.create_team(TeamNewBody(team_alias="owned")) + bound.proxy.create_model("owned", LiteLLMParamsBody(model="synthetic")) + assert tuple(received.get_nowait() for _ in range(4)) == ("Bearer caller",) * 4 + assert received.empty() + + def test_expired_shaped_token_is_sent_once_without_renewal(self) -> None: + with caller_boundary(status=401) as (bootstrap, received): + bound: Final = bootstrap.with_caller( + Caller(credential="expired.payload.signature", kind="direct_jwt", role="internal_user") + ) + result: Final = bound.key_info_as("owned") + assert not isinstance(result, Success) + assert received.get_nowait() == "Bearer expired.payload.signature" + assert received.empty() + + +@pytest.mark.parametrize("operation", ("server", "toolset")) +def test_partial_updates_preserve_explicit_null_at_the_http_boundary(operation: str) -> None: + bodies: Final[SimpleQueue[bytes]] = SimpleQueue() + with caller_boundary(status=401, bodies=bodies) as (bootstrap, _): + try: + if operation == "server": + bootstrap.update_mcp_server(McpServerUpdateBody(server_id="owned", alias=None)) + else: + bootstrap.proxy.update_toolset(ToolsetUpdateBody(toolset_id="owned", description=None)) + except AssertionError: + pass + expected: Final = ( + {"server_id": "owned", "alias": None} + if operation == "server" + else {"toolset_id": "owned", "description": None} + ) + assert json.loads(bodies.get_nowait()) == expected + assert bodies.empty() diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index e8caa801467..037db0c340f 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -7,11 +7,9 @@ client touches requests.* or builds raw dicts; they pass pydantic models here. from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Protocol -from pydantic import BaseModel - import e2e_http from e2e_http import ( URL, @@ -21,6 +19,7 @@ from e2e_http import ( Result, StreamingResponse, ) +from pydantic import BaseModel class Transport(Protocol): @@ -85,7 +84,7 @@ class Transport(Protocol): self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R] ) -> Result[R]: ... - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: ... + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: ... def upload[R: BaseModel]( self, @@ -113,7 +112,7 @@ class Transport(Protocol): @dataclass(frozen=True, slots=True) class HttpTransport: base_url: str - master_key: str + master_key: str = field(repr=False) request_timeout: float = 60.0 def _url(self, path: str) -> URL: @@ -245,10 +244,10 @@ class HttpTransport: timeout=self.request_timeout, ) - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), - headers=self.master, + headers=self.master if headers is None else headers, params=params, timeout=self.request_timeout, ) @@ -434,8 +433,8 @@ class SplitTransport: path, headers=headers, json=json, params=params, stream=stream ) - def probe(self, path: str, *, params: BaseModel) -> ProbeResult: - return self._route(path).probe(path, params=params) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: + return self._route(path).probe(path, params=params, headers=headers) def upload[R: BaseModel]( self, diff --git a/tests/e2e/ui/oidcSetup.ts b/tests/e2e/ui/oidcSetup.ts new file mode 100644 index 00000000000..943fa23cab3 --- /dev/null +++ b/tests/e2e/ui/oidcSetup.ts @@ -0,0 +1,30 @@ +import { chromium, expect } from "@playwright/test"; +import * as fs from "fs"; +import * as path from "path"; + +export default async function oidcSetup() { + const baseURL = process.env.E2E_OIDC_UI_URL; + const issuer = process.env.JWT_ISSUER; + const username = process.env.E2E_OIDC_USERNAME; + const password = process.env.E2E_OIDC_PASSWORD; + if (!baseURL || !issuer || !username || !password) { + throw new Error("The OIDC setup requires a running stack, issuer, and provisioned actor credentials"); + } + const artifactDir = process.env.E2E_UI_ARTIFACT_DIR || "."; + fs.mkdirSync(artifactDir, { recursive: true }); + const browser = await chromium.launch(); + try { + const page = await browser.newPage(); + await page.goto(`${baseURL.replace(/\/$/, "")}/sso/key/generate`); + await expect(page).toHaveURL(new RegExp(`^${issuer.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}/`)); + await page.getByLabel("Username or email").fill(username); + await page.getByLabel("Password", { exact: true }).fill(password); + await page.getByRole("button", { name: "Sign In", exact: true }).click(); + await page.waitForURL((url) => url.origin === new URL(baseURL).origin && url.pathname.startsWith("/ui")); + const statePath = path.join(artifactDir, "oidc.storageState.json"); + await page.context().storageState({ path: statePath }); + fs.chmodSync(statePath, 0o600); + } finally { + await browser.close(); + } +} diff --git a/tests/e2e/ui/playwright.oidc.config.ts b/tests/e2e/ui/playwright.oidc.config.ts new file mode 100644 index 00000000000..0fbe77e9bd2 --- /dev/null +++ b/tests/e2e/ui/playwright.oidc.config.ts @@ -0,0 +1,22 @@ +import { defineConfig, devices } from "@playwright/test"; +import * as path from "path"; + +const baseURL = process.env.E2E_OIDC_UI_URL; +if (!baseURL) throw new Error("E2E_OIDC_UI_URL must point to the running OIDC stack"); + +export default defineConfig({ + testDir: ".", + testMatch: "oidc/**/*.spec.ts", + retries: 0, + workers: 1, + outputDir: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc", "test-results"), + globalSetup: require.resolve("./oidcSetup"), + use: { + ...devices["Desktop Chrome"], + baseURL, + storageState: path.join(process.env.E2E_UI_ARTIFACT_DIR || ".", "oidc.storageState.json"), + trace: "off", + screenshot: "off", + video: "off", + }, +}); From c8bb54993e8b6db4eda84819b0015a1d9a85ca99 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:49:49 -0700 Subject: [PATCH 02/30] test: enforce isolated actors and stop OIDC process groups --- .../e2e/coverage_registry/management_cases.py | 2 +- tests/e2e/idp.py | 41 +++++++++++++++---- tests/e2e/management/jwt_actors.py | 2 +- tests/e2e/management/management_client.py | 17 ++++---- .../e2e/management/test_jwt_management_e2e.py | 20 ++++----- tests/e2e/test_idp.py | 25 ++++++++--- tests/e2e/test_proxy_client.py | 7 ++++ 7 files changed, 80 insertions(+), 34 deletions(-) diff --git a/tests/e2e/coverage_registry/management_cases.py b/tests/e2e/coverage_registry/management_cases.py index 15dc7d333c5..812dbfe5b8d 100644 --- a/tests/e2e/coverage_registry/management_cases.py +++ b/tests/e2e/coverage_registry/management_cases.py @@ -63,7 +63,7 @@ MANAGEMENT_CASES: Final = tuple( node=f"{JWT_CLASS}::test_admin_creates_reads_updates_clears_and_deletes_a_key[virtual_key]", credential_kind="virtual_key", actor="proxy_admin", - profile="database_role", + profile="group_scoped", method="POST", path="/key/generate", operation_family="key_lifecycle", diff --git a/tests/e2e/idp.py b/tests/e2e/idp.py index 12db91bbd88..2dc7c2ad71b 100644 --- a/tests/e2e/idp.py +++ b/tests/e2e/idp.py @@ -8,6 +8,7 @@ import secrets import signal import subprocess import sys +import time import warnings from collections.abc import Callable from contextlib import ExitStack @@ -426,6 +427,36 @@ def token_claims(token: str) -> TokenClaims: return TokenClaims.model_validate_json(base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))) +def _signal_process_group(process_id: int, signum: int) -> bool: + try: + os.killpg(process_id, signum) + except ProcessLookupError: + return False + return True + + +def _stop_process_group(child: subprocess.Popen[bytes]) -> None: + _signal_process_group(child.pid, signal.SIGTERM) + deadline: Final = time.monotonic() + 5 + while _process_group_exists(child.pid): + child.poll() + if time.monotonic() >= deadline: + _signal_process_group(child.pid, signal.SIGKILL) + break + time.sleep(0.05) + child.wait() + + +def _process_group_exists(process_id: int) -> bool: + try: + os.killpg(process_id, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + def run_oidc_profile(proxy_url: str, command: list[str]) -> int: idp: Final = keycloak_from_env().with_strict_cleanup() with ExitStack() as cleanup: @@ -441,17 +472,11 @@ def run_oidc_profile(proxy_url: str, command: list[str]) -> int: client: Final = idp.browser_client(callback_url=f"{proxy_url.rstrip('/')}/sso/callback", defer=defer) environment: Final = {**os.environ, **client.environment(idp.discovery()), "PROXY_BASE_URL": proxy_url} - with subprocess.Popen(command, env=environment) as child: + with subprocess.Popen(command, env=environment, start_new_session=True) as child: try: return child.wait() finally: - if child.poll() is None: - child.terminate() - try: - child.wait(timeout=5) - except subprocess.TimeoutExpired: - child.kill() - child.wait() + _stop_process_group(child) if __name__ == "__main__": diff --git a/tests/e2e/management/jwt_actors.py b/tests/e2e/management/jwt_actors.py index 909d1652ada..2d23549fe71 100644 --- a/tests/e2e/management/jwt_actors.py +++ b/tests/e2e/management/jwt_actors.py @@ -84,7 +84,7 @@ class ActorFactory: ) ) ) - self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key)) + self.resources.defer(lambda: self.bootstrap.delete_key_strict(created.key, missing_ok=True)) return created def tenant(self) -> Tenant: diff --git a/tests/e2e/management/management_client.py b/tests/e2e/management/management_client.py index a0243e868e2..8470d318db8 100644 --- a/tests/e2e/management/management_client.py +++ b/tests/e2e/management/management_client.py @@ -167,17 +167,18 @@ class ManagementClient: response_type=KeyInfoResponse, ) - def delete_key_strict(self, key: str, *, caller_key: str | None = None) -> None: + def delete_key_strict(self, key: str, *, caller_key: str | None = None, missing_ok: bool = False) -> None: """Strict delete for the act phase of a test: a failed delete is a hard failure, unlike the warn-only ProxyClient.delete_key used at teardown.""" - _ = unwrap( - self.proxy.transport.post( - "/key/delete", - headers=self.proxy.management_headers(caller_key), - json=KeyDeleteBody(keys=[key]), - response_type=NoBody, - ) + result = self.proxy.transport.post( + "/key/delete", + headers=self.proxy.management_headers(caller_key), + json=KeyDeleteBody(keys=[key]), + response_type=NoBody, ) + if missing_ok and isinstance(result, UnknownApiError) and result.status_code == 404: + return + _ = unwrap(result) def delete_model_strict(self, model_id: str) -> None: """Strict delete for the act phase of a test: a failed delete is a hard diff --git a/tests/e2e/management/test_jwt_management_e2e.py b/tests/e2e/management/test_jwt_management_e2e.py index 0d23f954158..5898073a4e6 100644 --- a/tests/e2e/management/test_jwt_management_e2e.py +++ b/tests/e2e/management/test_jwt_management_e2e.py @@ -102,31 +102,29 @@ class TestJwtManagement: @pytest.mark.parametrize("credential_kind", ("direct_jwt", "virtual_key")) def test_admin_creates_reads_updates_clears_and_deletes_a_key( self, - client: ManagementClient, - idp: Keycloak, - jwt_identity: Identity, - resources: ResourceManager, actor_factory: ActorFactory, credential_kind: Literal["direct_jwt", "virtual_key"], ) -> None: - actor: Final = actor_factory.create("proxy_admin") + tenant: Final = actor_factory.tenant() + actor: Final = actor_factory.create("proxy_admin", tenants=(tenant,), profile="group_scoped") virtual_key: Final = ( actor_factory.key(user_id=actor.identity.user_id).key if credential_kind == "virtual_key" else None ) - admin: Final = ( - virtual_key if virtual_key is not None else idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID) + admin: Final = virtual_key if virtual_key is not None else actor.mint_caller(actor_factory.idp).credential + bound: Final = actor_factory.bootstrap.with_caller( + Caller(credential=admin, kind=credential_kind, role="proxy_admin") ) - bound: Final = client.with_caller(Caller(credential=admin, kind=credential_kind, role="proxy_admin")) + assert bound.user_info().user_id == actor.identity.user_id alias: Final = f"e2e-jwt-key-{unique_marker()}" created: Final = unwrap( bound.generate_key( - KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]), + KeyGenerateBody(key_alias=alias, team_id=tenant.team_id, models=[CHEAP_OPENAI_MODEL]), ) ) - resources.defer(lambda: client.proxy.delete_key(created.key)) + actor_factory.resources.defer(lambda: actor_factory.bootstrap.delete_key_strict(created.key, missing_ok=True)) original: Final = unwrap(bound.key_info_as(created.key)).info - assert original.key_alias == alias and original.team_id == jwt_identity.group + assert original.key_alias == alias and original.team_id == tenant.team_id assert original.models == [CHEAP_OPENAI_MODEL] updated_alias: Final = f"{alias}-updated" diff --git a/tests/e2e/test_idp.py b/tests/e2e/test_idp.py index 8a92dfa527c..cf2d4f3118a 100644 --- a/tests/e2e/test_idp.py +++ b/tests/e2e/test_idp.py @@ -6,6 +6,7 @@ from __future__ import annotations import os import signal +import socket import subprocess import sys import time @@ -148,16 +149,27 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No assert deletions.empty() -@pytest.mark.parametrize("exit_mode", ("normal", "parent", "group")) +@pytest.mark.parametrize( + ("exit_mode", "ignore_termination"), (("normal", False), ("parent", False), ("group", False), ("parent", True)) +) def test_oidc_launcher_removes_client_on_exit_and_termination( - tmp_path: Path, exit_mode: Literal["normal", "parent", "group"] + tmp_path: Path, exit_mode: Literal["normal", "parent", "group"], ignore_termination: bool ) -> None: ready: Final = tmp_path / "ready" + descendant_command: Final = ( + "import signal,socket,time; from pathlib import Path; " + + ("signal.signal(signal.SIGTERM, signal.SIG_IGN); " if ignore_termination else "") + + "listener=socket.socket(); listener.bind(('127.0.0.1',0)); listener.listen(); " + f"Path({str(ready)!r}).write_text(str(listener.getsockname()[1])); time.sleep(120)" + ) child_command: Final = ( - "import os,time; from pathlib import Path; " + "import os,subprocess,sys,time; from pathlib import Path; " 'assert os.environ["GENERIC_CLIENT_SECRET"]; ' 'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; ' - f"Path({str(ready)!r}).touch(); " + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") + f"subprocess.Popen([sys.executable, '-c', {descendant_command!r}]); " + f"ready=Path({str(ready)!r})\n" + "while not ready.exists(): time.sleep(0.05)\n" + + ("raise SystemExit(7)" if exit_mode == "normal" else "time.sleep(120)") ) with _idp_server() as (idp, deletions): with subprocess.Popen( @@ -187,7 +199,10 @@ def test_oidc_launcher_removes_client_on_exit_and_termination( process.terminate() elif exit_mode == "group": os.killpg(process.pid, signal.SIGTERM) - assert process.wait(timeout=10) == (7 if exit_mode == "normal" else 143) + assert process.wait(timeout=15) == (7 if exit_mode == "normal" else 143) + with socket.socket() as connection: + connection.settimeout(1) + assert connection.connect_ex(("127.0.0.1", int(ready.read_text()))) != 0 finally: if process.poll() is None: os.killpg(process.pid, signal.SIGKILL) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 879bd88980c..0c4aed5bd65 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -121,6 +121,13 @@ def caller_boundary( class TestBoundManagementCaller: + def test_strict_key_cleanup_accepts_missing_only_when_requested(self) -> None: + with caller_boundary(delete_status=404) as (bootstrap, received), without_retries(): + with pytest.raises(AssertionError): + bootstrap.delete_key_strict("owned") + bootstrap.delete_key_strict("owned", missing_ok=True) + assert (received.get_nowait(), received.get_nowait()) == ("Bearer bootstrap", "Bearer bootstrap") + def test_actor_key_cleanup_reports_failure_and_continues(self) -> None: with caller_boundary(delete_status=500) as (bootstrap, received), without_retries(): resources: Final = ResourceManager(client=bootstrap.proxy, strict_cleanup=True) From 073d4fe2b01500526829523f6596a60428fbead9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:11:06 -0700 Subject: [PATCH 03/30] fix(responses): route mid-stream error events through exception_type so content_policy_fallbacks fire Mid-stream error events on the streaming Responses API were all raised as APIError, so a content_policy_violation event never matched the router's content-policy fallback dispatch and the client got the raw error instead of the fallback model's answer. Map each error event's code and status through the existing exception_type mapping, matching the non-streaming path, and unwrap the typed ContentPolicyViolationError and ContextWindowExceededError so the router routes them to the configured content_policy_fallbacks and context_window_fallbacks. --- litellm/responses/streaming_iterator.py | 48 ++++-- litellm/router.py | 10 +- .../test_streaming_iterator_error_events.py | 157 ++++++++++++++++-- tests/test_litellm/test_router.py | 105 ++++++++++++ 4 files changed, 288 insertions(+), 32 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 40ff88fc557..b3426fbbfef 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -221,6 +222,13 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None ) +def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: + if isinstance(mapped_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError)): + return True + status_code: Final = getattr(mapped_exception, "status_code", None) + return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 + + class BaseResponsesAPIStreamingIterator: """ Base class for streaming iterators that process responses from the Responses API. @@ -521,15 +529,8 @@ class BaseResponsesAPIStreamingIterator: getattr(self.completed_response, "response", None) if self.completed_response else None ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None - error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) - exception: Final = litellm.APIError( - status_code=_status_code_for_error_fields(error_type, error_code), - message=error_message, - llm_provider=self.custom_llm_provider or "", - model=self.model or "", - ) - self._handle_failure(exception) + self._handle_failure(self._map_error_event_exception(error_info)) def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: @@ -551,6 +552,26 @@ class BaseResponsesAPIStreamingIterator: self.logging_obj._response_cost_calculator(result=response_obj) or 0.0 ) + def _map_error_event_exception(self, error_obj: object) -> Exception: + error_message, error_type, error_code = _error_event_fields(error_obj) + status_code: Final = _status_code_for_error_fields(error_type, error_code) + error_body: Final = {"message": error_message, "type": error_type, "code": error_code} + provider_exception: Final = BaseLLMException( + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {error_body}}}", + body=error_body, + ) + try: + return litellm.exception_type( + model=self.model or "", + custom_llm_provider=self.custom_llm_provider or "", + original_exception=provider_exception, + completion_kwargs={}, + extra_kwargs={}, + ) + except Exception as mapped_exception: + return mapped_exception + def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) if chunk_type not in ("error", "response.failed"): @@ -562,15 +583,8 @@ class BaseResponsesAPIStreamingIterator: else getattr(result, "error", None) ) - error_message, error_type, error_code = _error_event_fields(error_obj) - status_code: Final = _status_code_for_error_fields(error_type, error_code) - mapped_exception: Final = litellm.APIError( - status_code=status_code, - message=error_message, - llm_provider=self.custom_llm_provider or "", - model=self.model or "", - ) - if 400 <= status_code < 500 and status_code != 429: + mapped_exception: Final = self._map_error_event_exception(error_obj) + if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception raise MidStreamFallbackError( message=str(mapped_exception), diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..79854a11150 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3271,8 +3271,16 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + fallback_trigger: Final[Exception] = ( + e.original_exception + if isinstance( + e.original_exception, + (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), + ) + else e + ) fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e, + e=fallback_trigger, disable_fallbacks=False, fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index ad74861c096..73afbb5e63a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -1,9 +1,14 @@ """ Regression: in-stream error events (type="error", type="response.failed") must raise instead of being returned as benign chunks, mirroring chat streaming -semantics (_handle_stream_fallback_error): non-retriable 4xx (except 429) -raise litellm.APIError directly; 429 and 5xx are wrapped in -MidStreamFallbackError so the Router's mid-stream fallback machinery fires. +semantics (_handle_stream_fallback_error). The event's code, type and status go +through litellm.exception_type, so each event raises the same typed exception +the non-streaming path raises for that provider error: non-retriable 4xx +(except 429) raise that typed exception directly, while 429, 5xx, +ContentPolicyViolationError and ContextWindowExceededError are wrapped in +MidStreamFallbackError so the Router's mid-stream fallback machinery fires and +its content_policy_fallbacks / context_window_fallbacks dispatch sees the +trigger it matches on. Status mapping must consider both the OpenAI error `type` (e.g. "invalid_request_error") and `code` (e.g. "invalid_prompt", @@ -66,12 +71,12 @@ def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback( with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 500 - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.InternalServerError) assert exc_info.value.original_exception.status_code == 500 def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fallback(): - """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped APIError.""" + """429 is retriable: it must be wrapped so the Router can fall back, carrying the mapped RateLimitError.""" iterator = _make_iterator() chunk = _make_error_chunk("tokens", "rate_limit_exceeded", "Too many requests") with pytest.raises(MidStreamFallbackError) as exc_info: @@ -79,15 +84,15 @@ def test_maybe_raise_for_error_event_maps_rate_limit_code_to_429_mid_stream_fall assert exc_info.value.status_code == 429 assert exc_info.value.generated_content == "" assert exc_info.value.is_pre_first_chunk is True - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) assert exc_info.value.original_exception.status_code == 429 def test_maybe_raise_for_error_event_maps_invalid_request_type_to_400(): - """Client errors classified via the `type` field must raise APIError directly (no fallback).""" + """Client errors classified via the `type` field must raise BadRequestError directly (no fallback).""" iterator = _make_iterator() chunk = _make_error_chunk("invalid_request_error", "invalid_prompt", "bad request") - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) @@ -99,12 +104,84 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): chunk = Mock() chunk.type = "error" chunk.error = {"code": "context_length_exceeded", "message": "too long"} - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) +def test_maybe_raise_for_error_event_wraps_context_window_exceeded_for_context_window_fallbacks(): + """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming + path and is wrapped so the Router's context_window_fallbacks dispatch fires mid-stream.""" + iterator = _make_iterator() + chunk = _make_error_chunk( + "invalid_request_error", + "context_length_exceeded", + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContextWindowExceededError) + assert exc_info.value.status_code == 400 + + +CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "azure"]) +def test_maybe_raise_for_error_event_wraps_content_policy_violation_for_content_policy_fallbacks( + custom_llm_provider: str, +): + """Regression: a content_policy_violation error event used to raise a bare APIError, so the Router's + content_policy_fallbacks never fired. It must map to ContentPolicyViolationError (the same exception the + non-streaming path raises) and be wrapped so the Router's mid-stream fallback catches it.""" + iterator = _make_iterator() + iterator.custom_llm_provider = custom_llm_provider + chunk = _make_error_chunk("invalid_request_error", "content_policy_violation", CONTENT_POLICY_MESSAGE) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + assert exc_info.value.original_exception.status_code == 400 + assert exc_info.value.status_code == 400 + assert exc_info.value.is_pre_first_chunk is True + assert CONTENT_POLICY_MESSAGE in str(exc_info.value.original_exception) + + +def test_maybe_raise_for_response_failed_event_wraps_content_policy_violation(): + iterator = _make_iterator() + chunk = _make_failed_chunk( + {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE} + ) + with pytest.raises(MidStreamFallbackError) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + + +@pytest.mark.parametrize( + "error_type,error_code,expected_exception", + [ + ("invalid_request_error", "content_policy_violation", litellm.ContentPolicyViolationError), + ("tokens", "rate_limit_exceeded", litellm.RateLimitError), + ("invalid_request_error", "insufficient_quota", litellm.RateLimitError), + ("server_error", "internal_error", litellm.InternalServerError), + ("invalid_request_error", "invalid_prompt", litellm.BadRequestError), + ("invalid_request_error", "model_not_found", litellm.NotFoundError), + ("server_error", "vector_store_timeout", litellm.Timeout), + ], +) +def test_error_event_raises_the_same_typed_exception_as_the_non_streaming_path( + error_type: str, error_code: str, expected_exception: type[Exception] +): + iterator = _make_iterator() + chunk = _make_error_chunk(error_type, error_code, "provider message") + with pytest.raises((MidStreamFallbackError, expected_exception)) as exc_info: + iterator._maybe_raise_for_error_event(chunk) + raised = exc_info.value + typed_exception = raised.original_exception if isinstance(raised, MidStreamFallbackError) else raised + assert type(typed_exception) is expected_exception + assert "provider message" in str(typed_exception) + + def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): """OpenAI returns HTTP 429 for insufficient_quota; it must not map to 400 even though its type is invalid_request_error-adjacent, and it must be wrapped for fallback.""" @@ -113,6 +190,7 @@ def test_maybe_raise_for_error_event_maps_insufficient_quota_to_429(): with pytest.raises(MidStreamFallbackError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 429 + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) def test_maybe_raise_for_error_event_passes_through_normal_chunk(): @@ -186,10 +264,43 @@ async def test_async_iterator_raises_mid_stream_fallback_on_rate_limit_error_eve assert exc_info.value.status_code == 429 assert exc_info.value.is_pre_first_chunk is True assert exc_info.value.generated_content == "" - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) assert exc_info.value.original_exception.status_code == 429 +@pytest.mark.asyncio +async def test_async_iterator_content_policy_violation_after_first_chunk_carries_generated_content(): + """The customer's case: text streams, then the provider halts the stream with a + content_policy_violation error event. The iterator must surface ContentPolicyViolationError + inside MidStreamFallbackError, together with the text already streamed.""" + iterator = _make_async_iterator_with_events( + [ + {"type": "response.output_text.delta", "delta": "partial "}, + { + "type": "error", + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": CONTENT_POLICY_MESSAGE, + }, + }, + ] + ) + + chunks = [] + + async def _drain(): + async for chunk in iterator: + chunks.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _drain() + assert len(chunks) == 1 + assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) + assert exc_info.value.is_pre_first_chunk is False + assert exc_info.value.generated_content == "partial " + + @pytest.mark.asyncio async def test_async_iterator_error_after_first_chunk_carries_generated_content(): """An error after streamed output must expose the accumulated text so the router's @@ -265,7 +376,7 @@ def test_handle_logging_failed_response_maps_rate_limit_to_429(): ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] - assert isinstance(logged_exception, litellm.APIError) + assert isinstance(logged_exception, litellm.RateLimitError) assert logged_exception.status_code == 429 assert "throttled" in str(logged_exception) @@ -282,10 +393,28 @@ def test_handle_logging_failed_response_maps_type_field_to_400(): ): iterator._handle_logging_failed_response() logged_exception = mock_run_async.call_args.kwargs["exception"] - assert isinstance(logged_exception, litellm.APIError) + assert isinstance(logged_exception, litellm.BadRequestError) assert logged_exception.status_code == 400 +def test_handle_logging_failed_response_logs_content_policy_violation(): + """Failure logging must record the same typed exception the stream raises, so logging + integrations see a content policy violation instead of a generic APIError.""" + iterator = _make_iterator() + iterator.completed_response = _make_failed_chunk( + {"type": "invalid_request_error", "code": "content_policy_violation", "message": CONTENT_POLICY_MESSAGE} + ) + with ( + patch.object(import_module("litellm.responses.streaming_iterator"), "run_async_function") as mock_run_async, + patch.object(import_module("litellm.responses.streaming_iterator"), "executor"), + ): + iterator._handle_logging_failed_response() + logged_exception = mock_run_async.call_args.kwargs["exception"] + assert isinstance(logged_exception, litellm.ContentPolicyViolationError) + assert logged_exception.status_code == 400 + assert CONTENT_POLICY_MESSAGE in str(logged_exception) + + def test_handle_logging_failed_response_records_usage_and_cost(): """Usage on a response.failed event must reach failure spend accounting via combined_usage_object.""" iterator = _make_iterator() @@ -357,7 +486,7 @@ def test_sync_iterator_raises_mid_stream_fallback_on_rate_limit_error_event(): for _ in iterator: pass assert exc_info.value.status_code == 429 - assert isinstance(exc_info.value.original_exception, litellm.APIError) + assert isinstance(exc_info.value.original_exception, litellm.RateLimitError) def test_every_openai_sdk_response_error_code_has_explicit_status_mapping(): @@ -413,7 +542,7 @@ def test_maybe_raise_for_response_failed_event_maps_image_code_to_400(): chunk = Mock() chunk.type = "response.failed" chunk.response = mock_response_obj - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index b8a0d70f5bc..d74074dd149 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3654,6 +3654,111 @@ async def test_aresponses_streaming_iterator_fallback(): assert call_kwargs["disable_fallbacks"] is False +@pytest.mark.asyncio +async def test_aresponses_streaming_content_policy_error_event_routes_to_content_policy_fallback(): + """Regression: a mid-stream content_policy_violation error event never reached + content_policy_fallbacks. The iterator raised a bare APIError the wrapper does not + catch, and even once wrapped, the MidStreamFallbackError envelope was handed to the + fallback dispatch, whose isinstance branch on ContentPolicyViolationError never matched. + The stream below is the customer's shape: a raw OpenAI error event with code + content_policy_violation, transformed by the real OpenAI config, and the router must + call the content_policy_fallbacks target, not the general fallbacks one.""" + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + router = litellm.Router( + model_list=[ + {"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k1"}}, + { + "model_name": "content-fallback", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k2"}, + }, + {"model_name": "general-fallback", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["general-fallback"]}], + content_policy_fallbacks=[{"primary": ["content-fallback"]}], + ) + error_event = { + "type": "error", + "sequence_number": 2, + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": "This content was flagged for possible cybersecurity risk. The response was halted mid-stream.", + "param": None, + }, + } + + async def aiter_bytes(): + yield f"data: {json.dumps(error_event)}\n\n".encode() + + raw_response = MagicMock() + raw_response.headers = {} + raw_response.aiter_bytes = aiter_bytes + logging_obj = MagicMock(spec=LiteLLMLogging) + logging_obj.model_call_details = {"litellm_params": {}} + logging_obj.completion_start_time = None + source = ResponsesAPIStreamingIterator( + response=raw_response, + model="gpt-5.4", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + fallback_chunks = [MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed")] + fallback_call = AsyncMock(return_value=_AsyncList(fallback_chunks)) + + wrapped = await router._aresponses_streaming_iterator( + response=source, + initial_kwargs={ + "model": "primary", + "stream": True, + "input": "Hi", + "original_generic_function": fallback_call, + }, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == fallback_chunks + fallback_call.assert_awaited_once() + assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" + + +@pytest.mark.asyncio +async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): + """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the + ContentPolicyViolationError carried inside MidStreamFallbackError, not the envelope.""" + router = _make_router_with_fallback("openai/gpt-5.4", "openai/gpt-5-mini") + content_policy_error = litellm.ContentPolicyViolationError( + message="flagged mid-stream", llm_provider="openai", model="openai/gpt-5.4" + ) + src = _make_responses_iterator( + chunks=[MagicMock(type="response.created")], + error=MidStreamFallbackError( + message=str(content_policy_error), + model="openai/gpt-5.4", + llm_provider="openai", + original_exception=content_policy_error, + is_pre_first_chunk=True, + ), + model="openai/gpt-5.4", + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AsyncList([MagicMock(type="response.completed")])), + ) as mock_fallback_utils: + wrapped = await router._aresponses_streaming_iterator( + response=src, + initial_kwargs={"model": "openai/gpt-5.4", "stream": True, "input": "Hi"}, + ) + [chunk async for chunk in wrapped] + + mock_fallback_utils.assert_awaited_once() + assert mock_fallback_utils.await_args.kwargs["e"] is content_policy_error + + @pytest.mark.asyncio @pytest.mark.parametrize( "fallback_headers", From c2463728593b935448bcb09f54ef4fd070e9f8bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:35:53 -0700 Subject: [PATCH 04/30] fix(responses): import BaseLLMException lazily and collect stream chunks via anext Move the BaseLLMException import into _map_error_event_exception so the module no longer imports it at load time, clearing the module-level cyclic import CodeQL flagged. The class is used only on the cold error path. Replace the mutable list-append test collector with aiter/anext so the regression tests read the stream immutably. --- litellm/responses/streaming_iterator.py | 3 ++- .../test_streaming_iterator_error_events.py | 22 ++++++++----------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index b3426fbbfef..a15571acb7c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,7 +31,6 @@ from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -553,6 +552,8 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: + from litellm.llms.base_llm.chat.transformation import BaseLLMException + error_message, error_type, error_code = _error_event_fields(error_obj) status_code: Final = _status_code_for_error_fields(error_type, error_code) error_body: Final = {"message": error_message, "type": error_type, "code": error_code} diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 73afbb5e63a..2f4fba45cee 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -287,15 +287,12 @@ async def test_async_iterator_content_policy_violation_after_first_chunk_carries ] ) - chunks = [] - - async def _drain(): - async for chunk in iterator: - chunks.append(chunk) + stream = aiter(iterator) + first_chunk = await anext(stream) + assert first_chunk is not None with pytest.raises(MidStreamFallbackError) as exc_info: - await _drain() - assert len(chunks) == 1 + await anext(stream) assert isinstance(exc_info.value.original_exception, litellm.ContentPolicyViolationError) assert exc_info.value.is_pre_first_chunk is False assert exc_info.value.generated_content == "partial " @@ -316,14 +313,13 @@ async def test_async_iterator_error_after_first_chunk_carries_generated_content( ] ) - chunks = [] - async def _drain(): - async for chunk in iterator: - chunks.append(chunk) + stream = aiter(iterator) + first_chunk = await anext(stream) + second_chunk = await anext(stream) + assert first_chunk is not None and second_chunk is not None with pytest.raises(MidStreamFallbackError) as exc_info: - await _drain() - assert len(chunks) == 2 + await anext(stream) assert exc_info.value.status_code == 500 assert exc_info.value.is_pre_first_chunk is False assert exc_info.value.generated_content == "hello world" From fff7a2cecfb113082681665d236525f2690aae45 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:18:59 -0700 Subject: [PATCH 05/30] fix(responses): keep context-window events out of mid-stream fallback and fix stale exception assertions --- litellm/responses/streaming_iterator.py | 2 +- litellm/router.py | 7 +++---- .../test_openai_responses_api.py | 9 +++++---- ...est_router_aresponses_streaming_fallback.py | 2 +- .../test_streaming_iterator_error_events.py | 18 ++++++++++-------- 5 files changed, 20 insertions(+), 18 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index a15571acb7c..b39e130242d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -222,7 +222,7 @@ def _status_code_for_error_fields(error_type: str | None, error_code: str | None def _mid_stream_fallback_eligible(mapped_exception: Exception) -> bool: - if isinstance(mapped_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError)): + if isinstance(mapped_exception, litellm.ContentPolicyViolationError): return True status_code: Final = getattr(mapped_exception, "status_code", None) return not isinstance(status_code, int) or status_code >= 500 or status_code == 429 diff --git a/litellm/router.py b/litellm/router.py index 79854a11150..a9a8f3e2739 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3271,12 +3271,11 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. fallback_trigger: Final[Exception] = ( e.original_exception - if isinstance( - e.original_exception, - (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), - ) + if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e ) fallback_response = await self.async_function_with_fallbacks_common_utils( diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index 05bb9113835..c7712d96969 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1627,9 +1627,10 @@ async def test_openai_responses_api_token_limit_error(): Parsing the in-stream ErrorEvent must not raise "pydantic_core._pydantic_core.ValidationError: 3 validation errors for ErrorEvent". - The iterator now surfaces the event as litellm.APIError with status 400 - (invalid_request_error is a non-retriable client error, so no - MidStreamFallbackError wrapping) carrying the provider's message. + The iterator routes the event through litellm.exception_type, so it surfaces as + the typed 400 client error the non-streaming path raises (litellm.BadRequestError) + carrying the provider's message. invalid_request_error is a non-retriable client + error, so there is no MidStreamFallbackError wrapping. """ litellm._turn_on_debug() @@ -1644,7 +1645,7 @@ async def test_openai_responses_api_token_limit_error(): async for event in response: print(event) - with pytest.raises(litellm.APIError) as exc_info: + with pytest.raises(litellm.BadRequestError) as exc_info: await _drain() assert exc_info.value.status_code == 400 diff --git a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py index ee4750e9db8..0d33435cf7a 100644 --- a/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py +++ b/tests/router_unit_tests/test_router_aresponses_streaming_fallback.py @@ -372,7 +372,7 @@ async def test_aresponses_fallback_on_in_stream_error_event(): raised = mock_fallback.await_args.kwargs["e"] assert isinstance(raised, MidStreamFallbackError) assert raised.status_code == 429 - assert isinstance(raised.original_exception, litellm.APIError) + assert isinstance(raised.original_exception, litellm.RateLimitError) assert raised.original_exception.status_code == 429 assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question" diff --git a/tests/test_litellm/responses/test_streaming_iterator_error_events.py b/tests/test_litellm/responses/test_streaming_iterator_error_events.py index 2f4fba45cee..3d7c220804a 100644 --- a/tests/test_litellm/responses/test_streaming_iterator_error_events.py +++ b/tests/test_litellm/responses/test_streaming_iterator_error_events.py @@ -4,11 +4,11 @@ raise instead of being returned as benign chunks, mirroring chat streaming semantics (_handle_stream_fallback_error). The event's code, type and status go through litellm.exception_type, so each event raises the same typed exception the non-streaming path raises for that provider error: non-retriable 4xx -(except 429) raise that typed exception directly, while 429, 5xx, -ContentPolicyViolationError and ContextWindowExceededError are wrapped in +(except 429) raise that typed exception directly, so a context-length event +surfaces as ContextWindowExceededError(400) with no MidStreamFallbackError +wrapping, while 429, 5xx and ContentPolicyViolationError are wrapped in MidStreamFallbackError so the Router's mid-stream fallback machinery fires and -its content_policy_fallbacks / context_window_fallbacks dispatch sees the -trigger it matches on. +its content_policy_fallbacks dispatch sees the trigger it matches on. Status mapping must consider both the OpenAI error `type` (e.g. "invalid_request_error") and `code` (e.g. "invalid_prompt", @@ -110,19 +110,21 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): assert not isinstance(exc_info.value, MidStreamFallbackError) -def test_maybe_raise_for_error_event_wraps_context_window_exceeded_for_context_window_fallbacks(): +def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(): """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming - path and is wrapped so the Router's context_window_fallbacks dispatch fires mid-stream.""" + path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream + fallback, preserving the direct-SDK 400 contract from issue #15785.""" iterator = _make_iterator() chunk = _make_error_chunk( "invalid_request_error", "context_length_exceeded", "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", ) - with pytest.raises(MidStreamFallbackError) as exc_info: + with pytest.raises(litellm.ContextWindowExceededError) as exc_info: iterator._maybe_raise_for_error_event(chunk) - assert isinstance(exc_info.value.original_exception, litellm.ContextWindowExceededError) assert exc_info.value.status_code == 400 + assert not isinstance(exc_info.value, MidStreamFallbackError) + assert "maximum context length" in str(exc_info.value) CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." From 03e6dd051cdc2d4fc307783b3a11ab618bc70216 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 06:59:10 +0000 Subject: [PATCH 06/30] fix(cli): drop enum.StrEnum so the CLI imports on Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/client/cli/commands/pi.py | 4 ++-- tests/test_litellm/proxy/client/cli/test_pi.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 9810e81ae36..5c749959638 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -10,7 +10,7 @@ import os import tempfile from collections.abc import Callable, Mapping from dataclasses import dataclass -from enum import StrEnum +from enum import Enum from pathlib import Path from types import MappingProxyType from typing import Annotated, Final @@ -25,7 +25,7 @@ LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" _REJECTED_STATUSES: Final = frozenset((401, 403)) -class ListingFailure(StrEnum): +class ListingFailure(str, Enum): """Why a proxy could not be listed, decided once where the HTTP outcome is classified. `unreachable` means no response at all; the other kinds prove the proxy answered, so callers diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 03e5d7dd197..3c343c9e268 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -1,3 +1,4 @@ +import inspect import json import os import stat @@ -7,6 +8,7 @@ from pathlib import Path import pytest import requests +from litellm.proxy.client.cli.commands import pi from litellm.proxy.client.cli.commands.pi import ( ListingFailure, ModelLimits, @@ -20,6 +22,18 @@ from litellm.proxy.client.cli.commands.pi import ( ) +def test_pi_module_has_no_strenum_import(): + """enum.StrEnum is Python 3.11+; pyproject allows 3.10, so pi.py must not import it.""" + assert "StrEnum" not in inspect.getsource(pi) + + +def test_listing_failure_is_str_enum(): + assert issubclass(ListingFailure, str) + assert ListingFailure.REJECTED.value == "rejected" + assert ListingFailure("rejected") is ListingFailure.REJECTED + assert str(ListingFailure.REJECTED.value) == "rejected" + + class _FakeResponse: def __init__(self, status_code, payload=None): self.status_code = status_code From 0679d799d483dabe74887a6a54b25b737676af8f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 07:07:06 +0000 Subject: [PATCH 07/30] ci: exercise the lite CLI on the Python 3.10 import smoke job Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-code-quality.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 9c7e0db7065..987f66773f2 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -178,7 +178,7 @@ jobs: version: "0.10.9" - name: Install dependencies - run: uv sync --frozen --extra proxy --python 3.10 + run: uv sync --frozen --extra proxy --extra cli --python 3.10 - run: uv run --no-sync python --version @@ -187,3 +187,6 @@ jobs: - name: Check litellm CLI run: uv run --no-sync litellm --version + + - name: Check lite CLI + run: uv run --no-sync lite version From 19c43eb875944e326d3c98743f4bcd0835298761 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 07:08:24 +0000 Subject: [PATCH 08/30] test(cli): drop structural StrEnum source check; smoke job covers the 3.10 import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/client/cli/test_pi.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index 3c343c9e268..6bb49566580 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/tests/test_litellm/proxy/client/cli/test_pi.py @@ -1,4 +1,3 @@ -import inspect import json import os import stat @@ -8,7 +7,6 @@ from pathlib import Path import pytest import requests -from litellm.proxy.client.cli.commands import pi from litellm.proxy.client.cli.commands.pi import ( ListingFailure, ModelLimits, @@ -22,11 +20,6 @@ from litellm.proxy.client.cli.commands.pi import ( ) -def test_pi_module_has_no_strenum_import(): - """enum.StrEnum is Python 3.11+; pyproject allows 3.10, so pi.py must not import it.""" - assert "StrEnum" not in inspect.getsource(pi) - - def test_listing_failure_is_str_enum(): assert issubclass(ListingFailure, str) assert ListingFailure.REJECTED.value == "rejected" From d4f2119b03faa175e790dd86cb3c8aa46f546293 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:41:46 +0000 Subject: [PATCH 09/30] fix(cost): bill gemini-embedding-2 per token and stop double charging audio Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/utils.py | 14 ++++- .../batch_embed_content_transformation.py | 60 ++----------------- ...odel_prices_and_context_window_backup.json | 15 ++--- model_prices_and_context_window.json | 15 ++--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 33 ++++++++++ ...test_batch_embed_content_transformation.py | 42 +++++-------- 6 files changed, 75 insertions(+), 104 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 8fc428b38ae..a004f46b291 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -956,12 +956,17 @@ def _calculate_input_cost( ) ### AUDIO COST - if prompt_tokens_details["audio_tokens"]: + if prompt_tokens_details["audio_tokens"] and not ( + prompt_tokens_details["audio_length_seconds"] + and model_info.get("input_cost_per_audio_per_second") is not None + ): audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier) prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"]) ### IMAGE TOKEN COST - if prompt_tokens_details["image_tokens"]: + if prompt_tokens_details["image_tokens"] and not ( + prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None + ): # For image token costs: # First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token. image_token_cost_key = "input_cost_per_image_token" @@ -970,7 +975,10 @@ def _calculate_input_cost( prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]) ### VIDEO TOKEN COST - if prompt_tokens_details["video_tokens"]: + if prompt_tokens_details["video_tokens"] and not ( + prompt_tokens_details["video_length_seconds"] + and model_info.get("input_cost_per_video_per_second") is not None + ): video_token_cost_key = "input_cost_per_video_token" if model_info.get(video_token_cost_key) is None: video_token_cost_key = "input_cost_per_token" diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index e7fd9a0d08b..8e120ab9fe6 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -297,9 +297,6 @@ def transform_openai_input_gemini_embed_content( return request_body -_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"}) -_VIDEO_TOKENS_PER_SECOND: Final = 258.0 -_AUDIO_TOKENS_PER_SECOND: Final = 32.0 _usage_metadata_adapter: Final = TypeAdapter(UsageMetadata) @@ -312,40 +309,6 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None: return None -def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: - if isinstance(input, str): - return (input,) - return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element])) - - -def _is_image_element( - element: str, - resolved_files: Mapping[str, Mapping[str, str]], -) -> bool: - if element.startswith("data:") and ";base64," in element: - try: - mime_type, _ = _parse_data_url(element) - except ValueError: - return False - return mime_type in _IMAGE_MIME_TYPES - if _is_gcs_url(element): - try: - return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES - except ValueError: - return False - if _is_file_reference(element): - file_info: Final = resolved_files.get(element) - return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES - return False - - -def _count_input_images( - input: GeminiEmbeddingInput, - resolved_files: Mapping[str, Mapping[str, str]], -) -> int: - return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) - - def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: return sum(detail["tokenCount"] for detail in details if detail["modality"] == modality) @@ -362,7 +325,6 @@ def _usage_from_embed_content_response( input: GeminiEmbeddingInput, model: str, raw_usage_metadata: object, - resolved_files: Mapping[str, Mapping[str, str]], ) -> Usage: usage_metadata: Final = _parse_usage_metadata(raw_usage_metadata) if usage_metadata is None: @@ -374,28 +336,17 @@ def _usage_from_embed_content_response( details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or () text_tokens: Final = _tokens_for_modality(details, "TEXT") audio_tokens: Final = _tokens_for_modality(details, "AUDIO") + image_tokens: Final = _tokens_for_modality(details, "IMAGE") video_tokens: Final = _tokens_for_modality(details, "VIDEO") - image_count: Final = _count_input_images(input, resolved_files) - - video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0 - audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0 - - # generic_cost_per_token rewrites text_tokens to the full prompt minus - # other modalities when both text_tokens and image_count are zero. For - # video, that misallocates video tokens to text; a 1-token floor sidesteps - # the rewrite and keeps billing on input_cost_per_video_per_second. - needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0 - resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens return Usage( prompt_tokens=prompt_tokens, total_tokens=total_tokens, prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=resolved_text_tokens, + text_tokens=text_tokens, audio_tokens=audio_tokens, - image_count=image_count, - video_length_seconds=video_length_seconds, - audio_length_seconds=audio_length_seconds, + image_tokens=image_tokens, + video_tokens=video_tokens, ), ) @@ -415,8 +366,6 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint - resolved_files: Mapping of file references (files/abc) to {mime_type, uri}, - used to bill resolved image references at the per-image rate Returns: EmbeddingResponse with single embedding @@ -438,7 +387,6 @@ def process_embed_content_response( input=input, model=model, raw_usage_metadata=response_json.get("usageMetadata"), - resolved_files=resolved_files or {}, ) return model_response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0c60ff26635..74ca23fcc8a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25615,13 +25615,11 @@ "uses_embed_content": true }, "gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25648,13 +25646,11 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25705,11 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0c60ff26635..74ca23fcc8a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25615,13 +25615,11 @@ "uses_embed_content": true }, "gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25648,13 +25646,11 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, "input_cost_per_audio_token": 6.5e-06, - "input_cost_per_image": 0.00012, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25705,11 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_token_batches": 1e-07, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 2289de9a951..4a02bb1a638 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -74,6 +74,39 @@ def test_missing_cache_read_policy_preserves_billing(prompt_tokens, read_rate, s assert prompt_cost == pytest.approx((prompt_tokens - 100) * billed[0] + 100 * billed[4]) +def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: + model_info: ModelInfo = { + "key": "gemini-embedding-2", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": 2e-7, + "input_cost_per_audio_token": 6.5e-6, + "input_cost_per_audio_per_second": 0.00016, + "output_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "embedding", + } + usage = Usage( + prompt_tokens=64, + completion_tokens=0, + total_tokens=64, + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=64, + audio_length_seconds=2, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(2 * 0.00016) + + def test_missing_cache_read_uses_off_peak_input_rate(): from datetime import datetime, timezone diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 86b3f0976ab..fbf86105e71 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -22,7 +22,6 @@ from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation from litellm.types.llms.vertex_ai import VertexAIBatchEmbeddingsResponseObject from litellm.types.utils import EmbeddingResponse - IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII" GCS_URL = "gs://my-bucket/image.png" @@ -324,7 +323,7 @@ class TestProcessEmbedContentResponseUsage: ) assert result.usage.prompt_tokens == 258 assert result.usage.total_tokens == 258 - assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.image_tokens == 258 prompt_cost, _ = generic_cost_per_token( model=self.MODEL, @@ -358,7 +357,7 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost > 0 - def test_video_modality_derives_seconds_and_text_floor(self): + def test_video_modality_preserves_token_count(self): response_json = { "embedding": {"values": [0.1]}, "usageMetadata": { @@ -374,10 +373,8 @@ class TestProcessEmbedContentResponseUsage: response_json=response_json, ) assert result.usage.prompt_tokens == 516 - assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( - 2.0 - ) - assert result.usage.prompt_tokens_details.text_tokens == 1 + assert result.usage.prompt_tokens_details.video_tokens == 516 + assert result.usage.prompt_tokens_details.text_tokens == 0 def test_missing_usage_metadata_does_not_estimate_from_base64(self): response_json = {"embedding": {"values": [0.1, 0.2]}} @@ -400,8 +397,7 @@ class TestProcessEmbedContentResponseUsage: ) assert result.usage.prompt_tokens > 0 - def test_file_reference_image_billed_per_image_not_text(self): - """files/... image refs must bill per-image, not at the text token rate.""" + def test_file_reference_image_billed_per_image_token_rate(self): response_json = { "embedding": {"values": [0.1, 0.2, 0.3]}, "usageMetadata": { @@ -422,7 +418,7 @@ class TestProcessEmbedContentResponseUsage: } }, ) - assert result.usage.prompt_tokens_details.image_count == 1 + assert result.usage.prompt_tokens_details.image_tokens == 258 assert result.usage.prompt_tokens_details.text_tokens == 0 prompt_cost, _ = generic_cost_per_token( @@ -430,10 +426,10 @@ class TestProcessEmbedContentResponseUsage: usage=result.usage, custom_llm_provider="vertex_ai", ) - assert prompt_cost == pytest.approx(0.00012) + assert prompt_cost == pytest.approx(258 * 4.5e-7) def test_file_reference_non_image_not_counted_as_image(self): - """A files/... ref resolving to a non-image mime must not be image-counted.""" + """A files/... ref resolving to a non-image mime keeps audio token billing.""" response_json = { "embedding": {"values": [0.1, 0.2]}, "usageMetadata": { @@ -454,21 +450,18 @@ class TestProcessEmbedContentResponseUsage: } }, ) - assert result.usage.prompt_tokens_details.image_count == 0 assert result.usage.prompt_tokens_details.audio_tokens == 64 - assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.image_tokens == 0 prompt_cost, _ = generic_cost_per_token( model=self.MODEL, usage=result.usage, custom_llm_provider="vertex_ai", ) - assert prompt_cost == pytest.approx(2.0 * 0.00016) + assert prompt_cost == pytest.approx(64 * 6.5e-6) def test_video_plus_audio_does_not_double_bill_text(self): - """Video+audio responses must not get video tokens reassigned to text.""" + """Video and audio responses are billed from their respective token counts.""" response_json = { "embedding": {"values": [0.1]}, "usageMetadata": { @@ -486,18 +479,13 @@ class TestProcessEmbedContentResponseUsage: model=self.MODEL, response_json=response_json, ) - assert result.usage.prompt_tokens_details.text_tokens == 1 - assert result.usage.prompt_tokens_details.video_length_seconds == pytest.approx( - 2.0 - ) - assert result.usage.prompt_tokens_details.audio_length_seconds == pytest.approx( - 2.0 - ) + assert result.usage.prompt_tokens_details.text_tokens == 0 + assert result.usage.prompt_tokens_details.video_tokens == 516 + assert result.usage.prompt_tokens_details.audio_tokens == 64 prompt_cost, _ = generic_cost_per_token( model=self.MODEL, usage=result.usage, custom_llm_provider="vertex_ai", ) - # 1 floor text token at 2e-7 + 2s of video at 7.9e-4 + 2s of audio at 1.6e-4 - assert prompt_cost == pytest.approx(1 * 2e-7 + 2 * 0.00079 + 2 * 0.00016) + assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6) From e26a4970dd0ba5efe277a5f42b51854daaf5da6d Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:42:21 +0000 Subject: [PATCH 10/30] fix(test): complete synthetic model metadata Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 4a02bb1a638..97d04a03a78 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -86,6 +86,7 @@ def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: "output_cost_per_token": 0.0, "litellm_provider": "vertex_ai", "mode": "embedding", + "supported_openai_params": None, } usage = Usage( prompt_tokens=64, From 0c91d9157c43ba7728b58393b6088641aa824367 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:45:08 +0000 Subject: [PATCH 11/30] refactor(vertex): drop unused resolved_files from embed response parsing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gemini_embeddings/batch_embed_content_handler.py | 2 -- .../batch_embed_content_transformation.py | 3 +-- .../test_batch_embed_content_transformation.py | 12 ------------ 3 files changed, 1 insertion(+), 16 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index f81d4ca777e..e09622ba236 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -268,7 +268,6 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, - resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) @@ -372,7 +371,6 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, - resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 8e120ab9fe6..f2cce775f3f 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,7 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Final from pydantic import TypeAdapter, ValidationError @@ -356,7 +356,6 @@ def process_embed_content_response( model_response: EmbeddingResponse, model: str, response_json: dict, - resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index fbf86105e71..0251a799b66 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -411,12 +411,6 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, - resolved_files={ - "files/img123": { - "mime_type": "image/png", - "uri": "https://example.com/img123", - } - }, ) assert result.usage.prompt_tokens_details.image_tokens == 258 assert result.usage.prompt_tokens_details.text_tokens == 0 @@ -443,12 +437,6 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, - resolved_files={ - "files/clip1": { - "mime_type": "audio/mpeg", - "uri": "https://example.com/clip1", - } - }, ) assert result.usage.prompt_tokens_details.audio_tokens == 64 assert result.usage.prompt_tokens_details.image_tokens == 0 From 6c9fe65608997dbfa85c35d01ad196f8b97c9a9d Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:47:30 +0000 Subject: [PATCH 12/30] style(cost): apply ruff formatting to modality guards Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index a004f46b291..88ea4b602cc 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -957,8 +957,7 @@ def _calculate_input_cost( ### AUDIO COST if prompt_tokens_details["audio_tokens"] and not ( - prompt_tokens_details["audio_length_seconds"] - and model_info.get("input_cost_per_audio_per_second") is not None + prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None ): audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier) prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"]) @@ -976,8 +975,7 @@ def _calculate_input_cost( ### VIDEO TOKEN COST if prompt_tokens_details["video_tokens"] and not ( - prompt_tokens_details["video_length_seconds"] - and model_info.get("input_cost_per_video_per_second") is not None + prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None ): video_token_cost_key = "input_cost_per_video_token" if model_info.get(video_token_cost_key) is None: From e5845c17ffde232ee1e460b648fb223ea4561348 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:51:49 +0000 Subject: [PATCH 13/30] fix(vertex): bill image inputs at the image rate when usage lacks modality details Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batch_embed_content_handler.py | 2 + .../batch_embed_content_transformation.py | 52 +++++++++++++++- ...test_batch_embed_content_transformation.py | 60 +++++++++++++++++++ 3 files changed, 113 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index e09622ba236..f81d4ca777e 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -268,6 +268,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) @@ -371,6 +372,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response=model_response, model=model, response_json=_json_response, + resolved_files=resolved_files, ) else: _predictions: Final = VertexAIBatchEmbeddingsResponseObject(**_json_response) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index f2cce775f3f..b61fb47cf5c 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -4,7 +4,7 @@ Transformation logic from OpenAI /v1/embeddings format to Google AI Studio /batc Why separate file? Make it easy to see how transformation works """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final from pydantic import TypeAdapter, ValidationError @@ -297,6 +297,7 @@ def transform_openai_input_gemini_embed_content( return request_body +_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"}) _usage_metadata_adapter: Final = TypeAdapter(UsageMetadata) @@ -309,6 +310,40 @@ def _parse_usage_metadata(raw_usage_metadata: object) -> UsageMetadata | None: return None +def _flatten_input(input: GeminiEmbeddingInput) -> tuple[str, ...]: + if isinstance(input, str): + return (input,) + return tuple(sub for element in input for sub in (element if isinstance(element, list) else [element])) + + +def _is_image_element( + element: str, + resolved_files: Mapping[str, Mapping[str, str]], +) -> bool: + if element.startswith("data:") and ";base64," in element: + try: + mime_type, _ = _parse_data_url(element) + except ValueError: + return False + return mime_type in _IMAGE_MIME_TYPES + if _is_gcs_url(element): + try: + return _infer_mime_type_from_gcs_url(element) in _IMAGE_MIME_TYPES + except ValueError: + return False + if _is_file_reference(element): + file_info: Final = resolved_files.get(element) + return file_info is not None and file_info.get("mime_type") in _IMAGE_MIME_TYPES + return False + + +def _count_input_images( + input: GeminiEmbeddingInput, + resolved_files: Mapping[str, Mapping[str, str]], +) -> int: + return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) + + def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: return sum(detail["tokenCount"] for detail in details if detail["modality"] == modality) @@ -325,6 +360,7 @@ def _usage_from_embed_content_response( input: GeminiEmbeddingInput, model: str, raw_usage_metadata: object, + resolved_files: Mapping[str, Mapping[str, str]], ) -> Usage: usage_metadata: Final = _parse_usage_metadata(raw_usage_metadata) if usage_metadata is None: @@ -334,6 +370,17 @@ def _usage_from_embed_content_response( total_tokens: Final = usage_metadata.get("totalTokenCount") or prompt_tokens details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or () + if not details: + image_tokens: Final = prompt_tokens if _count_input_images(input, resolved_files) else 0 + return Usage( + prompt_tokens=prompt_tokens, + total_tokens=total_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, + image_tokens=image_tokens, + ), + ) + text_tokens: Final = _tokens_for_modality(details, "TEXT") audio_tokens: Final = _tokens_for_modality(details, "AUDIO") image_tokens: Final = _tokens_for_modality(details, "IMAGE") @@ -356,6 +403,7 @@ def process_embed_content_response( model_response: EmbeddingResponse, model: str, response_json: dict, + resolved_files: Mapping[str, Mapping[str, str]] | None = None, ) -> EmbeddingResponse: """ Process Gemini embedContent response (single embedding for multimodal input). @@ -365,6 +413,7 @@ def process_embed_content_response( model_response: EmbeddingResponse to populate model: Model name response_json: Raw JSON response from embedContent endpoint + resolved_files: Mapping of file references to resolved metadata Returns: EmbeddingResponse with single embedding @@ -386,6 +435,7 @@ def process_embed_content_response( input=input, model=model, raw_usage_metadata=response_json.get("usageMetadata"), + resolved_files=resolved_files or {}, ) return model_response diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 0251a799b66..df5903b9285 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -411,6 +411,12 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, + resolved_files={ + "files/img123": { + "mime_type": "image/png", + "uri": "https://example.com/img123", + } + }, ) assert result.usage.prompt_tokens_details.image_tokens == 258 assert result.usage.prompt_tokens_details.text_tokens == 0 @@ -437,6 +443,12 @@ class TestProcessEmbedContentResponseUsage: model_response=EmbeddingResponse(), model=self.MODEL, response_json=response_json, + resolved_files={ + "files/clip1": { + "mime_type": "audio/mpeg", + "uri": "https://example.com/clip1", + } + }, ) assert result.usage.prompt_tokens_details.audio_tokens == 64 assert result.usage.prompt_tokens_details.image_tokens == 0 @@ -477,3 +489,51 @@ class TestProcessEmbedContentResponseUsage: custom_llm_provider="vertex_ai", ) assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6) + + def test_image_without_modality_details_uses_image_rate(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + }, + } + result = process_embed_content_response( + input=IMAGE_DATA_URI, + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.image_tokens == 258 + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(258 * 4.5e-7) + + def test_text_without_modality_details_uses_text_rate(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 12, + "totalTokenCount": 12, + }, + } + result = process_embed_content_response( + input="a short caption", + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.text_tokens == 0 + assert result.usage.prompt_tokens_details.image_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(12 * 2e-7) From e31c64d2e038fc0895cf02e45e2c6a798bba3cf3 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 00:54:58 +0000 Subject: [PATCH 14/30] fix(schema): sync model price schema with cost map Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- model_prices_and_context_window.schema.json | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index d1ac3e67b2b..eff3f192b3c 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -375,6 +375,10 @@ "minimum": 0, "description": "Rate applied once the prompt exceeds the token threshold in the field name." }, + "input_cost_per_video_token": { + "type": "number", + "minimum": 0 + }, "input_dbu_cost_per_token": { "type": "number", "minimum": 0 From ac8e1a355cf69830ec989fd19afb42a2bef78efd Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:00:32 +0000 Subject: [PATCH 15/30] test(vertex): load local pricing in embedding billing tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_batch_embed_content_transformation.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index df5903b9285..49f5167fd6b 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -10,6 +10,7 @@ Covers: import pytest +import litellm from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( _build_part_for_input, @@ -26,6 +27,15 @@ IMAGE_DATA_URI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+ GCS_URL = "gs://my-bucket/image.png" +@pytest.fixture(autouse=True) +def _local_model_cost_map(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + class TestIsMultimodalInput: def test_text_only_string(self): assert _is_multimodal_input("hello world") is False From 6cbed7b4c0ae643a429b0ebc7cb85a99e52b5b9e Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:10:27 +0000 Subject: [PATCH 16/30] fix(vertex): drop Final image_tokens redeclaration flagged by basedpyright Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../gemini_embeddings/batch_embed_content_transformation.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index b61fb47cf5c..b618f6e5165 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -371,13 +371,12 @@ def _usage_from_embed_content_response( details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or () if not details: - image_tokens: Final = prompt_tokens if _count_input_images(input, resolved_files) else 0 return Usage( prompt_tokens=prompt_tokens, total_tokens=total_tokens, prompt_tokens_details=PromptTokensDetailsWrapper( text_tokens=0, - image_tokens=image_tokens, + image_tokens=prompt_tokens if _count_input_images(input, resolved_files) else 0, ), ) From 6a18105275223a39170e24d8fec96d123af82b32 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:16:00 +0000 Subject: [PATCH 17/30] fix(vertex): only bill image rate without modality details when every input is an image Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batch_embed_content_transformation.py | 9 ++++---- ...test_batch_embed_content_transformation.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index b618f6e5165..d669acecfd9 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -337,11 +337,12 @@ def _is_image_element( return False -def _count_input_images( +def _is_image_only_input( input: GeminiEmbeddingInput, resolved_files: Mapping[str, Mapping[str, str]], -) -> int: - return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files)) +) -> bool: + elements: Final = _flatten_input(input) + return bool(elements) and all(_is_image_element(element, resolved_files) for element in elements) def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int: @@ -376,7 +377,7 @@ def _usage_from_embed_content_response( total_tokens=total_tokens, prompt_tokens_details=PromptTokensDetailsWrapper( text_tokens=0, - image_tokens=prompt_tokens if _count_input_images(input, resolved_files) else 0, + image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0, ), ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 49f5167fd6b..5dfeac6f469 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -524,6 +524,29 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(258 * 4.5e-7) + def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 270, + "totalTokenCount": 270, + }, + } + result = process_embed_content_response( + input=["a short caption", IMAGE_DATA_URI], + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + ) + assert result.usage.prompt_tokens_details.image_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(270 * 2e-7) + def test_text_without_modality_details_uses_text_rate(self): response_json = { "embedding": {"values": [0.1]}, From 4a8ec7b9d8c60896b28448b3bd87380617692d74 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:18:50 +0000 Subject: [PATCH 18/30] fix(cost): bill gemini-embedding-2-preview per token like the GA entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 18 +++++++-------- model_prices_and_context_window.json | 18 +++++++-------- ...test_batch_embed_content_transformation.py | 22 +++++++++++++++++++ 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 74ca23fcc8a..9299a1493f8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25601,10 +25601,10 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25631,10 +25631,10 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25689,10 +25689,10 @@ }, "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 74ca23fcc8a..9299a1493f8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25601,10 +25601,10 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models" }, "gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25631,10 +25631,10 @@ "uses_embed_content": true }, "vertex_ai/gemini-embedding-2-preview": { - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25689,10 +25689,10 @@ }, "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", - "input_cost_per_audio_per_second": 0.00016, - "input_cost_per_image": 0.00012, + "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, - "input_cost_per_video_per_second": 0.00079, + "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 5dfeac6f469..926570d7929 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -500,6 +500,28 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(516 * 1.2e-5 + 64 * 6.5e-6) + def test_preview_alias_bills_audio_per_token(self): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 64, + "totalTokenCount": 64, + "promptTokensDetails": [{"modality": "AUDIO", "tokenCount": 64}], + }, + } + result = process_embed_content_response( + input="audio", + model_response=EmbeddingResponse(), + model="gemini-embedding-2-preview", + response_json=response_json, + ) + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2-preview", + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + assert prompt_cost == pytest.approx(64 * 6.5e-6) + def test_image_without_modality_details_uses_image_rate(self): response_json = { "embedding": {"values": [0.1]}, From 27a486e4d320b4481c8aad6062f0c518d1c52ea4 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:34:17 +0000 Subject: [PATCH 19/30] test(cost): cover modality guards and image detection fallbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llm_cost_calc/test_llm_cost_calc_utils.py | 68 +++++++++++++++++++ ...test_batch_embed_content_transformation.py | 39 +++++++++++ 2 files changed, 107 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 97d04a03a78..854bc9bbb81 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -108,6 +108,74 @@ def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: assert prompt_cost == pytest.approx(2 * 0.00016) +def test_generic_cost_per_token_prefers_image_per_image_rate() -> None: + model_info: ModelInfo = { + "key": "gemini-embedding-2", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": 2e-7, + "input_cost_per_image_token": 4.5e-7, + "input_cost_per_image": 0.00012, + "output_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "embedding", + "supported_openai_params": None, + } + usage = Usage( + prompt_tokens=258, + completion_tokens=0, + total_tokens=258, + prompt_tokens_details=PromptTokensDetailsWrapper( + image_tokens=258, + image_count=1, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(0.00012) + + +def test_generic_cost_per_token_prefers_video_per_second_rate() -> None: + model_info: ModelInfo = { + "key": "gemini-embedding-2", + "max_tokens": None, + "max_input_tokens": None, + "max_output_tokens": None, + "input_cost_per_token": 2e-7, + "input_cost_per_video_token": 1.2e-5, + "input_cost_per_video_per_second": 0.00079, + "output_cost_per_token": 0.0, + "litellm_provider": "vertex_ai", + "mode": "embedding", + "supported_openai_params": None, + } + usage = Usage( + prompt_tokens=516, + completion_tokens=0, + total_tokens=516, + prompt_tokens_details=PromptTokensDetailsWrapper( + video_tokens=516, + video_length_seconds=2, + ), + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-embedding-2", + usage=usage, + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(2 * 0.00079) + + def test_missing_cache_read_uses_off_peak_input_rate(): from datetime import datetime, timezone diff --git a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index 926570d7929..fd8c2a9cf6a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -546,6 +546,45 @@ class TestProcessEmbedContentResponseUsage: ) assert prompt_cost == pytest.approx(258 * 4.5e-7) + @pytest.mark.parametrize( + "input_value,resolved_files,expected_image_tokens", + [ + (GCS_URL, {}, 258), + ("gs://my-bucket/clip.mp4", {}, 0), + ("gs://my-bucket/unknown.bin", {}, 0), + ("files/image-123", {"files/image-123": {"mime_type": "image/jpeg"}}, 258), + ("files/missing", {}, 0), + ("data:application/octet-stream;base64,abc", {}, 0), + ([[IMAGE_DATA_URI]], {}, 258), + ([], {}, 0), + ], + ) + def test_missing_modality_details_classifies_image_inputs(self, input_value, resolved_files, expected_image_tokens): + response_json = { + "embedding": {"values": [0.1]}, + "usageMetadata": { + "promptTokenCount": 258, + "totalTokenCount": 258, + }, + } + result = process_embed_content_response( + input=input_value, + model_response=EmbeddingResponse(), + model=self.MODEL, + response_json=response_json, + resolved_files=resolved_files, + ) + assert result.usage.prompt_tokens_details.image_tokens == expected_image_tokens + assert result.usage.prompt_tokens_details.text_tokens == 0 + + prompt_cost, _ = generic_cost_per_token( + model=self.MODEL, + usage=result.usage, + custom_llm_provider="vertex_ai", + ) + expected_rate = 4.5e-7 if expected_image_tokens else 2e-7 + assert prompt_cost == pytest.approx(258 * expected_rate) + def test_mixed_text_and_image_without_modality_details_not_billed_as_image(self): response_json = { "embedding": {"values": [0.1]}, From a28ea22ec131d1ce9f47af4dceb1faf7df7aa2f2 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:34:22 +0000 Subject: [PATCH 20/30] fix(cost): move gemini-embedding-2-preview to per-token rates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 3 +++ model_prices_and_context_window.json | 3 +++ tests/test_litellm/test_utils.py | 11 +++++++---- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9299a1493f8..29243832c45 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25604,6 +25604,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, @@ -25634,6 +25635,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, @@ -25692,6 +25694,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9299a1493f8..29243832c45 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25604,6 +25604,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, @@ -25634,6 +25635,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, @@ -25692,6 +25694,7 @@ "input_cost_per_audio_token": 6.5e-06, "input_cost_per_image_token": 4.5e-07, "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, "litellm_provider": "gemini", "max_input_tokens": 8192, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 04d2d35e05f..1da53fba923 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2946,7 +2946,7 @@ def test_model_info_for_openrouter_kimi_k2_5(): def test_gemini_embedding_2_ga_in_cost_map(): - """GA and Vertex preview gemini-embedding-2 entries align with multimodal unit pricing.""" + """GA and Vertex preview gemini-embedding-2 entries align with multimodal token pricing.""" import json from pathlib import Path @@ -2968,9 +2968,12 @@ def test_gemini_embedding_2_ga_in_cost_map(): assert info.get("mode") == "embedding" assert info.get("supports_multimodal") is True assert info.get("input_cost_per_token") == 2e-07 - assert info.get("input_cost_per_image") == 0.00012 - assert info.get("input_cost_per_audio_per_second") == 0.00016 - assert info.get("input_cost_per_video_per_second") == 0.00079 + assert info.get("input_cost_per_audio_token") == 6.5e-06 + assert info.get("input_cost_per_image_token") == 4.5e-07 + assert info.get("input_cost_per_video_token") == 1.2e-05 + assert "input_cost_per_image" not in info + assert "input_cost_per_audio_per_second" not in info + assert "input_cost_per_video_per_second" not in info if provider in ("vertex_ai-embedding-models", "vertex_ai"): assert ( info.get("uses_embed_content") is True From ce83fac3515c36c927ed133fe48abd7dc1a3ee74 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:50:21 +0000 Subject: [PATCH 21/30] fix(cost): bill batch embeddings per modality token rate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/generate_model_prices_schema.py | 3 ++ litellm/cost_calculator.py | 30 ++++++++++- ...odel_prices_and_context_window_backup.json | 18 +++++++ litellm/types/utils.py | 6 +++ litellm/utils.py | 3 ++ model_prices_and_context_window.json | 18 +++++++ model_prices_and_context_window.schema.json | 15 ++++++ tests/test_litellm/test_cost_calculator.py | 51 +++++++++++++++++++ tests/test_litellm/test_utils.py | 3 ++ 9 files changed, 146 insertions(+), 1 deletion(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ab29b70bdd4..ee0e7f22eb1 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -172,7 +172,10 @@ COST_DESCRIPTIONS: dict[str, str] = { ), "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", + "input_cost_per_audio_token_batches": "USD per audio prompt token via the provider's batch API.", + "input_cost_per_image_token_batches": "USD per image prompt token via the provider's batch API.", "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", + "input_cost_per_video_token_batches": "USD per video prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", } diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f5319776213..cbb9a45ada9 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2337,7 +2337,35 @@ def batch_cost_calculator( total_prompt_cost = 0.0 total_completion_cost = 0.0 if input_cost_per_token_batches is not None: - total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches + batch_details: Final = parse_prompt_tokens_details(usage) + audio_tokens, image_tokens, video_tokens = ( + batch_details["audio_tokens"], + batch_details["image_tokens"], + batch_details["video_tokens"], + ) + modality_rates: Final = ( + cast(float, model_info.get("input_cost_per_audio_token_batches")) + if model_info.get("input_cost_per_audio_token_batches") is not None + else input_cost_per_token_batches, + cast(float, model_info.get("input_cost_per_image_token_batches")) + if model_info.get("input_cost_per_image_token_batches") is not None + else input_cost_per_token_batches, + cast(float, model_info.get("input_cost_per_video_token_batches")) + if model_info.get("input_cost_per_video_token_batches") is not None + else input_cost_per_token_batches, + ) + total_prompt_cost = sum( + tokens * rate + for tokens, rate in zip( + ( + max(cast(int, usage.prompt_tokens) - audio_tokens - image_tokens - video_tokens, 0), + audio_tokens, + image_tokens, + video_tokens, + ), + (input_cost_per_token_batches, *modality_rates), + ) + ) elif input_cost_per_token: details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = details["cache_hit_tokens"] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 29243832c45..9f91cf82f41 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -25602,10 +25602,13 @@ }, "gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25617,10 +25620,13 @@ }, "gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25633,10 +25639,13 @@ }, "vertex_ai/gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25649,10 +25658,13 @@ }, "vertex_ai/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25692,10 +25704,13 @@ "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25724,13 @@ }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index fbfcc678de9..2e8b20edf7a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -283,8 +283,11 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_video_token: float | None # for gemini omni models with video input input_cost_per_audio_per_second: float | None # only for vertex ai models input_cost_per_video_per_second: float | None # only for vertex ai models + input_cost_per_audio_token_batches: ReadOnly[float | None] + input_cost_per_image_token_batches: ReadOnly[float | None] input_cost_per_second: float | None # for OpenAI Speech models input_cost_per_token_batches: float | None + input_cost_per_video_token_batches: ReadOnly[float | None] output_cost_per_token_batches: float | None output_cost_per_token: Required[float | None] output_cost_per_token_flex: float | None # OpenAI flex service tier pricing @@ -3583,7 +3586,10 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_video_per_second_above_128k_tokens: float | None = None input_cost_per_video_per_second_above_15s_interval: float | None = None input_cost_per_video_per_second_above_8s_interval: float | None = None + input_cost_per_audio_token_batches: float | None = None + input_cost_per_image_token_batches: float | None = None input_cost_per_token_batches: float | None = None + input_cost_per_video_token_batches: float | None = None output_cost_per_token_batches: float | None = None output_cost_per_token_flex: float | None = None output_cost_per_token_priority: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index d4e3d58ba9f..b2715b41739 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5923,10 +5923,13 @@ def _get_model_info_helper( input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), + input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None), + input_cost_per_image_token_batches=_model_info.get("input_cost_per_image_token_batches", None), input_cost_per_image=_model_info.get("input_cost_per_image", None), input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None), input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), input_cost_per_token_batches=_model_info.get("input_cost_per_token_batches"), + input_cost_per_video_token_batches=_model_info.get("input_cost_per_video_token_batches", None), output_cost_per_token_batches=_model_info.get("output_cost_per_token_batches"), output_cost_per_token=_output_cost_per_token, output_cost_per_token_flex=_model_info.get("output_cost_per_token_flex", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 29243832c45..9f91cf82f41 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -25602,10 +25602,13 @@ }, "gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25617,10 +25620,13 @@ }, "gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25633,10 +25639,13 @@ }, "vertex_ai/gemini-embedding-2-preview": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25649,10 +25658,13 @@ }, "vertex_ai/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25692,10 +25704,13 @@ "gemini/gemini-embedding-2-preview": { "deprecation_date": "2026-08-10", "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, @@ -25709,10 +25724,13 @@ }, "gemini/gemini-embedding-2": { "input_cost_per_audio_token": 6.5e-06, + "input_cost_per_audio_token_batches": 3.25e-06, "input_cost_per_image_token": 4.5e-07, + "input_cost_per_image_token_batches": 2.25e-07, "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "input_cost_per_video_token": 1.2e-05, + "input_cost_per_video_token_batches": 6e-06, "litellm_provider": "gemini", "max_input_tokens": 8192, "max_tokens": 8192, diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index eff3f192b3c..b7e0a9fd414 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -249,6 +249,11 @@ "type": "number", "minimum": 0 }, + "input_cost_per_audio_token_batches": { + "type": "number", + "minimum": 0, + "description": "USD per audio prompt token via the provider's batch API." + }, "input_cost_per_audio_token_priority": { "type": "number", "minimum": 0, @@ -276,6 +281,11 @@ "type": "number", "minimum": 0 }, + "input_cost_per_image_token_batches": { + "type": "number", + "minimum": 0, + "description": "USD per image prompt token via the provider's batch API." + }, "input_cost_per_pixel": { "type": "number", "minimum": 0 @@ -379,6 +389,11 @@ "type": "number", "minimum": 0 }, + "input_cost_per_video_token_batches": { + "type": "number", + "minimum": 0, + "description": "USD per video prompt token via the provider's batch API." + }, "input_dbu_cost_per_token": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 68e9b6143a0..8c3436d3108 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3909,6 +3909,57 @@ def _batch_cache_usage() -> Usage: ) +def test_batch_cost_calculator_prices_multimodal_tokens_at_modality_rates(): + from litellm.cost_calculator import batch_cost_calculator + + model_info: ModelInfo = { + "input_cost_per_token_batches": 1e-7, + "input_cost_per_audio_token_batches": 3.25e-6, + "input_cost_per_image_token_batches": 2.25e-7, + "input_cost_per_video_token_batches": 6e-6, + } + usage = Usage( + prompt_tokens=100, + completion_tokens=0, + total_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=64, + image_tokens=10, + video_tokens=6, + ), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-embedding-2", + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(20 * 1e-7 + 64 * 3.25e-6 + 10 * 2.25e-7 + 6 * 6e-6) + + +def test_batch_cost_calculator_falls_back_to_text_batch_rate_for_modalities(): + from litellm.cost_calculator import batch_cost_calculator + + model_info: ModelInfo = {"input_cost_per_token_batches": 1e-7} + usage = Usage( + prompt_tokens=100, + completion_tokens=0, + total_tokens=100, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=64), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-embedding-2", + custom_llm_provider="vertex_ai", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(100 * 1e-7) + + def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate(): """ LIT-4008 regression: anthropic batch usage is dominated by cache tokens. diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1da53fba923..02ffaee0543 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2971,6 +2971,9 @@ def test_gemini_embedding_2_ga_in_cost_map(): assert info.get("input_cost_per_audio_token") == 6.5e-06 assert info.get("input_cost_per_image_token") == 4.5e-07 assert info.get("input_cost_per_video_token") == 1.2e-05 + assert info.get("input_cost_per_audio_token_batches") == 3.25e-06 + assert info.get("input_cost_per_image_token_batches") == 2.25e-07 + assert info.get("input_cost_per_video_token_batches") == 6e-06 assert "input_cost_per_image" not in info assert "input_cost_per_audio_per_second" not in info assert "input_cost_per_video_per_second" not in info From 2b32f586c087ca94c3427c31eb83513c2a03c599 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:51:41 +0000 Subject: [PATCH 22/30] refactor(cost): extract batch modality rate lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cbb9a45ada9..dce6b7299a2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2278,6 +2278,19 @@ def default_video_cost_calculator( return 0.0 +def _batch_rate( + model_info: ModelInfo, + key: Literal[ + "input_cost_per_audio_token_batches", + "input_cost_per_image_token_batches", + "input_cost_per_video_token_batches", + ], + fallback: float, +) -> float: + rate: Final = model_info.get(key) + return fallback if rate is None else cast(float, rate) + + def batch_cost_calculator( usage: Usage, model: str, @@ -2344,15 +2357,9 @@ def batch_cost_calculator( batch_details["video_tokens"], ) modality_rates: Final = ( - cast(float, model_info.get("input_cost_per_audio_token_batches")) - if model_info.get("input_cost_per_audio_token_batches") is not None - else input_cost_per_token_batches, - cast(float, model_info.get("input_cost_per_image_token_batches")) - if model_info.get("input_cost_per_image_token_batches") is not None - else input_cost_per_token_batches, - cast(float, model_info.get("input_cost_per_video_token_batches")) - if model_info.get("input_cost_per_video_token_batches") is not None - else input_cost_per_token_batches, + _batch_rate(model_info, "input_cost_per_audio_token_batches", input_cost_per_token_batches), + _batch_rate(model_info, "input_cost_per_image_token_batches", input_cost_per_token_batches), + _batch_rate(model_info, "input_cost_per_video_token_batches", input_cost_per_token_batches), ) total_prompt_cost = sum( tokens * rate From c0c5044c45bac0a696cd270c442b00d29cb2756e Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 01:57:50 +0000 Subject: [PATCH 23/30] fix(batches): keep modality token details in raw vertex batch usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 34 +++++++++++++++++-- .../test_litellm/batches/test_batch_utils.py | 32 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +++++++ 3 files changed, 76 insertions(+), 2 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 87f8fd3946e..397bc0a35a2 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -3,14 +3,14 @@ from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum -from typing import Any, Final, Literal +from typing import Any, Final, Literal, cast import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import ModelInfo, Usage +from litellm.types.utils import ModelInfo, PromptTokensDetailsWrapper, Usage from litellm.utils import token_counter @@ -310,6 +310,35 @@ def _aggregate_batch_cost_usage_models( ) +def _vertex_prompt_tokens_details( + usage_metadata: Mapping[str, object], +) -> PromptTokensDetailsWrapper | None: + raw_details: Final = usage_metadata.get("promptTokensDetails") + if not isinstance(raw_details, list): + return None + + raw_list: Final = cast(list[object], raw_details) + if not all(isinstance(detail, Mapping) for detail in raw_list): + return None + + details: Final = tuple(cast(Mapping[str, object], detail) for detail in raw_list) + normalized: Final = tuple( + (modality.upper(), token_count) + for detail in details + if isinstance(modality := detail.get("modality"), str) + and isinstance(token_count := detail.get("tokenCount"), int) + ) + if len(normalized) != len(details): + return None + + return PromptTokensDetailsWrapper( + text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), + audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), + image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), + video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + ) + + def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: list[dict], model_name: str | None = None, @@ -356,6 +385,7 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=_prompt, completion_tokens=_completion, total_tokens=_total, + prompt_tokens_details=_vertex_prompt_tokens_details(cast(Mapping[str, object], usage_metadata)), ) try: diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 768ea332677..8b04d7af70a 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -695,6 +695,38 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch): assert result.failed_requests == 0 +def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/gemini-embedding-2", + { + "input_cost_per_token_batches": 1e-7, + "input_cost_per_audio_token_batches": 3.25e-6, + "input_cost_per_image_token_batches": 2.25e-7, + "input_cost_per_video_token_batches": 6e-6, + }, + ) + responses = [ + { + "response": { + "usageMetadata": { + "promptTokenCount": 84, + "candidatesTokenCount": 0, + "totalTokenCount": 84, + "promptTokensDetails": [ + {"modality": "AUDIO", "tokenCount": 64}, + {"modality": "TEXT", "tokenCount": 20}, + ], + } + } + } + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2") + + assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7) + + def test_vertex_cost_skips_none_response_body(monkeypatch): import litellm.cost_calculator as cc diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8aa05cf8c7c..73245806bb2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29857,6 +29857,8 @@ export interface components { input_cost_per_audio_per_second_above_128k_tokens?: number | null; /** Input Cost Per Audio Token */ input_cost_per_audio_token?: number | null; + /** Input Cost Per Audio Token Batches */ + input_cost_per_audio_token_batches?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -29867,6 +29869,8 @@ export interface components { input_cost_per_image_above_128k_tokens?: number | null; /** Input Cost Per Image Token */ input_cost_per_image_token?: number | null; + /** Input Cost Per Image Token Batches */ + input_cost_per_image_token_batches?: number | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; /** Input Cost Per Query */ @@ -29909,6 +29913,8 @@ export interface components { input_cost_per_video_per_second_above_8s_interval?: number | null; /** Input Cost Per Video Token */ input_cost_per_video_token?: number | null; + /** Input Cost Per Video Token Batches */ + input_cost_per_video_token_batches?: number | null; /** Itpm */ itpm?: number | null; /** Keepalive Seconds */ @@ -40071,6 +40077,8 @@ export interface components { input_cost_per_audio_per_second_above_128k_tokens?: number | null; /** Input Cost Per Audio Token */ input_cost_per_audio_token?: number | null; + /** Input Cost Per Audio Token Batches */ + input_cost_per_audio_token_batches?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -40081,6 +40089,8 @@ export interface components { input_cost_per_image_above_128k_tokens?: number | null; /** Input Cost Per Image Token */ input_cost_per_image_token?: number | null; + /** Input Cost Per Image Token Batches */ + input_cost_per_image_token_batches?: number | null; /** Input Cost Per Pixel */ input_cost_per_pixel?: number | null; /** Input Cost Per Query */ @@ -40123,6 +40133,8 @@ export interface components { input_cost_per_video_per_second_above_8s_interval?: number | null; /** Input Cost Per Video Token */ input_cost_per_video_token?: number | null; + /** Input Cost Per Video Token Batches */ + input_cost_per_video_token_batches?: number | null; /** Itpm */ itpm?: number | null; /** Keepalive Seconds */ From ca7364fb0568a1d7f6b085529d45da2487fcb62c Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:00:48 +0000 Subject: [PATCH 24/30] fix(batches): avoid strict lint violation in usage parser Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 397bc0a35a2..acc0036f27d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -3,7 +3,7 @@ from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger @@ -317,18 +317,18 @@ def _vertex_prompt_tokens_details( if not isinstance(raw_details, list): return None - raw_list: Final = cast(list[object], raw_details) - if not all(isinstance(detail, Mapping) for detail in raw_list): - return None + def _normalize(detail: object) -> tuple[str, int] | None: + if not isinstance(detail, Mapping): + return None + modality: Final = detail.get("modality") + token_count: Final = detail.get("tokenCount") + if not isinstance(modality, str) or not isinstance(token_count, int): + return None + return modality.upper(), token_count - details: Final = tuple(cast(Mapping[str, object], detail) for detail in raw_list) - normalized: Final = tuple( - (modality.upper(), token_count) - for detail in details - if isinstance(modality := detail.get("modality"), str) - and isinstance(token_count := detail.get("tokenCount"), int) - ) - if len(normalized) != len(details): + parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) + normalized: Final = tuple(detail for detail in parsed_details if detail is not None) + if len(normalized) != len(parsed_details): return None return PromptTokensDetailsWrapper( @@ -385,7 +385,7 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=_prompt, completion_tokens=_completion, total_tokens=_total, - prompt_tokens_details=_vertex_prompt_tokens_details(cast(Mapping[str, object], usage_metadata)), + prompt_tokens_details=_vertex_prompt_tokens_details(usage_metadata), ) try: From a5fc880c907356a79843987eecc6aa170263271c Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:03:32 +0000 Subject: [PATCH 25/30] refactor(vertex): move batch usage modality parsing under llms Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 34 ++----------------- .../llms/vertex_ai/batches/transformation.py | 32 ++++++++++++++++- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index acc0036f27d..26b4318da2d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -9,8 +9,9 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details +from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import ModelInfo, PromptTokensDetailsWrapper, Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -310,35 +311,6 @@ def _aggregate_batch_cost_usage_models( ) -def _vertex_prompt_tokens_details( - usage_metadata: Mapping[str, object], -) -> PromptTokensDetailsWrapper | None: - raw_details: Final = usage_metadata.get("promptTokensDetails") - if not isinstance(raw_details, list): - return None - - def _normalize(detail: object) -> tuple[str, int] | None: - if not isinstance(detail, Mapping): - return None - modality: Final = detail.get("modality") - token_count: Final = detail.get("tokenCount") - if not isinstance(modality, str) or not isinstance(token_count, int): - return None - return modality.upper(), token_count - - parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) - normalized: Final = tuple(detail for detail in parsed_details if detail is not None) - if len(normalized) != len(parsed_details): - return None - - return PromptTokensDetailsWrapper( - text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), - audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), - image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), - video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), - ) - - def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: list[dict], model_name: str | None = None, @@ -385,7 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=_prompt, completion_tokens=_completion, total_tokens=_total, - prompt_tokens_details=_vertex_prompt_tokens_details(usage_metadata), + prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata), ) try: diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index e63c80dd3cf..f5f1ab2068a 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final from urllib.parse import unquote @@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.utils import LiteLLMBatch +from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper + + +def vertex_prompt_tokens_details( + usage_metadata: Mapping[str, object], +) -> PromptTokensDetailsWrapper | None: + raw_details: Final = usage_metadata.get("promptTokensDetails") + if not isinstance(raw_details, list): + return None + + def _normalize(detail: object) -> tuple[str, int] | None: + if not isinstance(detail, Mapping): + return None + modality: Final = detail.get("modality") + token_count: Final = detail.get("tokenCount") + if not isinstance(modality, str) or not isinstance(token_count, int): + return None + return modality.upper(), token_count + + parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) + normalized: Final = tuple(detail for detail in parsed_details if detail is not None) + if len(normalized) != len(parsed_details): + return None + + return PromptTokensDetailsWrapper( + text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), + audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), + image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), + video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + ) class VertexAIBatchTransformation: From bf1bdb3045670322350e1789f210da8e41c5a8e9 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:07:44 +0000 Subject: [PATCH 26/30] fix(ci): keep cost map schema generated by the base branch generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ci_cd/generate_model_prices_schema.py | 3 --- model_prices_and_context_window.schema.json | 9 +++------ 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ee0e7f22eb1..ab29b70bdd4 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -172,10 +172,7 @@ COST_DESCRIPTIONS: dict[str, str] = { ), "cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.", "cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.", - "input_cost_per_audio_token_batches": "USD per audio prompt token via the provider's batch API.", - "input_cost_per_image_token_batches": "USD per image prompt token via the provider's batch API.", "input_cost_per_token_batches": "USD per prompt token via the provider's batch API.", - "input_cost_per_video_token_batches": "USD per video prompt token via the provider's batch API.", "output_cost_per_token_batches": "USD per generated token via the provider's batch API.", } diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index b7e0a9fd414..c2490041cf7 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -251,8 +251,7 @@ }, "input_cost_per_audio_token_batches": { "type": "number", - "minimum": 0, - "description": "USD per audio prompt token via the provider's batch API." + "minimum": 0 }, "input_cost_per_audio_token_priority": { "type": "number", @@ -283,8 +282,7 @@ }, "input_cost_per_image_token_batches": { "type": "number", - "minimum": 0, - "description": "USD per image prompt token via the provider's batch API." + "minimum": 0 }, "input_cost_per_pixel": { "type": "number", @@ -391,8 +389,7 @@ }, "input_cost_per_video_token_batches": { "type": "number", - "minimum": 0, - "description": "USD per video prompt token via the provider's batch API." + "minimum": 0 }, "input_dbu_cost_per_token": { "type": "number", From a5f00b9189fa1dad1515800b0e0797a2fa411099 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:15:22 +0000 Subject: [PATCH 27/30] fix(cost): drop unnecessary cast in batch rate lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index dce6b7299a2..39af725a231 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2288,7 +2288,7 @@ def _batch_rate( fallback: float, ) -> float: rate: Final = model_info.get(key) - return fallback if rate is None else cast(float, rate) + return fallback if rate is None else rate def batch_cost_calculator( From 168d0d3d997e11b39240479cf47ddd1675ba22c2 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:26:19 +0000 Subject: [PATCH 28/30] fix(cost): drop remaining unnecessary cast in batch cost calculator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 39af725a231..3dc6d81256b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2365,7 +2365,7 @@ def batch_cost_calculator( tokens * rate for tokens, rate in zip( ( - max(cast(int, usage.prompt_tokens) - audio_tokens - image_tokens - video_tokens, 0), + max((usage.prompt_tokens or 0) - audio_tokens - image_tokens - video_tokens, 0), audio_tokens, image_tokens, video_tokens, From 7761d044509e0ab0dae37ed6ded7835d4d06d7e9 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:37:51 +0000 Subject: [PATCH 29/30] test(vertex): cover malformed batch usage details Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../vertex_ai/batches/test_transformation.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 232c6413e78..e6126b02790 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -17,9 +17,9 @@ from unittest.mock import patch import pytest - from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, + vertex_prompt_tokens_details, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 VertexAIError, @@ -41,6 +41,22 @@ ENDPOINT_INPUT_FILE = ( ) +def test_vertex_prompt_tokens_details_rejects_malformed_details(): + assert vertex_prompt_tokens_details({"promptTokensDetails": [1]}) is None + assert vertex_prompt_tokens_details({"promptTokensDetails": [{"modality": "AUDIO"}]}) is None + assert ( + vertex_prompt_tokens_details( + { + "promptTokensDetails": [ + {"modality": "AUDIO", "tokenCount": 1}, + "malformed", + ] + } + ) + is None + ) + + # =========================================================================== # # transform_openai_batch_request_to_vertex_ai_batch_request # =========================================================================== # From 954dfa6ba74882eecb3109438c993d769c636c92 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 15 Sep 2026 02:38:53 +0000 Subject: [PATCH 30/30] test(utils): allow modality batch cost fields in cost map schema test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 02ffaee0543..d5feda6f892 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -892,7 +892,10 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_video_per_second_above_8s_interval", "input_cost_per_video_per_second_above_15s_interval", "input_cost_per_video_per_second_above_128k_tokens", + "input_cost_per_audio_token_batches", + "input_cost_per_image_token_batches", "input_cost_per_token_batches", + "input_cost_per_video_token_batches", "output_cost_per_token_batches", "input_cost_per_token_cache_hit", "cache_creation_input_token_cost", @@ -1041,7 +1044,10 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_second": {"type": "number"}, "input_cost_per_token": {"type": "number"}, "input_cost_per_token_above_128k_tokens": {"type": "number"}, + "input_cost_per_audio_token_batches": {"type": "number"}, + "input_cost_per_image_token_batches": {"type": "number"}, "input_cost_per_token_batches": {"type": "number"}, + "input_cost_per_video_token_batches": {"type": "number"}, "input_cost_per_token_cache_hit": {"type": "number"}, "input_cost_per_video_per_second": {"type": "number"}, "input_cost_per_video_per_second_above_8s_interval": {"type": "number"},