Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_python_version_ci

This commit is contained in:
mateo-berri 2026-09-02 18:30:13 -07:00
commit 85961201e7
34 changed files with 1546 additions and 73 deletions

View file

@ -128,6 +128,9 @@ jobs:
- name: check_fastuuid_usage
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
- name: check_py310_typing_imports
run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py
- name: check_e2e_no_raw_requests
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
@ -145,3 +148,33 @@ jobs:
- name: documentation_test_api_docs
run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
python-310-import-smoke:
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.10"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: uv sync --frozen --extra proxy --python 3.10
- run: uv run --no-sync python --version
- name: Import litellm
run: uv run --no-sync python -c "import litellm"
- name: Check litellm CLI
run: uv run --no-sync litellm --version

View file

@ -1449,6 +1449,7 @@ RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
"Truncation is a DB storage safeguard. "

View file

@ -17,7 +17,10 @@ from typing import TYPE_CHECKING, Any, Final, Optional
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -313,9 +316,14 @@ class A2AGuardrailHandler(BaseTranslation):
return responses_so_far
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
_, valid_parsed = self._parse_streaming_responses(responses_so_far)
combined_text, _ = self._collect_text_from_parsed_chunks(valid_parsed)
return StreamingScanKey(texts=(combined_text,))
def _parse_streaming_responses(
self,
responses_so_far: list[object],
responses_so_far: Sequence[object],
) -> tuple[list[dict[str, object] | None], list[tuple[int, dict[str, object]]]]:
"""Parse JSON-RPC items, returning aligned parsed list and valid entries."""
parsed: Final[list[dict[str, object] | None]] = [None] * len(responses_so_far)

View file

