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

This commit is contained in:
mateo-berri 2026-09-15 01:17:31 -07:00
commit 1771255b32
143 changed files with 7177 additions and 1057 deletions

View file

@ -2915,6 +2915,25 @@ jobs:
exit 1
fi
provider_replay_harness:
docker:
- *python312_image
working_directory: ~/project
resource_class: medium
steps:
- setup_litellm_test_deps
- run:
name: Test provider replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
tests/code_coverage_tests/test_provider_replay_harness.py
- store_test_results:
path: test-results/provider-replay-harness
integration_contracts:
parameters:
suite:
@ -2967,6 +2986,7 @@ workflows:
only:
- main
- /litellm_.*/
- provider_replay_harness
- base_sdk_install:
filters: *main_branches
- local_testing_part1:

View file

@ -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

5
.github/e2e-stack/oidc-profile.sh vendored Executable file
View file

@ -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 "$@"

View file

@ -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$"

View file

@ -57,6 +57,7 @@ permissions:
env:
UV_PYTHON: "3.12"
LITELLM_LOCAL_MODEL_COST_MAP: "True"
jobs:
run:
@ -113,6 +114,7 @@ jobs:
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'

View file

@ -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[@]}"

View file

@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
allowed_fails: int = 3
allow_dynamic_callback_disabling: bool = True
num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
####### SECRET MANAGERS #####################
secret_manager_client: Optional[Any] = (
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.

View file

@ -9,6 +9,7 @@ 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, Usage
from litellm.utils import token_counter
@ -356,6 +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),
)
try:

View file

