diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 5addc994df2..9332137e3b1 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -898,6 +898,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def redacts_messages_itself(self) -> bool: return False + def exempt_from_message_redaction(self) -> bool: + """Router pre-call checks that key routing off the request and never log it return True, so + `turn_off_message_logging` leaves the success event they receive unredacted.""" + return False + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1d6250fc528..95323c2ef5e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3751,7 +3751,11 @@ class Logging(LiteLLMLoggingBaseClass): if isinstance(callback, CustomLogger): # custom logger class from litellm.integrations.shadow_eval_logger import ShadowEvalLogger - model_call_details: dict = self.model_call_details + model_call_details: dict = ( + self._model_call_details_with_request_messages() + if callback.exempt_from_message_redaction() + else self.model_call_details + ) ################################## # call redaction hook for custom logger model_call_details = callback.redact_standard_logging_payload_from_model_call_details( @@ -4289,6 +4293,21 @@ class Logging(LiteLLMLoggingBaseClass): ) -> None: executor.submit(self.success_handler, result, start_time, end_time, cache_hit) + def _model_call_details_with_request_messages(self) -> dict[str, object]: # mutable-ok: loggers take a dict + """The success-event payload with the request messages put back where `turn_off_message_logging` + redacted them in place, for a callback that routes on the request and never logs it.""" + request_messages: Final = cast(object, self.messages) # cast-ok: the request messages are only put back as-is + standard_logging_object: Final = self.model_call_details.get("standard_logging_object") + return { + **self.model_call_details, + "messages": request_messages, + **( + {"standard_logging_object": {**standard_logging_object, "messages": request_messages}} + if isinstance(standard_logging_object, dict) + else {} + ), + } + def _should_run_sync_callbacks_for_async_calls(self) -> bool: """ Returns: diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index bf33070faa6..b4df0624e68 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1,3 +1,4 @@ +import inspect import time import types from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping @@ -351,6 +352,11 @@ _EXTRA_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str] | None) _EXTRA_QUERY_ADAPTER: Final = TypeAdapter(dict[str, object] | None) _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body")) +_CLIENT_API_KEY_ONLY_AUTH: Final[Mapping[str, Mapping[str, bool]]] = ( + types.MappingProxyType({"security": types.MappingProxyType({"bearer_auth": True})}) + if "security" in inspect.signature(make_request_options).parameters + else types.MappingProxyType({}) +) def _raw_chat_response( @@ -370,6 +376,7 @@ def _request_without_sdk_transform( extra_query=_EXTRA_QUERY_ADAPTER.validate_python(data.get("extra_query")), extra_body=data.get("extra_body"), timeout=timeout, + **_CLIENT_API_KEY_ONLY_AUTH, # pyright: ignore[reportArgumentType] # openai<3 make_request_options has no security kwarg ) return body, options diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index e9ec040023c..abbeab7f89d 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -55,3 +55,10 @@ class ToolDiscoveryQueue: """Return and clear all pending rows.""" items, self._pending = list(self._pending.values()), {} return items + + def requeue(self, items: Iterable[ToolDiscoveryQueueItem]) -> None: + """Put rows from a failed write back, ahead of and merged with anything queued since the flush.""" + since_flush: Final = self._pending + self._pending = {} + for item in (*items, *since_flush.values()): + self.add_update(item) diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 669df7454e3..23d83829d8c 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -31,6 +31,24 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( "too many connections for database", "remaining connection slots are reserved", ) +_NOT_ACCEPTING_CONNECTIONS_SQLSTATE: Final = "57P03" +_NOT_ACCEPTING_CONNECTIONS_PHRASES: Final = ( + "the database system is starting up", + "the database system is shutting down", + "the database system is not yet accepting connections", + "the database system is not accepting connections", + "the database system is in recovery mode", +) +_SERVER_TERMINATED_SQLSTATES: Final = frozenset({"57P01", "57P02", "57P05"}) +_SERVER_TERMINATED_PHRASES: Final = ( + "terminating connection due to administrator command", + "terminating connection due to immediate shutdown command", + "terminating connection because of unexpected sigquit signal", + "terminating connection because of crash of another server process", + "terminating connection due to idle-session timeout", +) +_UNREACHABLE_PRISMA_CODE: Final = "P1001" +_UNREACHABLE_PHRASES: Final = ("can't reach database server", "cannot reach database server") _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" @@ -74,6 +92,22 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]: return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException)) +def _prisma_error_code(e: Exception) -> str | None: + import prisma.engine.errors + + if isinstance(e, _exception_types(prisma.errors.DataError)): + code: Final[object] = getattr(e, "code", None) + return code if isinstance(code, str) else None + if not isinstance(e, _exception_types(prisma.engine.errors.EngineRequestError)): + return None + try: + body: Final = _DATABASE_ERROR_META.validate_json(str(e).partition(": ")[2]) + except ValidationError: + return None + engine_code: Final = body.get("error_code") + return engine_code if isinstance(engine_code, str) else None + + class PrismaDBExceptionHandler: """ Class to handle DB Exceptions or Connection Errors @@ -123,6 +157,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection @staticmethod @@ -271,6 +307,39 @@ class PrismaDBExceptionHandler: error_message: Final = str(e).lower() return any(phrase in error_message for phrase in _CONNECTION_CAPACITY_PHRASES) + @staticmethod + def is_database_not_accepting_connections_error(e: Exception) -> bool: + """True iff Postgres refused the connection because it is starting up, shutting + down or in crash recovery (SQLSTATE 57P03). The refusal comes before any statement + runs, so nothing the write carried was applied.""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) == _NOT_ACCEPTING_CONNECTIONS_SQLSTATE: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _NOT_ACCEPTING_CONNECTIONS_PHRASES) + + @staticmethod + def is_database_lifecycle_error(e: Exception) -> bool: + """True iff Postgres failed the statement because of a server lifecycle event: + a refused connection during startup, shutdown or recovery (57P03), or a session + the server terminated (57P01, 57P02, 57P05). The server comes back on its own, + so the rows are not at fault. A terminated session may have committed before it + was cut, so only ``is_database_not_accepting_connections_error`` proves a + non-idempotent write never applied.""" + import prisma + + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) in _SERVER_TERMINATED_SQLSTATES: + return True + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _SERVER_TERMINATED_PHRASES) + @staticmethod def is_lock_timeout_error(e: Exception) -> bool: """True iff Postgres cancelled a statement for exceeding ``lock_timeout`` @@ -388,6 +457,8 @@ class PrismaDBExceptionHandler: return True if PrismaDBExceptionHandler.is_database_capacity_error(e): return True + if PrismaDBExceptionHandler.is_database_lifecycle_error(e): + return True if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): return True if "cached plan must not change result type" in str(e).lower(): @@ -662,3 +733,18 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool: return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES + + +def db_write_never_sent(e: Exception) -> bool: + """Whether the database could not take the write at all, because it was unreachable (P1001), + refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again.""" + if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): + return True + if PrismaDBExceptionHandler.is_database_capacity_error(e): + return True + if _prisma_error_code(e) == _UNREACHABLE_PRISMA_CODE: + return True + if not PrismaDBExceptionHandler.is_prisma_error(e): + return False + error_message: Final = str(e).lower() + return any(phrase in error_message for phrase in _UNREACHABLE_PHRASES) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index f7ec88945a1..9600fb40ad3 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -23,7 +23,11 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, aggregate_tool_discoveries, ) -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.exception_handler import ( + call_with_db_reconnect_retry, + db_write_failure_is_requeue_safe, + db_write_never_sent, +) from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions @@ -134,8 +138,19 @@ async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discove return try: await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail - verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + except Exception as e: # noqa: BLE001 # the rows go back to the queue or are dropped, the scheduler job must not fail + if db_write_never_sent(e): + tool_discovery_queue.requeue(items) + verbose_proxy_logger.warning( + "Tool discovery - %d tool rows went back to the queue for the next tick. Error: %s", len(items), str(e) + ) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) async def commit_tool_discoveries( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 486028297b3..8479f8e45b9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -198,6 +198,7 @@ from litellm.proxy.db.db_url_settings import ( from litellm.proxy.db.exception_handler import ( PrismaDBExceptionHandler, call_with_db_reconnect_retry, + db_write_never_sent, ) from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, @@ -8500,9 +8501,20 @@ async def update_daily_tag_spend( def _rollup_batch_never_applied(e: Exception) -> bool: """True when a drained rollup batch provably never reached its rows: Postgres had no - connection for it (53300) or cancelled it under ``lock_timeout`` before it took the row - lock (55P03). Either way re-sending the batch cannot double-count, so it is requeued.""" - return PrismaDBExceptionHandler.is_database_capacity_error(e) or PrismaDBExceptionHandler.is_lock_timeout_error(e) + connection for it (53300), was unreachable or refusing connections (P1001, 57P03), or + cancelled it under ``lock_timeout`` before it took the row lock (55P03). In each case + re-sending the batch cannot double-count, so it is requeued.""" + return ( + PrismaDBExceptionHandler.is_database_capacity_error(e) + or db_write_never_sent(e) + or PrismaDBExceptionHandler.is_lock_timeout_error(e) + ) + + +def _rollup_failure_stops_job(e: Exception) -> bool: + """True when the database cannot take any write, so the job stops instead of draining + the queues after this one into the same failure.""" + return PrismaDBExceptionHandler.is_database_capacity_error(e) or db_write_never_sent(e) async def update_spend_logs_job( @@ -8593,12 +8605,12 @@ async def _run_spend_logs_job( *prisma_client.tool_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during tool usage " + "Spend tracking - database unavailable or rollup row lock timed out during tool usage " "flush, requeued %s tool usage transactions for the next flush: %s", len(tool_usage_to_process), tool_tracking_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(tool_tracking_err): + if _rollup_failure_stops_job(tool_tracking_err): raise else: verbose_proxy_logger.error( @@ -8625,12 +8637,12 @@ async def _run_spend_logs_job( *prisma_client.model_usage_transactions, ] verbose_proxy_logger.warning( - "Spend tracking - database out of connections or rollup row lock timed out during model usage " + "Spend tracking - database unavailable or rollup row lock timed out during model usage " "flush, requeued %s model usage transactions for the next flush: %s", len(model_usage_to_process), model_usage_err, ) - if PrismaDBExceptionHandler.is_database_capacity_error(model_usage_err): + if _rollup_failure_stops_job(model_usage_err): raise else: verbose_proxy_logger.error( @@ -8807,6 +8819,7 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool: return ( PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) + or PrismaDBExceptionHandler.is_database_lifecycle_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e) ) diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index b2f0c658161..435f9b9c5a4 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -138,6 +138,9 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments + def exempt_from_message_redaction(self) -> bool: + return True + async def async_log_success_event( self, kwargs: _PromptCachingSuccessEventKwargs, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 77160e0adf1..c6c6cad6f72 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -653,6 +653,14 @@ class CallTypes(str, Enum): retrieve_fine_tuning_job = "retrieve_fine_tuning_job" responses = "responses" aresponses = "aresponses" + get_responses = "get_responses" + aget_responses = "aget_responses" + delete_responses = "delete_responses" + adelete_responses = "adelete_responses" + cancel_responses = "cancel_responses" + acancel_responses = "acancel_responses" + compact_responses = "compact_responses" + acompact_responses = "acompact_responses" alist_input_items = "alist_input_items" llm_passthrough_route = "llm_passthrough_route" allm_passthrough_route = "allm_passthrough_route" @@ -754,6 +762,8 @@ CallTypesLiteral = Literal[ "send_message", "aresponses", "responses", + "compact_responses", + "acompact_responses", "acreate_skill", "acreate_realtime_client_secret", "arealtime_calls", diff --git a/tests/integration/_support/redaction_probe.py b/tests/integration/_support/redaction_probe.py index 2a344daa73c..83d1cbd9cba 100644 --- a/tests/integration/_support/redaction_probe.py +++ b/tests/integration/_support/redaction_probe.py @@ -19,6 +19,7 @@ def _payload(kwargs: Mapping[str, object], response_obj: object) -> str: "request": request, "response": response_obj.model_dump() if isinstance(response_obj, BaseModel) else response_obj, "standard_response": standard.get("response") if isinstance(standard, Mapping) else None, + "standard": standard, }, default=str, ) diff --git a/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py b/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py new file mode 100644 index 00000000000..3227caaa7f5 --- /dev/null +++ b/tests/integration/authorization/test_responses_lifecycle_during_call_guardrail.py @@ -0,0 +1,114 @@ +"""GET, compact, DELETE and cancel on /v1/responses answer 200 when a during_call guardrail is on by default +(LIT-9770): the hook used to raise `'acompact_responses' is not a valid CallTypes` and the proxy returned 500""" + +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, string_value +from integration._support.process import owned_proxy +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.shares_worker_state + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MODEL: Final = "openai/gpt-5.4-mini" + + +def _input() -> list[JsonValue]: + return [ + {"role": "user", "content": f"Remember the number {uuid.uuid4().hex}."}, + {"role": "assistant", "content": "Got it."}, + ] + + +def _response(response_id: str) -> Reply: + body: Final[dict[str, JsonValue]] = { + "id": response_id, + "object": "response", + "created_at": 1, + "status": "completed", + "model": _MODEL.removeprefix("openai/"), + "output": [ + { + "type": "message", + "id": f"msg_{uuid.uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Got it.", "annotations": []}], + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } + return Reply(body=json.dumps(body).encode()) + + +def _deleted(response_id: str) -> Reply: + return Reply(body=json.dumps({"id": response_id, "object": "response.deleted", "deleted": True}).encode()) + + +def _guarded_config(path: Path) -> Path: + base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + guardrail: Final[dict[str, JsonValue]] = { + "guardrail_name": f"content-filter-{uuid.uuid4().hex}", + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "during_call", + "default_on": True, + "blocked_words": [{"keyword": "forbiddenword", "action": "BLOCK"}], + }, + } + path.write_text(yaml.safe_dump({**base, "guardrails": [guardrail]})) + return path + + +@pytest.fixture(scope="module") +def guarded(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("during-call-guardrail") + with owned_proxy(session_gateway, directory, {}, config=_guarded_config(directory / "config.yaml")) as candidate: + yield candidate + + +def _create(guarded: Gateway, provider: SharedProvider, upstream_id: str) -> str: + provider.expect(_response(upstream_id)) + created: Final = guarded.request("POST", "/v1/responses", {"model": _MODEL, "input": _input(), "store": True}) + assert created.status_code == 200, created.text + (upstream,) = provider.received() + assert upstream.target.endswith("/responses"), upstream.target + return string_value(_JSON_OBJECT.validate_json(created.content)["id"]) + + +def _lifecycle(guarded: Gateway, provider: SharedProvider, operation: str) -> tuple[httpx.Response, str]: + if operation == "compact": + provider.expect(_response(f"resp_{uuid.uuid4().hex}")) + return guarded.request( + "POST", "/v1/responses/compact", {"model": _MODEL, "input": _input()} + ), "/responses/compact" + upstream_id: Final = f"resp_{uuid.uuid4().hex}" + managed_id: Final = _create(guarded, provider, upstream_id) + provider.expect(_deleted(upstream_id) if operation == "delete" else _response(upstream_id)) + if operation == "get": + return guarded.request("GET", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}" + if operation == "delete": + return guarded.request("DELETE", f"/v1/responses/{managed_id}"), f"/responses/{upstream_id}" + return guarded.request("POST", f"/v1/responses/{managed_id}/cancel"), f"/responses/{upstream_id}/cancel" + + +@pytest.mark.parametrize("operation", ["get", "compact", "delete", "cancel"]) +def test_responses_lifecycle_answers_200_with_a_during_call_guardrail_on( + guarded: Gateway, provider: SharedProvider, operation: str +) -> None: + response, upstream_target = _lifecycle(guarded, provider, operation) + assert response.status_code == 200, response.text + (upstream,) = provider.received() + assert upstream.target.endswith(upstream_target), upstream.target diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index e82430c3add..7b0149fd2dd 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -68,6 +68,7 @@ def pytest_addoption(parser: pytest.Parser) -> None: def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line("markers", "integration: owned real-service integration contracts") config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced") + config.addinivalue_line("markers", "shares_worker_state: owned proxies here share the worker DB and Redis") config.stash[REPORTS] = [] config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files"))) config.pluginmanager.register(IntegrationReportPlugin(config)) diff --git a/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py new file mode 100644 index 00000000000..13bf67582b3 --- /dev/null +++ b/tests/integration/routing/test_prompt_cache_pin_with_message_logging_off.py @@ -0,0 +1,95 @@ +import collections +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import eventually, gateway_from_environment, object_value +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.wire import Reply, Wire, wire_server +from integration.routing._prompt_cache_pin_support import ( + DEFAULT_MARK, + anthropic_deployment, + ask_turn, + chat_body, + claude_peer, + new_marker, + one_new_pin, + pin_config, + post_chat, + read_pins, + system_turn, +) +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +from litellm.constants import REDACTED_BY_LITELLM + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +_GROUP: Final = "pin-logging-off" +_LEFT: Final = "pin-logging-off-left" +_RIGHT: Final = "pin-logging-off-right" +_FOLLOW_UPS: Final = 20 +_REDACTED_MESSAGES: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}] +_LOGGED_EVENT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + redis: "Redis[bytes]" + logger_sink: Wire + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("prompt-cache-pin-message-logging-off") + with ( + gateway_from_environment() as gateway, + owned_redis(directory) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + wire_server(lambda request: Reply()) as logger_sink, + ): + config: Final = pin_config( + directory, + [anthropic_deployment(_GROUP, _LEFT, left.url), anthropic_deployment(_GROUP, _RIGHT, right.url)], + litellm_settings={ + "turn_off_message_logging": True, + "callbacks": ["integration._support.redaction_probe.probe"], + "cache": False, + "enable_redis_auth_cache": False, + }, + ) + environment: Final = { + "REDIS_HOST": cache.host, + "REDIS_PORT": str(cache.port), + "REDACTION_PROBE_URL": logger_sink.url, + } + with ( + owned_proxy_process(gateway, directory, environment, config=config, workers=1) as owned, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + yield _Rig(owned, redis, logger_sink) + + +def test_follow_ups_stay_on_the_first_deployment_while_the_logger_sees_redacted_messages(rig: _Rig) -> None: + before: Final = frozenset(read_pins(rig.redis)) + body: Final = chat_body(_GROUP, [system_turn(new_marker(), DEFAULT_MARK), ask_turn(new_marker())]) + + _, first_served = post_chat(rig.owned.gateway, body) + _, pin = one_new_pin(rig.redis, before) + assert pin.model_id == first_served, (pin, first_served) + + served: Final = tuple(post_chat(rig.owned.gateway, body)[1] for _ in range(_FOLLOW_UPS)) + assert served == (first_served,) * _FOLLOW_UPS, collections.Counter(served) + + eventually(rig.logger_sink.received.qsize, lambda count: count >= 1 + _FOLLOW_UPS, seconds=30) + logged: Final = tuple( + object_value(_LOGGED_EVENT.validate_json(request.body)["standard"])["messages"] + for request in rig.logger_sink.drain() + ) + assert len(logged) == 1 + _FOLLOW_UPS, len(logged) + assert all(messages == _REDACTED_MESSAGES for messages in logged), logged diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index b277cac1f29..98e3b84ff24 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -456,6 +456,10 @@ def test_default_max_retries_env_var_reaches_azure_sdk_client(): "add_message", "arun_thread_stream", "aresponses", + "aget_responses", + "adelete_responses", + "acancel_responses", + "acompact_responses", "aresponses_websocket", "alist_input_items", "acreate_fine_tuning_job", diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 749dc823f20..961b6c88012 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -1,4 +1,6 @@ import asyncio +import contextlib +import inspect import importlib import json import os @@ -1937,3 +1939,49 @@ async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_head assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True} assert sent.headers["x-caller"] == "unit" assert sent.headers["authorization"] == "Bearer sk-test" + + +_SDK_READS_OPENAI_ADMIN_KEY: Final = "admin_api_key" in inspect.signature(AsyncOpenAI.__init__).parameters + + +@pytest.mark.skipif(not _SDK_READS_OPENAI_ADMIN_KEY, reason="openai<3 has no OPENAI_ADMIN_KEY fallback to leak") +@pytest.mark.asyncio +async def test_acompletion_never_sends_the_openai_admin_key_for_a_deployment_without_a_key( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_ADMIN_KEY", "sk-admin-only") + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + outbound: Final = asyncio.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + outbound.put_nowait(request.headers.get("authorization")) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-no-admin-key", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + monkeypatch.setattr(litellm, "aclient_session", http_client) + with contextlib.suppress(litellm.AuthenticationError, litellm.InternalServerError): + await asyncio.wait_for( + litellm.acompletion( + model="openai/gpt-5.4-mini", + api_key="", + messages=[{"role": "user", "content": "who signs this request"}], + num_retries=0, + max_retries=0, + ), + timeout=10, + ) + + sent: Final = tuple(outbound.get_nowait() for _ in range(outbound.qsize())) + assert all(authorization is None for authorization in sent), sent diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index 0a015255f2e..814fdcb4eda 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -2,7 +2,6 @@ Unit tests for ToolDiscoveryQueue. """ - import pytest from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( @@ -63,3 +62,19 @@ def test_aggregate_sums_call_counts_across_workers_flushes(): {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, {"tool_name": "search", "call_count": 1}, ) + + +def test_requeued_rows_go_ahead_of_and_merge_with_sightings_since_the_flush(queue: ToolDiscoveryQueue) -> None: + queue.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + failed_write = queue.flush() + queue.add_update({"tool_name": "translate", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + + queue.requeue(failed_write) + + assert queue.flush() == [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 1}, + {"tool_name": "search", "origin": "user_defined", "call_count": 2}, + {"tool_name": "translate", "origin": "user_defined", "call_count": 1}, + ] diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index a4c16c0949b..6bf5f6b94eb 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -22,7 +22,7 @@ from prisma.errors import ( import litellm from litellm.proxy._types import ProxyErrorTypes, ProxyException -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_never_sent # Test is_database_connection_error method @@ -966,3 +966,137 @@ def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None ) def test_is_statement_timeout_error(error: Exception, expected: bool) -> None: assert PrismaDBExceptionHandler.is_statement_timeout_error(error) is expected + + +def _connector_error(message: str) -> DataError: + return DataError( + data={"user_facing_error": {"message": f"Error in connector: Error querying the database: FATAL: {message}"}} + ) + + +def _batched_postgres_error(sqlstate: str, message: str) -> DataError: + return DataError( + data={ + "user_facing_error": { + "message": "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: " + f'None, kind: QueryError(PostgresError {{ code: "{sqlstate}", message: "{message}", severity: "FATAL" }}) }})' + } + } + ) + + +def _localized_error(sqlstate: str) -> DataError: + return DataError(data={"user_facing_error": {"message": "FATAL: Datenbank", "meta": {"code": sqlstate}}}) + + +_NOT_ACCEPTING_CONNECTIONS: Final = ( + _connector_error("the database system is starting up"), + _connector_error("the database system is shutting down"), + _connector_error("the database system is not yet accepting connections"), + _connector_error("the database system is not accepting connections"), + _connector_error("the database system is in recovery mode"), + _localized_error("57P03"), + _batched_postgres_error("57P03", "the database system is starting up"), + EngineRequestError( + MagicMock(status=500), + '{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: the database system is ' + 'starting up","backtrace":null}', + ), + RawQueryError( + data={ + "user_facing_error": { + "message": "Raw query failed. Code: `57P03`. Message: `db error: FATAL: the database system is starting up`", + "meta": {"code": "57P03", "message": "FATAL: the database system is starting up"}, + "error_code": "P2010", + } + } + ), +) +_SERVER_TERMINATED: Final = ( + _connector_error("terminating connection due to administrator command"), + _connector_error("terminating connection due to immediate shutdown command"), + _connector_error("terminating connection because of unexpected SIGQUIT signal"), + _connector_error("terminating connection because of crash of another server process"), + _connector_error("terminating connection due to idle-session timeout"), + _localized_error("57P01"), + _localized_error("57P02"), + _localized_error("57P05"), + _batched_postgres_error("57P01", "terminating connection due to administrator command"), +) + + +@pytest.mark.parametrize( + ("error", "connection_error"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + ], +) +def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Exception, connection_error: bool) -> None: + """A server that is restarting, shutting down or cutting sessions is a transient outage: + auth answers 503 rather than 401, and the spend-log writer must not take the rows for poison. + Only a server refusing connections opens the ``allow_requests_on_db_unavailable`` fallback, + one terminated session does not prove the database is down.""" + assert { + "lifecycle": PrismaDBExceptionHandler.is_database_lifecycle_error(error), + "service_unavailable": PrismaDBExceptionHandler.is_database_service_unavailable_error(error), + "connection_error": PrismaDBExceptionHandler.is_database_connection_error(error), + "permanent": PrismaDBExceptionHandler.is_permanent_database_fault(error), + } == { + "lifecycle": True, + "service_unavailable": True, + "connection_error": connection_error, + "permanent": False, + } + + +@pytest.mark.parametrize( + ("error", "never_sent"), + [ + *((error, True) for error in _NOT_ACCEPTING_CONNECTIONS), + *((error, False) for error in _SERVER_TERMINATED), + (DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}), True), + (DataError(data={"user_facing_error": {"error_code": "P1001", "message": "unreachable"}}), True), + (_capacity_error("sorry, too many clients already"), True), + (_localized_error("53300"), True), + ( + DataError( + data={ + "user_facing_error": { + "error_code": "P2024", + "message": "Timed out fetching a new connection from the connection pool", + } + } + ), + True, + ), + (httpx.ReadTimeout("read timed out"), False), + (_localized_error("57014"), False), + ], +) +def test_db_write_never_sent_only_for_failures_before_any_statement_ran(error: Exception, never_sent: bool) -> None: + """A refused, unreachable or full server provably applied nothing, while a terminated session + or a timeout may have committed before the failure, so only the former lets a non-idempotent + write be replayed.""" + assert db_write_never_sent(error) is never_sent + + +@pytest.mark.parametrize( + "error", + [ + _localized_error("57P04"), + _localized_error("57014"), + _localized_error("23505"), + DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}), + RuntimeError("the database system is starting up"), + ], +) +def test_other_failures_are_not_postgres_lifecycle_errors(error: Exception) -> None: + assert PrismaDBExceptionHandler.is_database_lifecycle_error(error) is False + assert db_write_never_sent(error) is False + + +def test_lifecycle_error_is_seen_through_a_wrapping_exception() -> None: + wrapped: Final = RuntimeError("spend flush failed") + wrapped.__cause__ = _connector_error("the database system is starting up") + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(wrapped) is True diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index 76898b0eeab..a5ad897913e 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -69,9 +69,7 @@ def _make_prisma( prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) prisma.db.execute_raw = AsyncMock(return_value=1) - prisma.db.litellm_tooltable.find_many = AsyncMock( - return_value=find_many_rows if find_many_rows is not None else [] - ) + prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else []) prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row) prisma.db.litellm_verificationtoken = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows)) @@ -227,9 +225,7 @@ async def test_get_tool_found(): result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" - prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio @@ -266,9 +262,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): updated_by="admin", ) prisma = _make_prisma(find_unique_row=row) - result = await update_tool_policy( - prisma, "my_tool", updated_by="admin", input_policy="blocked" - ) + result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked") assert result is not None assert result.input_policy == "blocked" prisma.db.litellm_tooltable.upsert.assert_awaited_once() @@ -567,3 +561,42 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (_postgres_rejection("57P03"), True), + (_postgres_rejection("53300"), True), + (_postgres_rejection("57P01"), False), + (_postgres_rejection("54000"), False), + (_postgres_rejection("23502"), False), + ], + ids=[ + "restarting_server_keeps_the_rows", + "full_server_keeps_the_rows", + "terminated_session_may_have_counted_so_the_rows_are_dropped", + "oversized_value_is_dropped", + "rejected_row_is_dropped", + ], +) +async def test_failed_write_without_redis_keeps_the_rows_in_the_queue_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=_InMemoryRedis()), + pod_lock_manager=PodLockManager(redis_cache=_InMemoryRedis()), + use_redis_buffer=False, + ) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + expected = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2 if requeued else 1}] + assert queue.flush() == expected diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 3efbe6475b2..1417b87417c 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3053,3 +3053,21 @@ async def test_pre_call_hook_skips_attachment_scan_when_guardrail_has_none(): ) assert len(guardrail.apply_calls) == 1 + + +@pytest.mark.asyncio +async def test_moderation_hook_lets_responses_compact_through_without_a_translation() -> None: + handler: Final = UnifiedLLMGuardrails() + guardrail: Final = RecordingGuardrail() + data: Final = { + "guardrail_to_apply": guardrail, + "model": "gpt-5.4-mini", + "input": [{"role": "user", "content": "Remember the number 42."}], + } + + result: Final = await handler.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), call_type="acompact_responses" + ) + + assert result is data + assert guardrail.apply_calls == [] diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py index fc7ae531c34..293f2d49f43 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -23,9 +23,9 @@ import pytest from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, + drain_spend_logs_queue, monitor_spend_logs_queue, raise_failed_update_spend_exception, - drain_spend_logs_queue, recover_parked_spend_logs, update_daily_tag_spend, update_spend, @@ -985,6 +985,85 @@ async def test_update_spend_logs_job_requeues_whole_batch_when_postgres_is_out_o } == {"create_many_calls": 1, "queue_after": ["r1", "r2", "r3"], "backoff_sleeps": []} +def _postgres_restarting() -> Exception: + """prisma's shape for Postgres SQLSTATE 57P03 while the server runs crash recovery after + a restart: a base ``DataError`` carrying only the connector message.""" + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: the database system is starting up", + "backtrace": None, + } + } + ) + + +def _postgres_unreachable() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Can't reach database server at `db.internal`:`5432`", + "error_code": "P1001", + } + } + ) + + +def _postgres_terminated_session() -> Exception: + from prisma.errors import DataError + + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": "Error in connector: Error querying the database: FATAL: terminating connection due to " + "administrator command", + "meta": {"code": "57P01", "message": "terminating connection due to administrator command"}, + } + } + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_terminated_session()], ids=["57P03", "57P01"]) +async def test_update_spend_logs_job_keeps_every_row_while_postgres_restarts( + error: Exception, + mock_prisma_client: MagicMock, + make_spend_log_row: Callable[..., dict[str, object]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A restarting server refuses or cuts every statement, so the batch is not bisected + into per-row drops: it is retried whole and then goes back to the queue head.""" + monkeypatch.setattr(asyncio, "sleep", AsyncMock()) + proxy_logging: Final = MagicMock() + proxy_logging.failure_handler = AsyncMock() + mock_prisma_client.spend_log_transactions = [ + make_spend_log_row(request_id="r1"), + make_spend_log_row(request_id="r2"), + make_spend_log_row(request_id="r3"), + make_spend_log_row(request_id="r4"), + ] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=error) + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert { + "create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count, + "queue_after": [row["request_id"] for row in mock_prisma_client.spend_log_transactions], + } == {"create_many_calls": 4, "queue_after": ["r1", "r2", "r3", "r4"]} + + def _tool_usage_transaction(request_id: str) -> object: from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -1177,6 +1256,53 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out } == {"queue_after": [first, second], "autorouter_queue_after": 1} +@pytest.mark.asyncio +@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"]) +async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes( + error: Exception, + mock_prisma_client: MagicMock, +) -> None: + """A server that is restarting or unreachable never saw the rollup batch, so the tool + usage increments are kept for the next flush, and the job stops before draining the + auto-router queue into the same outage.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=error) + mock_prisma_client.spend_log_transactions = [] + first, second = _tool_usage_transaction("r1"), _tool_usage_transaction("r2") + mock_prisma_client.tool_usage_transactions = [first, second] + mock_prisma_client.autorouter_turn_transactions = [MagicMock(name="autorouter_turn")] + + with suppress(Exception): + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert { + "queue_after": mock_prisma_client.tool_usage_transactions, + "autorouter_queue_after": len(mock_prisma_client.autorouter_turn_transactions), + } == {"queue_after": [first, second], "autorouter_queue_after": 1} + + +@pytest.mark.asyncio +async def test_rollups_drop_a_batch_whose_session_the_server_terminated( + mock_prisma_client: MagicMock, +) -> None: + """A terminated session (57P01) may have committed before it was cut, and the rollup + increments are not idempotent, so the batch is dropped rather than counted twice.""" + mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=_postgres_terminated_session()) + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [_tool_usage_transaction("r1")] + + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert mock_prisma_client.tool_usage_transactions == [] + + @pytest.mark.asyncio async def test_model_usage_flush_still_drops_a_batch_that_may_have_applied( mock_prisma_client: MagicMock, diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 5f018ab5dda..5b31e0e5e3b 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1236,3 +1236,58 @@ async def test_router_with_prompt_caching(anthropic_messages): print("response=", response) assert response._hidden_params["model_id"] == initial_model_id + + +class _LoggedMessagesProbe(CustomLogger): + def __init__(self) -> None: + self.logged_messages: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_messages.append((kwargs["standard_logging_object"]["messages"], kwargs["messages"])) + + +@pytest.mark.asyncio +async def test_pin_is_written_with_message_logging_off_while_loggers_still_see_redacted_messages( + monkeypatch, local_model_cost_map +): + """ + `turn_off_message_logging` redacted the shared success-event payload before the check read it, so the + pin was keyed off `redacted-by-litellm`, never written, and a cacheable conversation bounced across the + group. The pin must land whatever the flag says, while a real logger alongside still sees only the + redacted messages. + """ + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + probe = _LoggedMessagesProbe() + monkeypatch.setattr(litellm, "callbacks", [probe]) + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + for model_id in ("dep-1", "dep-2") + ], + optional_pre_call_checks=["prompt_caching"], + ) + messages = _messages(word_count=5000) + + first = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok") + served_by = first._hidden_params["model_id"] + + affinity_key = PromptCachingCache.get_prompt_caching_cache_key(messages, None) + pin = await _eventually(lambda: router.prompt_caching_pin_cache.get_cache(key=affinity_key)) + assert pin == {"model_id": served_by}, pin + + subsequent = [ + (await router.acompletion(model=MODEL_GROUP_ALIAS, messages=messages, mock_response="ok"))._hidden_params[ + "model_id" + ] + for _ in range(4) + ] + assert subsequent == [served_by] * 4 + + await _eventually(lambda: len(probe.logged_messages) == 5) + redacted = [{"role": "user", "content": "redacted-by-litellm"}] + assert probe.logged_messages == [(redacted, redacted)] * 5, probe.logged_messages diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0a408e7c8ab..a9e3aa7da5d 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -98,6 +98,7 @@ from litellm.utils import ( _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, async_post_call_success_deployment_hook, + async_pre_call_deployment_hook, calculate_max_parallel_requests, client, get_non_default_completion_params, @@ -10052,3 +10053,25 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ ) assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} + + +class _RecordingPreCallDeploymentHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen_call_types: tuple[CallTypes | None, ...] = () + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.seen_call_types = (*self.seen_call_types, call_type) + return None + + +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_hands_responses_compact_its_call_type(monkeypatch: pytest.MonkeyPatch) -> None: + hook: Final = _RecordingPreCallDeploymentHook() + monkeypatch.setattr(litellm, "callbacks", [hook]) + + await async_pre_call_deployment_hook({"model": "gpt-5.4-mini"}, "acompact_responses") + + assert hook.seen_call_types == (CallTypes.acompact_responses,) diff --git a/tests/unit/types/test_utils.py b/tests/unit/types/test_utils.py index 32479c62372..bdbcbb698d4 100644 --- a/tests/unit/types/test_utils.py +++ b/tests/unit/types/test_utils.py @@ -1,7 +1,12 @@ import pytest from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_excluded_services_value -from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VALUES, CaptureMessageContent, captures_span_content +from litellm.types.utils import ( + CAPTURE_MESSAGE_CONTENT_VALUES, + CallTypes, + CaptureMessageContent, + captures_span_content, +) @pytest.mark.parametrize("value", ["", " ", "auth", "none,redis", "redis,", ",postgres"]) @@ -45,3 +50,8 @@ def test_capture_message_content_constructor_accepts_any_string() -> None: ) def test_captures_span_content(mode: str | None, captures: bool) -> None: assert captures_span_content(mode) is captures + + +@pytest.mark.parametrize("call_type", ["compact_responses", "acompact_responses"]) +def test_responses_compact_call_types_round_trip_their_public_names(call_type: str) -> None: + assert CallTypes(call_type).value == call_type diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 89cc9157359..13b84896a84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28370,7 +28370,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "decisions" | "adecisions" | "systemone" | "asystemone" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "get_responses" | "aget_responses" | "delete_responses" | "adelete_responses" | "cancel_responses" | "acancel_responses" | "compact_responses" | "acompact_responses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ diff --git a/uv.lock b/uv.lock index 674ffa049ef..ba0f16119ff 100644 --- a/uv.lock +++ b/uv.lock @@ -6194,11 +6194,11 @@ wheels = [ [[package]] name = "oauthlib" -version = "4.0.0" +version = "4.0.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7a/d8/a1bcc8ba112a627f8ffbdc212a78ce18d3ac07e91a5ca65d27918eee25a1/oauthlib-4.0.0.tar.gz", hash = "sha256:efb274799819440f95b4ab3b818869f1ce9ae26c5beacba0201d1a1b76b54f86", size = 187232, upload-time = "2026-09-28T06:01:18.77Z" } +sdist = { url = "https://files.pythonhosted.org/packages/33/17/2bb072de0b990303c985f26869fd47bb179024a051f47e7f9d208eb92f8e/oauthlib-4.0.1.tar.gz", hash = "sha256:3d73febdb6faa00f277c81fb04be7b7586fb96fc3b7f0b9c772dada69a25ede7", size = 191219, upload-time = "2026-10-31T09:52:58.008Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d9/f4/78229a1066068ca14fc60fb26cf7381cabe4382261392b90e5f9552722d4/oauthlib-4.0.0-py3-none-any.whl", hash = "sha256:624c28c13a0a59cabf9747dfa52af63be3e512a7f2714df16e91b5b3a145e6cd", size = 159715, upload-time = "2026-09-28T06:01:17.008Z" }, + { url = "https://files.pythonhosted.org/packages/df/ef/3b8adde1b32a6f540758bff734d3442c621734c83708e5f02b0ba2f43b84/oauthlib-4.0.1-py3-none-any.whl", hash = "sha256:498a298876e996301950a5b053d12c31628bb7b895db03666430c69ce182eb61", size = 161934, upload-time = "2026-10-31T09:52:56.206Z" }, ] [[package]]