@ -26,7 +26,10 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
LiteLLMAnthropicMessagesAdapter,
is_provider_native_tool_dict,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
anthropic_tool_name,
anthropic_tool_names,
@ -36,6 +39,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
merge_guardrailed_scoped_messages,
merge_returned_tools_into_request_tools,
scoped_structured_message_indices,
stream_item_fingerprint,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -1176,6 +1180,25 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (),
stream_ended=stream_ended,
)
@classmethod
def _streamed_tool_use_fingerprints(cls, responses_so_far: Sequence[object]) -> tuple[str, ...]:
return tuple(
stream_item_fingerprint(block)
for item in responses_so_far
for event in cls._iter_sse_events(item)
if event.get("type") == "content_block_start"
and isinstance(block := event.get("content_block"), Mapping)
and block.get("type") == "tool_use"
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Parse streaming responses and extract accumulated text content.

View file

@ -35,6 +35,22 @@ class StreamTransformSink:
holdback_per_choice: dict[int, int] = field(default_factory=dict)
@dataclass(frozen=True, slots=True)
class StreamingScanKey:
"""What a streaming guardrail round would hand to ``apply_guardrail``. Two keys
compare equal when the round would scan the same content again; ``stream_ended``
stays out of the comparison and only says whether the handler is on its
end-of-stream path, where an empty payload is still scanned today."""
texts: tuple[str, ...]
tool_calls: tuple[str, ...] = ()
stream_ended: bool = field(default=False, compare=False)
@property
def has_nothing_to_scan(self) -> bool:
return not self.stream_ended and not any(self.texts) and not self.tool_calls
class BaseTranslation(ABC):
@staticmethod
def transform_user_api_key_dict_to_metadata(
@ -151,6 +167,9 @@ class BaseTranslation(ABC):
"""
return responses_so_far
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
return None
def build_block_sse_chunks(
self,
exc: "ModifyResponseException",

View file

@ -4,6 +4,8 @@ import json
from collections.abc import Callable, Iterator, Sequence
from typing import Any, Final, TypeVar
from pydantic import BaseModel
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
@ -130,6 +132,16 @@ def stream_item_field(item: object, field: str) -> object | None:
return getattr(item, field, None)
def stream_item_fingerprint(item: object) -> str:
plain: Final = item.model_dump() if isinstance(item, BaseModel) else item
return json.dumps(plain, sort_keys=True, default=str)
def stream_item_items(item: object, field: str) -> tuple[object, ...]:
value: Final = stream_item_field(item, field)
return tuple(value) if isinstance(value, (list, tuple)) else ()
def blocked_chat_stream_usage(original_response: object) -> tuple[int, int]:
"""
``(prompt_tokens, completion_tokens)`` for a synthetic guardrail-blocked

View file

@ -2,6 +2,7 @@ from typing import Final
from httpx import Headers
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
@ -16,16 +17,18 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
"""
Session id to send as `x-session-affinity`, or None when the caller gave none.
Deliberately does not fall back to `litellm_trace_id`: that is generated per
request (`str(uuid.uuid4())` when absent), so using it pins every request to a
different Fireworks node and prompt caching never hits.
Deliberately does not fall back to `litellm_trace_id`, and ignores session ids the
proxy generated for a request that had none: both are per request, so using them
pins every request to a different Fireworks node and prompt caching never hits.
"""
params: Final = litellm_params
metadata: Final = params.get("metadata")
if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return None
for key in ("litellm_session_id", "session_id"):
value = params.get(key)
if value:
return str(value)
metadata: Final = params.get("metadata")
if isinstance(metadata, dict):
value = metadata.get("session_id")
if value:

View file

@ -26,6 +26,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
@ -39,6 +40,8 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
role_out_of_guardrail_scope,
scoped_structured_message_indices,
stream_item_field,
stream_item_fingerprint,
stream_item_items,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -503,12 +506,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
"""Block-only streaming path: run the guardrail so an in-flight BLOCK can
terminate the stream. Text rewrites are not propagated to the client here
(see ``_process_streaming_transform`` for the incremental_diff path)."""
# check if the stream has ended
has_stream_ended = False
for chunk in responses_so_far:
if chunk.choices and chunk.choices[0].finish_reason is not None:
has_stream_ended = True
break
has_stream_ended: Final = self._first_choice_has_finished(responses_so_far)
if has_stream_ended:
# convert to model response
@ -706,8 +704,33 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback)
}
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream))
stream_ended: Final = self._first_choice_has_finished(responses_so_far)
return StreamingScanKey(
texts=tuple(self._combine_streaming_texts(chunks).values()),
tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (),
stream_ended=stream_ended,
)
@staticmethod
def _streamed_tool_call_fingerprints(responses_so_far: Sequence[object]) -> tuple[str, ...]:
return tuple(
stream_item_fingerprint(tool_call)
for chunk in responses_so_far
for choice in _stream_chunk_choices(chunk)
for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls")
)
@staticmethod
def _first_choice_has_finished(responses_so_far: Sequence[object]) -> bool:
first_choices: Final = tuple(
choices[0] for choices in (_stream_chunk_choices(chunk) for chunk in responses_so_far) if choices
)
return any(stream_item_field(choice, "finish_reason") is not None for choice in first_choices)
def _combine_streaming_texts(
self, responses_so_far: list["ModelResponseStream"]
self, responses_so_far: Sequence["ModelResponseStream"]
) -> dict[tuple[int, int | None], str]:
"""
Combine all streaming chunks into complete text per choice.

View file

@ -44,10 +44,15 @@ from litellm._logging import verbose_proxy_logger
from litellm.completion_extras.litellm_responses_transformation.transformation import (
OpenAiResponsesToChatCompletionStreamIterator,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_responses_stream_usage,
stream_item_field,
stream_item_fingerprint,
stream_item_items,
)
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
from litellm.responses.litellm_completion_transformation.transformation import (
@ -593,18 +598,55 @@ class OpenAIResponsesHandler(BaseTranslation):
)
return responses_so_far
def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool:
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if the streaming has ended.
"""
if not responses_so_far:
return False
terminal_types: Final = {
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
}
return responses_so_far[-1].get("type") in terminal_types
terminal_types: Final = frozenset(
(
ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value,
ResponsesAPIStreamEvents.RESPONSE_FAILED.value,
ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value,
)
)
return stream_item_field(responses_so_far[-1], "type") in terminal_types
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
if not responses_so_far or not hasattr(responses_so_far[-1], "get"):
return None
last_event: Final = responses_so_far[-1]
last_event_type: Final = stream_item_field(last_event, "type")
if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value:
return None
if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value:
return self._completed_response_scan_key(stream_item_field(last_event, "response"))
return StreamingScanKey(
texts=(self.get_streaming_string_so_far(responses_so_far),),
stream_ended=self._check_streaming_has_ended(responses_so_far),
)
@staticmethod
def _completed_response_scan_key(response: object) -> StreamingScanKey:
output_items: Final = stream_item_items(response, "output")
message_items: Final = tuple(
item for item in output_items if stream_item_field(item, "type") != "function_call"
)
return StreamingScanKey(
texts=tuple(
text
for item in message_items
for part in stream_item_items(item, "content")
if isinstance(text := stream_item_field(part, "text"), str) and text
),
tool_calls=tuple(
stream_item_fingerprint(item)
for item in output_items
if stream_item_field(item, "type") == "function_call"
),
stream_ended=True,
)
def build_stream_error_items(
self,
@ -629,7 +671,7 @@ class OpenAIResponsesHandler(BaseTranslation):
),
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Get the string so far from the responses so far.
@ -641,12 +683,16 @@ class OpenAIResponsesHandler(BaseTranslation):
"""
keyed_events: Final = tuple(
(
(event.get("item_id"), event.get("output_index"), event.get("content_index")),
event.get("text"),
event.get("delta"),
(
stream_item_field(event, "item_id"),
stream_item_field(event, "output_index"),
stream_item_field(event, "content_index"),
),
stream_item_field(event, "text"),
stream_item_field(event, "delta"),
)
for event in responses_so_far
if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str)
if isinstance(stream_item_field(event, "text"), str) or isinstance(stream_item_field(event, "delta"), str)
)
def part_text(part_key: tuple[object, object, object]) -> str:

View file

@ -2594,6 +2594,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.",
)
missing_session_id: Literal["generate", "reject"] | None = Field(
None,
description="What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.",
)
enable_public_model_hub: bool = Field(
default=False,
description="Public model hub for users to see what models they have access to, supported openai params, etc.",

View file

@ -514,12 +514,26 @@ async def update_guardrail(
guardrail_name: Final = result.get("guardrail_name", "Unknown")
try:
IN_MEMORY_GUARDRAIL_HANDLER.update_in_memory_guardrail(
guardrail_id=guardrail_id, guardrail=cast(Guardrail, result)
)
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=cast(Guardrail, result))
verbose_proxy_logger.info(
"Immediate sync: Successfully updated guardrail '%s' (ID: %s)", guardrail_name, guardrail_id
)
except (ValueError, TypeError) as update_error:
# The new config is invalid (a raising guardrail __init__):
# reinitialize_guardrail already restored the previous live instance, but
# update_guardrail_in_db above already persisted the rejected config to
# the DB. Roll that back too, so the DB and the live guardrail never
# disagree about what's actually enforcing, and surface the rejection to
# the caller instead of a misleading 200.
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
guardrail=existing_guardrail,
prisma_client=prisma_client,
)
raise HTTPException(
status_code=422,
detail=f"Invalid guardrail configuration, update rejected: {update_error}",
) from update_error
except Exception as update_error:
verbose_proxy_logger.warning(
"Immediate sync: Failed to update '%s' (ID: %s) in memory: %s",

View file

@ -36,6 +36,7 @@ if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
)
# Call types that stream JSON-RPC events (A2A); guardrail HTTPException is emitted as in-stream error
@ -54,6 +55,9 @@ class _EndpointTranslation(Protocol):
@property
def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ...
@property
def get_streaming_scan_key(self) -> "Callable[[Sequence[object]], StreamingScanKey | None]": ...
@property
def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
@ -70,6 +74,12 @@ def _chunk_choices(item: object) -> Sequence[object]:
return choices
def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
if scan_key is None:
return False
return scan_key == last_scan_key or scan_key.has_nothing_to_scan
class _StreamTerminated(Exception):
"""Internal signal that the incremental transform stream has already emitted
its terminal chunks (block message or in-stream error) and must stop."""
@ -1011,6 +1021,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# Drives how a block terminates the stream: continue the in-progress
# message (True) vs emit a standalone block message (False, buffered).
chunks_yielded = False
last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round
async for item in response:
chunk_counter += 1
@ -1052,6 +1063,19 @@ class UnifiedLLMGuardrails(CustomLogger):
# Process chunk based on sampling rate
if chunk_counter % sampling_rate == 0:
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
if _is_redundant_scan(scan_key, last_scan_key):
verbose_proxy_logger.debug(
"Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round",
chunk_counter,
guardrail_to_apply.guardrail_name,
)
chunks_yielded = True
responses_yielded.append(item)
yield item
continue
verbose_proxy_logger.debug(
"Processing streaming chunk %s (sampling_rate=%s) with guardrail %s",
chunk_counter,
@ -1067,8 +1091,6 @@ class UnifiedLLMGuardrails(CustomLogger):
# string, permanently losing this chunk's content.
original_item = copy.deepcopy(item)
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
try:
await endpoint_translation.process_output_streaming_response(
responses_so_far=responses_so_far,
@ -1110,6 +1132,8 @@ class UnifiedLLMGuardrails(CustomLogger):
):
yield error_item
return
if scan_key is not None:
last_scan_key = scan_key
chunks_yielded = True
responses_yielded.append(original_item)
yield original_item
@ -1136,6 +1160,18 @@ class UnifiedLLMGuardrails(CustomLogger):
# preserve the list, not clone every chunk (deepcopy would double
# peak memory for large responses).
buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None
end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far)
if _is_redundant_scan(end_scan_key, last_scan_key):
verbose_proxy_logger.debug(
"Skipping end-of-stream scan for guardrail %s: the last sampled round already scanned it all",
guardrail_to_apply.guardrail_name,
)
for buffered_item in buffered_items or ():
yield buffered_item
for pending_item in pending_end_of_stream_items:
responses_yielded.append(pending_item)
yield pending_item
return
try:
await endpoint_translation.process_output_streaming_response(

View file

@ -826,11 +826,12 @@ class InMemoryGuardrailHandler:
Removes old callback from litellm.callbacks and creates fresh instance.
If the new config fails to initialize (e.g. an invalid on_flagged
combination), the previous instance is restored rather than left
deleted: initialize_guardrail's own ValueError/TypeError propagate
uncaught, so a caller reaching this point after already deleting the
old instance would otherwise leave the guardrail providing no
protection at all, not merely "still enforcing the old config."
combination or an invalid regex), the previous instance is restored
rather than left deleted, and the failure is re-raised as ValueError so
every init failure reaches callers as one exception type: a caller
reaching this point after already deleting the old instance would
otherwise leave the guardrail providing no protection at all, not
merely "still enforcing the old config."
"""
guardrail_id: Final = guardrail.get("guardrail_id")
if not guardrail_id:
@ -849,7 +850,7 @@ class InMemoryGuardrailHandler:
# that was enforcing must never fail open because an update was bad.
try:
return self.initialize_guardrail(guardrail=guardrail, config_file_path=config_file_path, source=source)
except Exception:
except Exception as init_error:
if previous_guardrail is not None:
verbose_proxy_logger.exception(
"Reinitializing guardrail %s with updated params failed; restoring the previous configuration",
@ -861,7 +862,7 @@ class InMemoryGuardrailHandler:
)
except Exception: # noqa: BLE001 # the original failure must propagate even if the restore breaks
verbose_proxy_logger.exception("Restoring previous guardrail %s also failed", guardrail_id)
raise
raise ValueError(f"Guardrail initialization failed: {init_error}") from init_error
def sync_guardrail_from_db(self, guardrail: Guardrail, config_file_path: str | None = None) -> Guardrail | None:
"""

View file

@ -16,6 +16,7 @@ from starlette.datastructures import Headers
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm._uuid import uuid
from litellm.constants import (
CONSUMED_REQUEST_TAGS_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
@ -23,6 +24,7 @@ from litellm.constants import (
OTEL_SERVICE_NAME_METADATA_KEYS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
@ -40,6 +42,7 @@ from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LiteLLMRoutes,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
@ -47,6 +50,8 @@ from litellm.proxy._types import (
TeamCallbackMetadata,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_utils import get_request_route
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.common_utils.callback_utils import (
decrypt_callback_vars,
get_metadata_variable_name_from_kwargs,
@ -715,6 +720,50 @@ def _get_anthropic_session_id_from_metadata(metadata: object) -> str | None:
return session_id
def _is_llm_inference_route(request: Request) -> bool:
route: Final = get_request_route(request)
return RouteChecks.is_llm_api_route(route=route) and not RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
def apply_missing_session_id_policy(
data: dict[str, object], # mutable-ok: stamps session ids in place on the request body the pipeline threads through
_metadata_variable_name: str,
general_settings: Mapping[str, object] | None,
request: Request,
) -> None:
policy: Final = general_settings.get("missing_session_id") if general_settings else None
if policy is None or not _is_llm_inference_route(request):
return
metadata: Final = data.get(_metadata_variable_name)
if not isinstance(metadata, dict):
return
if data.get("litellm_session_id") or metadata.get("session_id"):
return
match policy:
case "generate":
session_id: Final = str(data.get("litellm_trace_id") or metadata.get("trace_id") or uuid.uuid4())
data["litellm_session_id"] = session_id # rebind-ok: data is an out-param
data.setdefault("litellm_trace_id", session_id)
metadata["session_id"] = session_id
metadata[SESSION_ID_GENERATED_METADATA_KEY] = True
case "reject":
raise ProxyException(
message=(
"Request has no session id. Send an `x-litellm-session-id` header or `metadata.session_id`. "
"Required by `general_settings.missing_session_id: reject`."
),
type=ProxyErrorTypes.bad_request_error,
param="session_id",
code=400,
)
case _:
verbose_proxy_logger.warning(
"Ignoring unknown general_settings.missing_session_id=%r; expected 'generate' or 'reject'", policy
)
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code identifies itself as ``claude-cli/<version> ...``; the IDE
extensions and the Agent SDK run through the same CLI and share that prefix."""
@ -1818,6 +1867,12 @@ async def add_litellm_data_to_request(
data=data,
_metadata_variable_name=_metadata_variable_name,
)
apply_missing_session_id_policy(
data=data,
_metadata_variable_name=_metadata_variable_name,
general_settings=general_settings,
request=request,
)
# Expose request headers under the metadata field for guardrails (fixes #17477)
if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict):