@ -205,21 +205,35 @@ def _extract_anthropic_tool_exchange_spans(
return spans, None
def _message_has_cache_control(message: Mapping[str, object]) -> bool:
if message.get("cache_control") is not None:
return True
content: Final = message.get("content")
if isinstance(content, list):
return any(isinstance(part, Mapping) and part.get("cache_control") is not None for part in content)
return False
def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
"""
Return indices of messages that must never be compressed:
- All system messages
- The last user message
- The last assistant message
- Any message carrying an Anthropic cache_control breakpoint
The last user message is what the model is being asked to act on right now,
so compressing it replaces the live instruction with a marker. Compression
guardrails share this policy; see the Headroom guardrail.
guardrails share this policy; see the Headroom guardrail. A cache_control
breakpoint pins the provider's prompt-cache prefix to that row's exact
bytes, so rewriting a marked row anywhere in history turns the next
request's cache read into a cache write.
"""
system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
return system_indices + last_user + last_assistant
assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg))
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices))
def _combine_scores(
@ -421,7 +435,7 @@ def compress(
combined_scores = bm25_scores
# Protected messages are never compressed
protected_indices: Final = get_protected_indices(normalized_messages)
protected_indices: Final = get_protected_indices(original_messages)
kept_indices: set[int] = set(protected_indices)
tool_exchange_spans: list[set[int]] = []

View file

@ -1976,6 +1976,8 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
STRINGIFIED_NONE: Final[str] = "None"
# A retrieved response replays the usage of the call that created it, so pricing these
# read/management routes like inference bills the same tokens twice.
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(

View file

@ -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 rate
def batch_cost_calculator(
usage: Usage,
model: str,
@ -2337,7 +2350,29 @@ 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 = (
_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
for tokens, rate in zip(
(
max((usage.prompt_tokens or 0) - 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"]

View file

@ -309,8 +309,8 @@ def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_re
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
if not isinstance(metadata, Mapping):
return False
attempted_retries: Final = metadata.get("attempted_retries")
return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries
retry_count: Final = metadata.get("request_retry_count")
return type(retry_count) is int and 0 < retry_count and num_retries_per_request <= retry_count
def get_or_create_metadata_bucket(

View file

@ -1,6 +1,9 @@
import inspect
import json
import re
import traceback
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, Protocol, cast
import httpx
@ -202,11 +205,17 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None
return _response_headers
def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]:
accepted: Final = inspect.signature(exception_class).parameters
return MappingProxyType({name: value for name, value in candidates.items() if name in accepted})
def extract_and_raise_litellm_exception(
response: Any | None,
error_str: str,
model: str,
custom_llm_provider: str,
body: object | None = None,
):
"""
Covers scenario where litellm sdk calling proxy.
@ -216,32 +225,19 @@ def extract_and_raise_litellm_exception(
Relevant Issue: https://github.com/BerriAI/litellm/issues/7259
"""
pattern: Final = r"litellm\.\w+Error"
# Search for the exception in the error string
match: Final = re.search(pattern, error_str)
# Extract the exception if found
if match:
exception_name = match.group(0)
exception_name = exception_name.strip().replace("litellm.", "")
raised_exception_obj: Final = getattr(litellm, exception_name, None)
if raised_exception_obj:
# Try with response parameter first, fall back to without it
# Some exceptions (e.g., APIConnectionError) don't accept response param
try:
raise raised_exception_obj(
message=error_str,
llm_provider=custom_llm_provider,
model=model,
response=response,
)
except TypeError:
# Exception doesn't accept response parameter
raise raised_exception_obj(
message=error_str,
llm_provider=custom_llm_provider,
model=model,
)
if match is None:
return
exception_name: Final = match.group(0).removeprefix("litellm.")
raised_exception_obj: Final = getattr(litellm, exception_name, None)
if not raised_exception_obj:
return
raise raised_exception_obj(
message=error_str,
llm_provider=custom_llm_provider,
model=model,
**_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})),
)
class _ProviderHTTPException(Protocol):
@ -254,6 +250,23 @@ class _ProviderHTTPException(Protocol):
llm_provider: str
def _litellm_proxy_response(
original_exception: _ProviderHTTPException, custom_llm_provider: str
) -> httpx.Response | None:
response: Final = getattr(original_exception, "response", None)
if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers:
return response
headers: Final = getattr(original_exception, "headers", None)
if not isinstance(headers, Mapping) or not headers:
return response
pairs: Final = headers.multi_items() if isinstance(headers, httpx.Headers) else headers.items()
return httpx.Response(
status_code=response.status_code,
headers=[(str(k), str(v)) for k, v in pairs],
request=getattr(original_exception, "request", None),
)
def _map_openai_exception(
*,
model: str,
@ -264,6 +277,7 @@ def _map_openai_exception(
exception_provider: str,
extra_information: str,
) -> None:
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
# custom_llm_provider is openai, make it OpenAI
message = get_error_message(error_obj=original_exception)
if message is None:
@ -292,14 +306,14 @@ def _map_openai_exception(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
)
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
raise ContextWindowExceededError(
message=f"ContextWindowExceededError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif "invalid_request_error" in error_str and "model_not_found" in error_str:
@ -307,7 +321,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif "A timeout occurred" in error_str:
@ -326,8 +340,9 @@ def _map_openai_exception(
message=f"ContentPolicyViolationError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str:
helpful_message: Final = (
@ -345,7 +360,7 @@ def _map_openai_exception(
message=helpful_message,
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
@ -354,7 +369,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
@ -372,7 +387,7 @@ def _map_openai_exception(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif (
@ -383,7 +398,7 @@ def _map_openai_exception(
message=f"AuthenticationError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif "Mistral API raised a streaming error" in error_str:
@ -402,15 +417,16 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif original_exception.status_code == 401:
raise AuthenticationError(
message=f"AuthenticationError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 404:
@ -418,7 +434,7 @@ def _map_openai_exception(
message=f"NotFoundError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 408:
@ -433,7 +449,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
@ -442,7 +458,7 @@ def _map_openai_exception(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 500:
@ -450,7 +466,7 @@ def _map_openai_exception(
message=f"InternalServerError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 502:
@ -458,7 +474,7 @@ def _map_openai_exception(
message=f"BadGatewayError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 503:
@ -466,7 +482,7 @@ def _map_openai_exception(
message=f"ServiceUnavailableError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 504: # gateway timeout error
@ -2423,10 +2439,11 @@ def exception_type(
custom_llm_provider == "litellm_proxy"
): # handle special case where calling litellm proxy + exception str contains error message
extract_and_raise_litellm_exception(
response=getattr(original_exception, "response", None),
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
error_str=error_str,
model=model,
custom_llm_provider=custom_llm_provider,
body=getattr(original_exception, "body", None),
)
if (
custom_llm_provider == "openai"

View file

@ -484,11 +484,12 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None,
def _apply_off_peak_to_base_costs(
model_info: ModelInfo,
current_time: datetime | None,
base_costs: tuple[float, float, float, float, float],
base_costs: tuple[float, float, float, float | None, float],
) -> tuple[float, float, float, float, float]:
"""Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
produced them. The one-hour cache-creation rate passes through untouched, since
off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate.
produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a
present one passes through untouched and an absent one resolves to the applied
cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate.
"""
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
rates: Final = apply_off_peak_pricing(
@ -506,7 +507,7 @@ def _apply_off_peak_to_base_costs(
rates.input_rate,
rates.output_rate,
rates.cache_creation_rate,
cache_creation_above_1hr,
rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr,
rates.cache_read_rate,
)
@ -532,6 +533,11 @@ def _get_token_base_cost(
`missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
An absent cache-creation rate always resolves to the resolved input rate, the way the
tiered table and custom deployment pricing already do, since a provider that publishes
no write price bills cache writes as ordinary input. An absent 1h write rate resolves
to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for both.
Returns:
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
"""
@ -554,10 +560,9 @@ def _get_token_base_cost(
output_image_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_image_token", None)
if output_image_cost is not None:
completion_base_cost = cast(float, output_image_cost)
cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key))
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None)
cache_creation_cost_above_1hr = _get_cost_per_unit(
model_info, "cache_creation_input_token_cost_above_1hr", default_value=None
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
@ -639,22 +644,10 @@ def _get_token_base_cost(
else f"cache_read_input_token_cost_above_{threshold_str}_tokens"
)
cache_creation_cost = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_tiered_key,
cache_creation_cost,
),
)
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost)
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_1hr_tiered_key,
cache_creation_cost_above_1hr,
),
cache_creation_cost_above_1hr = _get_cost_per_unit(
model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
@ -665,16 +658,16 @@ def _get_token_base_cost(
except Exception:
continue
input_rate_for_missing_cache_rates: Final = _off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if cache_read_cost is None:
cache_read_cost = (
_off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if missing_cache_read_uses_input
else 0.0
)
cache_read_cost = input_rate_for_missing_cache_rates if missing_cache_read_uses_input else 0.0
resolved_cache_creation_cost: Final = (
input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost
)
return _apply_off_peak_to_base_costs(
model_info,
@ -682,7 +675,7 @@ def _get_token_base_cost(
(
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
resolved_cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
),
@ -956,12 +949,16 @@ 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 +967,9 @@ 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"

View file

@ -20,6 +20,7 @@ from itertools import chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_proxy_logger
@ -44,6 +45,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
scoped_structured_message_indices,
stream_item_field,
stream_item_fingerprint,
unappliable_request_rewrite,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -103,9 +105,24 @@ class ToolResultBlockTextTarget:
block_idx: int
InputWriteBackTarget = (
MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
)
@dataclass(frozen=True, slots=True)
class SystemStringTarget:
pass
@dataclass(frozen=True, slots=True)
class SystemBlockTextTarget:
block_idx: int
@dataclass(frozen=True, slots=True)
class ToolUseInputTarget:
msg_idx: int
content_idx: int
MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
@ -146,10 +163,17 @@ class ScannedText:
target: InputWriteBackTarget
@dataclass(frozen=True, slots=True)
class ScannedToolCall:
tool_call: ChatCompletionToolCallChunk
target: ToolUseInputTarget
@dataclass(frozen=True, slots=True)
class ExtractedInput:
scanned: tuple[ScannedText, ...]
images: tuple[str, ...]
tool_calls: tuple[ScannedToolCall, ...] = ()
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
@ -161,6 +185,74 @@ class _ToolCallShape:
arguments: str
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
return (
block.get("type") == "tool_use"
and isinstance(block.get("id"), str)
and isinstance(block.get("name"), str)
and isinstance(block.get("input"), dict)
)
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
if not isinstance(system, list):
return
text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text")
if block_idx < len(text_blocks):
text_blocks[block_idx]["text"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None:
content: Final = message.get("content", None)
if content is None:
return
match target:
case MessageContentTarget():
if isinstance(content, str):
message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place
case ContentBlockTextTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["text"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultStringTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["content"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
if isinstance(content, list):
content[content_idx]["content"][block_idx]["text"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case _:
assert_never(target)
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None:
try:
return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments)
except ValidationError:
return None
def _write_back_tool_use(
message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object]
) -> None:
content: Final = message.get("content", None)
block: Final = content[target.content_idx] if isinstance(content, list) else None
if not isinstance(block, dict):
return
block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place
if shape.name is not None and shape.name != block.get("name"):
block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place
@dataclass(frozen=True, slots=True)
class _SSEFieldRewrite:
"""One field of one nested section of a buffered SSE event, rewritten."""
@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation):
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
# Exclude only the trusted top-level prompt. In-sequence system entries are untrusted
# and must stay aligned with texts_to_check for positional masking. When the top-level
# prompt is included, the pre-existing count mismatch disables positional masking.
# The top-level prompt is translated on its own below so it can be hoisted in front of
# any mid-turn system entries and scanned first, aligned with that structured position.
translation_source: Final = { # mutable-ok: API message payload
key: value for key, value in data.items() if key != "system"
}
@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation):
]
)
# Step 1: Extract all text content and images
# Step 1: Extract all text content, images, and tool calls
top_level_system_scanned: Final = (
()
if hoisted_system_message is None or scan_only_tool_results
else self._extract_top_level_system_text(hoisted_system_message)
)
extracted: Final = tuple(
self._extract_input_text_and_images(
message=message,
@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation):
)
for msg_idx, message in enumerate(messages)
)
scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned)
scanned: Final = (
*top_level_system_scanned,
*(item for one_message in extracted for item in one_message.scanned),
)
texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
images_to_check: Final = [
image for one_message in extracted for image in one_message.images
] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls)
tool_calls_to_check: Final = [
item.tool_call for item in scanned_tool_calls
] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk]
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
# Step 2: Apply guardrail to all texts and tool calls in batch
if texts_to_check or tool_calls_to_check:
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
if tools_to_check:
inputs["tools"] = tools_to_check
original_structured_messages: Final = structured_messages
@ -570,9 +676,18 @@ class AnthropicMessagesHandler(BaseTranslation):
preserve_system_messages=has_midturn_system_message,
)
else:
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
self._apply_guardrail_tool_calls_to_input(
messages=messages,
scanned_tool_calls=scanned_tool_calls,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
guardrail_name=guardrail_to_apply.guardrail_name,
)
# Step 3: Map guardrail responses back to original message structure
await self._apply_guardrail_responses_to_input(
messages=messages,
data=data,
responses=guardrailed_texts,
scanned=scanned,
)
@ -598,6 +713,19 @@ class AnthropicMessagesHandler(BaseTranslation):
hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload
return hoisted[0] if hoisted else None
@staticmethod
def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]:
content: Final = hoisted_system_message.get("content")
if isinstance(content, str):
return (ScannedText(content, SystemStringTarget()),)
if not isinstance(content, list):
return ()
return tuple(
ScannedText(text_str, SystemBlockTextTarget(block_idx))
for block_idx, block in enumerate(content)
if isinstance(block, dict) and isinstance(text_str := block.get("text"), str)
)
@staticmethod
def _openai_system_message_to_anthropic(
message: Mapping[str, object],
@ -852,9 +980,25 @@ class AnthropicMessagesHandler(BaseTranslation):
for content_idx, content_item in enumerate(content)
if isinstance(content_item, dict)
)
tool_use_blocks: Final = (
()
if scan_only_tool_results
else tuple(
(content_idx, content_item)
for content_idx, content_item in enumerate(content)
if isinstance(content_item, dict) and _is_client_tool_use(content_item)
)
)
return ExtractedInput(
scanned=tuple(item for block in blocks for item in block.scanned),
images=tuple(image for block in blocks for image in block.images),
tool_calls=tuple(
ScannedToolCall(
tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx),
target=ToolUseInputTarget(msg_idx, content_idx),
)
for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks)
),
)
@classmethod
@ -940,43 +1084,59 @@ class AnthropicMessagesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input(
self,
messages: Sequence[_WritableMessage],
responses: list[str],
data: dict[str, object], # mutable-ok: API message payload
responses: Sequence[str],
scanned: tuple[ScannedText, ...],
) -> None:
"""
Apply guardrail responses back to input messages.
Apply guardrail responses back to the top-level system prompt and the input messages.
"""
raw_messages: Final = data.get("messages")
messages: Final[Sequence[_WritableMessage]] = raw_messages if isinstance(raw_messages, list) else ()
for item, guardrail_response in zip(scanned, responses):
target = item.target
message = messages[target.msg_idx]
content = message.get("content", None)
if content is None:
continue
match target:
case MessageContentTarget():
if isinstance(content, str):
message["content"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ContentBlockTextTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["text"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultStringTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["content"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
if isinstance(content, list):
content[content_idx]["content"][block_idx]["text"] = (
match item.target:
case SystemStringTarget():
if isinstance(data.get("system"), str):
data["system"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case SystemBlockTextTarget(block_idx=block_idx):
_write_back_system_block(data.get("system"), block_idx, guardrail_response)
case (
MessageContentTarget()
| ContentBlockTextTarget()
| ToolResultStringTarget()
| ToolResultBlockTextTarget() as message_target
):
_write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response)
case _:
assert_never(target)
assert_never(item.target)
@staticmethod
def _apply_guardrail_tool_calls_to_input(
messages: Sequence[_WritableMessage],
scanned_tool_calls: tuple[ScannedToolCall, ...],
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
returned_tool_calls: Sequence[object] | None,
guardrail_name: str | None,
) -> None:
post_guardrail_tool_calls: Final = _tool_call_shapes(
returned_tool_calls
if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
else tuple(item.tool_call for item in scanned_tool_calls)
)
rewritten: Final = tuple(
(item, after, _rewritten_tool_use_input(after.arguments))
for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if before != after
)
applicable: Final = tuple(
(item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None
)
if len(applicable) != len(rewritten):
raise unappliable_request_rewrite(guardrail_name)
for item, after, rewritten_input in applicable:
_write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input)
async def process_output_response(
self,

View file

@ -1,8 +1,8 @@
from __future__ import annotations
import json
from collections.abc import Callable, Iterator, Sequence
from typing import Final, TypeVar
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
from pydantic import BaseModel
@ -364,3 +364,67 @@ def merge_guardrailed_scoped_messages(
yield from appended
return list(_merged())
def _content_part_text(part: object) -> str | None:
if not isinstance(part, Mapping):
return None
text: Final = part.get("text")
return text if isinstance(text, str) else None
def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]:
content: Final = message.get("content")
if isinstance(content, str):
return (content,)
if isinstance(content, list):
return tuple(text for part in content if (text := _content_part_text(part)) is not None)
return ()
def message_text_slot_count(message: AllMessageValues) -> int:
return len(message_slot_texts(message))
def _part_with_text(part: object, text: str) -> object:
if not isinstance(part, Mapping):
return part
return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts
def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]:
remaining_texts: Final = iter(texts)
return [ # mutable-ok: message content stays a JSON list
_part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part
for part in content
]
def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None:
"""Swap one rewritten text into each text slot of a chat row, in order.
A slot is a string ``content`` or one list part carrying a string ``text``;
images and other parts ride along untouched. Returns None unless the counts
line up exactly, so a rewrite never lands on the wrong slot.
"""
if message_text_slot_count(message) != len(texts):
return None
content: Final = message.get("content")
if not isinstance(content, (str, list)):
return message
rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts)
rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts
return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped
class UnappliableRequestRewrite(Exception):
def __init__(self, guardrail_name: str) -> None:
super().__init__(
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
"so the request was rejected rather than sent unrewritten"
)
self.guardrail_name: Final = guardrail_name
def unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite:
return UnappliableRequestRewrite(guardrail_name or "unknown")

View file

@ -17,7 +17,7 @@ BaseAWSLLM._sign_request after the request body is finalized.
import json
from collections.abc import Mapping
from typing import Any, Final
from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
import httpx
from typing_extensions import ReadOnly, TypedDict
@ -32,6 +32,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
BedrockMantleAuthMixin,
)
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
ResponseInputParam,
@ -58,8 +59,6 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset(
_BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"})
_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message"
_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction"
_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call"
@ -233,17 +232,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
normalized_input: Final = self._normalize_codex_input_items(remaining_input)
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
)
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
normalized_input: Final = self._normalize_codex_input_items(hoisted.input)
request_params: Final = (
{
**response_api_optional_request_params,
"tools": [
*(response_api_optional_request_params.get("tools") or []),
*hoisted_tools,
],
}
if hoisted_tools
self._params_with_hoisted_tools(params, hoisted)
if hoisted.hoisted
else response_api_optional_request_params
)
return super().transform_responses_api_request(
@ -254,41 +250,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
headers=headers,
)
@staticmethod
def _is_codex_additional_tools_item(item: Any) -> bool:
return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
@staticmethod
def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]":
tools: Final = item.get("tools")
return tools if isinstance(tools, list) else []
@classmethod
def _hoist_codex_additional_tools(
cls,
input: "str | ResponseInputParam",
) -> "tuple[str | ResponseInputParam, list[Any]]":
"""Codex's "responses lite" wire mode ships tool definitions inside
`input` as {"type": "additional_tools", "role": "developer",
"tools": [...]} items. api.openai.com accepts that item type; Mantle
rejects the whole request with 400 "Invalid 'input': value did not
match any expected variant" but accepts the same tools at the top
level, so move them there and strip the items from `input`.
"""
if not isinstance(input, list):
return input, []
additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)]
if not additional_tools_items:
return input, []
remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)]
hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)]
verbose_logger.debug(
"Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) "
"into the top-level tools param (Mantle rejects that input item type).",
len(hoisted_tools),
len(additional_tools_items),
)
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
def _params_with_hoisted_tools(
cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools
) -> dict[str, object]:
supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools))
if supported_tools:
return {**params, "tools": supported_tools}
return {key: value for key, value in params.items() if key != "tools"}
@staticmethod
def _agent_message_text(item: "Mapping[str, object]") -> str:

View file

@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_field,
stream_item_fingerprint,
stream_item_items,
unappliable_request_rewrite,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -196,6 +197,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
else:
# Step 3: Map guardrail responses back to original message structure
if guardrailed_texts and texts_to_check:
if len(guardrailed_texts) != len(text_task_mappings):
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
await self._apply_guardrail_responses_to_input_texts(
messages=messages,
responses=guardrailed_texts,
@ -210,6 +213,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
task_mappings=tool_call_task_mappings,
)
elif (
not images_to_check
and not guardrail_to_apply.records_own_guardrail_information
and (not_run_reason := self._not_run_reason(messages)) is not None
):
guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=not_run_reason,
request_data=data,
guardrail_status="not_run",
)
verbose_proxy_logger.debug(
"OpenAI Chat Completions: Processed input messages: %s",
data.get("messages"),
@ -217,6 +231,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return data
def _not_run_reason(
self,
messages: Sequence[dict[str, Any]], # mutable-ok: raw request messages consumed by _extract_inputs
) -> str | None:
"""Why nothing was scanned, or None when the only unscoped content is images, which this handler never scans."""
texts: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
images: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
tool_calls: Final[list[ChatCompletionToolParam]] = [] # mutable-ok: filled by _extract_inputs
for msg_idx, message in enumerate(messages):
self._extract_inputs(
message=message,
msg_idx=msg_idx,
texts_to_check=texts,
images_to_check=images,
tool_calls_to_check=tool_calls,
text_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
tool_call_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
)
if texts or tool_calls:
return "no scannable content after message scoping"
return None if images else "no scannable content"
def extract_request_tool_names(self, data: dict) -> list[str]:
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
names: Final[list[str]] = []

View file

@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_field,
stream_item_fingerprint,
stream_item_items,
unappliable_request_rewrite,
)
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
from litellm.responses.litellm_completion_transformation.transformation import (
@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation):
data["instructions"] = written_back.instructions # rebind-ok: data is an out-param
elif isinstance(input_data, str):
guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
if len(guardrailed_texts) > 1:
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param
else:
rewritten_texts: Final = guardrailed_inputs.get("texts") or ()
if len(rewritten_texts) != len(extracted.task_mappings):
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown")
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
await self._apply_guardrail_responses_to_input(
messages=input_data,
responses=rewritten_texts,

View file

@ -6,8 +6,10 @@ from typing import Final, TypeAlias
from pydantic import BaseModel, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix
from litellm.responses.litellm_completion_transformation.transformation import (
NAMESPACE_DESCRIPTION_SEPARATOR,
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS,
LiteLLMCompletionResponsesConfig,
)
@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]:
return tuple(tool for tool in validated if tool is not None)
def _is_function(tool: Tool) -> bool:
return tool.get("type") == "function"
def _has_chat_tool(member: Tool) -> bool:
return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS
def _chat_tool_key(tool: Tool) -> str:
@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool:
return function if function is not None else MappingProxyType({})
def _without_namespace_prefix(key: str, value: object, prefix: str) -> object:
if key != "description" or not isinstance(value, str) or not value.startswith(prefix):
def _member_description(key: str, value: object, prefix: str, suffix: str) -> object:
if key != "description" or not isinstance(value, str):
return value
return value[len(prefix) :]
return value.replace(prefix, "", 1).replace(suffix, "", 1)
def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool:
flattened_function: Final = _function_fields(flattened)
prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else ""
suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else ""
changed_function: Final = MappingProxyType(
{
key: _without_namespace_prefix(key, value, prefix)
key: _member_description(key, value, prefix, suffix)
for key, value in _function_fields(guardrailed).items()
if flattened_function.get(key) != value
}
@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_
return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType
def _rebuilt_function_members(
function_members: Sequence[Tool],
def _rebuilt_flattened_members(
flattened_members: Sequence[Tool],
flattened_group: Sequence[Tool],
group_keys: Sequence[IndexedKey],
guardrailed_by_key: Mapping[IndexedKey, Tool],
@ -106,7 +109,7 @@ def _rebuilt_function_members(
else member
if guardrailed_by_key[key] == flattened
else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description)
for member, flattened, key in zip(function_members, flattened_group, group_keys)
for member, flattened, key in zip(flattened_members, flattened_group, group_keys)
)
@ -118,9 +121,9 @@ def _rebuilt_namespace(
guardrailed_by_key: Mapping[IndexedKey, Tool],
) -> tuple[Tool, ...]:
namespace_description: Final = str(original.get("description") or "")
rebuilt_functions: Final = iter(
_rebuilt_function_members(
tuple(member for member in members if _is_function(member)),
rebuilt_flattened: Final = iter(
_rebuilt_flattened_members(
tuple(member for member in members if _has_chat_tool(member)),
flattened_group,
group_keys,
guardrailed_by_key,
@ -129,7 +132,7 @@ def _rebuilt_namespace(
)
rebuilt_members: Final = tuple(
rebuilt
for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members)
for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members)
if rebuilt is not None
)
if not rebuilt_members:
@ -149,7 +152,7 @@ def _merged_original(
if guardrailed_group == tuple(flattened_group):
return (original,)
members: Final = _namespace_members(original) if original.get("type") == "namespace" else ()
if members and sum(map(_is_function, members)) == len(flattened_group):
if members and sum(map(_has_chat_tool, members)) == len(flattened_group):
return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key)
if not guardrailed_group:
return ()

View file

@ -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:

View file

@ -298,8 +298,6 @@ def transform_openai_input_gemini_embed_content(
_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)
@ -339,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:
@ -372,30 +371,29 @@ 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:
return Usage(
prompt_tokens=prompt_tokens,
total_tokens=total_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=0,
image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0,
),
)
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 +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 (files/abc) to {mime_type, uri},
used to bill resolved image references at the per-image rate
resolved_files: Mapping of file references to resolved metadata
Returns:
EmbeddingResponse with single embedding

View file

@ -25601,10 +25601,14 @@
"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_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_video_per_second": 0.00079,
"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,
@ -25615,13 +25619,14 @@
"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_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_per_second": 0.00079,
"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 +25638,14 @@
"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_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_video_per_second": 0.00079,
"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,
@ -25648,13 +25657,14 @@
"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_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_per_second": 0.00079,
"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,
@ -25693,10 +25703,14 @@
},
"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_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_video_per_second": 0.00079,
"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 +25723,14 @@
"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_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_video_per_second": 0.00079,
"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,

View file

@ -4,11 +4,12 @@ import os
from collections.abc import Callable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple
import httpx
from pydantic import (
BaseModel,
BeforeValidator,
ConfigDict,
Field,
Json,
@ -47,6 +48,7 @@ from litellm.types.proxy.carried_budget_state import (
)
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.router_weights import validate_router_settings_dict
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CallTypes,
@ -1981,8 +1983,14 @@ class OrgMember(MemberBase):
from litellm.models.team import TeamBase as TeamBase # noqa: E402
RouterSettingsDict = Annotated[
dict[str, object],
BeforeValidator(validate_router_settings_dict, json_schema_input_type=UpdateRouterConfig),
]
class NewTeamRequest(TeamBase):
router_settings: RouterSettingsDict | None = None
model_aliases: dict | None = None
tags: list | None = None
guardrails: list[str] | None = None
@ -2080,7 +2088,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
enforced_batch_output_expires_after: dict | None = None
enforced_file_expires_after: dict | None = None
router_settings: dict | None = None
router_settings: RouterSettingsDict | None = None
access_group_ids: list[str] | None = None
budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows
default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members

View file

@ -249,6 +249,7 @@ async def authenticate_user(
if os.getenv("DATABASE_URL") is not None:
response = await generate_key_helper_fn(
llm_router=None,
request_type="key",
**{
"user_role": LitellmUserRoles.PROXY_ADMIN,
@ -324,6 +325,7 @@ async def authenticate_user(
await _rehash_password_if_needed(_user_row.user_id, password, _password)
if os.getenv("DATABASE_URL") is not None:
response = await generate_key_helper_fn(
llm_router=None,
request_type="key",
**{
"user_role": user_role,

View file

@ -878,6 +878,7 @@ async def _auto_register_jwt_mapping(
# the NOT NULL @id constraint. Every successful key-creation caller (e.g.
# /key/generate) passes table_name="key" explicitly.
key_data: Final = await generate_key_helper_fn(
llm_router=None,
request_type="key",
table_name="key",
team_id=team_id,

View file

@ -580,8 +580,8 @@ What the command changed is recorded in `~/.litellm/claude_configure_state.json`
`lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline:
```
claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14
Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
claude-auto ████████░░░░░░░░░░░░░░░░ $0.14
Claude Opus 5 ████████████████████████ $0.38
```

View file

@ -28,6 +28,7 @@ import os
import sys
import tempfile
import time
import unicodedata
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping
@ -42,7 +43,6 @@ FETCH_TIMEOUT_SECONDS: Final = 3
BAR_WIDTH: Final = 24
BAR_FULL: Final = "\u2588"
BAR_EMPTY: Final = "\u2591"
SEPARATOR: Final = " \u00b7 "
TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024
CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",)
CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
@ -50,7 +50,6 @@ CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",)
CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",)
CODEX_STOP_EVENT: Final = "Stop"
SYNTHETIC_MODEL: Final = "<synthetic>"
LITELLM_LABEL: Final = "LiteLLM"
RESET: Final = "\033[0m"
BOLD: Final = "\033[1m"
DIM: Final = "\033[90m"
@ -302,31 +301,37 @@ def _bar(fraction: float, color: str, width: int, use_color: bool) -> str:
return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}"
def _display_width(label: str) -> int:
return sum(
2 if unicodedata.east_asian_width(character) in ("W", "F") else 1
for character in label
if unicodedata.category(character) not in ("Mn", "Me")
)
def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str:
def paint(code: str, text: str) -> str:
return f"{code}{text}{RESET}" if use_color else text
routed: Final = paint(BOLD, f"Routed to: {model}")
if session is None:
if session is None or session.baseline_model is None or session.baseline_spend <= 0:
return routed
header: Final = f"{session.router_name}{SEPARATOR}{routed}"
if session.baseline_model is None or session.baseline_spend <= 0:
return header
reference: Final = baseline_label(session.baseline_model, config_dir)
pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100
delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}")
peak: Final = max(session.spend, session.baseline_spend)
label_width: Final = max(len(LITELLM_LABEL), len(reference))
label_width: Final = max(_display_width(session.router_name), _display_width(reference))
rows: Final = (
(LITELLM_LABEL, session.spend, LITELLM_COLOR),
(session.router_name, session.spend, LITELLM_COLOR),
(reference, session.baseline_spend, BASELINE_COLOR),
)
lines: Final = (
f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} "
f"{paint(DIM, label + ' ' * (label_width - _display_width(label)))} "
f"{_bar(amount / peak, color, bar_width, use_color)} "
f"{paint(DIM, f'${amount:.2f}')}"
for label, amount, color in rows
)
return "\n".join((f"{header} {delta}", *lines))
return "\n".join((f"{routed} {delta}", *lines))
def color_enabled(env: Mapping[str, str]) -> bool:

View file

@ -14,6 +14,7 @@ import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import ValidationError
from starlette.types import Receive, Scope, Send
import litellm
@ -76,6 +77,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.router_weights import validate_router_weights
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
_LlmCallT = TypeVar("_LlmCallT")
@ -1939,6 +1941,13 @@ class ProxyBaseLLMRequestProcessing:
# This avoids expensive Router instantiation on each request
if router_settings is not None:
self.data["router_settings_override"] = router_settings
try:
self.data["_router_weights"] = validate_router_weights(router_settings.get("weights"))
except ValidationError:
self.data["_router_weights"] = None
verbose_proxy_logger.warning(
"Ignoring invalid saved router weights; update team/key router_settings"
)
alias_target: Final = await _resolve_per_request_model_group_alias(
requested_model=self.data.get("model"),
router_settings=router_settings,

View file

@ -8,6 +8,8 @@ from typing import Final
from fastapi import status
from litellm.constants import STRINGIFIED_NONE
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
{
status.HTTP_401_UNAUTHORIZED: "authentication_error",
@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str:
"""OpenAI types ``error.type`` as a required string, so an exception carrying none
falls back to the type its status code stands for."""
carried: Final = attribute_of(exc, "type")
if isinstance(carried, str):
if isinstance(carried, str) and carried != STRINGIFIED_NONE:
return carried
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
if mapped is not None:
@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None:
"""OpenAI types ``error.param`` as nullable, so an exception carrying none
serializes as JSON ``null``."""
carried: Final = attribute_of(exc, "param")
return carried if isinstance(carried, str) else None
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None

View file

@ -26,7 +26,7 @@ class ComplianceChecker:
def __init__(self, data: ComplianceCheckRequest):
self.data = data
self.guardrails = data.guardrail_information or []
self.guardrails = tuple(g for g in data.guardrail_information or () if g.get("guardrail_status") != "not_run")
def _get_guardrails_by_mode(self, mode: str) -> list[dict]:
"""

View file

@ -7,7 +7,7 @@
import fnmatch
import os
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
import httpx
@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIMetadata,
GenericGuardrailAPIRequest,
@ -150,6 +150,26 @@ def _extract_inbound_headers(
return None
def _structured_rows_to_write_back(
original_rows: Sequence[AllMessageValues] | None,
shown_rows: Sequence[AllMessageValues] | None,
returned_rows: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...] | None:
"""The request model drops row keys its message types do not declare, so a
row the server echoes back verbatim is restored to the original row object.
A server that echoes every row back unchanged has not rewritten anything
per row, so its answer is read from texts, as it was before rows could be
returned at all."""
if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows):
return tuple(returned_rows)
if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)):
return None
return tuple(
original if returned == shown else returned
for original, shown, returned in zip(original_rows, shown_rows, returned_rows)
)
class GenericGuardrailAPI(CustomGuardrail):
"""
Generic Guardrail API integration for LiteLLM.
@ -322,6 +342,8 @@ class GenericGuardrailAPI(CustomGuardrail):
texts: list,
images: list[str] | None,
tools: list[ChatCompletionToolParam] | None,
structured_messages: Sequence[AllMessageValues] | None,
shown_messages: Sequence[AllMessageValues] | None,
guardrail_response: GenericGuardrailAPIResponse,
) -> GenericGuardrailAPIInputs:
# Action is NONE or no modifications needed
@ -336,6 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail):
return_inputs["tools"] = guardrail_response.tools
elif tools:
return_inputs["tools"] = tools
rows_to_write_back: Final = (
_structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages)
if guardrail_response.structured_messages
else None
)
if rows_to_write_back is not None:
return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list
if guardrail_response.stream_holdback_chars is not None:
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
return return_inputs
@ -473,6 +502,8 @@ class GenericGuardrailAPI(CustomGuardrail):
texts=texts,
images=images,
tools=tools,
structured_messages=structured_messages,
shown_messages=guardrail_request.structured_messages,
guardrail_response=guardrail_response,
)

View file

@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
Args:
texts: Flattened text entries from the framework.
messages: Original request messages (request_data["messages"]),
NOT structured_messages (which may have injected system content).
messages: The structured messages the framework flattened into ``texts``,
hoisted top-level system prompt included, so positions line up.
Returns a set of scannable indices, or None on count mismatch or no user/developer
message (safety fallback to existing role-filter behavior).
@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail):
structured_messages: Final = inputs.get("structured_messages")
if structured_messages:
# For Anthropic /v1/messages: default to latest-user-only scanning.
# Uses request_data["messages"] (original format), NOT structured_messages
# (which has injected system content from adapter translation).
if self._use_latest_user_only(request_data, logging_obj):
original_messages: Final = request_data.get("messages")
if original_messages:
scannable_indices = self._get_latest_user_text_indices(texts, original_messages)
scannable_indices = self._get_latest_user_text_indices(texts, structured_messages)
# Fall through to existing role filtering if:
# - not Anthropic, OR flag explicitly False, OR
# - no original messages, OR
# - latest-user extraction returned None (no user / count mismatch)
if scannable_indices is None:
scannable_indices = self._get_scannable_text_indices(texts, structured_messages)

View file

@ -2,6 +2,7 @@ import asyncio
import base64
import os
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Optional
import httpx
@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -28,12 +31,36 @@ if TYPE_CHECKING:
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
class PromptSecurityGuardrailMissingSecrets(Exception):
pass
def _inputs_with_structured_messages(
inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
) -> GenericGuardrailAPIInputs:
if rewritten_messages is None:
return inputs
patched: Final[GenericGuardrailAPIInputs] = {
**inputs,
"structured_messages": list(rewritten_messages), # mutable-ok: the TypedDict field is declared as a list
}
return patched
def _inputs_with_modifications(
inputs: GenericGuardrailAPIInputs,
modified_texts: list[str],
rewritten_messages: Sequence[AllMessageValues] | None,
) -> GenericGuardrailAPIInputs:
if not modified_texts:
return _inputs_with_structured_messages(inputs, rewritten_messages)
with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts}
return _inputs_with_structured_messages(with_texts, rewritten_messages)
class _ProtectVerdict(TypedDict, total=False):
"""One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
@ -276,14 +303,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
)
elif action == "modify":
# Extract modified texts from modified_messages
modified_messages: Final = result.get("modified_messages", [])
modified_texts: Final = self._extract_texts_from_messages(modified_messages)
if modified_texts:
inputs["texts"] = modified_texts
return _inputs_with_modifications(
inputs,
self._extract_texts_from_messages(modified_messages),
self._structured_messages_with_modifications(structured_messages, modified_messages),
)
return inputs
def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool:
return self.check_tool_results or message.get("role") in _PROTECT_ROLES
def _structured_messages_with_modifications(
self,
structured_messages: Sequence[AllMessageValues],
modified_messages: Sequence[Mapping[str, object]],
) -> tuple[AllMessageValues, ...] | None:
sent_indices: Final = tuple(
index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message)
)
if not sent_indices or len(sent_indices) != len(modified_messages):
return None
rewritten: Final = tuple(
message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,)))
for index, modified in zip(sent_indices, modified_messages)
)
replacements: Final = MappingProxyType(
{index: message for index, message in zip(sent_indices, rewritten) if message is not None}
)
if len(replacements) != len(sent_indices):
return None
return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages))
async def _apply_guardrail_on_response(
self,
inputs: GenericGuardrailAPIInputs,
@ -347,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
return inputs
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
"""Extract text content from messages."""
texts: Final = []
for message in messages:
content = message.get("content")
if isinstance(content, str):
texts.append(content)
elif isinstance(content, list):
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
text = item.get("text")
if text:
texts.append(text)
return texts
return [text for message in messages for text in message_slot_texts(message)]
async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None:
"""Process standalone images from inputs (data URLs)."""
@ -681,14 +721,13 @@ class PromptSecurityGuardrail(CustomGuardrail):
This allows checking tool results for indirect prompt injection when enabled.
"""
supported_roles: Final = ["system", "user", "assistant"]
filtered_messages: Final = []
transformed_count = 0
filtered_count = 0
for message in messages:
role = message.get("role", "")
if role in supported_roles:
if role in _PROTECT_ROLES:
filtered_messages.append(message)
else:
if self.check_tool_results:

View file

@ -42,7 +42,7 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2})
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"not_run": 0, "passed": 1, "flagged": 2, "blocked": 3})
_T = TypeVar("_T")
@ -325,7 +325,7 @@ class UsageDetailResponse(BaseModel):
class UsageLogEntry(BaseModel):
id: str
timestamp: str
action: str # blocked | passed | flagged
action: str # blocked | passed | flagged | not_run
score: float | None
latency_ms: float | None
model: str | None

View file

@ -193,10 +193,12 @@ async def _upsert_rows_with_retry(
def guardrail_status_to_action(status: str | None) -> str:
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
"""Map StandardLogging guardrail_status to blocked/passed/flagged/not_run."""
if not status:
return "passed"
s: Final = (status or "").lower()
if s == "not_run":
return "not_run"
if "intervened" in s or "block" in s:
return "blocked"
if "flagged" in s or "fail" in s or "error" in s:
@ -354,37 +356,49 @@ async def process_spend_logs_guardrail_usage(
"flagged_count": 0,
}
)
index_rows: Final[list[dict[str, object]]] = []
index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {}
for payload in logs_to_process:
request_id = payload.get("request_id")
start_time = _parse_payload_start_time(payload)
if not request_id or start_time is None:
if not isinstance(request_id, str) or not request_id or start_time is None:
continue
date_key = _date_str(start_time)
for entry in _parse_guardrail_info_from_payload(payload):
guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or ""
if not guardrail_id:
entries = _parse_guardrail_info_from_payload(payload)
ids_by_name = MappingProxyType(
{
e["guardrail_name"]: e["guardrail_id"]
for e in entries
if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) and e["guardrail_name"]
}
)
for entry in entries:
raw_name = entry.get("guardrail_name")
guardrail_name = raw_name if isinstance(raw_name, str) else ""
guardrail_id = entry.get("guardrail_id") or ids_by_name.get(guardrail_name) or guardrail_name
if not isinstance(guardrail_id, str) or not guardrail_id:
continue
key = _MetricsKey(guardrail_id, date_key)
daily_guardrail[key]["requests_evaluated"] += 1
action = guardrail_status_to_action(entry.get("guardrail_status"))
if action == "passed":
daily_guardrail[key]["passed_count"] += 1
elif action == "blocked":
daily_guardrail[key]["blocked_count"] += 1
else:
daily_guardrail[key]["flagged_count"] += 1
if action != "not_run":
key = _MetricsKey(guardrail_id, date_key)
daily_guardrail[key]["requests_evaluated"] += 1
if action == "passed":
daily_guardrail[key]["passed_count"] += 1
elif action == "blocked":
daily_guardrail[key]["blocked_count"] += 1
else:
daily_guardrail[key]["flagged_count"] += 1
policy_id = entry.get("policy_id")
index_rows.append(
{
prior = index_rows_by_key.get((request_id, guardrail_id))
if prior is None or (prior["policy_id"] is None and policy_id is not None):
index_rows_by_key[(request_id, guardrail_id)] = {
"request_id": request_id,
"guardrail_id": guardrail_id,
"policy_id": policy_id,
"start_time": start_time,
}
)
index_rows: Final = tuple(index_rows_by_key.values())
async with pending.lock:
pending_metrics: Final = pending.metrics

View file

@ -221,6 +221,8 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
)
_UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"weights",
"_router_weights",
"proxy_server_request",
"standard_logging_object",
"secret_fields",
@ -334,7 +336,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
# and read by spend logs as fact; a client value has no legitimate meaning and no
# key or team setting keeps it, so the strip is never gated.
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
{"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY}
)
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"

View file

@ -567,7 +567,7 @@ async def new_user(
teams = check_if_default_team_set()
organization_ids: Final = cast(list[str] | None, data_json.pop("organizations", None))
response: Final = await generate_key_helper_fn(request_type="user", **data_json)
response: Final = await generate_key_helper_fn(request_type="user", **data_json, llm_router=None)
# Admin UI Logic
# Add User to Team and Organization
# if team_id passed add this user to the team

View file

@ -96,6 +96,7 @@ from litellm.proxy.management_endpoints.common_utils import (
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
from litellm.proxy.management_helpers.access_group_key_sync import (
sync_key_access_group_membership,
sync_key_regeneration_access_group_membership,
@ -201,6 +202,10 @@ class _KeyUpdateResult(TypedDict):
data: ReadOnly[Mapping[str, object]]
class _StoredKeyRouterSettings(BaseModel):
router_settings: Mapping[str, object] | None = None
class _KeyRowWhere(TypedDict):
token: ReadOnly[str]
@ -1330,7 +1335,7 @@ async def _common_key_generation_helper(
prisma_client=prisma_client,
)
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key")
response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key", llm_router=llm_router)
response["soft_budget"] = data.soft_budget # include the user-input soft budget in the response
@ -2234,7 +2239,26 @@ async def _update_key_row_with_soft_budget(
async def prepare_key_update_data(
data: UpdateKeyRequest | RegenerateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
*,
prisma_client: PrismaClient | None = None,
llm_router: Router | None = None,
):
if data.router_settings is not None or (
"router_settings" not in data.model_fields_set
and "team_id" in data.model_fields_set
and data.team_id != existing_key_row.team_id
):
effective_settings: Final = (
data.router_settings
if data.router_settings is not None
else _StoredKeyRouterSettings.model_validate(existing_key_row, from_attributes=True).router_settings
)
await validate_router_settings_weights(
effective_settings,
team_id=data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id,
prisma_client=prisma_client,
llm_router=llm_router,
)
data_json: Final[dict] = data.model_dump(exclude_unset=True)
data_json.pop("key", None)
data_json.pop("new_key", None)
@ -2575,7 +2599,9 @@ async def _process_single_key_update(
)
# Prepare update data
non_default_values = await prepare_key_update_data(data=update_key_request, existing_key_row=existing_key_row)
non_default_values = await prepare_key_update_data(
data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
)
# Update key in database
if prisma_client is None:
@ -3093,7 +3119,9 @@ async def update_key_fn(
# Enforce upperbound key params on update (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
non_default_values: Final = await prepare_key_update_data(data=data, existing_key_row=existing_key_row)
non_default_values: Final = await prepare_key_update_data(
data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router
)
# Only validate key_alias format if it's actually being changed
new_key_alias: Final = non_default_values.get("key_alias", None)
@ -4137,15 +4165,24 @@ async def generate_key_helper_fn(
object_permission: LiteLLM_ObjectPermissionBase | None = None,
auto_rotate: bool | None = None,
rotation_interval: str | None = None,
router_settings: dict | None = None,
router_settings: dict[str, object] | None = None,
access_group_ids: list[str] | None = None,
budget_limits: list | None = None, # multiple concurrent budget windows
*,
llm_router: Router | None = None,
):
from litellm.proxy.proxy_server import premium_user, prisma_client
if prisma_client is None:
raise Exception("Connect Proxy to database to generate keys - https://docs.litellm.ai/docs/proxy/virtual_keys ")
await validate_router_settings_weights(
router_settings,
team_id=team_id,
prisma_client=prisma_client,
llm_router=llm_router,
)
if token is None:
if key is not None:
token = key
@ -5070,6 +5107,7 @@ async def _insert_deprecated_key(
async def _execute_virtual_key_regeneration(
*,
prisma_client: PrismaClient,
llm_router: Router | None = None,
key_in_db: LiteLLM_VerificationToken,
hashed_api_key: str,
key: str,
@ -5129,7 +5167,9 @@ async def _execute_virtual_key_regeneration(
if data is not None:
# Enforce upperbound key params on regenerate (don't fill defaults)
_enforce_upperbound_key_params(data, fill_defaults=False)
non_default_values = await prepare_key_update_data(data=data, existing_key_row=key_in_db)
non_default_values = await prepare_key_update_data(
data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router
)
# Only validate key_alias format if it's actually being changed
new_key_alias: Final = non_default_values.get("key_alias")
if new_key_alias != key_in_db.key_alias:
@ -5268,6 +5308,7 @@ async def regenerate_key_fn(
try:
from litellm.proxy.proxy_server import (
hash_token,
llm_router,
master_key,
premium_user,
prisma_client,
@ -5456,6 +5497,7 @@ async def regenerate_key_fn(
return await _execute_virtual_key_regeneration(
prisma_client=prisma_client,
llm_router=llm_router,
key_in_db=_key_in_db,
hashed_api_key=hashed_api_key,
key=key,

View file

@ -0,0 +1,129 @@
from abc import abstractmethod
from collections.abc import Mapping
from typing import Annotated, Final, Protocol
from fastapi import HTTPException
from pydantic import BaseModel, BeforeValidator, ValidationError
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.router_weights import RouterWeights
class _StoredModel(Protocol):
@property
@abstractmethod
def model_id(self) -> str:
pass
class _ModelDb(Protocol):
@property
@abstractmethod
def litellm_proxymodeltable(self) -> TableActions[_StoredModel]:
pass
class _PrismaClient(Protocol):
@property
@abstractmethod
def db(self) -> _ModelDb:
pass
class _Router(Protocol):
@abstractmethod
def get_deployment(self, model_id: str) -> object | None:
pass
class _RouterWeightSettings(BaseModel):
weights: RouterWeights | None = None
class _RouterWeightModelInfo(BaseModel):
team_id: str | None = None
db_model: bool | None = None
team_public_model_name: str | None = None
def _router_weight_model_info(value: object) -> _RouterWeightModelInfo:
if isinstance(value, str):
return _RouterWeightModelInfo.model_validate_json(value)
return _RouterWeightModelInfo.model_validate(value or {}, from_attributes=True)
class _RouterWeightDeployment(BaseModel):
model_name: str
model_info: Annotated[_RouterWeightModelInfo, BeforeValidator(_router_weight_model_info)]
def _validate_router_weight_reference(
model_group: str,
deployment_id: str,
team_id: str | None,
stored: _RouterWeightDeployment | None,
configured: object | None,
) -> None:
reference: Final = (
stored
if stored is not None
else (
_RouterWeightDeployment.model_validate(configured, from_attributes=True) if configured is not None else None
)
)
if (
reference is None
or (stored is None and reference.model_info.db_model)
or (reference.model_info.team_id is not None and reference.model_info.team_id != team_id)
):
raise HTTPException(status_code=400, detail=f"Unknown deployment ID in router weights: {deployment_id}")
canonical_group: Final = (
reference.model_info.team_public_model_name if reference.model_info.team_id is not None else None
) or reference.model_name
if model_group != canonical_group:
raise HTTPException(
status_code=400,
detail=f"Deployment {deployment_id} does not belong to model group {model_group}",
)
async def validate_router_settings_weights(
router_settings: BaseModel | Mapping[str, object] | None,
*,
team_id: str | None,
prisma_client: _PrismaClient | None,
llm_router: _Router | None,
) -> None:
try:
weights: Final = (
_RouterWeightSettings.model_validate(router_settings, from_attributes=True).weights
if router_settings is not None
else None
)
except ValidationError:
raise HTTPException(
status_code=400,
detail="Invalid router weights. Replace or clear router_settings.weights.",
) from None
if not weights:
return
deployment_ids: Final = frozenset(deployment_id for group in weights.values() for deployment_id in group)
if not deployment_ids:
return
if prisma_client is None:
raise HTTPException(status_code=503, detail="Database unavailable while validating router weights")
stored_models: Final = await prisma_client.db.litellm_proxymodeltable.find_many(
where={"model_id": {"in": list(deployment_ids)}}
)
stored_by_id: Final = {
row.model_id: _RouterWeightDeployment.model_validate(row, from_attributes=True) for row in stored_models
}
for model_group, group_weights in weights.items():
for deployment_id in group_weights:
_validate_router_weight_reference(
model_group,
deployment_id,
team_id,
stored_by_id.get(deployment_id),
llm_router.get_deployment(model_id=deployment_id) if llm_router is not None else None,
)

View file

@ -112,6 +112,7 @@ from litellm.proxy.management_endpoints.common_utils import (
from litellm.proxy.management_endpoints.organization_endpoints import (
add_member_to_organization,
)
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
@ -1288,6 +1289,7 @@ async def new_team(
create_audit_log_for_update,
general_settings,
litellm_proxy_admin_name,
llm_router,
prisma_client,
user_api_key_cache,
)
@ -1462,6 +1464,13 @@ async def new_team(
user_api_key_dict=user_api_key_dict,
)
await validate_router_settings_weights(
data.router_settings,
team_id=data.team_id,
prisma_client=prisma_client,
llm_router=llm_router,
)
## ADD TO MODEL TABLE
_model_id = None
if data.model_aliases is not None and isinstance(data.model_aliases, dict):
@ -2075,6 +2084,13 @@ async def update_team(
user_api_key_dict=user_api_key_dict,
)
await validate_router_settings_weights(
data.router_settings,
team_id=data.team_id,
prisma_client=prisma_client,
llm_router=llm_router,
)
_existing_team_metadata: Final[object] = getattr(existing_team_row, "metadata", None)
enforce_output_token_estimates_are_admin_only(
data=data,

View file

@ -3592,6 +3592,7 @@ class SSOAuthenticationHandler:
verbose_proxy_logger.info("user_defined_values for creating ui key: %s", user_defined_values)
response: Final = await generate_key_helper_fn(
llm_router=None,
request_type="key",
duration=LITELLM_UI_SESSION_DURATION,
key_max_budget=litellm.max_ui_session_budget,

View file

@ -58,15 +58,6 @@ class UndeliverableStreamRewrite(Exception):
self.guardrail_name: Final = guardrail_name
class UnappliableRequestRewrite(Exception):
def __init__(self, guardrail_name: str) -> None:
super().__init__(
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
"so the request was rejected rather than sent unrewritten"
)
self.guardrail_name: Final = guardrail_name
def _tool_call_shape(tool_call: object) -> tuple[object, object]:
plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call
function: Final = plain.get("function") if isinstance(plain, Mapping) else None

View file

@ -9546,6 +9546,7 @@ class ProxyStartupEvent:
gate the first duration window.
"""
await generate_key_helper_fn(
llm_router=llm_router,
request_type="user",
table_name="user",
user_id=LITELLM_PROXY_BUDGET_NAME,
@ -16290,6 +16291,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
global master_key, general_settings
response: Final = await generate_key_helper_fn(
llm_router=llm_router,
request_type="key",
**{
"user_role": user_obj.user_role,

View file

@ -0,0 +1,65 @@
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final, cast # noqa: TID251 # validating the openai tool union strips vendor keys from raw tools
from pydantic import BaseModel, ValidationError
from litellm._logging import verbose_logger
from litellm.types.llms.openai import ALL_RESPONSES_API_TOOL_PARAMS, ResponseInputParam
ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
class _InputItemType(BaseModel):
type: str = ""
class _AdditionalToolsItem(BaseModel):
tools: tuple[dict[str, object], ...] = ()
@dataclass(frozen=True, slots=True)
class HoistedAdditionalTools:
input: str | ResponseInputParam
tools: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]
hoisted: tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]
def _is_additional_tools_item(item: object) -> bool:
try:
return _InputItemType.model_validate(item).type == ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
except ValidationError:
return False
def _tools_of_item(item: object) -> tuple[ALL_RESPONSES_API_TOOL_PARAMS, ...]:
try:
parsed: Final = _AdditionalToolsItem.model_validate(item)
except ValidationError:
return ()
return tuple(
cast(
"ALL_RESPONSES_API_TOOL_PARAMS", tool
) # cast-ok: nested tools carry the same raw tool JSON as top-level tools
for tool in parsed.tools
)
def hoist_additional_tools(
input: str | ResponseInputParam,
tools: Sequence[ALL_RESPONSES_API_TOOL_PARAMS] | None,
) -> HoistedAdditionalTools:
existing: Final = tuple(tools or ())
if isinstance(input, str):
return HoistedAdditionalTools(input=input, tools=existing, hoisted=())
items: Final = tuple(item for item in input if _is_additional_tools_item(item))
if not items:
return HoistedAdditionalTools(input=input, tools=existing, hoisted=())
hoisted: Final = tuple(tool for item in items for tool in _tools_of_item(item))
verbose_logger.debug(
"Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) into the top-level tools param.",
len(hoisted),
len(items),
)
remaining_input: Final = [item for item in input if not _is_additional_tools_item(item)]
return HoistedAdditionalTools(input=remaining_input, tools=(*existing, *hoisted), hoisted=hoisted)

View file

@ -39,15 +39,38 @@ def openai_shaped_tool_call_item_id(item_type: str, tool_id: str) -> str:
return f"{prefix}_{tool_id}"
class _ToolNameFields(BaseModel):
type: str = ""
name: str = ""
tools: tuple[object, ...] = ()
def _tool_name_fields_of(tool: object) -> _ToolNameFields | None:
try:
return _ToolNameFields.model_validate(tool)
except ValidationError:
return None
def _custom_tool_name_of(tool: object) -> str | None:
parsed: Final = _tool_name_fields_of(tool)
if parsed is None or parsed.type != "custom" or not parsed.name:
return None
return parsed.name
def _nested_tools_of(tool: object) -> tuple[object, ...]:
parsed: Final = _tool_name_fields_of(tool)
if parsed is None or parsed.type != "namespace":
return ()
return parsed.tools
def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
"""Extract names of tools originally defined as ``type: "custom"``."""
if not tools:
return set()
names: Final[set[str]] = set()
for tool in tools:
if isinstance(tool, dict) and tool.get("type") == "custom" and "name" in tool:
names.add(tool["name"])
return names
"""Extract names of ``type: "custom"`` tools, at the top level or one level inside a ``namespace`` tool."""
top_level: Final = tuple(tools or ())
nested: Final = tuple(nested_tool for tool in top_level for nested_tool in _nested_tools_of(tool))
return {name for tool in (*top_level, *nested) if (name := _custom_tool_name_of(tool)) is not None}
def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool:
@ -143,7 +166,7 @@ def validated_allowed_callers(value: object) -> list[str] | None:
raise ValueError("allowed_callers must be a list of strings") from exc
def _grammar_suffix(fmt: object) -> str:
def custom_tool_grammar_suffix(fmt: object) -> str:
try:
parsed: Final = _CustomToolFormat.model_validate(fmt)
except ValidationError:
@ -167,7 +190,9 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp
raw_name: Final = tool.get("name")
name: Final = raw_name if isinstance(raw_name, str) else ""
raw_description: Final = tool.get("description")
description = (raw_description if isinstance(raw_description, str) else "") + _grammar_suffix(tool.get("format"))
description: Final = (raw_description if isinstance(raw_description, str) else "") + custom_tool_grammar_suffix(
tool.get("format")
)
allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers"))
function_chunk: Final = ChatCompletionToolParamFunctionChunk(
name=name,

View file

@ -6,6 +6,7 @@ from collections.abc import Coroutine, Mapping
from typing import Final
import litellm
from litellm.responses.additional_tools import hoist_additional_tools
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
@ -37,11 +38,16 @@ class LiteLLMCompletionTransformationHandler:
| BaseResponsesAPIStreamingIterator
| Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]
):
hoisted: Final = hoist_additional_tools(input, responses_api_request.get("tools"))
bridged_input: Final = hoisted.input
bridged_request: Final[ResponsesAPIOptionalRequestParams] = (
{**responses_api_request, "tools": list(hoisted.tools)} if hoisted.hoisted else responses_api_request
)
litellm_completion_request: Final[dict] = (
LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
model=model,
input=input,
responses_api_request=responses_api_request,
input=bridged_input,
responses_api_request=bridged_request,
custom_llm_provider=custom_llm_provider,
stream=stream,
extra_headers=extra_headers,
@ -52,8 +58,8 @@ class LiteLLMCompletionTransformationHandler:
if _is_async:
return self.async_response_api_handler(
litellm_completion_request=litellm_completion_request,
request_input=input,
responses_api_request=responses_api_request,
request_input=bridged_input,
responses_api_request=bridged_request,
**kwargs,
)
@ -70,8 +76,8 @@ class LiteLLMCompletionTransformationHandler:
responses_api_response: Final[ResponsesAPIResponse] = (
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
chat_completion_response=litellm_completion_response,
request_input=input,
responses_api_request=responses_api_request,
request_input=bridged_input,
responses_api_request=bridged_request,
)
)
@ -81,8 +87,8 @@ class LiteLLMCompletionTransformationHandler:
return LiteLLMCompletionStreamingIterator(
model=model,
litellm_custom_stream_wrapper=litellm_completion_response,
request_input=input,
responses_api_request=responses_api_request,
request_input=bridged_input,
responses_api_request=bridged_request,
custom_llm_provider=custom_llm_provider,
litellm_metadata=kwargs.get("litellm_metadata", {}),
)

View file

@ -8,6 +8,7 @@ from litellm.main import stream_chunk_builder
from litellm.responses.litellm_completion_transformation.custom_tools import (
build_tool_call_item_kwargs,
extract_custom_tool_names,
is_custom_tool_call,
serialize_tool_call_arguments,
)
from litellm.responses.litellm_completion_transformation.transformation import (
@ -166,6 +167,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
return tool_name, namespace
return fn_name, None
def _tool_call_item_kwargs(self, call_id: str, fn_name: str, arguments: str, status: str) -> dict[str, str]:
item_kwargs: Final = build_tool_call_item_kwargs(call_id, fn_name, arguments, status, self._custom_tool_names)
if is_custom_tool_call(fn_name, self._custom_tool_names):
return item_kwargs
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
namespace_kwargs: Final = {"namespace": tool_namespace} if tool_namespace else {}
return {**item_kwargs, "name": tool_name, **namespace_kwargs}
def _is_reasoning_end(self, chunk):
delta: Final = chunk.choices[0].delta
@ -244,17 +253,13 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
else:
fn_name = str(getattr(fn, "name", "") or "")
fn_args_delta = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
output_index = self._get_or_assign_tool_output_index(call_id)
if call_id not in self._tool_args_by_call_id:
self._tool_args_by_call_id[call_id] = ""
self._sequence_number += 1
names = self._custom_tool_names
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
if tool_namespace:
item_kwargs["namespace"] = tool_namespace
event = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=output_index,
@ -315,7 +320,6 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
else:
fn_name = str(getattr(fn, "name", "") or "")
fn_args = serialize_tool_call_arguments(getattr(fn, "arguments", ""))
tool_name, tool_namespace = self._responses_namespace_tool_call_fields(fn_name)
web_search_call = self._web_search_calls.get(call_id)
if web_search_call is not None:
if call_id not in self._queued_web_search_call_ids:
@ -330,11 +334,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if is_new_tool_call:
self._tool_args_by_call_id[call_id] = ""
self._sequence_number += 1
names = self._custom_tool_names
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, "", "in_progress", names)
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, "", "in_progress")
self._tool_item_id_by_call_id[call_id] = item_kwargs["id"]
if tool_namespace:
item_kwargs["namespace"] = tool_namespace
event = OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=output_index,
@ -376,11 +377,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self._pending_tool_events.append(done_event)
self._sequence_number += 1
names = self._custom_tool_names
item_kwargs = build_tool_call_item_kwargs(call_id, tool_name, final_args, "completed", names)
item_kwargs = self._tool_call_item_kwargs(call_id, fn_name, final_args, "completed")
item_kwargs["id"] = self._tool_item_id_by_call_id.setdefault(call_id, item_kwargs["id"])
if tool_namespace:
item_kwargs["namespace"] = tool_namespace
item_done_event = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=output_index,

View file

@ -110,6 +110,7 @@ NamespaceTool: TypeAlias = Mapping[str, object]
ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None
ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool
NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n"
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"})
@dataclass(frozen=True, slots=True)
@ -1891,9 +1892,21 @@ class LiteLLMCompletionResponsesConfig:
namespace_tool: NamespaceTool,
nested: bool,
) -> ChatCompletionToolParam | None:
if nested and namespace_tool.get("type") != "function":
tool_type: Final = namespace_tool.get("type")
if nested and tool_type not in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS:
return None
raw_description: Final = str(namespace_tool.get("description") or "")
description: Final = (
f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}"
if nested and namespace_description and raw_description
else namespace_description
if nested and namespace_description
else raw_description
)
if nested and tool_type == "custom":
return convert_custom_tool_to_function_tool({**namespace_tool, "description": description})
raw_parameters: Final = namespace_tool.get("parameters")
parameters: Final = (
MappingProxyType(raw_parameters) if isinstance(raw_parameters, Mapping) else MappingProxyType({})
@ -1902,14 +1915,6 @@ class LiteLLMCompletionResponsesConfig:
parameters if parameters and "type" in parameters else MappingProxyType({**parameters, "type": "object"})
)
tool_name: Final = str(namespace_tool.get("name") or "")
raw_description: Final = str(namespace_tool.get("description") or "")
description: Final = (
f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}{raw_description}"
if nested and namespace_description and raw_description
else namespace_description
if nested and namespace_description
else raw_description
)
chat_tool_name: Final = f"{namespace}__{tool_name}" if nested else tool_name
function: Final = ChatCompletionToolParamFunctionChunk(
name=chat_tool_name,

View file

@ -4849,6 +4849,7 @@ class Router:
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
data: Final = deployment["litellm_params"].copy()
@ -5163,13 +5164,11 @@ class Router:
return healthy_deployments[0]
# Use simple_shuffle for weighted selection
return cast(
GuardrailTypedDict,
simple_shuffle(
llm_router_instance=self,
healthy_deployments=healthy_deployments,
model=guardrail_name,
),
return simple_shuffle(
resolve_model_alias=self._get_model_from_alias,
healthy_deployments=healthy_deployments,
model=guardrail_name,
request_kwargs=None,
)
async def _ageneric_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
@ -8378,7 +8377,8 @@ class Router:
def log_retry(self, kwargs: dict, e: Exception) -> dict:
"""
When a retry or fallback happens, record which model group, deployment and attempt just failed and why
When a retry or fallback happens, record which model group, deployment and attempt just failed and why,
and count it toward the request-wide num_retries_per_request cap
"""
from litellm.types.router import RetryAttemptRecord
@ -8402,7 +8402,10 @@ class Router:
else ()
)
breadcrumbs: Final = (*kept_breadcrumbs, attempt_record)
earlier: Final = request_metadata.get("request_retry_count")
request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1
kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict
kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap
return kwargs
def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int:
@ -13045,9 +13048,10 @@ class Router:
start_time: Final = time.time()
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
resolve_model_alias=self._get_model_from_alias,
healthy_deployments=healthy_deployments,
model=model,
request_kwargs=request_kwargs,
)
deployment: Final = await self._select_deployment_async(
strategy=strategy,
@ -13190,9 +13194,10 @@ class Router:
start_time: Final = time.perf_counter()
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
resolve_model_alias=self._get_model_from_alias,
healthy_deployments=pass_through_deployments,
model=model,
request_kwargs=request_kwargs,
)
deployment: Final = await self._select_deployment_async(
strategy=strategy,
@ -13888,9 +13893,10 @@ class Router:
# if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm
############## Check 'weight' param set for weighted pick #################
return simple_shuffle(
llm_router_instance=self,
resolve_model_alias=self._get_model_from_alias,
healthy_deployments=healthy_deployments,
model=model,
request_kwargs=request_kwargs,
)
deployment: Final = self._select_deployment_sync(
strategy=strategy,
@ -13958,6 +13964,7 @@ class Router:
messages=messages,
input=input,
specific_deployment=specific_deployment,
request_kwargs=request_kwargs,
)
strategy, strategy_selector = self._get_routing_context(model, request_kwargs)
@ -14040,9 +14047,10 @@ class Router:
# 6. Apply load balancing strategy
if strategy == "simple-shuffle":
return simple_shuffle(
llm_router_instance=self,
resolve_model_alias=self._get_model_from_alias,
healthy_deployments=pass_through_deployments,
model=model,
request_kwargs=request_kwargs,
)
deployment: Final = self._select_deployment_sync(
strategy=strategy,

View file

@ -1,71 +1,67 @@
"""
Returns a random deployment from the list of healthy deployments.
"""Choose among eligible deployments using request weights, then global metrics."""
If weights are provided, it will return a deployment based on the weights.
"""
from __future__ import annotations
import logging
import random
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Callable, Mapping, Sequence
from itertools import chain
from typing import Final, TypeVar
from litellm._logging import verbose_router_logger
from litellm.types.router_weights import validate_router_weights
if TYPE_CHECKING:
from litellm.router import Router as _Router
_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object])
_ROUTER_LOGGER: Final = logging.getLogger("LiteLLM Router")
LitellmRouter = _Router
else:
LitellmRouter = Any
def _metric_weight(deployment: Mapping[str, object], metric: str) -> float:
params: Final = deployment.get("litellm_params")
value: Final = params.get(metric) if isinstance(params, Mapping) else None
if value is None:
return 0.0
if isinstance(value, (int, float)):
return float(value)
raise TypeError(f"Deployment {metric} must be numeric")
def _scoped_weights(
deployments: Sequence[Mapping[str, object]],
model: str,
request_kwargs: Mapping[str, object] | None,
) -> tuple[float, ...]:
settings: Final = validate_router_weights((request_kwargs or {}).get("_router_weights"))
model_weights: Final = settings.get(model) if settings is not None else None
if not model_weights:
return ()
return tuple(
model_weights.get(str(info.get("id")), 0.0) if isinstance(info, Mapping) else 0.0
for deployment in deployments
for info in (deployment.get("model_info"),)
)
def simple_shuffle(
llm_router_instance: LitellmRouter,
healthy_deployments: list[Any] | dict[Any, Any],
resolve_model_alias: Callable[[str], str | None],
healthy_deployments: Sequence[_DeploymentT],
model: str,
) -> dict:
"""
Returns a random deployment from the list of healthy deployments.
If weights are provided, it will return a deployment based on the weights.
If users pass `rpm` or `tpm`, we do a random weighted pick - based on `rpm`/`tpm`.
Args:
llm_router_instance: LitellmRouter instance
healthy_deployments: List of healthy deployments
model: Model name
Returns:
Dict: A single healthy deployment
"""
############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick #################
for weight_by in ["weight", "rpm", "tpm"]:
if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments):
weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments]
verbose_router_logger.debug("\nweight %s", weights)
total_weight = sum(weights)
if total_weight <= 0:
# All remaining candidates have weight 0 for this metric (e.g.
# after a weighted-failover exclusion left only zero-weight
# backups). Skip to the next metric (rpm/tpm) which may still
# provide a meaningful weighted pick; if none do, we fall
# through to the uniform random pick at the end.
continue
weights = [weight / total_weight for weight in weights]
verbose_router_logger.debug("\n weights %s by %s", weights, weight_by)
# Perform weighted random pick
selected_index = random.choices(range(len(weights)), weights=weights)[0]
verbose_router_logger.debug("\n selected index, %s", selected_index)
deployment = healthy_deployments[selected_index]
verbose_router_logger.info(
"get_available_deployment for model: %s, Selected deployment: %s for model: %s",
model,
llm_router_instance.print_deployment(deployment) or deployment[0],
model,
)
return deployment or deployment[0]
############## No RPM/TPM passed, we do a random pick #################
item: Final = random.choice(healthy_deployments)
return item or item[0]
request_kwargs: Mapping[str, object] | None,
) -> _DeploymentT:
resolved_model: Final = resolve_model_alias(model) or model
weight_sets: Final = chain(
(_scoped_weights(healthy_deployments, resolved_model, request_kwargs),),
(
tuple(_metric_weight(deployment, metric) for deployment in healthy_deployments)
for metric in ("weight", "rpm", "tpm")
),
)
for weights in weight_sets:
largest = max(weights, default=0.0)
if largest <= 0:
continue
normalized = tuple(weight / largest for weight in weights)
if sum(normalized) <= 0:
continue
selected = random.choices(healthy_deployments, weights=normalized)[0]
_ROUTER_LOGGER.info("Selected deployment for model %s: %s", model, selected.get("model_info"))
return selected
return random.choice(healthy_deployments)

View file

@ -48,6 +48,7 @@ from litellm.exceptions import (
ServiceUnavailableError,
)
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.prompt_templates.common_utils import (
encrypted_content_of_block,
strip_encrypted_reasoning_from_messages,
@ -215,11 +216,11 @@ class EncryptedContentAffinityCheck(CustomLogger):
@staticmethod
def _encryption_boundary_key(
litellm_params: object,
) -> tuple | None:
) -> tuple[object, object] | None:
"""
``(api_base, api_key)`` pair identifying an Azure resource. Two
deployments sharing both are interchangeable for ``encrypted_content``
follow-ups; Azure rejects content produced by any other resource.
``(api_base, api_key)`` identifies an upstream encryption boundary.
The values are resolved from the deployment and its named credential
without modifying the deployment.
Accepts any object exposing dict-style ``.get(key, default)``: plain
dicts (the common case in ``healthy_deployments``) as well as
@ -234,9 +235,25 @@ class EncryptedContentAffinityCheck(CustomLogger):
return None
api_base: Final = getter("api_base")
api_key: Final = getter("api_key")
if not api_base or not api_key:
credential_name: Final = getter("litellm_credential_name")
credential_values: Final[Mapping[str, object] | None] = (
CredentialAccessor.get_credential_values(credential_name)
if isinstance(credential_name, str) and credential_name
else None
)
effective_api_base: Final = (
credential_values.get("api_base")
if credential_values is not None and "api_base" in credential_values
else api_base
)
effective_api_key: Final = (
credential_values.get("api_key")
if credential_values is not None and "api_key" in credential_values
else api_key
)
if not effective_api_base or not effective_api_key:
return None
return (api_base, api_key)
return (effective_api_base, effective_api_key)
def _find_deployments_on_same_encryption_boundary(
self,

View file

@ -1,4 +1,5 @@
from typing import Any, Final, Literal
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal, cast # noqa: TID251 # JSON chat rows have no typed constructor across roles
from pydantic import BaseModel, ConfigDict, Field
from typing_extensions import TypedDict
@ -158,12 +159,21 @@ def coerce_stream_holdback_value(value: Any) -> int:
return 0
def structured_messages_from_response(value: object) -> Sequence[AllMessageValues] | None:
if not isinstance(value, list):
return None
if not all(isinstance(message, Mapping) and isinstance(message.get("role"), str) for message in value):
return None
return cast("Sequence[AllMessageValues]", value) # cast-ok: JSON rows checked for a role, the same trust texts get
class GenericGuardrailAPIResponse:
"""Response model for the Generic Guardrail API"""
texts: list[str] | None
images: list[str] | None
tools: list[GuardrailToolParam] | None
structured_messages: Sequence[AllMessageValues] | None
action: str
blocked_reason: str | None
stream_holdback_chars: list[int] | None
@ -176,12 +186,14 @@ class GenericGuardrailAPIResponse:
images: list[str] | None = None,
tools: list[GuardrailToolParam] | None = None,
stream_holdback_chars: list[int] | None = None,
structured_messages: Sequence[AllMessageValues] | None = None,
) -> None:
self.action = action
self.blocked_reason = blocked_reason
self.texts = texts
self.images = images
self.tools = tools
self.structured_messages = structured_messages
# Number of trailing chars, indexed the same as ``texts``, that the
# framework must withhold from streaming emission until the next
# processing round (word-boundary safety for text transformations).
@ -200,4 +212,5 @@ class GenericGuardrailAPIResponse:
images=data.get("images"),
tools=data.get("tools"),
stream_holdback_chars=stream_holdback_chars,
structured_messages=structured_messages_from_response(data.get("structured_messages")),
)

View file

@ -15,6 +15,7 @@ from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_c
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.types.router_weights import RouterWeights
if TYPE_CHECKING:
from litellm.router import Router
@ -146,6 +147,7 @@ class UpdateRouterConfig(BaseModel):
context_window_fallbacks: list[dict] | None = None
model_group_alias: dict[str, str | dict] | None = {}
enable_tag_filtering: bool | None = None
weights: RouterWeights | None = None
tag_routing_prefix: str | None = None
optional_pre_call_checks: OptionalPreCallChecks | None = None

View file

@ -0,0 +1,30 @@
from collections.abc import Mapping
from typing import Annotated, Final
from pydantic import AfterValidator, Field, TypeAdapter
def _validate_positive_router_weights(weights: Mapping[str, Mapping[str, float]]) -> Mapping[str, Mapping[str, float]]:
if any(group and not any(weight > 0 for weight in group.values()) for group in weights.values()):
raise ValueError("Each nonempty weights group must contain at least one positive weight")
return weights
RouterWeightIdentifier = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")]
RouterWeight = Annotated[float, Field(strict=True, ge=0, allow_inf_nan=False)]
RouterWeights = Annotated[
dict[RouterWeightIdentifier, dict[RouterWeightIdentifier, RouterWeight]],
AfterValidator(_validate_positive_router_weights),
]
_ROUTER_WEIGHTS_ADAPTER: Final[TypeAdapter[RouterWeights | None]] = TypeAdapter(RouterWeights | None)
_ROUTER_SETTINGS_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
def validate_router_weights(value: object) -> RouterWeights | None:
return _ROUTER_WEIGHTS_ADAPTER.validate_python(value)
def validate_router_settings_dict(value: object) -> dict[str, object]:
settings: Final = _ROUTER_SETTINGS_DICT_ADAPTER.validate_python(value)
validate_router_weights(settings.get("weights"))
return settings

View file

@ -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
@ -3776,6 +3782,7 @@ all_litellm_params = (
"id",
"fallbacks",
"routing_strategy",
"_router_weights",
"azure",
"headers",
"model_list",

View file

@ -1208,30 +1208,13 @@ def _dispatch_success_logging(
is_litellm_internal_call: bool,
) -> None:
if not is_litellm_internal_call:
if getattr(logging_obj, "_defer_async_logging", False):
def _enqueue_deferred_logging() -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
logging_obj._enqueue_deferred_logging = _enqueue_deferred_logging
else:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
_schedule_async_success_logging(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=result,
@ -1240,6 +1223,43 @@ def _dispatch_success_logging(
)
def _schedule_async_success_logging(
logging_obj: LiteLLMLoggingObject,
result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
is_completion_with_fallbacks: bool,
) -> None:
"""Fire the async success log for ``result`` now, or park it on the logging object while
the proxy defers logging past its post-call guardrails.
Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses
bridge) each exit through here with the same logging object and their own shape of the same
response. The immediate path already logs one request once, since the first task marks
``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same
first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log
reads usage from, and a later wrapper never swaps in its client-shaped translation.
"""
def _enqueue_async_logging() -> None:
asyncio.create_task(
_client_async_logging_helper(
logging_obj=logging_obj,
result=result,
start_time=start_time,
end_time=end_time,
is_completion_with_fallbacks=is_completion_with_fallbacks,
)
)
if not getattr(logging_obj, "_defer_async_logging", False):
_enqueue_async_logging()
return
if getattr(logging_obj, "_enqueue_deferred_logging", None) is not None:
return
logging_obj._enqueue_deferred_logging = _enqueue_async_logging
async def _client_async_logging_helper(
logging_obj: LiteLLMLoggingObject,
result,
@ -5923,10 +5943,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),

View file

@ -25601,10 +25601,14 @@
"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_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_video_per_second": 0.00079,
"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,
@ -25615,13 +25619,14 @@
"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_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_per_second": 0.00079,
"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 +25638,14 @@
"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_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_video_per_second": 0.00079,
"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,
@ -25648,13 +25657,14 @@
"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_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_per_second": 0.00079,
"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,
@ -25693,10 +25703,14 @@
},
"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_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_video_per_second": 0.00079,
"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 +25723,14 @@
"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_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_video_per_second": 0.00079,
"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,

View file

@ -249,6 +249,10 @@
"type": "number",
"minimum": 0
},
"input_cost_per_audio_token_batches": {
"type": "number",
"minimum": 0
},
"input_cost_per_audio_token_priority": {
"type": "number",
"minimum": 0,
@ -276,6 +280,10 @@
"type": "number",
"minimum": 0
},
"input_cost_per_image_token_batches": {
"type": "number",
"minimum": 0
},
"input_cost_per_pixel": {
"type": "number",
"minimum": 0
@ -375,6 +383,14 @@
"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_cost_per_video_token_batches": {
"type": "number",
"minimum": 0
},
"input_dbu_cost_per_token": {
"type": "number",
"minimum": 0

View file

@ -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",

View file

@ -0,0 +1,329 @@
from __future__ import annotations
import json
import os
import subprocess
import sys
import threading
from pathlib import Path
from typing import Final
import pytest
from fixture_bundle import BundleRecorder, LoadedBundle, load_bundle, prepare_bundle
from fixture_mode import current_test_key
from fixture_profile import MatchProfile
from provider_edge import REPLAY_MISS_STATUS, RecordEdge, ReplayEdge, ReplaySource
from test_provider_edge import (
CHAT_PATH,
SSE_CHUNKS,
STREAM_BODY,
UPLOAD_PATH,
call_edge,
chunked_provider,
fake_provider,
json_object,
provider_url,
raw_stream_post,
running_edge,
this_tests_files,
)
class TestStrictIdentity:
@pytest.mark.parametrize("path", [CHAT_PATH, "/anthropic/v1/messages"])
def test_roundtrip_rejects_semantic_changes(self, tmp_path: Path, path: str) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
original: Final = (
b'{"model":"synthetic","messages":[{"role":"user",'
b'"content":"2031-04-05 00000000-0000-0000-0000-000000000001"}],"options":[1,2]}'
)
headers: Final = {
"content-type": "application/json",
"accept": "application/json",
"anthropic-version": "2023-06-01",
"anthropic-beta": "feature-a",
"openai-beta": "feature-b",
"authorization": "Bearer synthetic-secret-one",
}
query: Final = "?part=one&part=two&blank="
with fake_provider() as provider:
mounts: Final = {"openai": provider_url(provider), "anthropic": provider_url(provider)}
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
captured: Final = call_edge(edge, "POST", path + query, body=original, headers=headers)
assert captured.status_code == 200
assert json_object(captured.body)["echo"] == original.decode()
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
assert isinstance(loaded, LoadedBundle)
assert loaded.manifest.match_profile == "stateless_v1"
source: Final = ReplaySource(loaded)
with running_edge(ReplayEdge(source), mounts) as edge:
cases: Final = (
(original.replace(b"2031-04-05", b"2032-06-07"), headers, query, "body"),
(original.replace(b"000000000001", b"000000000002"), headers, query, "body"),
(original.replace(b"[1,2]", b"[2,1]"), headers, query, "body"),
(original.replace(b"synthetic", b"other"), headers, query, "body"),
(original, headers, "?part=three&part=two&blank=", "query"),
(original, headers, "?part=two&part=one&blank=", "query"),
*(
(original, {k: v for k, v in headers.items() if k != name}, query, "headers")
for name in ("accept", "anthropic-version", "anthropic-beta", "openai-beta")
),
*(
(original, {**headers, name: value}, query, "headers")
for name in ("accept", "anthropic-version", "anthropic-beta", "openai-beta")
for value in ("different", "")
),
(original, {**headers, "authorization": "Basic synthetic-secret-two"}, query, "auth"),
(original, {k: v for k, v in headers.items() if k != "authorization"}, query, "auth"),
)
for rejected, reason in (
(call_edge(edge, "POST", path + changed_query, body=body, headers=changed_headers), reason)
for body, changed_headers, changed_query, reason in cases
):
assert rejected.status_code == REPLAY_MISS_STATUS
assert reason in rejected.body.decode()
assert b"synthetic-secret" not in rejected.body
reordered: Final = json.dumps(dict(reversed(list(json_object(original).items())))).encode()
accepted: Final = call_edge(
edge, "POST", path + query, body=reordered, headers={k.upper(): v for k, v in headers.items()}
)
assert accepted.status_code == 200
assert accepted.body == captured.body
assert source.leftover_error(current_test_key()) is None
assert len(provider.hits) == 1
assert "synthetic-secret" not in "".join(file.read_text() for file in recorder.root.rglob("*.json"))
@pytest.mark.parametrize(
"body",
[
b'{"value":null}',
b'{"value":""}',
b'{"value":false}',
b'{"value":0}',
b'{"value":[]}',
b'{"value":{}}',
b'{"value":0.123456789012345678901}',
b'{"value":0.123456789012345678902}',
b'{"value":1e400}',
b'{"value":1}',
b'{"value":1e0}',
b'{"value":-0}',
b'{"value":1e9999999999999999999}',
],
)
def test_json_values_remain_distinct(self, tmp_path: Path, body: bytes) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
with fake_provider() as provider:
mounts: Final = {"openai": provider_url(provider)}
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
assert (
call_edge(
edge, "POST", CHAT_PATH, body=body, headers={"content-type": "application/json"}
).status_code
== 200
)
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
assert isinstance(loaded, LoadedBundle)
with running_edge(ReplayEdge(ReplaySource(loaded)), mounts) as edge:
values: Final = (
b"{}",
b'{"value":null}',
b'{"value":""}',
b'{"value":false}',
b'{"value":0}',
b'{"value":[]}',
b'{"value":{}}',
b'{"value":0.123456789012345678901}',
b'{"value":0.123456789012345678902}',
b'{"value":1e400}',
b'{"value":1}',
b'{"value":1e0}',
b'{"value":-0}',
b'{"value":1e9999999999999999999}',
)
for rejected in (
call_edge(edge, "POST", CHAT_PATH, body=value, headers={"content-type": "application/json"})
for value in values
if value != body
):
assert rejected.status_code == REPLAY_MISS_STATUS
assert b"body" in rejected.body
assert (
call_edge(
edge, "POST", CHAT_PATH, body=body, headers={"content-type": "application/json"}
).status_code
== 200
)
assert len(provider.hits) == 1
@pytest.mark.parametrize(
"path,body,headers",
[
(UPLOAD_PATH, b"{}", {"content-type": "application/json"}),
(CHAT_PATH + "?part=%FF", b"{}", {"content-type": "application/json"}),
(CHAT_PATH + "?part=%FE", b"{}", {"content-type": "application/json"}),
(CHAT_PATH, b"opaque", {"content-type": "application/octet-stream"}),
(CHAT_PATH, b"--boundary", {"content-type": "multipart/form-data; boundary=boundary"}),
(CHAT_PATH, b'{"x":1,"x":2}', {"content-type": "application/json"}),
(CHAT_PATH, b"{}", {"content-type": "application/json", "x-custom-behavior": "synthetic-private-value"}),
],
)
def test_ineligible_capture_never_calls_provider(
self, tmp_path: Path, path: str, body: bytes, headers: dict[str, str]
) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
with fake_provider() as provider:
with running_edge(RecordEdge(recorder, threading.Lock()), {"openai": provider_url(provider)}) as edge:
result: Final = call_edge(edge, "POST", path, body=body, headers=headers)
assert result.status_code == REPLAY_MISS_STATUS
assert b"eligibility error" in result.body
assert b"synthetic-private-value" not in result.body
assert provider.hits == []
assert this_tests_files(recorder.root) == []
def test_destination_is_part_of_actual_http_identity(self, tmp_path: Path) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
with fake_provider() as provider:
with running_edge(RecordEdge(recorder, threading.Lock()), {"openai": provider_url(provider)}) as edge:
assert (
call_edge(
edge, "POST", CHAT_PATH, body=b"{}", headers={"content-type": "application/json"}
).status_code
== 200
)
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
assert isinstance(loaded, LoadedBundle)
with running_edge(ReplayEdge(ReplaySource(loaded)), {"openai": provider_url(provider) + "/other"}) as edge:
result: Final = call_edge(
edge, "POST", CHAT_PATH, body=b"{}", headers={"content-type": "application/json"}
)
assert result.status_code == REPLAY_MISS_STATUS
assert b"upstream" in result.body
assert len(provider.hits) == 1
def test_credentials_are_not_identity_and_fresh_process_replays(self, tmp_path: Path) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
headers: Final = {
"content-type": "application/json",
"authorization": "bEaReR synthetic-token",
"x-api-key": "synthetic-api-key",
"cookie": "synthetic-cookie",
}
path: Final = CHAT_PATH + "?api_key=synthetic-query-secret&part=one&part=two"
body: Final = b'{"model":"synthetic","messages":[]}'
with fake_provider(echo_request=False) as provider:
mounts: Final = {"openai": provider_url(provider)}
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
captured: Final = call_edge(edge, "POST", path, body=body, headers=headers)
assert captured.status_code == 200
seen_headers, seen_body = provider.requests[0]
assert {k.lower(): v for k, v in seen_headers.items()}.items() >= headers.items()
assert seen_body == body
assert provider.hits == ["POST " + path.removeprefix("/openai")]
artifacts: Final = "".join(file.read_text() for file in recorder.root.rglob("*.json"))
for secret in ("synthetic-token", "synthetic-api-key", "synthetic-cookie", "synthetic-query-secret"):
assert secret not in artifacts
child: Final = subprocess.run(
[
sys.executable,
"-c",
"""
import json, sys
from pathlib import Path
from fixture_bundle import LoadedBundle, load_bundle
from provider_edge import ProviderRequestObservation, observed_provider_edge, replay_leftover_error
from test_provider_edge import call_edge
from fixture_profile import MatchProfile
from fixture_mode import current_test_key
loaded = load_bundle(Path(sys.argv[1]), profile="stateless_v1")
assert isinstance(loaded, LoadedBundle)
with observed_provider_edge(ProviderRequestObservation("synthetic"), mode_raw="replay", bundle_dir=Path(sys.argv[1]), bind_host="127.0.0.1", advertise_host="127.0.0.1", mounts={"openai": sys.argv[2]}) as edge:
response = call_edge(edge, "POST", sys.argv[3], body=sys.argv[4].encode(), headers=json.loads(sys.argv[5]))
assert response.status_code == 200
print(response.body.decode())
assert replay_leftover_error(mode_raw="replay", bundle_dir=Path(sys.argv[1]), test_key=current_test_key()) is None
""",
str(recorder.root),
provider_url(provider),
path.replace("synthetic-query-secret", "new-query-credential"),
body.decode(),
json.dumps({**headers, "authorization": "Bearer another-credential", "x-api-key": "another-key"}),
],
env={
**os.environ,
"PYTHONPATH": str(Path(__file__).resolve().parents[1] / "e2e"),
"E2E_REPLAY_MATCH_PROFILE": "stateless_v1",
},
capture_output=True,
text=True,
timeout=30,
)
assert child.returncode == 0, child.stderr
assert child.stdout.strip().encode() == captured.body
assert len(provider.hits) == 1
@pytest.mark.parametrize("profile,other", [("legacy", "stateless_v1"), ("stateless_v1", "legacy")])
def test_profiles_cannot_load_each_others_bundles(
self, tmp_path: Path, profile: MatchProfile, other: MatchProfile
) -> None:
from fixture_bundle import UnreadableBundle
recorder: Final = prepare_bundle(tmp_path / profile, profile=profile)
assert isinstance(recorder, BundleRecorder)
mismatch: Final = load_bundle(recorder.root, profile=other)
assert isinstance(mismatch, UnreadableBundle)
assert "profile mismatch" in mismatch.reason
assert "re-record" in mismatch.reason
@pytest.mark.parametrize("abort_after", [None, 2])
def test_strict_stream_preserves_chunks_and_truncation(self, tmp_path: Path, abort_after: int | None) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
with chunked_provider(abort_after=abort_after) as provider:
mounts: Final = {"anthropic": provider_url(provider)}
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
_, captured, captured_ending = raw_stream_post(edge.port, "/anthropic/v1/messages", STREAM_BODY)
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
assert isinstance(loaded, LoadedBundle)
source: Final = ReplaySource(loaded)
with running_edge(ReplayEdge(source), mounts) as edge:
_, replayed, ending = raw_stream_post(edge.port, "/anthropic/v1/messages", STREAM_BODY)
assert captured == replayed == list(SSE_CHUNKS[:abort_after])
assert ending == captured_ending
assert (ending == "terminated") == (abort_after is None)
assert source.leftover_error(current_test_key()) is None
assert len(provider.hits) == 1
def test_auth_scheme_survives_missing_credentials(self, tmp_path: Path) -> None:
recorder: Final = prepare_bundle(tmp_path / "strict", profile="stateless_v1")
assert isinstance(recorder, BundleRecorder)
headers: Final = {"content-type": "application/json", "authorization": "Bearer"}
with fake_provider() as provider:
mounts: Final = {"openai": provider_url(provider)}
with running_edge(RecordEdge(recorder, threading.Lock()), mounts) as edge:
assert call_edge(edge, "POST", CHAT_PATH, body=b"{}", headers=headers).status_code == 200
loaded: Final = load_bundle(recorder.root, profile="stateless_v1")
assert isinstance(loaded, LoadedBundle)
with running_edge(ReplayEdge(ReplaySource(loaded)), mounts) as edge:
for result in (
call_edge(edge, "POST", CHAT_PATH, body=b"{}", headers={**headers, "authorization": scheme})
for scheme in ("Basic", "Digest")
):
assert result.status_code == REPLAY_MISS_STATUS
assert b"auth" in result.body
assert (
call_edge(
edge,
"POST",
CHAT_PATH,
body=b"{}",
headers={**headers, "authorization": "bEaReR synthetic-token"},
).status_code
== 200
)
assert len(provider.hits) == 1

View file

@ -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" <server-command>`. 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`:
@ -232,3 +236,15 @@ Before you push
4. Capture screenshots of the test run and attach them to the PR as proof
5. If a test fails because it surfaced a real issue in the product, flag that explicitly in the PR rather than reworking the test until it passes
### Strict stateless replay matching
Set `E2E_REPLAY_MATCH_PROFILE=stateless_v1` for both recording and replay to bind OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` requests to their upstream destination, ordered query pairs, semantic headers and literal JSON content. The default remains `legacy`. Strict bundles use format 5 and cannot load as legacy bundles; select the matching profile or re-record with `E2E_FIXTURE_MODE=record`. Missing profile metadata never enrolls a legacy bundle in strict matching
Strict matching preserves dates, UUIDs, hashes, model names, tool arguments, array order and omitted/null/empty/false/zero values. JSON object key order and header name casing may change. The strict body uses tagged JSON values so number precision and JSON types survive persistence, including exact numeric spelling and numbers larger than a floating-point value. Invalid UTF-8 query values fail eligibility. Duplicate JSON keys, unsupported endpoints, non-JSON bodies and unknown semantic headers fail eligibility before contacting a provider
The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthropic-beta` and `openai-beta`, including missing versus present values. Authorization records presence and the case-insensitive scheme; `x-api-key` records presence only. Credential values and cookies are excluded. Credential query values are redacted while their position and field name remain in the identity. Never use real customer inputs in fixture qualification
Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the existing legacy harness files with `--noconftest -o pythonpath=tests/e2e`; they need only synthetic HTTP providers and temporary fixture storage

View file

@ -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="group_scoped",
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),
)

View file

@ -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"}

View file

@ -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):
@ -168,6 +170,7 @@ class StreamingResponse(BaseModel):
# the consumed body is elided, so this is the only place they surface.
stream_error: str | None = None
stream_done: bool = False
stream_done_positions: tuple[int, ...] = ()
@property
def ok(self) -> bool:
@ -292,6 +295,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 +338,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 +433,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 +442,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:
@ -628,6 +648,7 @@ def streaming_outcome(
stream_events=[payload for payload, _ in events],
stream_event_arrivals=[arrived for _, arrived in events],
stream_done=any(payload == _SSE_DONE for payload, _ in payloads),
stream_done_positions=tuple(index for index, (payload, _) in enumerate(payloads) if payload == _SSE_DONE),
stream_error=next(
(line.decode(errors="replace")[:300] for line, _ in stamped if _is_stream_error_line(line)),
None,

View file

@ -30,9 +30,11 @@ from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Annotated, Final, Literal
from fixture_profile import MatchProfile, StrictIdentity
from pydantic import BaseModel, Field, JsonValue
BUNDLE_FORMAT_VERSION: Final = 4
STRICT_BUNDLE_FORMAT_VERSION: Final = 5
MAX_BUNDLE_AGE: Final = timedelta(days=7)
MANIFEST_FILENAME: Final = "manifest.json"
@ -41,6 +43,7 @@ class Manifest(BaseModel):
format_version: int
recorded_at: datetime
harness_version: str
match_profile: MatchProfile = "legacy"
class RecordedRequest(BaseModel):
@ -69,6 +72,7 @@ class RecordedRequest(BaseModel):
file_name: str | None = None
file_sha256: str | None = None
file_bytes: int | None = None
strict_identity: StrictIdentity | None = None
class RecordedHttpResponse(BaseModel):
@ -100,9 +104,7 @@ class RecordedStreamedResponse(BaseModel):
truncated: str | None = None
type RecordedResponse = Annotated[
RecordedHttpResponse | RecordedStreamedResponse, Field(discriminator="kind")
]
type RecordedResponse = Annotated[RecordedHttpResponse | RecordedStreamedResponse, Field(discriminator="kind")]
class Interaction(BaseModel):
@ -152,6 +154,7 @@ class BundleRecorder:
manifest, so record mode never reads (or merges into) an existing bundle."""
root: Path
profile: MatchProfile = "legacy"
_ordinals: dict[str, int] = field(default_factory=dict)
def record(self, *, test_key: str, request: RecordedRequest, response: RecordedResponse) -> None:
@ -162,7 +165,12 @@ class BundleRecorder:
directory.mkdir(parents=True, exist_ok=True)
interaction = Interaction(request=request, response=response)
target = directory / interaction_filename(ordinal, request)
target.write_text(interaction.model_dump_json(indent=2), encoding="utf-8")
target.write_text(
interaction.model_dump_json(
indent=2, exclude={"request": {"strict_identity"}} if self.profile == "legacy" else None
),
encoding="utf-8",
)
@dataclass(frozen=True, slots=True)
@ -171,7 +179,7 @@ class UnsafeBundleDir:
reason: str
def prepare_bundle(root: Path) -> BundleRecorder | UnsafeBundleDir:
def prepare_bundle(root: Path, *, profile: MatchProfile = "legacy") -> BundleRecorder | UnsafeBundleDir:
"""Start a fresh bundle at ``root`` for record mode: wipe whatever bundle is
there and write a new manifest. Refuses to wipe a directory that is neither
empty nor a bundle (no manifest.json), so a mistyped E2E_FIXTURE_DIR can
@ -188,12 +196,15 @@ def prepare_bundle(root: Path) -> BundleRecorder | UnsafeBundleDir:
shutil.rmtree(root)
root.mkdir(parents=True)
manifest = Manifest(
format_version=BUNDLE_FORMAT_VERSION,
format_version=BUNDLE_FORMAT_VERSION if profile == "legacy" else STRICT_BUNDLE_FORMAT_VERSION,
match_profile=profile,
recorded_at=datetime.now(timezone.utc),
harness_version=harness_version(),
)
(root / MANIFEST_FILENAME).write_text(manifest.model_dump_json(indent=2), encoding="utf-8")
return BundleRecorder(root=root)
(root / MANIFEST_FILENAME).write_text(
manifest.model_dump_json(indent=2, exclude={"match_profile"} if profile == "legacy" else None), encoding="utf-8"
)
return BundleRecorder(root=root, profile=profile)
@dataclass(frozen=True, slots=True)
@ -226,25 +237,30 @@ def _read_manifest(root: Path) -> Manifest | UnreadableBundle:
return UnreadableBundle(reason=f"{MANIFEST_FILENAME} is invalid: {exc}")
def _supported_manifest(root: Path) -> Manifest | UnreadableBundle:
def _supported_manifest(root: Path, profile: MatchProfile = "legacy") -> Manifest | UnreadableBundle:
"""The manifest, refused when it was written under a different format version.
A bundle is atomic (record wipes and rewrites the whole directory and never
merges), so a foreign version is a hard reject rather than a partial read."""
manifest = _read_manifest(root)
if isinstance(manifest, UnreadableBundle):
return manifest
if manifest.format_version != BUNDLE_FORMAT_VERSION:
expected_version: Final = BUNDLE_FORMAT_VERSION if profile == "legacy" else STRICT_BUNDLE_FORMAT_VERSION
if manifest.match_profile != profile:
return UnreadableBundle(
reason="match profile mismatch; select the recorded E2E_REPLAY_MATCH_PROFILE or re-record"
)
if manifest.format_version != expected_version:
return UnreadableBundle(
reason=(
f"format_version {manifest.format_version} != supported {BUNDLE_FORMAT_VERSION}; "
f"format_version {manifest.format_version} != supported {expected_version}; "
"re-record with E2E_FIXTURE_MODE=record"
)
)
return manifest
def check_freshness(root: Path, *, now: datetime) -> BundleFreshness:
manifest = _supported_manifest(root)
def check_freshness(root: Path, *, now: datetime, profile: MatchProfile = "legacy") -> BundleFreshness:
manifest = _supported_manifest(root, profile)
if isinstance(manifest, UnreadableBundle):
return manifest
recorded_at = (
@ -269,16 +285,27 @@ class LoadedBundle:
interactions: dict[str, tuple[Interaction, ...]]
def load_bundle(root: Path) -> LoadedBundle | UnreadableBundle:
manifest = _supported_manifest(root)
def load_bundle(root: Path, *, profile: MatchProfile = "legacy") -> LoadedBundle | UnreadableBundle:
manifest = _supported_manifest(root, profile)
if isinstance(manifest, UnreadableBundle):
return manifest
interactions = {
directory.name: tuple(
Interaction.model_validate_json(file.read_text(encoding="utf-8"))
for file in sorted(directory.glob("*.json"))
)
for directory in sorted(root.iterdir())
if directory.is_dir()
}
try:
interactions = {
directory.name: tuple(
Interaction.model_validate_json(file.read_text(encoding="utf-8"))
for file in sorted(directory.glob("*.json"))
)
for directory in sorted(root.iterdir())
if directory.is_dir()
}
except (ValueError, OSError):
if profile == "legacy":
raise
return UnreadableBundle(reason="invalid stateless_v1 interaction; re-record with the selected profile")
if any(
(item.request.strict_identity is not None) != (profile == "stateless_v1")
for items in interactions.values()
for item in items
):
return UnreadableBundle(reason="request identity/profile mismatch; re-record with the selected profile")
return LoadedBundle(manifest=manifest, interactions=interactions)

View file

@ -23,9 +23,8 @@ from dataclasses import dataclass
from functools import reduce
from typing import Final
from pydantic import JsonValue
from fixture_bundle import RecordedRequest
from pydantic import JsonValue
VOLATILE_HEADER_NAMES: Final[frozenset[str]] = frozenset(
{
@ -123,6 +122,12 @@ class CanonicalRequest:
def canonicalize(request: RecordedRequest) -> CanonicalRequest:
if request.strict_identity is not None:
return CanonicalRequest(
method=request.method,
path=request.path,
content=json.dumps(request.strict_identity.model_dump(mode="json"), sort_keys=True, separators=(",", ":")),
)
file_identity: Final[JsonValue | None] = (
None
if request.file_name is None and request.file_sha256 is None

View file

@ -26,6 +26,7 @@ from fixture_bundle import (
check_freshness,
format_age,
)
from fixture_profile import match_profile
type FixtureMode = Literal["live", "record", "replay"]
@ -82,6 +83,7 @@ def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datet
Called at collection time (conftest pytest_sessionstart) so a stale or missing
bundle fails the whole run up front, naming the bundle age, instead of failing
every test individually."""
match_profile()
mode = parse_fixture_mode(mode_raw)
match mode:
case InvalidFixtureMode(value=value):
@ -89,7 +91,7 @@ def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datet
case "live" | "record":
return None
case "replay":
freshness = check_freshness(bundle_dir, now=now)
freshness = check_freshness(bundle_dir, now=now, profile=match_profile())
match freshness:
case FreshBundle():
return None
@ -110,6 +112,7 @@ def fixture_mode_collection_error(mode_raw: str, bundle_dir: Path, *, now: datet
def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> list[str]:
"""pytest report-header lines; empty in live mode so an unset
E2E_FIXTURE_MODE keeps today's output byte-identical."""
match_profile()
mode = parse_fixture_mode(mode_raw)
match mode:
case InvalidFixtureMode() | "live":
@ -117,7 +120,7 @@ def fixture_report_lines(mode_raw: str, bundle_dir: Path, *, now: datetime) -> l
case "record":
return [f"e2e fixture mode: record -> {bundle_dir}"]
case "replay":
freshness = check_freshness(bundle_dir, now=now)
freshness = check_freshness(bundle_dir, now=now, profile=match_profile())
match freshness:
case FreshBundle(manifest=manifest):
return [

View file

@ -0,0 +1,179 @@
from __future__ import annotations
import json
import os
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal
from urllib.parse import parse_qsl, urlsplit
from pydantic import BaseModel, JsonValue, TypeAdapter
type MatchProfile = Literal["legacy", "stateless_v1"]
@dataclass(frozen=True, slots=True)
class NumberToken:
literal: str
type ExactJson = dict[str, ExactJson] | list[ExactJson] | str | bool | NumberToken | None
SEMANTIC_HEADERS: Final = frozenset({"content-type", "accept", "anthropic-version", "anthropic-beta", "openai-beta"})
AUTH_HEADERS: Final = frozenset({"authorization", "x-api-key"})
EXCLUDED_HEADERS: Final = frozenset(
{
"host",
"content-length",
"transfer-encoding",
"connection",
"accept-encoding",
"user-agent",
"traceparent",
"tracestate",
"x-request-id",
"x-client-request-id",
"cookie",
}
)
CREDENTIAL_QUERY: Final = frozenset(
{
"api_key",
"api-key",
"apikey",
"key",
"token",
"access_token",
"signature",
"password",
"secret",
"credentials",
"authorization",
"sig",
"client_secret",
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
}
)
JSON_VALUE: Final[TypeAdapter[ExactJson]] = TypeAdapter(ExactJson)
def match_profile() -> MatchProfile:
raw: Final = os.environ.get("E2E_REPLAY_MATCH_PROFILE", "legacy")
if raw in ("legacy", "stateless_v1"):
return raw
raise ValueError("E2E_REPLAY_MATCH_PROFILE must be legacy or stateless_v1")
class StrictIdentity(BaseModel):
upstream: str
mount: str
query: tuple[tuple[str, str], ...]
headers: dict[str, str]
auth: dict[str, str]
body_present: bool
body: JsonValue
@dataclass(frozen=True, slots=True)
class IneligibleRequest:
reason: str
def _unique_object(pairs: list[tuple[str, ExactJson]]) -> dict[str, ExactJson]:
if len({key for key, _ in pairs}) != len(pairs):
raise ValueError("duplicate JSON object keys")
return dict(pairs)
def _invalid_constant(value: str) -> ExactJson:
raise ValueError("nonfinite JSON number")
def _exact_value(value: ExactJson) -> JsonValue:
match value:
case dict():
return {"object": {key: _exact_value(item) for key, item in value.items()}}
case list():
return {"array": [_exact_value(item) for item in value]}
case bool():
return {"boolean": value}
case NumberToken(literal=literal):
return {"number": literal}
case str():
return {"string": value}
case None:
return None
def strict_identity(
*,
method: str,
path: str,
query: str,
headers: Mapping[str, str],
body: bytes | None,
mount: str,
upstream_base: str,
) -> StrictIdentity | IneligibleRequest:
if (mount, path, method.upper()) not in {
("openai", "/openai/v1/chat/completions", "POST"),
("anthropic", "/anthropic/v1/messages", "POST"),
}:
return IneligibleRequest("unsupported endpoint or method")
lowered: Final = {key.lower(): value for key, value in headers.items()}
if len(lowered) != len(headers):
return IneligibleRequest("duplicate header names")
if any(
key not in SEMANTIC_HEADERS | AUTH_HEADERS | EXCLUDED_HEADERS and not key.startswith("x-stainless-")
for key in lowered
):
return IneligibleRequest("unsupported semantic header")
if "transfer-encoding" in lowered:
return IneligibleRequest("unsupported request transfer-encoding; send a content-length framed JSON body")
authorization: Final = lowered.get("authorization")
if authorization is not None and authorization.partition(" ")[0].lower() not in {"bearer", "basic", "digest"}:
return IneligibleRequest("unsupported authorization scheme")
destination: Final = urlsplit(upstream_base)
if destination.username or destination.password or destination.query or destination.fragment:
return IneligibleRequest("upstream destination contains credentials, query or fragment")
if destination.scheme not in ("http", "https") or not destination.netloc:
return IneligibleRequest("unsupported upstream destination")
if body and lowered.get("content-type", "").split(";", 1)[0].strip().lower() != "application/json":
return IneligibleRequest("unsupported body content-type; stateless_v1 requires JSON")
try:
parsed: Final = (
JSON_VALUE.validate_python(
json.loads(
body,
object_pairs_hook=_unique_object,
parse_constant=_invalid_constant,
parse_float=NumberToken,
parse_int=NumberToken,
)
)
if body
else None
)
except (ValueError, UnicodeError):
return IneligibleRequest("invalid JSON or duplicate JSON object keys")
if body and not isinstance(parsed, dict):
return IneligibleRequest("stateless inference requires a JSON object")
try:
query_pairs: Final = tuple(parse_qsl(query, keep_blank_values=True, errors="strict"))
except UnicodeError:
return IneligibleRequest("invalid UTF-8 query encoding")
return StrictIdentity(
upstream=upstream_base,
mount=mount,
query=tuple((key, "<credential>" if key.lower() in CREDENTIAL_QUERY else value) for key, value in query_pairs),
headers={key: value for key, value in lowered.items() if key in SEMANTIC_HEADERS},
auth={
key: (value.partition(" ")[0].lower() if key == "authorization" else "present")
for key, value in lowered.items()
if key in AUTH_HEADERS
},
body_present=bool(body),
body=_exact_value(parsed),
)

View file

@ -2,11 +2,18 @@
from __future__ import annotations
import base64
import os
import secrets
import signal
import subprocess
import sys
import time
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 +21,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 +57,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 +76,7 @@ class GroupCreateBody(BaseModel):
class PasswordCredential(BaseModel):
type: Literal["password"] = "password"
value: str
value: str = Field(repr=False)
temporary: bool = False
@ -101,8 +114,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 +136,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 +179,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 +189,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 +202,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 +233,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 +275,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 +349,137 @@ 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 _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:
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, start_new_session=True) as child:
try:
return child.wait()
finally:
_stop_process_group(child)
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:]))

View file

@ -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:

View file

@ -1,51 +1,90 @@
"""Vendor §12.3: chat completions streaming SSE contract (LIT-4778).
Asserts a streamed /chat/completions response is SSE, carries content chunks,
and terminates with the OpenAI [DONE] sentinel.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_config import provider_edge_base, unique_marker
from e2e_http import require_successful_call
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, LiteLLMParamsBody
from models import ChatBody, ChatMessage, ChatStreamOptions, LiteLLMParamsBody, Usage
from proxy_client import ProxyClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
class _Delta(BaseModel):
content: str | None = None
class _Choice(BaseModel):
index: int
delta: _Delta
finish_reason: str | None = None
class _Chunk(BaseModel):
choices: tuple[_Choice, ...]
usage: Usage | None = None
class TestChatStreamContract:
@pytest.mark.covers("llm.chat_completions.openai.basic.stream.works")
def test_chat_stream_is_sse_and_ends_with_done(self, proxy: ProxyClient, resources: ResourceManager) -> None:
model = f"e2e-chat-stream-{unique_marker()}"
model_id = proxy.create_model(
model: Final = f"e2e-chat-stream-{unique_marker()}"
base: Final = provider_edge_base("openai")
model_id: Final = proxy.create_model(
model,
LiteLLMParamsBody(model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY"),
LiteLLMParamsBody(
model="openai/gpt-5.6",
api_key="os.environ/OPENAI_API_KEY",
api_base=f"{base}/v1" if base else None,
),
)
resources.defer(lambda: proxy.delete_model(model_id))
key = resources.key()
result = proxy.chat_stream(
key: Final = resources.key()
expected: Final = "The amber kite crosses the quiet lake."
result: Final = proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[
ChatMessage(
role="user",
content=f"Reply with the single word ok. {unique_marker()}",
role="user", content=f"Repeat exactly this sentence, with no additional text: {expected}"
)
],
stream=True,
max_completion_tokens=32,
temperature=0.0,
stream_options=ChatStreamOptions(include_usage=True),
max_completion_tokens=256,
reasoning_effort="none",
),
)
require_successful_call(result)
assert result.is_streaming, f"expected SSE content-type, got {result.content_type!r}"
assert result.stream_events, "stream returned no data events"
assert result.stream_done, (
f"stream must terminate with [DONE]; "
f"chunks={result.chunks} done={result.stream_done} events={len(result.stream_events)}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_done, "stream must terminate with [DONE]"
assert result.stream_done_positions == (len(result.stream_events),), "[DONE] must occur once after all events"
chunks: Final = tuple(_Chunk.model_validate_json(event) for event in result.stream_events)
text_positions: Final = tuple(
i for i, chunk in enumerate(chunks) if any(c.delta.content for c in chunk.choices)
)
terminal_positions: Final = tuple(
i for i, chunk in enumerate(chunks) if any(c.finish_reason is not None for c in chunk.choices)
)
assert text_positions, "stream completed without meaningful text"
assert len(terminal_positions) == 1, "expected exactly one terminal choice"
assert text_positions[0] < terminal_positions[0], "meaningful text must arrive before termination"
assert text_positions[-1] <= terminal_positions[0], "text arrived after termination"
assert all(c.index == 0 for chunk in chunks for c in chunk.choices)
assert tuple(c.finish_reason for c in chunks[terminal_positions[0]].choices) == ("stop",)
text: Final = "".join(c.delta.content or "" for chunk in chunks for c in chunk.choices)
assert text.strip() == expected, f"streamed answer was altered or incomplete: {text!r}"
usage_positions: Final = tuple(i for i, chunk in enumerate(chunks) if chunk.usage is not None)
assert usage_positions == (len(chunks) - 1,), "expected one final usage chunk"
assert terminal_positions[0] < usage_positions[0], "usage must follow the terminal choice"
usage: Final = chunks[-1].usage
assert usage is not None
assert usage.prompt_tokens is not None and usage.prompt_tokens > 0
assert usage.completion_tokens is not None and usage.completion_tokens > 0
assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens

View file

@ -21,7 +21,12 @@ from e2e_http import assert_client_error, require_successful_call, unwrap
from endpoints_client import EndpointsClient, MessagesResult
from lifecycle import ResourceManager
from models import (
AnthropicAssistantTurn,
AnthropicContentBlock,
AnthropicCustomTool,
AnthropicToolChoice,
AnthropicToolResultBlock,
AnthropicToolResultTurn,
AnthropicMessagesBody,
ChatMessage,
JsonSchemaProperty,
@ -29,7 +34,7 @@ from models import (
SpendLogRow,
ToolInputSchema,
)
from pydantic import BaseModel
from pydantic import BaseModel, ConfigDict
pytestmark = [pytest.mark.e2e, pytest.mark.replayable]
@ -284,8 +289,139 @@ class TestAnthropicMessages:
result = endpoints_client.proxy.transport.send(
"/v1/messages",
headers=endpoints_client.proxy.transport.bearer(key),
json=_OptionalMessagesBody(
messages=[ChatMessage(role="user", content="hi")], max_tokens=50
),
json=_OptionalMessagesBody(messages=[ChatMessage(role="user", content="hi")], max_tokens=50),
)
assert_client_error(result, "messages missing model")
class _BridgeDelta(BaseModel):
type: str | None = None
partial_json: str | None = None
stop_reason: str | None = None
class _BridgeEvent(BaseModel):
type: str
index: int | None = None
content_block: AnthropicContentBlock | None = None
delta: _BridgeDelta | None = None
class _ParcelInput(BaseModel):
model_config = ConfigDict(extra="forbid", strict=True)
parcel: str
shelf: int
def _tool_from_stream(events: tuple[_BridgeEvent, ...]) -> AnthropicContentBlock:
starts: Final = tuple(
event
for event in events
if event.type == "content_block_start"
and event.content_block is not None
and event.content_block.type == "tool_use"
)
assert len(starts) == 1, "expected exactly one tool call"
start: Final = starts[0]
block: Final = start.content_block
assert block is not None and block.id and start.index is not None
fragments: Final = tuple(
event
for event in events
if event.type == "content_block_delta" and event.delta is not None and event.delta.type == "input_json_delta"
)
assert fragments, "tool stream contained no argument fragments"
assert all(event.index == start.index for event in fragments), "tool fragments changed index"
positions: Final = tuple(i for i, event in enumerate(events) if event in fragments)
stops: Final = tuple(
i for i, event in enumerate(events) if event.type == "content_block_stop" and event.index == start.index
)
assert len(stops) == 1 and events.index(start) < positions[0] <= positions[-1] < stops[0]
assert tuple(
event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None
) == ("tool_use",)
terminal_positions: Final = tuple(i for i, event in enumerate(events) if event.type == "message_delta")
assert len(terminal_positions) == 1 and stops[0] < terminal_positions[0] < len(events) - 1
assert tuple(i for i, event in enumerate(events) if event.type == "message_stop") == (len(events) - 1,), (
"tool stream did not terminate exactly once"
)
arguments: Final = _ParcelInput.model_validate_json(
"".join(event.delta.partial_json or "" for event in fragments if event.delta is not None)
)
return AnthropicContentBlock(type="tool_use", id=block.id, name=block.name, input=arguments.model_dump())
def _parcel_result(tool: AnthropicContentBlock, result: AnthropicToolResultBlock) -> AnthropicToolResultTurn:
assert tool.id and result.tool_use_id == tool.id, "tool result ID does not match the emitted call"
return AnthropicToolResultTurn(content=[result])
def _request_tool(
client: EndpointsClient, key: str, request: AnthropicMessagesBody, stream: bool
) -> AnthropicContentBlock:
if stream:
response: Final = client.proxy.messages_stream(key, request)
require_successful_call(response)
assert response.is_streaming and not response.stream_error
return _tool_from_stream(tuple(_BridgeEvent.model_validate_json(event) for event in response.stream_events))
response_body: Final = unwrap(client.proxy.messages(key, request))
blocks: Final = tuple(block for block in response_body.content or () if block.type == "tool_use")
assert len(blocks) == 1
return blocks[0]
class TestOpenAIMessagesToolContinuation:
@pytest.mark.parametrize("stream", [True, False], ids=["stream", "nonstream"])
def test_required_tool_arguments_and_correlated_result(
self, endpoints_client: EndpointsClient, resources: ResourceManager, stream: bool
) -> None:
model: Final = f"e2e-bridge-tool-{unique_marker()}"
base: Final = provider_edge_base("openai")
model_id: Final = endpoints_client.create_model(
model,
LiteLLMParamsBody(
model="openai/gpt-5.6", api_key="os.environ/OPENAI_API_KEY", api_base=f"{base}/v1" if base else None
),
)
resources.defer(lambda: endpoints_client.delete_model(model_id))
key: Final = resources.key(models=[model])
tool: Final = AnthropicCustomTool(
name="locate_parcel",
description="Look up the receipt for a parcel on a shelf. Return the receipt verbatim.",
input_schema=ToolInputSchema(
properties={"parcel": JsonSchemaProperty(type="string"), "shelf": JsonSchemaProperty(type="integer")},
required=["parcel", "shelf"],
),
)
question: Final = ChatMessage(
role="user",
content="Call locate_parcel with parcel exactly amber-kite and shelf exactly 7. After the tool result, reply with only the receipt returned by the tool.",
)
request: Final = AnthropicMessagesBody(
model=model,
max_tokens=2048,
messages=[question],
tools=[tool],
tool_choice=AnthropicToolChoice(type="tool", name=tool.name),
stream=stream,
)
emitted: Final = _request_tool(endpoints_client, key, request, stream)
assert emitted.id and emitted.name == "locate_parcel"
assert emitted.input == {"parcel": "amber-kite", "shelf": 7}, "required tool arguments were lost or changed"
receipt: Final = f"receipt-{unique_marker()}"
result_turn: Final = _parcel_result(emitted, AnthropicToolResultBlock(tool_use_id=emitted.id, content=receipt))
continuation: Final = unwrap(
endpoints_client.proxy.messages(
key,
AnthropicMessagesBody(
model=model,
max_tokens=2048,
tools=[tool],
tool_choice=AnthropicToolChoice(type="none"),
messages=[question, AnthropicAssistantTurn(content=[emitted]), result_turn],
),
)
)
answer: Final = "".join(block.text or "" for block in continuation.content or ())
assert answer.strip() == receipt, "continuation did not consume the correlated tool result"
assert all(block.type != "tool_use" for block in continuation.content or ())

View file

@ -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()

View file

@ -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, missing_ok=True))
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)

View file

@ -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,25 +159,26 @@ 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,
)
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.transport.master if caller_key is None else self.proxy.transport.bearer(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
@ -179,7 +186,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 +197,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 +207,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 +216,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 +226,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 +235,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 +273,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 +283,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 +296,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 +307,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 +316,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 +328,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 +336,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 +353,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 +387,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 +397,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 +407,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 +418,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 +427,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 +436,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 +445,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 +456,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 +476,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 +486,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 +497,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 +507,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 +516,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 +525,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 +551,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 +561,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 +572,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 +584,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 +595,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,
)

View file

@ -2,60 +2,247 @@
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.covers("mgmt.key.jwt.lifecycle")
def test_admin_creates_reads_updates_clears_and_deletes_a_key(
self, client: ManagementClient, idp: Keycloak, jwt_identity: Identity, resources: ResourceManager
) -> None:
admin: Final = idp.access_token(jwt_identity, client_id=ADMIN_CLIENT_ID)
alias: Final = f"e2e-jwt-key-{unique_marker()}"
created: Final = unwrap(
client.generate_key(
KeyGenerateBody(key_alias=alias, team_id=jwt_identity.group, models=[CHEAP_OPENAI_MODEL]),
caller_key=admin,
@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,
)
)
resources.defer(lambda: client.proxy.delete_key(created.key))
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")}
original: Final = unwrap(client.key_info_as(created.key, caller_key=admin)).info
assert original.key_alias == alias and original.team_id == jwt_identity.group
@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,
actor_factory: ActorFactory,
credential_kind: Literal["direct_jwt", "virtual_key"],
) -> None:
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 actor.mint_caller(actor_factory.idp).credential
bound: Final = actor_factory.bootstrap.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=tenant.team_id, models=[CHEAP_OPENAI_MODEL]),
)
)
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 == tenant.team_id
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 +250,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 +275,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

View file

@ -283,10 +283,15 @@ class ChatToolResultTurn(BaseModel):
type ChatTurn = ChatMessage | ChatAssistantTurn | ChatToolResultTurn
class ChatStreamOptions(BaseModel):
include_usage: bool
class ChatBody(BaseModel):
model: str
messages: Sequence[ChatTurn]
stream: bool = False
stream_options: ChatStreamOptions | None = None
max_tokens: int | None = None
max_completion_tokens: int | None = None
temperature: float | None = None
@ -488,12 +493,18 @@ class AnthropicToolResultTurn(BaseModel):
type AnthropicMessage = ChatMessage | AnthropicAssistantTurn | AnthropicToolResultTurn
class AnthropicToolChoice(BaseModel):
type: Literal["auto", "any", "tool", "none"]
name: str | None = None
class AnthropicMessagesBody(BaseModel):
model: str
messages: list[AnthropicMessage]
max_tokens: int
stream: bool | None = None
tools: list[AnthropicTool] | None = None
tool_choice: AnthropicToolChoice | None = None
guardrails: list[str] | None = None
cache: dict[str, bool] | None = {"no-cache": True}
@ -1091,13 +1102,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 +1146,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 +1187,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 +1200,7 @@ class UserUpdateBody(BaseModel):
class UserInfoParams(BaseModel):
user_id: str
user_id: str | None = None
class UserData(BaseModel):
@ -1240,16 +1253,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) ----------

View file

@ -92,6 +92,7 @@ from fixture_mode import (
current_test_key,
parse_fixture_mode,
)
from fixture_profile import IneligibleRequest, MatchProfile, match_profile, strict_identity
from pydantic import JsonValue, TypeAdapter
EDGE_MOUNTS: Final[Mapping[str, str]] = MappingProxyType(
@ -404,6 +405,12 @@ def _miss_message(test_key: str, slug: str, canonical: CanonicalRequest, bundle:
f"under {slug}; re-record with E2E_FIXTURE_MODE=record"
)
closest, closest_file = _closest_recorded(canonical, recorded)
if bundle.manifest.match_profile == "stateless_v1":
expected: Final = _JSON.validate_json(closest.content)
actual: Final = _JSON.validate_json(canonical.content)
assert isinstance(expected, dict) and isinstance(actual, dict)
changed: Final = ", ".join(key for key in expected if expected[key] != actual.get(key))
return f"stateless_v1 replay mismatch: {changed or 'method/path'}; re-record with E2E_FIXTURE_MODE=record"
diff: Final = "\n".join(
islice(
difflib.unified_diff(
@ -785,11 +792,33 @@ def handle_edge_request(
mount, _, upstream_path = split.path.lstrip("/").partition("/")
upstream_base: Final = mounts.get(mount)
if upstream_base is None:
return _text_reply(
404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}"
return _text_reply(404, f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(mounts))}")
profile: Final = (
backend.recorder.profile
if isinstance(backend, RecordEdge)
else backend.source.bundle.manifest.match_profile
if isinstance(backend, ReplayEdge)
else "legacy"
)
identity: Final = (
strict_identity(
method=method,
path=split.path,
query=split.query,
headers=headers,
body=body,
mount=mount,
upstream_base=upstream_base,
)
request: Final = edge_request(
method, split.path, split.query, body, _header_value(headers, "content-type")
if profile == "stateless_v1"
else None
)
if isinstance(identity, IneligibleRequest):
return _text_reply(REPLAY_MISS_STATUS, f"stateless_v1 eligibility error: {identity.reason}")
request: Final = (
RecordedRequest(method=method.lower(), path=split.path, headers={}, strict_identity=identity)
if identity is not None
else edge_request(method, split.path, split.query, body, _header_value(headers, "content-type"))
)
match backend:
case LiveEdge():
@ -837,6 +866,14 @@ class _EdgeHandler(BaseHTTPRequestHandler):
body: Final = self.rfile.read(length) if length else None
if edge_server.observation is not None:
edge_server.observation.observe(body)
strict: Final = (
isinstance(edge_server.backend, RecordEdge) and edge_server.backend.recorder.profile == "stateless_v1"
or isinstance(edge_server.backend, ReplayEdge)
and edge_server.backend.source.bundle.manifest.match_profile == "stateless_v1"
)
if strict and len({name.lower() for name in self.headers.keys()}) != len(self.headers):
self._write_reply(_text_reply(REPLAY_MISS_STATUS, "stateless_v1 eligibility error: duplicate headers"))
return
outcome: Final = handle_edge_request(
edge_server.backend,
edge_server.mounts,
@ -955,16 +992,16 @@ def start_provider_edge(
@functools.lru_cache(maxsize=8)
def _shared_recorder(root: Path) -> BundleRecorder:
prepared = prepare_bundle(root)
def _shared_recorder(root: Path, profile: MatchProfile = "legacy") -> BundleRecorder:
prepared = prepare_bundle(root, profile=profile)
if isinstance(prepared, UnsafeBundleDir):
raise ValueError(f"E2E_FIXTURE_DIR {prepared.path} {prepared.reason}")
return prepared
@functools.lru_cache(maxsize=8)
def _shared_replay_source(root: Path) -> ReplaySource:
loaded = load_bundle(root)
def _shared_replay_source(root: Path, profile: MatchProfile = "legacy") -> ReplaySource:
loaded = load_bundle(root, profile=profile)
if isinstance(loaded, UnreadableBundle):
raise ValueError(f"cannot replay from {root}: {loaded.reason}")
return ReplaySource(bundle=loaded)
@ -977,11 +1014,12 @@ def _shared_edge(
bind_host: str,
advertise_host: str,
forward_timeout: float,
profile: MatchProfile,
) -> ProviderEdge:
backend: Final[EdgeBackend] = (
RecordEdge(recorder=_shared_recorder(bundle_dir), lock=threading.Lock())
RecordEdge(recorder=_shared_recorder(bundle_dir, profile), lock=threading.Lock())
if mode == "record"
else ReplayEdge(source=_shared_replay_source(bundle_dir))
else ReplayEdge(source=_shared_replay_source(bundle_dir, profile))
)
return start_provider_edge(
backend,
@ -998,7 +1036,7 @@ def replay_leftover_error(*, mode_raw: str, bundle_dir: Path, test_key: str) ->
recording it no longer matches. Inert in every other mode."""
if parse_fixture_mode(mode_raw) != "replay":
return None
return _shared_replay_source(bundle_dir).leftover_error(test_key)
return _shared_replay_source(bundle_dir, match_profile()).leftover_error(test_key)
def provider_edge_api_base(
@ -1021,10 +1059,10 @@ def provider_edge_api_base(
return None
case "record" | "replay":
if mount not in EDGE_MOUNTS:
raise ValueError(
f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}"
)
return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout).api_base(mount)
raise ValueError(f"unknown provider mount {mount!r}; known mounts: {', '.join(sorted(EDGE_MOUNTS))}")
return _shared_edge(mode, bundle_dir, bind_host, advertise_host, forward_timeout, match_profile()).api_base(
mount
)
case _:
assert_never(mode)
@ -1037,9 +1075,9 @@ def _observed_backend(mode_raw: str, bundle_dir: Path) -> EdgeBackend:
case "live":
return LiveEdge()
case "record":
return RecordEdge(_shared_recorder(bundle_dir), threading.Lock())
return RecordEdge(_shared_recorder(bundle_dir, match_profile()), threading.Lock())
case "replay":
return ReplayEdge(_shared_replay_source(bundle_dir))
return ReplayEdge(_shared_replay_source(bundle_dir, match_profile()))
case _:
assert_never(mode)

View file

@ -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(

View file

@ -0,0 +1,113 @@
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from math import isclose
from typing import Final
from e2e_config import provider_edge_base, unique_marker
from e2e_http import unwrap
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, ChatResponse, KeyGenerateBody, LiteLLMParamsBody, TeamNewBody
from spend_e2e_client import SpendClient
INPUT_RATE: Final = 0.00004
OUTPUT_RATE: Final = 0.00008
@dataclass(frozen=True)
class TeamTraffic:
team_id: str
key: str
responses: tuple[ChatResponse, ...]
@property
def prompt_tokens(self) -> int:
return sum(response.usage.prompt_tokens or 0 for response in self.responses if response.usage)
@property
def completion_tokens(self) -> int:
return sum(response.usage.completion_tokens or 0 for response in self.responses if response.usage)
@property
def spend(self) -> float:
return self.prompt_tokens * INPUT_RATE + self.completion_tokens * OUTPUT_RATE
def create_traffic(client: SpendClient, resources: ResourceManager) -> tuple[TeamTraffic, ...]:
base: Final = provider_edge_base("openai")
model: Final = f"e2e-reconciliation-{unique_marker()}"
model_id: Final = client.proxy.create_model(
model,
LiteLLMParamsBody(
model="openai/gpt-5.6-luna",
api_key="os.environ/OPENAI_API_KEY",
api_base=None if base is None else f"{base}/v1",
input_cost_per_token=INPUT_RATE,
output_cost_per_token=OUTPUT_RATE,
),
)
resources.defer(lambda: client.proxy.delete_model(model_id))
def team_traffic() -> TeamTraffic:
team: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-spend-{unique_marker()}"))
resources.defer(lambda: client.proxy.delete_team(team))
key: Final = client.proxy.generate_key(KeyGenerateBody(team_id=team, models=[model]))
resources.defer(lambda: client.proxy.delete_key(key))
prompts: Final = tuple(f"Reply with one word. {index} {unique_marker()}" for index in range(7))
def call(index: int) -> ChatResponse:
response: Final = unwrap(
client.proxy.chat(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=prompts[index])],
max_completion_tokens=128,
),
)
)
assert response.id, "successful response must have an ID"
assert response.usage is not None, "successful response must have usage"
assert response.usage.prompt_tokens is not None and response.usage.prompt_tokens > 0
assert response.usage.completion_tokens is not None and response.usage.completion_tokens > 0
assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens
assert not response.usage.cache_creation_input_tokens
assert not response.usage.cache_read_input_tokens
assert not response.usage.prompt_tokens_details or not response.usage.prompt_tokens_details.cached_tokens
return response
sequential: Final = call(0)
with ThreadPoolExecutor(max_workers=6) as pool:
concurrent: Final = tuple(pool.map(call, range(1, 7)))
return TeamTraffic(team, key, (sequential, *concurrent))
return tuple(team_traffic() for _ in range(2))
def assert_logs_match(client: SpendClient, traffic: TeamTraffic) -> None:
expected_ids: Final = frozenset(response.id for response in traffic.responses)
assert len(expected_ids) == len(traffic.responses), "responses must have distinct IDs"
rows: Final = client.poll_logs_for_key(
traffic.key,
min_rows=len(traffic.responses),
predicate=lambda values: frozenset(row.request_id for row in values) == expected_ids,
)
assert frozenset(row.request_id for row in rows) == expected_ids, "stored IDs must equal returned response IDs"
assert len(rows) == len(traffic.responses), "expected exactly one scoped spend row per response"
by_id: Final = {row.request_id: row for row in rows}
def assert_response(response: ChatResponse) -> None:
row: Final = by_id[response.id]
usage: Final = response.usage
assert usage is not None and usage.prompt_tokens is not None and usage.completion_tokens is not None
assert row.team_id == traffic.team_id
assert row.status == "success"
assert row.cache_hit != "True"
assert row.prompt_tokens == usage.prompt_tokens
assert row.completion_tokens == usage.completion_tokens
assert row.total_tokens == usage.total_tokens
expected_cost: Final = usage.prompt_tokens * INPUT_RATE + usage.completion_tokens * OUTPUT_RATE
assert row.spend is not None and isclose(row.spend, expected_cost, rel_tol=1e-6, abs_tol=1e-9)
for response in traffic.responses:
assert_response(response)

View file

@ -17,13 +17,13 @@ fails the test; a pricing or token-count drift does not.
import time
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from math import isclose
from typing import Final
import pytest
from e2e_http import Result, Success
from e2e_http import Success
from lifecycle import ResourceManager
from models import ChatResponse, LiteLLMParamsBody, SpendLogs, SpendLogsParams
from models import LiteLLMParamsBody, SpendLogs, SpendLogsParams
from spend_e2e_client import SpendClient, SpendLogRow, is_ok, unique_marker, unwrap
pytestmark = pytest.mark.e2e
@ -280,51 +280,22 @@ def test_key_spend_equals_sum_of_logs(client: SpendClient, scoped_key: str) -> N
), f"key aggregate {key_spend} != sum of logs {logs_total}; rows: {_summarize(rows)}"
@pytest.mark.replayable
@pytest.mark.covers("quota_management.spend_tracking.concurrent_burst.loses_no_spend")
def test_burst_of_concurrent_calls_loses_no_spend(
client: SpendClient, scoped_key: str
client: SpendClient, resources: ResourceManager
) -> None:
"""Six concurrent calls on one key: every call lands its own spend row under a
distinct request_id and the key aggregate equals the sum of the rows.
Sequential accuracy is covered by test_key_spend_equals_sum_of_logs; this pins
the concurrent increment path (parallel writers racing on one key's counter),
where a lost update can never be reproduced by sequential calls."""
burst = 6
from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic
def call(idx: int) -> Result[ChatResponse]:
return client.chat(
scoped_key,
"gemini-2.5-flash",
f"burst call {idx} {unique_marker()}",
max_tokens=16,
)
traffic: Final = create_traffic(client, resources)
with ThreadPoolExecutor(max_workers=burst) as pool:
results = tuple(pool.map(call, range(burst)))
failed = [r for r in results if not is_ok(r)]
assert not failed, f"{len(failed)}/{burst} burst calls failed; first: {failed[0]}"
def assert_team(team: TeamTraffic) -> None:
assert_logs_match(client, team)
key_spend: Final = client.poll_key_spend(team.key, minimum=team.spend * 0.999999)
assert isclose(key_spend, team.spend, rel_tol=1e-6, abs_tol=1e-9)
rows = client.poll_logs_for_key(
scoped_key,
min_rows=burst,
predicate=lambda rs: len([r for r in rs if (r.spend or 0) > 0]) >= burst,
)
costed = [r for r in rows if (r.spend or 0) > 0]
assert len(costed) >= burst, (
f"only {len(costed)}/{burst} burst calls produced a costed row - "
f"rows lost under concurrency: {_summarize(rows)}"
)
request_ids = [r.request_id for r in costed]
assert len(set(request_ids)) == len(request_ids), (
f"concurrent rows collapsed onto shared request_ids: {_summarize(rows)}"
)
logs_total = sum((r.spend or 0) for r in rows)
key_spend = client.poll_key_spend(scoped_key, minimum=logs_total * 0.999)
assert _approx_equal(key_spend, logs_total), (
f"key aggregate {key_spend} != sum of {len(rows)} rows {logs_total} - "
f"spend increments lost under concurrency: {_summarize(rows)}"
)
for team in traffic:
assert_team(team)
@pytest.mark.covers("quota_management.spend_tracking.pagination.keeps_total")

View file

@ -7,13 +7,18 @@ missing start/end dates are rejected.
from __future__ import annotations
import time
from datetime import datetime, timedelta, timezone
from math import isclose
from typing import Final
import pytest
from e2e_http import ProbeResult
from models import DateRangeParams
from lifecycle import ResourceManager
from proxy_client import Converged, await_converged
from pydantic import BaseModel
from spend_e2e_client import SpendClient
from spend_reconciliation import TeamTraffic, assert_logs_match, create_traffic
pytestmark = pytest.mark.e2e
@ -24,22 +29,45 @@ class TeamDailyActivityParams(BaseModel):
start_date: str | None = None
end_date: str | None = None
page: int = 1
page_size: int = 1
team_ids: str | None = None
class TeamDailyActivityRow(BaseModel):
date: str
metrics: TeamDailyActivityMetrics
breakdown: TeamDailyActivityBreakdown
class TeamDailyActivityMetrics(BaseModel):
spend: float
total_tokens: int
prompt_tokens: int
completion_tokens: int
api_requests: int
successful_requests: int
failed_requests: int
class TeamDailyActivityEntity(BaseModel):
metrics: TeamDailyActivityMetrics
class TeamDailyActivityBreakdown(BaseModel):
entities: dict[str, TeamDailyActivityEntity]
class TeamDailyActivityMetadata(BaseModel):
page: int
total_pages: int
has_more: bool
total_spend: float
total_prompt_tokens: int
total_completion_tokens: int
total_tokens: int
total_api_requests: int
total_successful_requests: int
total_failed_requests: int
class TeamDailyActivityResponse(BaseModel):
@ -47,32 +75,128 @@ class TeamDailyActivityResponse(BaseModel):
metadata: TeamDailyActivityMetadata
def _range_days(days: int) -> DateRangeParams:
end = datetime.now(timezone.utc).date()
start = end - timedelta(days=days)
return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat())
def _probe(client: SpendClient, params: BaseModel) -> ProbeResult:
return client.proxy.transport.probe(ROUTE, params=params)
class TestTeamDailyActivity:
@pytest.mark.replayable
@pytest.mark.covers("mgmt.team.daily_activity.happy_path")
@pytest.mark.parametrize("days", [1, 7, 30])
def test_valid_date_range_returns_results_and_metadata(self, client: SpendClient, days: int) -> None:
result = _probe(client, _range_days(days))
assert result.status_code == 200, (
f"{ROUTE} range={days}d must be 200, got {result.status_code}: {result.body[:600]}"
def test_valid_date_range_returns_results_and_metadata(
self, client: SpendClient, resources: ResourceManager
) -> None:
started: Final = datetime.now(timezone.utc).date()
traffic: Final = create_traffic(client, resources)
for team in traffic:
assert_logs_match(client, team)
ended: Final = datetime.now(timezone.utc).date()
team_ids: Final = ",".join(team.team_id for team in traffic)
def fetch(
page: int, start: str = (started - timedelta(days=1)).isoformat(), end: str = ended.isoformat()
) -> TeamDailyActivityResponse:
result: Final = _probe(
client,
TeamDailyActivityParams(
start_date=start,
end_date=end,
page=page,
page_size=1,
team_ids=team_ids,
),
)
assert result.status_code == 200, f"daily activity failed: {result.status_code} {result.body[:300]}"
return TeamDailyActivityResponse.model_validate_json(result.body)
def pages() -> tuple[TeamDailyActivityResponse, ...]:
first: Final = fetch(1)
assert first.metadata.total_pages <= len(traffic) * 2, "unexpected extra scoped daily groups"
return (first, *(fetch(page) for page in range(2, first.metadata.total_pages + 1)))
outcome: Final = await_converged(
pages,
converged=lambda values: (
sum(page.metadata.total_api_requests for page in values) >= sum(len(team.responses) for team in traffic)
),
timeout=client.proxy.poll_timeout,
interval=client.proxy.poll_interval,
now=time.monotonic,
sleep=time.sleep,
)
parsed = TeamDailyActivityResponse.model_validate_json(result.body)
assert parsed.metadata.page == 1
assert parsed.metadata.total_pages >= 1
if parsed.results:
first = parsed.results[0]
assert first.date
assert first.metrics.spend >= 0
assert first.metrics.total_tokens >= 0
observed: Final = outcome.result if isinstance(outcome, Converged) else outcome.last_result
assert observed is not None, "daily aggregation must return a response before the deadline"
assert len(observed) >= 2, "two teams must exercise a page boundary"
def assert_page(index: int, page: TeamDailyActivityResponse) -> None:
assert page.metadata.page == index
assert page.metadata.total_pages == len(observed)
assert page.metadata.has_more == (index < len(observed))
assert len(page.results) == 1, "each fetched daily group must appear in results"
row: Final = page.results[0]
assert started <= datetime.fromisoformat(row.date).date() <= ended
assert len(row.breakdown.entities) == 1
assert row.metrics.total_tokens == page.metadata.total_tokens
assert row.metrics.prompt_tokens == page.metadata.total_prompt_tokens
assert row.metrics.completion_tokens == page.metadata.total_completion_tokens
assert row.metrics.api_requests == page.metadata.total_api_requests
assert row.metrics.successful_requests == page.metadata.total_successful_requests
assert row.metrics.failed_requests == page.metadata.total_failed_requests
assert isclose(row.metrics.spend, page.metadata.total_spend, rel_tol=1e-6, abs_tol=1e-9)
for index, page in enumerate(observed, 1):
assert_page(index, page)
entities: Final = tuple(
(team_id, entity.metrics)
for page in observed
for row in page.results
for team_id, entity in row.breakdown.entities.items()
)
assert frozenset(team_id for team_id, _ in entities) == frozenset(team.team_id for team in traffic)
def assert_team(team: TeamTraffic) -> None:
metrics: Final = tuple(metrics for team_id, metrics in entities if team_id == team.team_id)
assert sum(m.api_requests for m in metrics) == len(team.responses)
assert sum(m.successful_requests for m in metrics) == len(team.responses)
assert sum(m.failed_requests for m in metrics) == 0
assert sum(m.prompt_tokens for m in metrics) == team.prompt_tokens
assert sum(m.completion_tokens for m in metrics) == team.completion_tokens
assert sum(m.total_tokens for m in metrics) == team.prompt_tokens + team.completion_tokens
assert isclose(sum(m.spend for m in metrics), team.spend, rel_tol=1e-6, abs_tol=1e-9)
for team in traffic:
assert_team(team)
assert isclose(
sum(page.metadata.total_spend for page in observed),
sum(team.spend for team in traffic),
rel_tol=1e-6,
abs_tol=1e-9,
)
assert sum(page.metadata.total_tokens for page in observed) == sum(
team.prompt_tokens + team.completion_tokens for team in traffic
)
for days in (7, 30):
assert (
tuple(fetch(page, (started - timedelta(days=days)).isoformat()) for page in range(1, len(observed) + 1))
== observed
), f"{days}-day activity must preserve the same isolated groups and totals"
empty_date: Final = (started - timedelta(days=7)).isoformat()
empty: Final = fetch(1, empty_date, empty_date)
assert empty.results == []
assert empty.metadata.total_pages == 0
assert empty.metadata.page == 1
assert not empty.metadata.has_more
assert empty.metadata.total_spend == 0
assert empty.metadata.total_tokens == 0
assert empty.metadata.total_api_requests == 0
assert empty.metadata.total_prompt_tokens == 0
assert empty.metadata.total_completion_tokens == 0
assert empty.metadata.total_successful_requests == 0
assert empty.metadata.total_failed_requests == 0
@pytest.mark.covers("mgmt.team.daily_activity.missing_start_date_rejected")
def test_missing_start_date_is_rejected(self, client: SpendClient) -> None:

View file

@ -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

View file

@ -4,12 +4,20 @@ these carry no `e2e` marker and run everywhere."""
from __future__ import annotations
import os
import signal
import socket
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 +26,8 @@ from idp import (
KEYCLOAK_ADMIN_USER_ENV,
KEYCLOAK_REALM_ENV,
KEYCLOAK_URL_ENV,
BrowserClientBody,
Discovery,
Keycloak,
PasswordCredential,
UserCreateBody,
@ -60,24 +70,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 +149,68 @@ def test_partial_provisioning_removes_the_group_when_user_creation_fails() -> No
assert deletions.empty()
@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"], 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,subprocess,sys,time; from pathlib import Path; "
'assert os.environ["GENERIC_CLIENT_SECRET"]; '
'assert os.environ["GENERIC_CLIENT_USE_PKCE"] == "true"; '
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(
[
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=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)
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 +230,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()

View file

@ -31,6 +31,7 @@ from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from types import MappingProxyType
from typing import Final
import pytest
@ -81,9 +82,14 @@ def json_object(body: bytes) -> dict[str, object]:
class _FakeProvider(ThreadingHTTPServer):
daemon_threads = True
def __init__(self, bind: tuple[str, int]) -> None:
def __init__(self, bind: tuple[str, int], *, echo_request: bool = True) -> None:
super().__init__(bind, _FakeProviderHandler)
self.hits: list[str] = []
self.echo_request = echo_request
self.requests: tuple[tuple[Mapping[str, str], bytes], ...] = ()
def capture_request(self, headers: Mapping[str, str], body: bytes) -> None:
self.requests = (*self.requests, (MappingProxyType(dict(headers)), body))
class _FakeProviderHandler(BaseHTTPRequestHandler):
@ -101,8 +107,11 @@ class _FakeProviderHandler(BaseHTTPRequestHandler):
length = int(self.headers.get("content-length") or "0")
body = self.rfile.read(length) if length else b""
provider.hits.append(f"{self.command} {self.path}")
payload = json.dumps(
provider.capture_request(dict(self.headers.items()), body)
payload: Final = json.dumps(
{"echo": body.decode("utf-8"), "path": self.path, "hit": len(provider.hits)}
if provider.echo_request
else {"ok": True}
).encode()
self.send_response(200)
self.send_header("content-type", "application/json")
@ -117,8 +126,8 @@ class _FakeProviderHandler(BaseHTTPRequestHandler):
@contextmanager
def fake_provider() -> Generator[_FakeProvider]:
server = _FakeProvider(("127.0.0.1", 0))
def fake_provider(*, echo_request: bool = True) -> Generator[_FakeProvider]:
server = _FakeProvider(("127.0.0.1", 0), echo_request=echo_request)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:

View file

@ -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,115 @@ 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_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)
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 +418,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()

View file

@ -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,

30
tests/e2e/ui/oidcSetup.ts Normal file
View file

@ -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();
}
}

View file

@ -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",
},
});

View file

@ -11,7 +11,7 @@ import litellm
from unittest.mock import patch, MagicMock, AsyncMock
from create_mock_standard_logging_payload import create_standard_logging_payload
from litellm.types.utils import StandardLoggingPayload
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS
@ -630,10 +630,12 @@ def test_deployment_callback_respects_cooldown_time(model_list):
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_log_retry(model_list, metadata_key):
"""log_retry appends one flat record per failed attempt and copies neither the request kwargs nor
the request metadata into it"""
def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> None:
"""log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the
request metadata into it, counts every failed attempt of the request independently of the
per-hop attempted_retries, and never trusts a negative count planted before the first failure"""
router = Router(model_list=model_list)
rate_limit_error = litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo")
new_kwargs = router.log_retry(
kwargs={
"model": "gpt-3.5-turbo",
@ -641,7 +643,7 @@ def test_log_retry(model_list, metadata_key):
"messages": [{"role": "user", "content": "hi"}],
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
},
e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"),
e=rate_limit_error,
)
assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [
{
@ -652,6 +654,10 @@ def test_log_retry(model_list, metadata_key):
"attempted_retries": 2,
}
]
assert new_kwargs[metadata_key]["request_retry_count"] == 1
assert router.log_retry(kwargs=new_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 2
planted_kwargs = {"model": "gpt-3.5-turbo", metadata_key: {"request_retry_count": -100}}
assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1
def test_update_usage(model_list):

View file

@ -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

View file

@ -6,7 +6,8 @@ never rewrite. It is consumed by compress() and by the Headroom guardrail, so
the two agree on what "never compress this" means.
"""
from litellm.compression.compress import get_protected_indices
from litellm.compression.compress import compress, get_protected_indices
from litellm.types.utils import CallTypes
def test_protects_system_last_user_and_last_assistant():
@ -53,3 +54,94 @@ def test_every_system_row_is_protected():
def test_no_user_or_assistant_rows():
assert sorted(get_protected_indices([{"role": "system", "content": "sys"}])) == [0]
assert get_protected_indices([]) == ()
def test_mid_history_cache_control_part_is_protected():
# A large cached tool result from a few turns back, not the last user or
# last assistant row -- exactly the row a provider prompt-cache pins to
# exact bytes. Rewriting it (even leaving the marker on) changes those
# bytes and turns the next request's cache read into a cache write.
messages = [
{"role": "user", "content": "old question"},
{"role": "assistant", "content": "old answer"},
{
"role": "user",
"content": [
{"type": "text", "text": "a large cached tool result", "cache_control": {"type": "ephemeral"}},
],
},
{"role": "assistant", "content": "ack"},
{"role": "user", "content": "live instruction"},
]
# index 3 = last assistant, index 4 = last user (both protected by role
# regardless), index 2 = the cache_control-marked row itself.
assert sorted(get_protected_indices(messages)) == [2, 3, 4]
def test_cache_control_directly_on_message_is_protected():
messages = [
{"role": "user", "content": "old question", "cache_control": {"type": "ephemeral"}},
{"role": "assistant", "content": "old answer"},
{"role": "user", "content": "live instruction"},
]
assert sorted(get_protected_indices(messages)) == [0, 1, 2]
def test_cache_control_protection_does_not_duplicate_already_protected_rows():
# The last user row is already protected by role; marking it too must not
# produce a duplicate index.
messages = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "live", "cache_control": {"type": "ephemeral"}},
]
protected = get_protected_indices(messages)
assert sorted(protected) == [0, 1]
assert len(protected) == len(set(protected))
def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control():
# Defensive: a plain string content, or a list of non-dict items, must not
# raise or be misread as carrying a breakpoint.
messages = [
{"role": "assistant", "content": "plain string content"},
{"role": "user", "content": ["not", "a", "dict", "list"]},
{"role": "user", "content": "live instruction"},
]
assert sorted(get_protected_indices(messages)) == [0, 2]
def test_compress_keeps_part_level_cache_control_row_verbatim():
# compress() scores text-only copies of the rows, where a part-level marker
# is gone; protection has to read the original rows or the pinned row is stubbed.
stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]}
pinned = {
"role": "user",
"content": [
{"type": "text", "text": "cached tool result " * 2000, "cache_control": {"type": "ephemeral"}},
],
}
messages = [
stale_log,
{"role": "assistant", "content": "old answer"},
pinned,
{"role": "assistant", "content": "ack"},
{"role": "user", "content": "live instruction"},
]
result = compress(
messages,
model="gpt-4o",
call_type=CallTypes.anthropic_messages,
compression_trigger=1000,
compression_target=500,
)
assert len(result["messages"]) == len(messages)
assert result["messages"][2] == pinned
assert result["messages"][0] != stale_log
assert len(result["cache"]) >= 1

View file

@ -2,6 +2,7 @@ import json
from datetime import datetime, timezone
import pytest
from collections.abc import Mapping
from fastapi.testclient import TestClient
import litellm
@ -74,6 +75,108 @@ 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",
"supported_openai_params": None,
}
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_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
@ -4039,7 +4142,7 @@ def test_billed_token_rates_follow_the_token_tier_the_breakdown_bills_at(monkeyp
cache_read_input_token_cost=6e-7,
cache_read_input_audio_token_cost=6e-7,
cache_creation_input_token_cost=7.5e-6,
cache_creation_input_token_cost_above_1hr=0.0,
cache_creation_input_token_cost_above_1hr=7.5e-6,
output_cost_per_reasoning_token=3e-5,
)
assert breakdown.cache_read_cost == pytest.approx(200_000 * rates.cache_read_input_token_cost)
@ -5334,3 +5437,72 @@ def test_realtime_models_bill_cached_text_and_audio_at_their_cache_read_rates(
prompt_cost, _ = generic_cost_per_token(model=model, usage=usage, custom_llm_provider=custom_llm_provider)
assert prompt_cost == pytest.approx(expected_prompt_cost)
def test_generic_cost_per_token_bills_cache_creation_at_the_input_rate_without_a_write_price():
"""Azure and OpenAI publish no cache-write price and bill cache writes as ordinary input.
A deployment priced with only input, output, and cache-read rates must bill the creation
tokens the provider reports at the input rate, never at 0. The numbers are a cold 7,336-token
prompt on a deployment that reports all but 3 of them as cache creation."""
model_info = {
"input_cost_per_token": 2e-7,
"output_cost_per_token": 1.25e-6,
"cache_read_input_token_cost": 2e-8,
}
usage = Usage(
prompt_tokens=7336,
completion_tokens=23,
total_tokens=7359,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=0, cache_creation_tokens=7333),
)
prompt_cost, completion_cost = generic_cost_per_token(
model="custom-priced-deployment", usage=usage, custom_llm_provider="azure", model_info=model_info
)
assert prompt_cost == pytest.approx(7336 * 2e-7)
assert completion_cost == pytest.approx(23 * 1.25e-6)
@pytest.mark.parametrize(
("cache_rates", "current_time", "expected_creation", "expected_creation_1h"),
(
pytest.param({}, None, 2e-7, 2e-7, id="no-write-price-uses-the-input-rate"),
pytest.param({"cache_creation_input_token_cost": 2.5e-7}, None, 2.5e-7, 2.5e-7, id="no-1h-price-uses-the-write-price"),
pytest.param({"cache_creation_input_token_cost": 0.0}, None, 0.0, 0.0, id="explicit-zero-stays-zero"),
pytest.param(
{"off_peak_pricing": {"hours_utc": "00:00-23:59", "input_cost_per_token": 1e-7}},
datetime(2026, 9, 14, 12, tzinfo=timezone.utc),
1e-7,
1e-7,
id="no-write-price-uses-the-off-peak-input-rate",
),
pytest.param(
{
"off_peak_pricing": {
"hours_utc": "00:00-23:59",
"input_cost_per_token": 1e-7,
"cache_creation_input_token_cost": 3e-7,
}
},
datetime(2026, 9, 14, 12, tzinfo=timezone.utc),
3e-7,
3e-7,
id="no-1h-price-uses-the-off-peak-write-price",
),
),
)
def test_get_token_base_cost_resolves_missing_cache_write_rates_like_the_tiered_path(
cache_rates: Mapping[str, float | Mapping[str, float | str]],
current_time: datetime | None,
expected_creation: float,
expected_creation_1h: float,
) -> None:
model_info = {"input_cost_per_token": 2e-7, "output_cost_per_token": 1.25e-6, **cache_rates}
usage = Usage(prompt_tokens=10, completion_tokens=1, total_tokens=11)
_, _, creation, creation_1h, _ = _get_token_base_cost(model_info, usage, current_time=current_time)
assert creation == pytest.approx(expected_creation)
assert creation_1h == pytest.approx(expected_creation_1h)

View file

@ -1,4 +1,3 @@
import httpx
import openai
import pytest
@ -178,9 +177,7 @@ class TestExceptionCheckers:
]
for error_str in error_strings:
result = ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
)
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
assert result is True, f"Should detect policy violation in: {error_str}"
def test_is_azure_content_policy_violation_error_case_insensitive(self):
@ -194,12 +191,8 @@ class TestExceptionCheckers:
]
for error_str in error_strings:
result = ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
)
assert (
result is True
), f"Should detect policy violation in uppercase: {error_str}"
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
assert result is True, f"Should detect policy violation in uppercase: {error_str}"
def test_is_azure_content_policy_violation_error_with_non_policy_errors(self):
"""Test that non-policy violation errors are not detected as policy violations"""
@ -216,12 +209,8 @@ class TestExceptionCheckers:
]
for error_str in error_strings:
result = ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
)
assert (
result is False
), f"Should NOT detect policy violation in: {error_str}"
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
assert result is False, f"Should NOT detect policy violation in: {error_str}"
def test_is_azure_content_policy_violation_error_with_partial_matches(self):
"""Test that partial keyword matches work correctly"""
@ -234,9 +223,7 @@ class TestExceptionCheckers:
]
for error_str in positive_cases:
result = ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
)
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
assert result is True, f"Should detect policy violation in: {error_str}"
# These should not match even though they contain similar words
@ -248,12 +235,8 @@ class TestExceptionCheckers:
]
for error_str in negative_cases:
result = ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
)
assert (
result is False
), f"Should NOT detect policy violation in: {error_str}"
result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str)
assert result is False, f"Should NOT detect policy violation in: {error_str}"
gemini_context_window_test_cases = [
@ -271,12 +254,8 @@ gemini_context_window_test_cases = [
]
@pytest.mark.parametrize(
"error_message, should_raise_context_window", gemini_context_window_test_cases
)
def test_gemini_context_window_error_mapping(
error_message, should_raise_context_window
):
@pytest.mark.parametrize("error_message, should_raise_context_window", gemini_context_window_test_cases)
def test_gemini_context_window_error_mapping(error_message, should_raise_context_window):
"""
Tests that the exception_type function correctly maps Gemini's
context window exceeded errors to litellm.ContextWindowExceededError.
@ -421,9 +400,7 @@ vertex_rate_limit_test_cases = [
]
@pytest.mark.parametrize(
"error_message, should_raise_rate_limit", vertex_rate_limit_test_cases
)
@pytest.mark.parametrize("error_message, should_raise_rate_limit", vertex_rate_limit_test_cases)
def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_limit):
"""
Tests that the exception_type function correctly maps Vertex AI's
@ -458,10 +435,7 @@ class TestGetBodyErrorCode:
"""Unit tests for _get_body_error_code helper."""
def test_parses_int_code(self):
body = (
'{"error":{"message":"high demand","type":"upstream_error",'
'"param":"","code":429}}'
)
body = '{"error":{"message":"high demand","type":"upstream_error","param":"","code":429}}'
assert _get_body_error_code(body) == 429
def test_parses_string_code(self):
@ -498,8 +472,7 @@ gemini_body_code_429_test_cases = [
),
(
503,
'{"error":{"message":"upstream unavailable","type":"upstream_error",'
'"param":"","code":429}}',
'{"error":{"message":"upstream unavailable","type":"upstream_error","param":"","code":429}}',
litellm.RateLimitError,
"HTTP 503 envelope with body code:429 -> RateLimitError",
),
@ -769,9 +742,7 @@ class _UpstreamHTTPError(Exception):
self.message = "upstream failure"
self.status_code = status_code
self.request = httpx.Request("POST", "https://api.example.com/v1/chat/completions")
self.response = httpx.Response(
status_code=status_code, request=self.request, text="upstream failure"
)
self.response = httpx.Response(status_code=status_code, request=self.request, text="upstream failure")
UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503)
@ -892,15 +863,13 @@ PROVIDERS_WITHOUT_A_HANDLER = tuple(
MINIMAX_401_BODY = (
'{"type":"error","error":{"type":"authorized_error","message":"login fail: Please carry the API secret key '
"in the 'Authorization' field of the request header (1004)\",\"http_code\":\"401\"},"
'in the \'Authorization\' field of the request header (1004)","http_code":"401"},'
'"request_id":"06ddc9ba97ee6340e38f10e09787f547"}'
)
def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]:
return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(
status_code, OPENAI_SHAPED[status_code]
)
return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(status_code, OPENAI_SHAPED[status_code])
@pytest.fixture
@ -910,9 +879,7 @@ def quiet_exception_mapping(monkeypatch):
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
def test_an_upstream_status_maps_to_one_exception_per_provider(
provider, status_code, quiet_exception_mapping
):
def test_an_upstream_status_maps_to_one_exception_per_provider(provider, status_code, quiet_exception_mapping):
expected_class, expected_status = _expected_for(provider, status_code)
with pytest.raises(openai.APIError) as raised:
@ -928,9 +895,7 @@ def test_an_upstream_status_maps_to_one_exception_per_provider(
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(
provider, status_code, quiet_exception_mapping
):
def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(provider, status_code, quiet_exception_mapping):
with pytest.raises(openai.APIError) as raised:
exception_type(
model="test-model",
@ -943,12 +908,8 @@ def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
def test_an_already_mapped_litellm_exception_passes_through_untouched(
provider, quiet_exception_mapping
):
already_mapped = litellm.RateLimitError(
message="already mapped", llm_provider=provider, model="test-model"
)
def test_an_already_mapped_litellm_exception_passes_through_untouched(provider, quiet_exception_mapping):
already_mapped = litellm.RateLimitError(message="already mapped", llm_provider=provider, model="test-model")
returned = exception_type(
model="test-model",
@ -961,9 +922,7 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched(
@pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES)
@pytest.mark.parametrize("provider", PROVIDERS_WITHOUT_A_HANDLER)
def test_a_provider_without_a_handler_maps_by_the_upstream_status(
provider, status_code, quiet_exception_mapping
):
def test_a_provider_without_a_handler_maps_by_the_upstream_status(provider, status_code, quiet_exception_mapping):
expected_class, expected_status = STATUS_KEYED[status_code]
with pytest.raises(openai.APIError) as raised:
@ -1015,9 +974,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q
assert "boom" in raised.value.message
def _raise_and_map(
model: str | None, original_exception: Exception, custom_llm_provider: str | None
) -> None:
def _raise_and_map(model: str | None, original_exception: Exception, custom_llm_provider: str | None) -> None:
"""Calls exception_type() from inside the except block, as litellm/main.py does,
so traceback.format_exc() has a real stack."""
try:
@ -1058,9 +1015,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_message_keeps_traceback
CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens."
CONTENT_POLICY_MESSAGE = (
'{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
)
CONTENT_POLICY_MESSAGE = '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}'
TIMEOUT_MESSAGE = "Request timed out."
PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = (
@ -1103,15 +1058,11 @@ class _UpstreamErrorWithMessage(_UpstreamHTTPError):
super().__init__(status_code=status_code)
self.args = (message,)
self.message = message
self.response = httpx.Response(
status_code=status_code, request=self.request, text=message
)
self.response = httpx.Response(status_code=status_code, request=self.request, text=message)
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(
provider, quiet_exception_mapping
):
def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping):
if provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW:
expected_class, expected_status = litellm.ContextWindowExceededError, 400
else:
@ -1129,9 +1080,7 @@ def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(
provider, quiet_exception_mapping
):
def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping):
if provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK:
expected_class, expected_status = litellm.ContentPolicyViolationError, 400
else:
@ -1149,9 +1098,7 @@ def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(
@pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER)
def test_a_timed_out_request_is_a_timeout_for_every_provider(
provider, quiet_exception_mapping
):
def test_a_timed_out_request_is_a_timeout_for_every_provider(provider, quiet_exception_mapping):
with pytest.raises(litellm.Timeout) as raised:
exception_type(
model="test-model",
@ -1409,3 +1356,97 @@ def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code):
exception_headers = _get_response_headers(original_exception=exc_info.value)
assert exception_headers is not None
assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7
_GUARDRAIL_BLOCK_ERROR = {
"message": "Content blocked: secret_project_codename pattern detected",
"param": "None",
"code": "400",
"provider_specific_fields": {
"error": "Content blocked: secret_project_codename pattern detected",
"pattern": "secret_project_codename",
"guardrail_name": "block-secret-project",
"guardrail_mode": "pre_call",
},
}
def _openai_handler_error(
error_type: str,
headers: dict[str, str] | list[tuple[str, str]],
status_code: int = 400,
message: str = _GUARDRAIL_BLOCK_ERROR["message"],
) -> OpenAIError:
wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message}
return OpenAIError(
status_code=status_code,
message=f"Error code: {status_code} - {{'error': {wire_error}}}",
headers=httpx.Headers(headers),
body=wire_error,
)
_PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"}
@pytest.mark.parametrize(("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)])
def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int):
with pytest.raises(litellm.BadRequestError) as exc_info:
exception_type(
model="claude-haiku-4-5",
original_exception=_openai_handler_error(error_type, _PROXY_HEADERS, status_code=status_code),
custom_llm_provider="litellm_proxy",
completion_kwargs={},
extra_kwargs={},
)
assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project"
assert exc_info.value.body["type"] == error_type
assert dict(exc_info.value.response.headers) == _PROXY_HEADERS
@pytest.mark.parametrize("relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError])
def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]):
message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}"
with pytest.raises(relayed_class) as exc_info:
exception_type(
model="claude-haiku-4-5",
original_exception=_openai_handler_error("None", _PROXY_HEADERS, message=message),
custom_llm_provider="litellm_proxy",
completion_kwargs={},
extra_kwargs={},
)
assert type(exc_info.value) is relayed_class
assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project"
assert dict(exc_info.value.response.headers) == _PROXY_HEADERS
def test_openai_compatible_vendor_400_keeps_body_but_not_headers():
with pytest.raises(litellm.BadRequestError) as exc_info:
exception_type(
model="gpt-5.4-mini",
original_exception=_openai_handler_error("vendor_specific_error", {"openai-organization": "org-1"}),
custom_llm_provider="openai",
completion_kwargs={},
extra_kwargs={},
)
assert exc_info.value.body["type"] == "vendor_specific_error"
assert not exc_info.value.response.headers
def test_litellm_proxy_repeated_response_header_keeps_each_value():
repeated = [("x-litellm-call-id", "call-guardrail"), ("set-cookie", "a=1"), ("set-cookie", "b=2")]
with pytest.raises(litellm.BadRequestError) as exc_info:
exception_type(
model="claude-haiku-4-5",
original_exception=_openai_handler_error("None", repeated),
custom_llm_provider="litellm_proxy",
completion_kwargs={},
extra_kwargs={},
)
assert exc_info.value.response.headers.multi_items() == repeated

View file

@ -13,6 +13,7 @@ import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
AnthropicMessagesHandler,
@ -635,14 +636,19 @@ class TestAnthropicMessagesHandlerInputProcessing:
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.inputs is not None
assert guardrail.inputs["texts"] == ["safe text", "prohibited correction"]
assert guardrail.inputs["texts"] == [
"trusted top-level system prompt",
"safe text",
"prohibited correction",
]
structured = guardrail.inputs["structured_messages"]
assert [m["role"] for m in structured] == ["system", "user", "system"]
assert structured[0]["content"] == "trusted top-level system prompt"
assert data["system"] == "trusted top-level system prompt"
assert data["messages"][1]["content"] == "[MASKED]"
@pytest.mark.asyncio
async def test_bedrock_masking_slice_is_unavailable_when_top_level_system_is_included(
async def test_bedrock_masking_slice_lines_up_when_top_level_system_is_included(
self,
):
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
@ -668,25 +674,25 @@ class TestAnthropicMessagesHandlerInputProcessing:
structured = guardrail.inputs["structured_messages"]
bedrock = BedrockGuardrail(guardrailIdentifier="gi", guardrailVersion="1")
assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts) + 1
assert sum(bedrock._count_message_texts(m) for m in structured) == len(texts)
latest_user_index = bedrock._find_latest_message_index(structured, target_role="user")
assert (
bedrock._locate_message_texts_slice(
structured_messages=structured,
target_index=latest_user_index,
texts=texts,
)
is None
)
assert (
bedrock._merge_masked_texts(
masked_texts=["{MASKED}"],
texts=texts,
scanned_slice=None,
scanned_role_subset=True,
)
== texts
scanned_slice = bedrock._locate_message_texts_slice(
structured_messages=structured,
target_index=latest_user_index,
texts=texts,
)
assert scanned_slice == (3, 1)
assert bedrock._merge_masked_texts(
masked_texts=["{MASKED}"],
texts=texts,
scanned_slice=scanned_slice,
scanned_role_subset=True,
) == [
"trusted top-level system prompt",
"safe text",
"prohibited correction",
"{MASKED}",
]
@pytest.mark.asyncio
@pytest.mark.parametrize("skip_system_message_in_guardrail", [True, None])
@ -1611,7 +1617,8 @@ class TestAnthropicMessagesIncrementalScan:
)
assert mock_api.call_count == 1
assert [m["content"] for m in mock_api.call_args.kwargs["messages"]] == [
"What is the capital of France?"
"You are a helpful geography assistant.",
"What is the capital of France?",
]
mock_api.reset_mock()
await handler.process_input_messages(
@ -2150,6 +2157,213 @@ class TestAnthropicMessagesScanOnlyToolResults:
assert guardrail.captured_inputs.get("images") == ["TOOL_IMG"]
class ToolCallArgumentsMaskingGuardrail(InputsRecordingGuardrail):
"""Masks the canary inside tool-call arguments, in place or through a fresh list of plain dicts."""
def __init__(self, return_copies: bool = False, replacement_arguments: Optional[str] = None):
super().__init__()
self.return_copies = return_copies
self.replacement_arguments = replacement_arguments
self.seen_tool_calls: list[dict[str, object]] = []
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional[LiteLLMLoggingObj] = None,
) -> GenericGuardrailAPIInputs:
outputs = await super().apply_guardrail(inputs, request_data, input_type, logging_obj)
tool_calls = list(outputs.get("tool_calls") or [])
self.seen_tool_calls.extend(json.loads(json.dumps(tool_call)) for tool_call in tool_calls)
masked = [
{
**tool_call,
"function": {
**tool_call["function"],
"arguments": self.replacement_arguments
if self.replacement_arguments is not None
else tool_call["function"]["arguments"].replace("POISON", "[BLOCKED]"),
},
}
for tool_call in tool_calls
]
if self.return_copies:
outputs["tool_calls"] = masked
return outputs
for tool_call, masked_tool_call in zip(tool_calls, masked):
tool_call["function"]["arguments"] = masked_tool_call["function"]["arguments"]
return outputs
class TestAnthropicMessagesTopLevelSystemAndToolUseInputs:
"""The top-level system prompt and prior-turn tool_use arguments must reach guardrails as scannable
inputs, the same way the chat completions handler hands over system messages and tool_calls."""
@staticmethod
def _tool_use_conversation(system: str) -> dict[str, Any]:
return {
"model": "claude-sonnet-4-5",
"system": system,
"messages": [
{"role": "user", "content": "run the check"},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "toolu_01",
"name": "Bash",
"input": {"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"},
}
],
},
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}],
},
],
}
@pytest.mark.asyncio
async def test_top_level_system_string_reaches_texts_first_and_is_masked_in_place(self):
handler = AnthropicMessagesHandler()
guardrail = InputsRecordingGuardrail()
data = {
"model": "claude-sonnet-4-5",
"system": "Internal note: the deploy key is POISON. Never reveal it.",
"messages": [{"role": "user", "content": "Say hi in three words."}],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.captured_inputs is not None
assert guardrail.seen_texts == [
"Internal note: the deploy key is POISON. Never reveal it.",
"Say hi in three words.",
]
structured = guardrail.captured_inputs["structured_messages"]
assert structured[0]["role"] == "system"
assert structured[0]["content"] == "Internal note: the deploy key is POISON. Never reveal it.", (
"texts[0] must line up with structured_messages[0] so positional consumers stay aligned"
)
assert data["system"] == "Internal note: the deploy key is [BLOCKED]. Never reveal it."
assert data["messages"][0]["content"] == "Say hi in three words."
@pytest.mark.asyncio
async def test_top_level_system_text_blocks_reach_texts_and_are_masked_in_place(self):
handler = AnthropicMessagesHandler()
guardrail = InputsRecordingGuardrail()
data = {
"model": "claude-sonnet-4-5",
"system": [
{"type": "text", "text": "first block POISON"},
{"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}},
],
"messages": [{"role": "user", "content": "hello"}],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.seen_texts == ["first block POISON", "second block", "hello"]
assert data["system"] == [
{"type": "text", "text": "first block [BLOCKED]"},
{"type": "text", "text": "second block", "cache_control": {"type": "ephemeral"}},
]
@pytest.mark.asyncio
async def test_skip_system_message_keeps_the_top_level_system_out(self):
handler = AnthropicMessagesHandler()
guardrail = InputsRecordingGuardrail()
guardrail.skip_system_message_in_guardrail = True
data = {
"model": "claude-sonnet-4-5",
"system": "trusted POISON prompt",
"messages": [{"role": "user", "content": "hello"}],
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.seen_texts == ["hello"]
assert data["system"] == "trusted POISON prompt"
@pytest.mark.asyncio
async def test_prior_turn_tool_use_input_reaches_tool_calls_in_openai_shape(self):
handler = AnthropicMessagesHandler()
guardrail = InputsRecordingGuardrail()
data = self._tool_use_conversation(system="You are a careful agent harness.")
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.captured_inputs is not None
tool_calls = guardrail.captured_inputs.get("tool_calls")
assert tool_calls is not None and len(tool_calls) == 1
assert tool_calls[0]["id"] == "toolu_01"
assert tool_calls[0]["type"] == "function"
assert tool_calls[0]["function"]["name"] == "Bash"
assert json.loads(tool_calls[0]["function"]["arguments"]) == {
"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"
}
assert data["messages"][1]["content"][0]["input"] == {
"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"
}, "a guardrail that leaves tool_calls alone must leave the tool_use input alone"
@pytest.mark.asyncio
@pytest.mark.parametrize("return_copies", [False, True])
async def test_masked_tool_call_arguments_write_back_into_the_tool_use_input(self, return_copies: bool):
handler = AnthropicMessagesHandler()
guardrail = ToolCallArgumentsMaskingGuardrail(return_copies=return_copies)
data = self._tool_use_conversation(system="You are a careful agent harness.")
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert [tool_call["function"]["name"] for tool_call in guardrail.seen_tool_calls] == ["Bash"]
tool_use = data["messages"][1]["content"][0]
assert tool_use == {
"type": "tool_use",
"id": "toolu_01",
"name": "Bash",
"input": {"cmd": "AWS_ACCESS_KEY_ID=[BLOCKED] aws sts get-caller-identity"},
}
assert data["messages"][2]["content"][0]["tool_use_id"] == "toolu_01"
@pytest.mark.asyncio
async def test_non_json_rewritten_arguments_are_rejected_by_name(self):
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
handler = AnthropicMessagesHandler()
guardrail = ToolCallArgumentsMaskingGuardrail(replacement_arguments="[REDACTED]")
data = self._tool_use_conversation(system="Internal note: the deploy key is POISON. Never reveal it.")
data["messages"][2]["content"][0]["content"] = "fetched POISON page"
original = json.loads(json.dumps(data))
with pytest.raises(UnappliableRequestRewrite) as excinfo:
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert excinfo.value.guardrail_name == "scan-only-capture"
assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched"
assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched"
@pytest.mark.asyncio
async def test_scan_only_tool_results_keeps_system_and_tool_use_out(self):
handler = AnthropicMessagesHandler()
guardrail = InputsRecordingGuardrail()
guardrail.scan_only_tool_results = True
data = self._tool_use_conversation(system="trusted POISON prompt")
data["messages"][2]["content"][0]["content"] = "fetched POISON page"
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.seen_texts == ["fetched POISON page"]
assert guardrail.captured_inputs is not None
assert guardrail.captured_inputs.get("tool_calls") is None
assert data["system"] == "trusted POISON prompt"
assert data["messages"][1]["content"][0]["input"] == {
"cmd": "AWS_ACCESS_KEY_ID=POISON aws sts get-caller-identity"
}
assert data["messages"][2]["content"][0]["content"] == "fetched [BLOCKED] page"
class TestStructuredWriteBackKeepsToolResults:
"""A guardrail rewrite must never leave a tool_use without its tool_result (Claude Code ToolSearch, LIT-6103)."""
@ -2272,6 +2486,116 @@ class TestAnthropicMessagesHandlerStreamingScanKey:
assert ended_key != open_key
class PerRowTextGuardrail(CustomGuardrail):
"""Answers one redacted text per chat row it was shown, the way a guardrail
that scans per message does, and hands back only texts."""
def __init__(self):
super().__init__(guardrail_name="per-row-redactor")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
rows = inputs.get("structured_messages") or []
return {**inputs, "texts": [str(row.get("content")).replace("123-45-6789", "<US_SSN>") for row in rows]}
class PerSlotTextGuardrail(CustomGuardrail):
"""Answers one redacted text per text slot of every chat row it was shown, the
way a guardrail that counts slots per message does, and hands back only texts."""
def __init__(self):
super().__init__(guardrail_name="per-slot-redactor")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts
rows = inputs.get("structured_messages") or []
return {
**inputs,
"texts": [text.replace("123-45-6789", "<US_SSN>") for row in rows for text in message_slot_texts(row)],
}
class TestPerMessageTextWriteBack:
"""Texts that no longer pair one-to-one with what the handler extracted must be
rejected by name instead of sliding onto the wrong messages."""
@pytest.mark.asyncio
async def test_one_text_per_row_over_a_system_prompt_is_applied(self):
data = {
"model": "claude-sonnet-4-5",
"system": "Reply with exactly the SSN you were given.",
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
}
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
assert data["system"] == "Reply with exactly the SSN you were given."
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
@pytest.mark.asyncio
async def test_one_text_per_row_over_a_multi_block_system_prompt_is_rejected_by_name(self):
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
data = {
"model": "claude-sonnet-4-5",
"system": [
{"type": "text", "text": "Reply with exactly the SSN you were given."},
{"type": "text", "text": "Never apologize."},
],
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
}
original = json.loads(json.dumps(data))
with pytest.raises(UnappliableRequestRewrite) as excinfo:
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
assert excinfo.value.guardrail_name == "per-row-redactor"
assert data["system"] == original["system"], "a rejected rewrite must leave the request untouched"
assert data["messages"] == original["messages"], "a rejected rewrite must leave the request untouched"
@pytest.mark.asyncio
async def test_one_text_per_slot_over_a_system_prompt_with_an_empty_block_is_applied(self):
data = {
"model": "claude-sonnet-4-5",
"system": [
{"type": "text", "text": ""},
{"type": "text", "text": "Reply with exactly the SSN you were given."},
],
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
}
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerSlotTextGuardrail())
assert data["system"] == [
{"type": "text", "text": ""},
{"type": "text", "text": "Reply with exactly the SSN you were given."},
]
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
@pytest.mark.asyncio
async def test_one_text_per_row_without_a_system_prompt_is_applied(self):
data = {
"model": "claude-sonnet-4-5",
"messages": [{"role": "user", "content": "My SSN is 123-45-6789."}],
}
await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=PerRowTextGuardrail())
assert data["messages"] == [{"role": "user", "content": "My SSN is <US_SSN>."}]
class TestAnthropicMessagesHandlerPostCallHookResponse:
def test_openai_shaped_stream_assembly_reaches_the_hook_as_a_messages_response(self):
from litellm.types.utils import Choices, Message, ModelResponse, Usage

View file

@ -8,6 +8,8 @@ Without the fix, the AnthropicStreamWrapper silently dropped these
arguments, causing tool_use blocks to arrive with empty input {}.
"""
import json
from typing import List
from unittest.mock import MagicMock
@ -139,9 +141,7 @@ async def test_async_stream_emits_input_json_delta_for_bundled_tool_args():
# Verify the delta carries the tool arguments
delta_event = events[input_json_delta_idx]
assert delta_event["delta"][
"partial_json"
], "input_json_delta should have non-empty partial_json"
assert json.loads(delta_event["delta"]["partial_json"]) == {"location": "Boston"}
@pytest.mark.asyncio
@ -300,7 +300,7 @@ def test_sync_stream_emits_input_json_delta_for_bundled_tool_args():
assert (
input_json_delta_idx == tool_start_idx + 1
), "input_json_delta should immediately follow the tool_use content_block_start"
assert events[input_json_delta_idx]["delta"]["partial_json"]
assert json.loads(events[input_json_delta_idx]["delta"]["partial_json"]) == {"location": "Boston"}
def test_sync_stream_no_extra_delta_when_tool_args_empty():

View file

@ -773,6 +773,12 @@ class TestBedrockMantleCodexAdditionalTools:
assert body["input"] == codex_agentic_items
assert "tools" not in body
def test_input_without_additional_tools_sanitizes_tools_on_the_caller_params_object(self):
params = {"tools": [{"type": "function", "name": "wait", "parameters": '{"type": "object"}'}]}
body = self._transform(input=[self._USER_MESSAGE], params=params)
assert body["tools"][0]["parameters"] == {"type": "object"}
assert params["tools"][0]["parameters"] == {"type": "object"}
def test_malformed_additional_tools_item_without_tools_list_is_stripped(self):
body = self._transform(
input=[

View file

@ -1893,6 +1893,183 @@ class TestScanOnlyToolResults:
assert data["messages"][4]["content"] == "and then?"
class TestNoScannableContentRecordsNotRun:
"""LIT-6314: a guardrail whose scoping leaves nothing to scan must still persist an evaluation record"""
def _system_only_data(self) -> dict:
return {"messages": [{"role": "system", "content": "SYSTEM-PROMPT"}]}
def _recorded_entries(self, data: dict) -> list:
metadata = data.get("metadata") or data.get("litellm_metadata") or {}
return metadata.get("standard_logging_guardrail_information") or []
@pytest.mark.asyncio
async def test_skipped_scan_records_not_run_entry(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="skip-system-guardrail")
guardrail.skip_system_message_in_guardrail = True
data = self._system_only_data()
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.last_inputs is None, "nothing survived scoping, apply_guardrail must not run"
entries = self._recorded_entries(data)
assert len(entries) == 1
assert entries[0]["guardrail_name"] == "skip-system-guardrail"
assert entries[0]["guardrail_status"] == "not_run"
assert entries[0]["guardrail_response"] == "no scannable content after message scoping"
@pytest.mark.asyncio
@pytest.mark.parametrize("skip_system", [False, True])
async def test_empty_content_does_not_blame_scoping(self, skip_system: bool):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="unscoped-guardrail")
guardrail.skip_system_message_in_guardrail = skip_system
data = {"messages": [{"role": "user", "content": None}]}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.last_inputs is None
entries = self._recorded_entries(data)
assert len(entries) == 1
assert entries[0]["guardrail_status"] == "not_run"
assert entries[0]["guardrail_response"] == "no scannable content"
@pytest.mark.asyncio
async def test_self_recording_guardrail_is_left_alone(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="self-recording-guardrail")
guardrail.skip_system_message_in_guardrail = True
guardrail.records_own_guardrail_information = True
data = self._system_only_data()
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.last_inputs is None
assert self._recorded_entries(data) == []
@pytest.mark.asyncio
async def test_scannable_content_records_no_extra_entry(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="normal-guardrail")
data = {"messages": [{"role": "user", "content": "hello"}]}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.last_inputs is not None
assert all(e.get("guardrail_status") != "not_run" for e in self._recorded_entries(data))
@pytest.mark.asyncio
async def test_image_only_content_is_not_reported_as_not_run(self):
"""Images are only scanned alongside text, so an image-only request is a
pre-existing scan gap, not a message-scoping skip, and must not be labelled one"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="image-guardrail")
guardrail.skip_system_message_in_guardrail = True
data = {
"messages": [
{"role": "system", "content": "SYSTEM-PROMPT"},
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}],
},
]
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert self._recorded_entries(data) == []
@pytest.mark.asyncio
async def test_scoped_out_image_only_message_is_not_reported_as_not_run(self):
"""An image in a skipped role must behave like any other image-only request"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="image-guardrail")
guardrail.skip_system_message_in_guardrail = True
data = {
"messages": [
{
"role": "system",
"content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}],
},
]
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.last_inputs is None
assert self._recorded_entries(data) == []
@pytest.mark.asyncio
async def test_scoped_out_text_with_image_records_not_run(self):
"""Scoping removed text too, so the skip is recorded even though an image sat beside it"""
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="image-guardrail")
guardrail.skip_system_message_in_guardrail = True
data = {
"messages": [
{
"role": "system",
"content": [
{"type": "text", "text": "Describe this picture."},
{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}},
],
},
]
}
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
assert guardrail.last_inputs is None
entries = self._recorded_entries(data)
assert len(entries) == 1
assert entries[0]["guardrail_status"] == "not_run"
assert entries[0]["guardrail_response"] == "no scannable content after message scoping"
class ToolDroppingTextGuardrail(CustomGuardrail):
"""Answers one text per non-tool message it saw, the way a guardrail that
filters tool rows out before scanning does, and hands back only texts."""
def __init__(self):
super().__init__(guardrail_name="tool-dropping-redactor")
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional[Any] = None,
) -> GenericGuardrailAPIInputs:
kept = [m for m in inputs.get("structured_messages") or [] if m.get("role") != "tool"]
return {**inputs, "texts": [str(m.get("content")).replace("POISON", "[BLOCKED]") for m in kept]}
class TestPerMessageTextWriteBack:
"""Texts that no longer pair one-to-one with what the handler extracted must be
rejected by name instead of sliding onto the wrong messages."""
@pytest.mark.asyncio
async def test_fewer_texts_than_extracted_over_a_tool_message_is_rejected(self):
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
handler = OpenAIChatCompletionsHandler()
original_messages = [
{"role": "system", "content": "SYSTEM-PROMPT"},
{"role": "user", "content": "fetch the page"},
{"role": "assistant", "content": "fetching"},
{"role": "tool", "tool_call_id": "call_1", "content": "page says POISON here"},
{"role": "user", "content": "and then?"},
]
data = {"messages": json.loads(json.dumps(original_messages))}
with pytest.raises(UnappliableRequestRewrite) as excinfo:
await handler.process_input_messages(data=data, guardrail_to_apply=ToolDroppingTextGuardrail())
assert excinfo.value.guardrail_name == "tool-dropping-redactor"
assert data["messages"] == original_messages, "a rejected rewrite must leave the request untouched"
class TestBuildBlockSseChunks:
"""build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE chunks"""

View file

@ -8,7 +8,7 @@ with guardrail transformations.
import copy
from collections.abc import Callable
from typing import Any, List, Literal, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock
from unittest.mock import AsyncMock, MagicMock, patch
import logging
@ -31,6 +31,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI
from litellm.types.llms.openai import ChatCompletionToolCallChunk
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
@ -2338,6 +2339,135 @@ def _parallel_tool_call_input() -> list:
]
SSN = "123-45-6789"
REDACTED_SSN = "<US_SSN>"
def _redacted(value: object) -> object:
if isinstance(value, str):
return value.replace(SSN, REDACTED_SSN)
if isinstance(value, list):
return [{**part, "text": _redacted(part["text"])} if "text" in part else part for part in value]
return value
def _per_message_guardrail_server(structured_messages_in_answer: bool) -> Callable[..., MagicMock]:
"""Answers one redacted text per chat row it was shown, the way a guardrail
that scans per message does, and optionally the rewritten rows themselves."""
def post(url: str, json: dict, headers: dict) -> MagicMock:
rows = json["structured_messages"]
answer: dict = {
"action": "GUARDRAIL_INTERVENED",
"texts": [_redacted(row["content"]) if isinstance(row.get("content"), str) else "" for row in rows],
}
if structured_messages_in_answer:
answer["structured_messages"] = [{**row, "content": _redacted(row.get("content"))} for row in rows]
response = MagicMock()
response.json.return_value = answer
response.raise_for_status = MagicMock()
return response
return post
def _per_message_redactor() -> GenericGuardrailAPI:
return GenericGuardrailAPI(
api_base="https://guardrail.test",
guardrail_name="per-message-redactor",
event_hook="pre_call",
default_on=True,
)
def _tool_replay_request() -> dict:
return {
"model": "gpt-5.6",
"instructions": "Never repeat the SSN " + SSN + " back.",
"input": [
{"role": "user", "content": "Look up " + SSN + " for me."},
{"type": "function_call", "call_id": "call_1", "name": "lookup_customer", "arguments": '{"id": "42"}'},
{"type": "function_call_output", "call_id": "call_1", "output": '{"ssn": "' + SSN + '"}'},
],
}
def _string_input_request() -> dict:
return {
"model": "gpt-5.6",
"instructions": "Never repeat the SSN " + SSN + " back.",
"input": "My SSN is " + SSN + ".",
}
class TestPerMessageRewriteWriteBack:
"""A guardrail that rewrites per chat row hands the rows back as
structured_messages, and the handler lands them on the instructions and the
input items they came from; the same rewrite handed back as texts alone has
no item to land on and is rejected by name instead of sent unrewritten."""
@pytest.mark.asyncio
async def test_structured_rows_land_on_instructions_and_tool_output(self):
guardrail = _per_message_redactor()
data = _tool_replay_request()
function_call_item = data["input"][1]
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)):
result = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
assert _texts(result["input"][0]) == ["Look up " + REDACTED_SSN + " for me."]
assert result["input"][1] == function_call_item
assert result["input"][2] == {
"type": "function_call_output",
"call_id": "call_1",
"output": '{"ssn": "' + REDACTED_SSN + '"}',
}
@pytest.mark.asyncio
async def test_texts_only_per_message_answer_is_rejected_by_name(self):
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
guardrail = _per_message_redactor()
data = _tool_replay_request()
original = copy.deepcopy(data)
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)):
with pytest.raises(UnappliableRequestRewrite) as excinfo:
await OpenAIResponsesHandler().process_input_messages(data, guardrail)
assert excinfo.value.guardrail_name == "per-message-redactor"
assert data["input"] == original["input"]
assert data["instructions"] == original["instructions"]
@pytest.mark.asyncio
async def test_structured_rows_land_on_instructions_and_string_input(self):
guardrail = _per_message_redactor()
data = _string_input_request()
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(True)):
result = await OpenAIResponsesHandler().process_input_messages(data, guardrail)
assert result["instructions"] == "Never repeat the SSN " + REDACTED_SSN + " back."
assert [_texts(item) for item in result["input"]] == [["My SSN is " + REDACTED_SSN + "."]]
@pytest.mark.asyncio
async def test_texts_only_per_message_answer_over_a_string_input_is_rejected_by_name(self):
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
guardrail = _per_message_redactor()
data = _string_input_request()
original = copy.deepcopy(data)
with patch.object(guardrail.async_handler, "post", side_effect=_per_message_guardrail_server(False)):
with pytest.raises(UnappliableRequestRewrite) as excinfo:
await OpenAIResponsesHandler().process_input_messages(data, guardrail)
assert excinfo.value.guardrail_name == "per-message-redactor"
assert data["input"] == original["input"]
assert data["instructions"] == original["instructions"]
class TestProvenancePatching:
"""The O(n) provenance pass must keep patching rewritten rows in place for the
shapes real agent loops produce, and fall back safely everywhere else."""

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