View file

@ -26,7 +26,11 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
from pydantic import BaseModel, create_model
from litellm._logging import verbose_router_logger
from litellm.constants import EMPTY_MAPPING, RETURN_RAW_MODEL_NAME_METADATA_KEY
from litellm.constants import (
EMPTY_MAPPING,
RETURN_RAW_MODEL_NAME_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
@ -2712,7 +2716,7 @@ class ComplexityRouter(CustomLogger):
"""Resolve a client-supplied session_id."""
for metadata in ComplexityRouter._iter_metadata_dicts(request_kwargs):
session_id = metadata.get("session_id")
if session_id is not None:
if session_id is not None and not metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return str(session_id)
return None

View file

@ -21,7 +21,7 @@ from typing_extensions import TypedDict
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import AllMessageValues
@ -265,7 +265,7 @@ class DeploymentAffinityCheck(CustomLogger):
@staticmethod
def _get_session_id_from_metadata_dict(metadata: dict) -> str | None:
session_id: Final = metadata.get("session_id")
if session_id is None:
if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return None
return str(session_id)

View file

@ -0,0 +1,150 @@
import ast
import os
import sys
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
PY311_PLUS_TYPING_NAMES: Final[frozenset[str]] = frozenset(
{
"NotRequired",
"Required",
"Self",
"LiteralString",
"Never",
"assert_never",
"assert_type",
"reveal_type",
"TypeVarTuple",
"Unpack",
"dataclass_transform",
"override",
"TypeAliasType",
"get_original_bases",
"ReadOnly",
"TypeIs",
"NoDefault",
"get_protocol_members",
"is_protocol",
"evaluate_forward_ref",
"TypeForm",
}
)
@dataclass(frozen=True, slots=True)
class TypingImportViolation:
file: str
line: int
name: str
def _walk_with_ancestors(
node: ast.AST, ancestors: tuple[tuple[ast.AST, str], ...] = ()
) -> Iterator[tuple[ast.AST, tuple[tuple[ast.AST, str], ...]]]:
yield node, ancestors
for field_name, field_value in ast.iter_fields(node):
if isinstance(field_value, ast.AST):
yield from _walk_with_ancestors(field_value, (*ancestors, (node, field_name)))
elif isinstance(field_value, list):
for child in field_value:
if isinstance(child, ast.AST):
yield from _walk_with_ancestors(child, (*ancestors, (node, field_name)))
def _is_sys_version_info(node: ast.AST) -> bool:
return (
isinstance(node, ast.Attribute)
and isinstance(node.value, ast.Name)
and node.value.id == "sys"
and node.attr == "version_info"
)
def _is_version_guarded(ancestors: tuple[tuple[ast.AST, str], ...]) -> bool:
nearest_if: Final[tuple[ast.If, str] | None] = next(
(
(ancestor, field_name)
for ancestor, field_name in reversed(ancestors)
if isinstance(ancestor, ast.If)
),
None,
)
if nearest_if is None:
return False
enclosing_if, branch = nearest_if
test: Final[ast.expr] = enclosing_if.test
if not isinstance(test, ast.Compare) or len(test.ops) != 1 or not _is_sys_version_info(test.left):
return False
operator: Final[ast.cmpop] = test.ops[0]
return (isinstance(operator, (ast.Gt, ast.GtE)) and branch == "body") or (
isinstance(operator, (ast.Lt, ast.LtE)) and branch == "orelse"
)
def scan_file(file_path: str | os.PathLike[str]) -> tuple[TypingImportViolation, ...]:
path: Final[Path] = Path(file_path)
tree: Final[ast.Module] = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
return tuple(
violation
for node, ancestors in _walk_with_ancestors(tree)
if not _is_version_guarded(ancestors)
for violation in _violations_for_node(node, path)
)
def _violations_for_node(
node: ast.AST, path: Path
) -> tuple[TypingImportViolation, ...]:
if isinstance(node, ast.ImportFrom) and node.module == "typing":
return tuple(
TypingImportViolation(file=str(path), line=node.lineno, name=alias.name)
for alias in node.names
if alias.name in PY311_PLUS_TYPING_NAMES
)
if (
isinstance(node, ast.Attribute)
and isinstance(node.value, ast.Name)
and node.value.id == "typing"
and node.attr in PY311_PLUS_TYPING_NAMES
):
return (TypingImportViolation(file=str(path), line=node.lineno, name=node.attr),)
return ()
def scan_directory(base_dir: str | os.PathLike[str] = ".") -> tuple[TypingImportViolation, ...]:
base_path: Final[Path] = Path(base_dir)
return tuple(
violation
for directory in (
base_path / "litellm",
base_path / "enterprise",
base_path / "litellm-proxy-extras" / "litellm_proxy_extras",
)
if directory.exists()
for path in directory.rglob("*.py")
for violation in scan_file(path)
)
def main() -> None:
violations: Final[tuple[TypingImportViolation, ...]] = scan_directory()
if violations:
message: Final[str] = "\n".join(
(
"Python 3.10-incompatible typing imports found:",
*(
f"{violation.file}:{violation.line}: {violation.name} is unavailable in Python 3.10; "
"import it from typing_extensions instead because litellm supports Python 3.10"
for violation in violations
),
)
)
sys.stdout.write(f"{message}\n")
raise RuntimeError("Import Python 3.10-incompatible typing names from typing_extensions instead")
sys.stdout.write("No Python 3.10-incompatible typing imports found.\n")
if __name__ == "__main__":
main()

View file

@ -1,5 +1,5 @@
import httpx
from openai import OpenAI, BadRequestError, APIStatusError
from openai import OpenAI, BadRequestError, NotFoundError, APIStatusError
import pytest
@ -105,10 +105,9 @@ def test_streaming_response():
assert len(collected_chunks) > 0
def test_bad_request_error():
def test_model_not_found_error():
client = get_test_client()
with pytest.raises(BadRequestError):
# Trigger error with invalid model name
with pytest.raises(NotFoundError):
client.responses.create(model="non-existent-model", input="This should fail")

View file

@ -0,0 +1,44 @@
# Expected Structure
```text
tests/rust-python-harness/
├── __main__.py
├── strategies/
│ ├── e2e_parity/
│ │ ├── runner.py
│ │ ├── sdk/
│ │ │ ├── ocr/
│ │ │ ├── messages/
│ │ │ ├── chat_completions/
│ │ │ └── responses/
│ │ └── gateway/
│ │
│ ├── trace_parity/
│ │ ├── runner.py
│ │ ├── sdk/
│ │ └── gateway/
│ │
│ └── unit_tests/
│ ├── runner.py
│ ├── mapping_validator.py
│ ├── python_runner.py
│ └── rust_runner.py
└── shared/
├── parity/
├── tracing/
└── reporting/
```
- Run locally only; no CI integration
- `__main__.py` selects strategies and combines their reports; each strategy also runs independently
- `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses
- `trace_parity/` compares mapped operations, call counts, and required execution ordering
- E2E and trace runners share orchestration across `sdk/` and `gateway/`; surface-specific execution lives in those folders
- `unit_tests/runner.py` combines mapping validation, Python test runs, and native Rust test runs
- `mapping_validator.py` matches Python/Rust tests by agreed names or annotations and reports missing or ambiguous counterparts
- `python_runner.py` runs existing Python tests with Rust disabled and enabled in separate processes, verifies backend selection, and compares results
- `rust_runner.py` runs Cargo tests; native Rust unit tests stay beside their implementation
- `shared/` contains reusable parity, tracing, and reporting machinery
- Keep fixtures with their owning API and existing Python tests in their current locations

View file

@ -0,0 +1,35 @@
"""Tests for litellm/llms/a2a/chat/guardrail_translation/handler.py."""
import json
from litellm.llms.a2a.chat.guardrail_translation.handler import A2AGuardrailHandler
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
def _text_event(text: str) -> str:
return json.dumps(
{
"jsonrpc": "2.0",
"id": "req-1",
"result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": text}]},
}
)
def _status_event() -> str:
return json.dumps({"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "status-update", "status": {}}})
class TestA2AGuardrailHandlerStreamingScanKey:
def test_key_joins_the_text_of_every_message_event(self):
key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hello "), _text_event("world")])
assert key == StreamingScanKey(texts=("hello world",))
def test_events_without_text_leave_the_key_unchanged(self):
handler = A2AGuardrailHandler()
events = [_text_event("hello")]
assert handler.get_streaming_scan_key(events + [_status_event()]) == handler.get_streaming_scan_key(events)
def test_unparseable_items_are_ignored(self):
key = A2AGuardrailHandler().get_streaming_scan_key([_text_event("hi"), "not json", b"bytes"])
assert key.texts == ("hi",)

View file

@ -13,6 +13,7 @@ import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
AnthropicMessagesHandler,
)
@ -1991,3 +1992,56 @@ class TestStructuredWriteBackKeepsToolResults:
}
later_blocks = [b for m in messages[tool_use_index + 1 :] for b in self._blocks(m)]
assert {"type": "text", "text": "Now fetch the page."} in later_blocks
class TestAnthropicMessagesHandlerStreamingScanKey:
"""get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
@staticmethod
def _sse(event_type, data):
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
def _text_delta(self, text):
return self._sse(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}},
)
def test_key_is_empty_before_any_text_arrives(self):
head = self._sse("message_start", {"type": "message_start", "message": {"stop_reason": None}})
key = AnthropicMessagesHandler().get_streaming_scan_key([head])
assert key == StreamingScanKey(texts=("",))
def test_key_accumulates_text_deltas(self):
key = AnthropicMessagesHandler().get_streaming_scan_key([self._text_delta("hello "), self._text_delta("world")])
assert key.texts == ("hello world",)
assert key.stream_ended is False
def _stop(self, stop_reason):
return self._sse(
"message_delta",
{"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {}},
)
def test_stop_without_tool_use_scans_the_same_payload(self):
handler = AnthropicMessagesHandler()
open_key = handler.get_streaming_scan_key([self._text_delta("hi")])
ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), self._stop("end_turn")])
assert ended_key.stream_ended is True
assert ended_key == open_key
def test_tool_use_blocks_enter_the_key_once_the_stream_has_ended(self):
handler = AnthropicMessagesHandler()
tool_use = self._sse(
"content_block_start",
{
"type": "content_block_start",
"index": 1,
"content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}},
},
)
open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use])
ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")])
assert open_key == StreamingScanKey(texts=("hi",))
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key

View file

@ -8,6 +8,7 @@ import litellm
from litellm import get_model_info, supports_reasoning, supports_vision
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
from litellm.types.utils import (
ChatCompletionMessageToolCall,
@ -235,6 +236,21 @@ def test_get_fireworks_session_id_prefers_litellm_session_id_over_trace_id():
)
def test_get_fireworks_session_id_ignores_proxy_generated_session_id():
"""general_settings.missing_session_id: generate stamps a fresh id per request; sending it
as x-session-affinity would pin every request to a different node."""
assert (
get_fireworks_session_id(
{
"litellm_session_id": "generated-1",
"litellm_trace_id": "generated-1",
"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True},
}
)
is None
)
def test_handle_message_content_with_tool_calls():
config = FireworksAIConfig()
message = Message(

View file

@ -12,6 +12,7 @@ import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
from litellm.llms.openai.chat.guardrail_translation.handler import (
OpenAIChatCompletionsHandler,
)
@ -1643,3 +1644,74 @@ class TestCheckStreamingHasEnded:
)
]
assert handler._check_streaming_has_ended(chunks) is True
class TestStreamingScanKey:
"""get_streaming_scan_key identifies what a sampled round would scan so the
unified hook can skip rounds that would re-scan already-cleared text"""
@staticmethod
def _chunk(content, finish_reason=None, index=0):
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
return ModelResponseStream(
choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)]
)
def test_key_carries_accumulated_text_and_open_stream(self):
handler = OpenAIChatCompletionsHandler()
key = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")])
assert key == StreamingScanKey(texts=("hello",))
def test_chunks_without_text_leave_the_key_unchanged(self):
handler = OpenAIChatCompletionsHandler()
before = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo")])
after = handler.get_streaming_scan_key([self._chunk("hel"), self._chunk("lo"), self._chunk(None)])
assert after == before
def test_finish_chunk_without_tool_calls_scans_the_same_payload(self):
handler = OpenAIChatCompletionsHandler()
open_key = handler.get_streaming_scan_key([self._chunk("hi")])
ended_key = handler.get_streaming_scan_key([self._chunk("hi"), self._chunk(None, finish_reason="stop")])
assert open_key.stream_ended is False
assert ended_key.stream_ended is True
assert ended_key == open_key
def test_tool_calls_only_enter_the_key_once_the_stream_has_ended(self):
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
ModelResponseStream,
StreamingChoices,
)
handler = OpenAIChatCompletionsHandler()
tool_call = ChatCompletionDeltaToolCall(
id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}')
)
tool_chunk = ModelResponseStream(
choices=[StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason=None)]
)
open_key = handler.get_streaming_scan_key([self._chunk("hi"), tool_chunk])
ended_key = handler.get_streaming_scan_key(
[self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")]
)
assert open_key == StreamingScanKey(texts=("hi",))
assert ended_key.texts == ("hi",)
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key
def test_text_after_the_first_choice_finishes_still_changes_the_key(self):
handler = OpenAIChatCompletionsHandler()
first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)]
key_at_first_finish = handler.get_streaming_scan_key(first_done)
key_after_more_text = handler.get_streaming_scan_key(first_done + [self._chunk("y", index=1)])
assert key_at_first_finish.stream_ended is True
assert key_after_more_text.stream_ended is True
assert key_after_more_text != key_at_first_finish
def test_non_stream_items_are_ignored(self):
handler = OpenAIChatCompletionsHandler()
key = handler.get_streaming_scan_key([self._chunk("hi"), b"data: [DONE]"])
assert key.texts == ("hi",)

View file

@ -1731,3 +1731,93 @@ class TestBuildBlockSseChunks:
dones = [payload for payload in payloads if payload["type"] == "response.output_item.done"]
assert len(dones) == 1
assert dones[0]["item"]["content"][0]["text"] == "Blocked by policy."
class TestOpenAIResponsesHandlerStreamingScanKey:
"""get_streaming_scan_key mirrors what process_output_streaming_response would scan"""
@staticmethod
def _delta(sequence_number, text):
return {
"type": "response.output_text.delta",
"sequence_number": sequence_number,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": text,
}
def test_no_events_yields_no_key(self):
assert OpenAIResponsesHandler().get_streaming_scan_key([]) is None
def test_key_accumulates_deltas_while_the_stream_is_open(self):
from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey
key = OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hel"), self._delta(1, "lo")])
assert key == StreamingScanKey(texts=("hello",))
def test_typed_delta_events_accumulate_like_dicts(self):
from litellm.types.llms.openai import OutputTextDeltaEvent
events = [
OutputTextDeltaEvent(
type="response.output_text.delta",
item_id="msg_1",
output_index=0,
content_index=0,
delta=text,
sequence_number=i,
)
for i, text in enumerate(("hel", "lo"))
]
key = OpenAIResponsesHandler().get_streaming_scan_key(events)
assert key.texts == ("hello",)
assert key.stream_ended is False
def test_events_without_text_leave_the_key_unchanged(self):
handler = OpenAIResponsesHandler()
events = [self._delta(0, "hi")]
quiet = events + [{"type": "response.in_progress", "sequence_number": 1}]
assert handler.get_streaming_scan_key(quiet) == handler.get_streaming_scan_key(events)
@staticmethod
def _completed(sequence_number, output):
return {"type": "response.completed", "sequence_number": sequence_number, "response": {"output": output}}
def test_completed_event_keys_on_the_final_output_text(self):
handler = OpenAIResponsesHandler()
message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
open_key = handler.get_streaming_scan_key([self._delta(0, "hi")])
ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message])])
assert ended_key.stream_ended is True
assert ended_key == open_key
def test_completed_event_with_a_function_call_changes_the_key(self):
handler = OpenAIResponsesHandler()
message = {"type": "message", "content": [{"type": "output_text", "text": "hi"}]}
function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}
open_key = handler.get_streaming_scan_key([self._delta(0, "hi")])
ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), self._completed(1, [message, function_call])])
assert ended_key.texts == ("hi",)
assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0]
assert ended_key != open_key
def test_completed_event_reads_every_output_text_part(self):
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
item = GenericResponseOutputItem(
type="message",
id="msg_1",
status="completed",
role="assistant",
content=[
OutputText(type="output_text", text="one", annotations=[]),
OutputText(type="output_text", text="two", annotations=[]),
],
)
key = OpenAIResponsesHandler().get_streaming_scan_key([self._completed(0, [item])])
assert key.texts == ("one", "two")
def test_output_item_done_round_is_never_deduped(self):
done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}}
assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None

View file

@ -310,8 +310,8 @@ def _make_stream_chunk(content: str, finish_reason=None):
@pytest.mark.asyncio
async def test_openai_moderation_streaming_default_uses_sampled_cadence():
"""Default config samples every 5th streamed chunk and runs a final aggregate
pass after the stream ends. 10 chunks sampled at chunks 5 and 10 2 in-stream
calls, plus 1 final = 3 total.
pass after the stream ends. 10 chunks are sampled at 5 and 10; the end-of-stream
round is skipped because chunk 10 already scanned the full text, for 2 total calls
"""
import litellm
@ -370,8 +370,9 @@ async def test_openai_moderation_streaming_default_uses_sampled_cadence():
):
pass
assert patched_make_request.await_count == 3, (
f"Expected 3 moderation calls (2 sampled at chunks 5 / 10 + 1 final), "
assert patched_make_request.await_count == 2, (
f"Expected 2 moderation calls (2 sampled at chunks 5 / 10; "
f"the end-of-stream round is skipped because chunk 10 already scanned the full text), "
f"got {patched_make_request.await_count}"
)
@ -448,7 +449,8 @@ async def test_openai_moderation_streaming_end_of_stream_only_opt_in_calls_moder
@pytest.mark.asyncio
async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disabled():
"""With streaming_end_of_stream_only=False and streaming_sampling_rate=2,
moderation runs every 2nd chunk during the stream, plus once more at end.
moderation runs every 2nd chunk during the stream. The terminal chunk scan covers
the final aggregate, for 3 total calls
"""
import litellm
@ -509,9 +511,8 @@ async def test_openai_moderation_streaming_sampled_when_end_of_stream_only_disab
):
pass
# 6 chunks, sampling_rate=2 → in-stream calls at chunks 2, 4, 6 (3 calls),
# plus the final aggregate pass after the stream ends (1 call) = 4 total.
assert patched_make_request.await_count == 4, (
f"Expected 4 moderation calls (3 sampled + 1 final aggregate), "
assert patched_make_request.await_count == 3, (
f"Expected 3 moderation calls (3 sampled; the end-of-stream round is skipped "
f"because chunk 6 already scanned the full text), "
f"got {patched_make_request.await_count}"
)

View file

@ -1517,7 +1517,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
@pytest.mark.asyncio
async def test_streaming_default_uses_sampled_cadence(self):
"""Default samples every 5th chunk + final pass: 10 chunks → calls at 5, 10, and final = 3."""
"""Default samples every 5th chunk. For 10 chunks, sampled scans at 5 and 10
cover the full text, so the end-of-stream round is skipped and there are 2 calls
"""
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -1566,8 +1568,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
):
pass
assert mock_post.await_count == 3, (
f"Expected 3 guardrail calls (2 sampled at chunks 5 / 10 + 1 final), "
assert mock_post.await_count == 2, (
f"Expected 2 guardrail calls (2 sampled at chunks 5 / 10; "
f"the end-of-stream round is skipped because chunk 10 already scanned the full text), "
f"got {mock_post.await_count}"
)
for call in mock_post.await_args_list:
@ -1631,7 +1634,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
@pytest.mark.asyncio
async def test_streaming_sampling_rate_override(self):
"""sampling_rate=2 on 6 chunks → in-stream at 2,4,6 plus final = 4 calls."""
"""sampling_rate=2 on 6 chunks. Scans at 2, 4, and 6 cover the full text, so
the end-of-stream round is skipped and there are 3 calls
"""
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -1680,8 +1685,9 @@ class TestGenericGuardrailAPIStreamingViaUnified:
):
pass
assert mock_post.await_count == 4, (
f"Expected 4 guardrail calls (3 sampled + 1 final aggregate), "
assert mock_post.await_count == 3, (
f"Expected 3 guardrail calls (3 sampled; the end-of-stream round is skipped "
f"because chunk 6 already scanned the full text), "
f"got {mock_post.await_count}"
)

View file

@ -1971,3 +1971,271 @@ class TestStreamingGuardrailInformationBucket:
assert recorded[0]["guardrail_name"] == "audit-recorder"
assert recorded[0]["guardrail_status"] == "success"
assert request_data["metadata"]["user_api_key_user_id"] == "user-1"
class _ScanCountingGuardrail(CustomGuardrail):
"""Pass-through guardrail that records every response-side scan payload."""
def __init__(self, *, sampling_rate=5, end_of_stream_only=False, buffer_until_moderated=False):
super().__init__(guardrail_name="scan-counter")
self.streaming_sampling_rate = sampling_rate
self.streaming_end_of_stream_only = end_of_stream_only
self.streaming_buffer_until_moderated = buffer_until_moderated
self.guardrail_config = {}
self.scans: tuple[dict[str, object], ...] = ()
def should_run_guardrail(self, data, event_type): # type: ignore[override]
return True
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
self.scans = (
*self.scans,
{
"texts": list(inputs.get("texts") or []),
"tool_calls": list(inputs.get("tool_calls") or []),
"model": inputs.get("model"),
},
)
return inputs
def _responses_delta(sequence_number, text):
return {
"type": "response.output_text.delta",
"sequence_number": sequence_number,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": text,
}
def _responses_tail(sequence_number, text):
return [
{
"type": "response.output_text.done",
"sequence_number": sequence_number,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"text": text,
},
{
"type": "response.completed",
"sequence_number": sequence_number + 1,
"response": {
"model": "gpt-5.6",
"output": [{"type": "message", "content": [{"type": "output_text", "text": text}]}],
},
},
]
class TestStreamingScanDedup:
"""A sampled round whose scan payload matches the previous round (or carries
no text yet) is skipped, so a stream is never re-scanned for output the
guardrail already cleared. Regression for LIT-6692."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self, monkeypatch):
monkeypatch.setattr(
unified_module,
"endpoint_guardrail_translation_mappings",
load_guardrail_translation_mappings(),
)
@pytest.mark.asyncio
async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 3
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
@pytest.mark.asyncio
async def test_chat_round_with_unchanged_text_is_skipped(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [
_stream_chunk("a"),
_stream_chunk("b"),
_stream_chunk("c"),
_stream_chunk(None),
_stream_chunk(None),
_stream_chunk(None),
_stream_chunk("d", finish_reason="stop"),
]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 7
assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abcd"]]
@pytest.mark.asyncio
async def test_chat_finish_chunk_right_after_a_sampled_round_is_not_rescanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), _stream_chunk(None, finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 4
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]
@pytest.mark.asyncio
async def test_chat_finish_chunk_carrying_tool_calls_is_still_scanned(self):
from litellm.types.utils import ChatCompletionDeltaToolCall, Function
guardrail = _ScanCountingGuardrail(sampling_rate=3)
tool_call = ChatCompletionDeltaToolCall(
id="call_1", index=0, type="function", function=Function(name="get_weather", arguments='{"city": "Paris"}')
)
finish = ModelResponseStream(
choices=[
StreamingChoices(index=0, delta=Delta(content=None, tool_calls=[tool_call]), finish_reason="tool_calls")
]
)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c"), finish]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(out) == 4
assert [scan["texts"] for scan in guardrail.scans] == [["abc"], ["abc"]]
assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"]
@pytest.mark.asyncio
async def test_chat_second_choice_finishing_later_still_gets_the_end_scan(self):
guardrail = _ScanCountingGuardrail(sampling_rate=3)
chunks = [
_stream_chunk("a", index=0),
_stream_chunk("x", index=1),
_stream_chunk("b", finish_reason="stop", index=0),
_stream_chunk("y", index=1),
_stream_chunk("z", finish_reason="stop", index=1),
]
await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert len(guardrail.scans) == 2
assert any("yz" in text for text in guardrail.scans[-1]["texts"])
@pytest.mark.asyncio
async def test_responses_completed_event_on_sampled_index_is_scanned_once(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(8)]
full_text = "".join(f"t{i}" for i in range(8))
chunks = deltas + _responses_tail(8, full_text)
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 10
assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], [full_text]]
assert guardrail.scans[-1]["model"] == "gpt-5.6"
@pytest.mark.asyncio
async def test_responses_completed_right_after_a_sampled_round_is_not_rescanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
chunks = deltas + _responses_tail(5, "t0t1t2t3t4")
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 7
assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"]]
@pytest.mark.asyncio
async def test_responses_completed_carrying_a_function_call_is_still_scanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
completed = {
"type": "response.completed",
"sequence_number": 5,
"response": {
"model": "gpt-5.6",
"output": [
{"type": "message", "content": [{"type": "output_text", "text": "t0t1t2t3t4"}]},
{
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "get_weather",
"arguments": '{"city": "Paris"}',
"status": "completed",
},
],
},
}
chunks = deltas + [completed]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 6
assert [scan["texts"] for scan in guardrail.scans] == [["t0t1t2t3t4"], ["t0t1t2t3t4"]]
assert [call["function"]["name"] for call in guardrail.scans[1]["tool_calls"]] == ["get_weather"]
@pytest.mark.asyncio
async def test_responses_round_with_unchanged_text_is_skipped(self):
guardrail = _ScanCountingGuardrail(sampling_rate=5)
deltas = [_responses_delta(i, f"t{i}") for i in range(5)]
quiet = [{"type": "response.in_progress", "sequence_number": i} for i in range(5, 10)]
chunks = deltas + quiet + _responses_tail(10, "t0t1t2t3t4")
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 12
assert guardrail.scans == ({"texts": ["t0t1t2t3t4"], "tool_calls": [], "model": None},)
@pytest.mark.asyncio
async def test_responses_tool_call_done_event_is_still_scanned(self):
guardrail = _ScanCountingGuardrail(sampling_rate=2)
tool_call_done = {
"type": "response.output_item.done",
"sequence_number": 1,
"output_index": 1,
"item": {
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "get_weather",
"arguments": '{"city": "Paris"}',
"status": "completed",
},
}
chunks = [_responses_delta(0, "hi"), tool_call_done] + _responses_tail(2, "hi")
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses")
assert len(out) == 4
assert len(guardrail.scans) == 2
assert [call["function"]["name"] for call in guardrail.scans[0]["tool_calls"]] == ["get_weather"]
assert guardrail.scans[1]["texts"] == ["hi"]
@pytest.mark.asyncio
async def test_anthropic_skips_empty_round_and_terminal_duplicate(self):
guardrail = _ScanCountingGuardrail(sampling_rate=2)
chunks = _anthropic_message_chunks(["hello ", "world"])
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
assert out == chunks
assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]]
@pytest.mark.asyncio
async def test_end_of_stream_only_still_scans_exactly_once(self):
guardrail = _ScanCountingGuardrail(sampling_rate=2, end_of_stream_only=True)
chunks = _anthropic_message_chunks(["hello ", "world"])
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages")
assert out == chunks
assert [scan["texts"] for scan in guardrail.scans] == [["hello world"]]
@pytest.mark.asyncio
async def test_buffer_until_moderated_still_scans_exactly_once_and_releases_every_chunk(self):
guardrail = _ScanCountingGuardrail(sampling_rate=1, buffer_until_moderated=True)
chunks = [_stream_chunk("a"), _stream_chunk("b"), _stream_chunk("c", finish_reason="stop")]
out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks)
assert out == chunks
assert [scan["texts"] for scan in guardrail.scans] == [["abc"]]

View file

@ -104,7 +104,7 @@ def mock_in_memory_handler(mocker):
mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL
mock_handler.get_source.return_value = "config"
mock_handler.initialize_guardrail = mocker.Mock()
mock_handler.update_in_memory_guardrail = mocker.Mock()
mock_handler.sync_guardrail_from_db = mocker.Mock()
mock_handler.delete_in_memory_guardrail = mocker.Mock()
mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[])
return mock_handler
@ -1047,13 +1047,15 @@ async def test_create_guardrail_endpoint(
"scenario,expected_result,expected_exception",
[
("success_with_sync", "test-db-guardrail", None),
("success_sync_fails", "test-db-guardrail", None),
("success_sync_fails_unexpected_error", "test-db-guardrail", None),
("sync_fails_invalid_config", None, HTTPException),
("database_failure", None, HTTPException),
("no_prisma_client", None, HTTPException),
],
ids=[
"success_with_immediate_sync",
"success_but_sync_fails",
"success_but_sync_fails_with_unexpected_error",
"sync_rejects_invalid_config",
"database_error",
"missing_prisma_client",
],
@ -1073,6 +1075,7 @@ async def test_update_guardrail_endpoint(
mock_logger = None
if scenario == "success_with_sync":
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock()
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
@ -1083,10 +1086,13 @@ async def test_update_guardrail_endpoint(
mock_in_memory_handler,
)
elif scenario == "success_sync_fails":
elif scenario == "success_sync_fails_unexpected_error":
# A non-ValueError/TypeError failure is not a config-rejection signal,
# so it keeps the pre-existing swallow-and-warn behavior rather than
# rolling back the DB write.
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.update_in_memory_guardrail.side_effect = Exception(
"Sync failed"
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=Exception("Sync failed")
)
mock_logger = mocker.patch(
"litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger"
@ -1102,6 +1108,25 @@ async def test_update_guardrail_endpoint(
mock_in_memory_handler,
)
elif scenario == "sync_fails_invalid_config":
# Regression for the PUT half of the fix: a TypeError from the sync (the
# in-place update_in_memory_guardrail raised exactly this on every PUT)
# must roll back the DB write and surface a 422, not persist the
# rejected config with a 200.
mock_prisma_client = mocker.Mock()
mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(
side_effect=TypeError("vars() argument must have __dict__ attribute")
)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # test-quality-ok: reused pattern
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY",
mock_guardrail_registry,
)
mocker.patch( # test-quality-ok: reused pattern
"litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER",
mock_in_memory_handler,
)
elif scenario == "database_failure":
mock_prisma_client = mocker.Mock()
mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception(
@ -1130,6 +1155,16 @@ async def test_update_guardrail_endpoint(
assert "Database error" in str(exc_info.value.detail)
elif scenario == "no_prisma_client":
assert "Prisma client not initialized" in str(exc_info.value.detail)
elif scenario == "sync_fails_invalid_config":
assert exc_info.value.status_code == 422
assert "update rejected" in str(exc_info.value.detail)
# Rolled back: update_guardrail_in_db is called once for the
# rejected write and once more to restore the previous config.
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
assert (
mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"]
== MOCK_DB_GUARDRAIL
)
else:
result = await update_guardrail(
@ -1145,11 +1180,11 @@ async def test_update_guardrail_endpoint(
prisma_client=mocker.ANY,
)
mock_in_memory_handler.update_in_memory_guardrail.assert_called_once_with(
guardrail_id="test-guardrail-id", guardrail=mocker.ANY
mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(
guardrail=mocker.ANY
)
if scenario == "success_sync_fails":
if scenario == "success_sync_fails_unexpected_error":
assert mock_logger is not None
mock_logger.warning.assert_called_once()
assert "Failed to update" in str(mock_logger.warning.call_args)

View file

@ -913,3 +913,96 @@ def test_reinitialize_guardrail_restores_previous_on_failure():
assert restored.guardrail_name == "restore-me"
finally:
registry_module.guardrail_initializer_registry.pop("restore_test", None)
def test_reinitialize_guardrail_raises_value_error_for_non_value_error_init_failures():
"""Regression for the LIT-6479 fix's 422 path: a constructor failure that is not
already a ValueError/TypeError (re.error from an invalid regex has neither in its
MRO) must still surface as ValueError, so the PUT/PATCH endpoints' rollback+422
catch is exhaustive instead of warn-and-200 persisting a broken config."""
import re
from litellm.proxy.guardrails import guardrail_registry as registry_module
def _initializer(litellm_params, guardrail):
if litellm_params.api_key == "bad-regex":
re.compile("([")
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
registry_module.guardrail_initializer_registry["regex_test"] = _initializer
try:
handler = InMemoryGuardrailHandler()
created = handler.initialize_guardrail(
guardrail={
"guardrail_name": "regex-me",
"litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "ok"},
},
)
guardrail_id = created["guardrail_id"]
with pytest.raises(ValueError, match="Guardrail initialization failed") as excinfo:
handler.reinitialize_guardrail(
guardrail={
"guardrail_id": guardrail_id,
"guardrail_name": "regex-me",
"litellm_params": {"guardrail": "regex_test", "mode": "pre_call", "api_key": "bad-regex"},
},
)
assert isinstance(excinfo.value.__cause__, re.error)
assert guardrail_id in handler.IN_MEMORY_GUARDRAILS
restored = handler.guardrail_id_to_custom_guardrail[guardrail_id]
assert restored is not None and restored.guardrail_name == "regex-me"
finally:
registry_module.guardrail_initializer_registry.pop("regex_test", None)
def test_sync_guardrail_from_db_applies_db_dict_params_to_live_instance():
"""
Regression for PUT /guardrails/{id}: the DB row arrives with litellm_params as
a plain jsonb dict, and the in-place update_in_memory_guardrail cast it to
LitellmParams without constructing one, so vars() raised and the running proxy
kept enforcing the stale config forever. The PUT endpoint now routes through
sync_guardrail_from_db, which must rebuild the live instance from the dict:
new blocked words compiled in, old ones gone, and the event hook re-derived
from mode (the base-class setattr path wrote self.mode while dispatch reads
self.event_hook, so only a full re-init applies a mode change).
"""
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
handler = InMemoryGuardrailHandler()
gid = "66666666-6666-6666-6666-666666666666"
def db_guardrail(word: str, mode: str) -> Guardrail:
return Guardrail(
guardrail_id=gid,
guardrail_name="cf-put-sync",
litellm_params={
"guardrail": "litellm_content_filter",
"mode": mode,
"default_on": True,
"blocked_words": [{"keyword": word, "action": "BLOCK"}],
},
)
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler.sync_guardrail_from_db(db_guardrail("foobarblock", "pre_call"))
handler.sync_guardrail_from_db(db_guardrail("quxnewblock", "during_call"))
instance = handler.guardrail_id_to_custom_guardrail[gid]
assert isinstance(instance, ContentFilterGuardrail)
assert instance._check_blocked_words("hello QUXNEWBLOCK") is not None
assert instance._check_blocked_words("hello FOOBARBLOCK") is None
assert instance.event_hook == GuardrailEventHooks.during_call
assert instance.should_run_guardrail(data={}, event_type=GuardrailEventHooks.during_call) is True
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot

View file

@ -41,7 +41,9 @@ from litellm.litellm_core_utils.get_provider_specific_headers import (
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
TRUSTED_CALLBACK_VARS_FIELD,
)
from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id
from litellm.types.utils import CredentialItem
@ -7719,3 +7721,177 @@ def test_stamped_model_access_groups_survive_the_litellm_metadata_merge():
}
assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"]
def _request_for(path: str) -> MagicMock:
request = MagicMock(spec=Request)
request.scope = {"path": path}
request.url = MagicMock()
request.url.path = path
request.url.__str__.return_value = f"http://localhost{path}"
request.method = "POST"
request.query_params = {}
request.headers = {"Content-Type": "application/json"}
request.client = MagicMock()
request.client.host = "127.0.0.1"
return request
def _spend_log_session_id(data: dict[str, object]) -> str:
"""Resolve session_id the way LiteLLM_SpendLogs does: standard_logging_payload.trace_id."""
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.proxy.spend_tracking.spend_tracking_utils import _get_session_id_for_spend_log
metadata = data["metadata"]
assert isinstance(metadata, dict)
litellm_params = get_litellm_params(
litellm_session_id=str(data["litellm_session_id"]) if "litellm_session_id" in data else None,
litellm_trace_id=str(data["litellm_trace_id"]) if "litellm_trace_id" in data else None,
metadata=metadata,
)
trace_id = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id(
logging_obj=SimpleNamespace(litellm_trace_id="per-call-random-trace-id"),
litellm_params=litellm_params,
)
return _get_session_id_for_spend_log(kwargs={}, standard_logging_payload={"trace_id": trace_id})
@pytest.mark.asyncio
@pytest.mark.parametrize("request_correlation_in_logs", [False, True])
async def test_missing_session_id_generate_makes_spend_log_and_callback_session_ids_agree(
monkeypatch: pytest.MonkeyPatch, request_correlation_in_logs: bool
):
"""Without a session header, SpendLogs.session_id and the metadata.session_id that Langfuse logs
must be the same generated id, so cross-referencing the two by session_id works. The id is marked
as generated so affinity consumers (Fireworks x-session-affinity, router session pins) skip it."""
monkeypatch.setattr(litellm, "request_correlation_in_logs", request_correlation_in_logs)
data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
updated = await add_litellm_data_to_request(
data=data,
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "generate"},
)
callback_session_id = updated["metadata"]["session_id"]
assert isinstance(callback_session_id, str) and len(callback_session_id) == 36
assert _spend_log_session_id(updated) == callback_session_id
assert updated["metadata"][SESSION_ID_GENERATED_METADATA_KEY] is True
assert get_fireworks_session_id(
{"litellm_session_id": updated["litellm_session_id"], "metadata": updated["metadata"]}
) is None
@pytest.mark.asyncio
async def test_missing_session_id_unset_keeps_legacy_divergence():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={},
)
assert "session_id" not in updated["metadata"]
assert "litellm_session_id" not in updated
assert _spend_log_session_id(updated) == "per-call-random-trace-id"
@pytest.mark.asyncio
async def test_missing_session_id_generate_reuses_traceparent_trace_id():
"""A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it."""
request = _request_for("/v1/chat/completions")
request.headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "generate"},
)
assert updated["metadata"]["session_id"] == "4bf92f3577b34da6a3ce929d0e0e4736"
assert _spend_log_session_id(updated) == "4bf92f3577b34da6a3ce929d0e0e4736"
@pytest.mark.asyncio
@pytest.mark.parametrize("policy", ["generate", "reject"])
async def test_missing_session_id_policy_keeps_client_supplied_session_id(policy: str):
request = _request_for("/v1/chat/completions")
request.headers = {"x-litellm-session-id": "client-session-1"}
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=request,
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": policy},
)
assert updated["litellm_session_id"] == "client-session-1"
assert updated["metadata"]["session_id"] == "client-session-1"
assert _spend_log_session_id(updated) == "client-session-1"
assert SESSION_ID_GENERATED_METADATA_KEY not in updated["metadata"]
assert (
get_fireworks_session_id({"litellm_session_id": "client-session-1", "metadata": updated["metadata"]})
== "client-session-1"
)
@pytest.mark.asyncio
async def test_missing_session_id_reject_accepts_body_metadata_session_id():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": [], "metadata": {"session_id": "body-session-1"}},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "reject"},
)
assert updated["metadata"]["session_id"] == "body-session-1"
@pytest.mark.asyncio
async def test_missing_session_id_reject_returns_400_without_session_id():
with pytest.raises(ProxyException) as exc_info:
await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "reject"},
)
assert exc_info.value.code == "400"
assert exc_info.value.param == "session_id"
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/mcp/", "/mcp/tools", "/key/health"])
async def test_missing_session_id_policy_skips_non_inference_routes(path: str):
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o"},
request=_request_for(path),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "reject"},
)
assert "session_id" not in updated["metadata"]
@pytest.mark.asyncio
async def test_missing_session_id_unknown_value_is_ignored():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": []},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "typo"},
)
assert "session_id" not in updated["metadata"]

View file

@ -17,7 +17,7 @@ import litellm
from litellm import Router
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.router_strategy.complexity_router.complexity_router import (
_CLASSIFICATION_CURRENT_MESSAGE_ONLY,
_CLASSIFICATION_WITH_CONVERSATION,
@ -4294,6 +4294,26 @@ class TestSessionAffinity:
assert first.model == "o1-preview"
assert second.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config):
"""A session id the proxy generated for a request that had none is per request, so
it must not create a pin even with session_affinity enabled."""
mock_router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config=session_affinity_config,
)
request_kwargs = {"metadata": {"session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}}
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
)
second = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert first.model == "o1-preview"
assert second.model == "gpt-4o-mini"
@pytest.mark.asyncio
async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config):
"""Regression: session_affinity=True is the opt-in, so a shared session_id reuses the

View file

@ -7,7 +7,7 @@ import json
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
@ -180,6 +180,47 @@ async def test_async_session_id_affinity_priority_over_user_key():
assert filtered[0]["model_info"]["id"] == "deployment-2"
@pytest.mark.asyncio
async def test_proxy_generated_session_id_does_not_pin_a_deployment():
"""A session id the proxy generated for a request that had none is per request, so a
pin stored under it must be ignored and none must be written."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=123,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=True,
)
healthy_deployments = [
{"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}},
{"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}},
]
await cache.async_set_cache(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"),
{"model_id": "deployment-2"},
)
request_kwargs = {
"metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}
}
filtered = await callback.async_filter_deployments(
model="model_group", healthy_deployments=healthy_deployments, messages=[], request_kwargs=request_kwargs
)
await callback.async_pre_call_deployment_hook(
kwargs={
"metadata": {**request_kwargs["metadata"], "deployment_model_name": "model_group"},
"model_info": {"id": "deployment-1"},
},
call_type=None,
)
assert len(filtered) == 2
assert await cache.async_get_cache(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1")
) == {"model_id": "deployment-2"}
MOCK_RESPONSES_API_RESPONSE = {
"id": "resp_mock-resp-456",
"object": "response",

View file

@ -0,0 +1,86 @@
import sys
from pathlib import Path
from typing import Final
_CODE_COVERAGE_DIR: Final[Path] = Path(__file__).resolve().parents[1] / "code_coverage_tests"
sys.path.insert(0, str(_CODE_COVERAGE_DIR)) # test-quality-ok: required to import checker from its source directory
import check_py310_typing_imports as checker # noqa: E402 # load checker from its source directory
def _scan(tmp_path: Path, source: str) -> tuple[object, ...]:
file_path = tmp_path / "fixture.py"
file_path.write_text(source, encoding="utf-8")
return checker.scan_file(file_path)
def test_typing_import_flags_python_311_name(tmp_path: Path) -> None:
violations = _scan(tmp_path, "from typing import NotRequired, TypedDict\n")
assert tuple(violation.name for violation in violations) == ("NotRequired",)
def test_typing_extensions_import_passes(tmp_path: Path) -> None:
assert _scan(tmp_path, "from typing_extensions import NotRequired\n") == ()
def test_typing_attribute_flags_python_311_name(tmp_path: Path) -> None:
violations = _scan(tmp_path, "import typing\nx: typing.Self\n")
assert tuple(violation.name for violation in violations) == ("Self",)
def test_version_guarded_typing_import_passes(tmp_path: Path) -> None:
source = (
"import sys\n"
"if sys.version_info >= (3, 11):\n"
" from typing import NotRequired\n"
"else:\n"
" from typing_extensions import NotRequired\n"
)
assert _scan(tmp_path, source) == ()
def test_python_310_branch_flags_typing_import(tmp_path: Path) -> None:
source = (
"import sys\n"
"if sys.version_info >= (3, 11):\n"
" from typing_extensions import NotRequired\n"
"else:\n"
" from typing import NotRequired\n"
)
violations = _scan(tmp_path, source)
assert tuple(violation.name for violation in violations) == ("NotRequired",)
def test_python_310_branch_is_exempt_for_less_than_guard(tmp_path: Path) -> None:
source = (
"import sys\n"
"if sys.version_info < (3, 11):\n"
" from typing_extensions import NotRequired\n"
"else:\n"
" from typing import NotRequired\n"
)
assert _scan(tmp_path, source) == ()
def test_nearest_if_controls_version_guard(tmp_path: Path) -> None:
source = (
"if sys.version_info >= (3, 11):\n"
" from typing import Self\n"
" x = 1\n"
"if True:\n"
" from typing import Self\n"
)
violations = _scan(tmp_path, source)
assert tuple((violation.name, violation.line) for violation in violations) == (("Self", 5),)
def test_scan_directory_includes_proxy_extras(tmp_path: Path) -> None:
file_path = tmp_path / "litellm-proxy-extras" / "litellm_proxy_extras" / "m.py"
file_path.parent.mkdir(parents=True)
file_path.write_text("from typing import NotRequired\n", encoding="utf-8")
violations = checker.scan_directory(tmp_path)
assert tuple((violation.name, violation.file) for violation in violations) == (("NotRequired", str(file_path)),)
def test_python_310_typing_name_passes(tmp_path: Path) -> None:
assert _scan(tmp_path, "from typing import Optional\n") == ()

View file

@ -25772,6 +25772,11 @@ export interface components {
* @description Number of trusted reverse proxies/load balancers in front of the gateway that append to X-Forwarded-For. When set (and mcp_trusted_proxy_ranges validates the direct peer), the client IP for MCP access control is read this many entries from the right of the chain instead of the spoofable leftmost value, defeating append-style X-Forwarded-For forgery.
*/
mcp_xff_num_trusted_hops?: number | null;
/**
* Missing Session Id
* @description What to do with LLM API requests that carry no session id (x-litellm-session-id header, metadata.session_id, etc.). 'generate' stamps one id into litellm_session_id, litellm_trace_id and metadata.session_id so SpendLogs and logging callbacks agree; 'reject' returns 400. Unset keeps the legacy behavior where SpendLogs falls back to the trace id while callbacks get no session id.
*/
missing_session_id?: ("generate" | "reject") | null;
/**
* Model List Healthy Only
* @description When true, `/models`, `/v1/models/{id}` and `/model/info` hide models whose backing deployments are all unhealthy, for every caller, without needing `healthy_only=true` per request. Requires `background_health_checks: true`, and keeps deployment health state cached without turning on `enable_health_check_routing`, so routing is unaffected. With no health state nothing is hidden. Hiding is presentation-only, a hidden model can still be called.