mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix(passthrough): log upstream 4xx/5xx error bodies and carry them into the failure hook (#42695)
* test(integration): reproduce passthrough upstream error body missing from logs and spend row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): log upstream 4xx/5xx error bodies and carry them into the failure hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(error_normalization): let the passthrough prefix win over upstream body text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): honor message redaction for upstream error bodies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): bound the upstream error body read and sanitize it before logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(passthrough): use the Sequence import directly in the allowed-routes cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): rechunk the upstream error stream so the preview read stays bounded Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit matrix for passthrough upstream error visibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): drop the restating docstring on the upstream failure logger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop the retired covers markers from the passthrough error tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): keep the upstream status when the error body peek fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(passthrough): cover the relay aclose in the mid-read failure test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): relay decoded partial body on mid-read failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
a5b9b6da4d
commit
6dbd65b230
8 changed files with 1595 additions and 36 deletions
|
|
@ -1518,6 +1518,7 @@ PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int(
|
|||
CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60))
|
||||
MCP_TOOL_NAME_PREFIX: Final = "mcp_tool"
|
||||
MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100))
|
||||
PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096
|
||||
|
||||
# Headers to control callbacks
|
||||
X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
|
||||
|
|
|
|||
|
|
@ -60,13 +60,13 @@ class _HasProxyErrorType(Protocol):
|
|||
|
||||
|
||||
_MESSAGE_PATTERNS: Final[tuple[tuple[re.Pattern[str], str], ...]] = (
|
||||
(re.compile(r"upstream passthrough request failed", re.IGNORECASE), UPSTREAM_PASSTHROUGH),
|
||||
(
|
||||
re.compile(r"budget has been exceeded|max budget|crossed budget", re.IGNORECASE),
|
||||
BUDGET_EXCEEDED,
|
||||
),
|
||||
(re.compile(r"no healthy deployments?|no deployments available", re.IGNORECASE), NO_HEALTHY_DEPLOYMENTS),
|
||||
(re.compile(r"not allowed to access model due to tags configuration", re.IGNORECASE), MODEL_ACCESS_DENIED),
|
||||
(re.compile(r"upstream passthrough request failed", re.IGNORECASE), UPSTREAM_PASSTHROUGH),
|
||||
(re.compile(r"is not supported for provider|not implemented", re.IGNORECASE), UNSUPPORTED_OPERATION),
|
||||
(
|
||||
re.compile(r"context window|context length|(prompt|input) is too long|tokens? ?> ?\d+ ?maximum", re.IGNORECASE),
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import json
|
|||
import posixpath
|
||||
import traceback
|
||||
from base64 import b64encode
|
||||
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from itertools import count, groupby
|
||||
|
|
@ -41,6 +41,8 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS,
|
||||
REDACTED_BY_LITELLM,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
|
||||
)
|
||||
|
|
@ -56,6 +58,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.base_llm.managed_resources.utils import (
|
||||
resolve_passthrough_managed_id_provider,
|
||||
|
|
@ -849,23 +852,106 @@ def _resolve_team_callback_wiring(
|
|||
)
|
||||
|
||||
|
||||
def _truncate_upstream_error_body(body: str) -> str:
|
||||
if len(body) <= PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS:
|
||||
return body
|
||||
return (
|
||||
f"{body[:PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS]}... "
|
||||
f"(truncated at {PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS} chars)"
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_upstream_error_body(body: str) -> str:
|
||||
return " ".join("".join(char if char.isprintable() else " " for char in body).split())
|
||||
|
||||
|
||||
class _PrefixReplayStream(httpx.AsyncByteStream):
|
||||
def __init__(self, prefix: bytes, rest: AsyncIterator[bytes], upstream: httpx.Response) -> None:
|
||||
self._prefix: Final = prefix
|
||||
self._rest: Final = rest
|
||||
self._upstream: Final = upstream
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[bytes]:
|
||||
if self._prefix:
|
||||
yield self._prefix
|
||||
async for chunk in self._rest:
|
||||
yield chunk
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._upstream.aclose()
|
||||
|
||||
|
||||
async def _no_more_chunks() -> AsyncIterator[bytes]:
|
||||
return
|
||||
yield b""
|
||||
|
||||
|
||||
async def _read_error_body_preview(
|
||||
stream: AsyncIterator[bytes],
|
||||
) -> tuple[bytes, AsyncIterator[bytes]]:
|
||||
collected: Final[list[bytes]] = [] # mutable-ok: accumulated until the preview byte budget, then joined once
|
||||
total = 0 # rebind-ok: running byte count against the preview budget
|
||||
try:
|
||||
async for chunk in stream:
|
||||
collected.append(chunk)
|
||||
total += len(chunk)
|
||||
if total > PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS:
|
||||
break
|
||||
except httpx.HTTPError as err:
|
||||
partial: Final = b"".join(collected)
|
||||
verbose_proxy_logger.warning(
|
||||
"pass_through_endpoint: upstream error body read failed after %d bytes: %s",
|
||||
len(partial),
|
||||
type(err).__name__,
|
||||
)
|
||||
return partial, _no_more_chunks()
|
||||
return b"".join(collected), stream
|
||||
|
||||
|
||||
def _headers_without_body_framing(headers: httpx.Headers) -> httpx.Headers:
|
||||
return httpx.Headers(
|
||||
[(name, value) for name, value in headers.raw if name.lower() not in (b"content-encoding", b"content-length")]
|
||||
)
|
||||
|
||||
|
||||
async def _error_body_preview_and_relay(response: httpx.Response) -> tuple[str, httpx.Response]:
|
||||
if response.is_stream_consumed:
|
||||
return response.text, response
|
||||
body_iter: Final = response.aiter_bytes()
|
||||
prefix, rest = await _read_error_body_preview(body_iter)
|
||||
preview_text: Final = prefix.decode(response.encoding or "utf-8", errors="replace")
|
||||
return preview_text, httpx.Response(
|
||||
status_code=response.status_code,
|
||||
headers=_headers_without_body_framing(response.headers),
|
||||
stream=_PrefixReplayStream(prefix=prefix, rest=rest, upstream=response),
|
||||
request=response.request,
|
||||
extensions=response.extensions,
|
||||
)
|
||||
|
||||
|
||||
async def _log_passthrough_upstream_failure(
|
||||
response: httpx.Response,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
request_payload: dict,
|
||||
) -> None:
|
||||
"""Fire LiteLLM-side failure hooks (spend tracking, alerting callbacks) for
|
||||
an upstream 4xx/5xx passthrough response.
|
||||
|
||||
Passthrough must return the upstream status/body/headers to the client
|
||||
unchanged, so this never raises or transforms the response - it only
|
||||
mirrors the monitoring side effect that ``post_call_failure_hook`` would
|
||||
have received had the error originated inside LiteLLM.
|
||||
"""
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> httpx.Response:
|
||||
if response.status_code < 400:
|
||||
return
|
||||
return response
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
preview_text, relay_response = await _error_body_preview_and_relay(response)
|
||||
upstream_error_body: Final = (
|
||||
REDACTED_BY_LITELLM
|
||||
if should_redact_message_logging(logging_obj.model_call_details)
|
||||
else _truncate_upstream_error_body(_sanitize_upstream_error_body(preview_text))
|
||||
)
|
||||
verbose_proxy_logger.warning(
|
||||
"pass_through_endpoint: upstream %s %s returned %s: %s",
|
||||
response.request.method,
|
||||
response.url.copy_with(query=None, fragment=None),
|
||||
response.status_code,
|
||||
upstream_error_body,
|
||||
)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError:
|
||||
|
|
@ -878,7 +964,7 @@ async def _log_passthrough_upstream_failure(
|
|||
# rate-limit errors already are.
|
||||
synthetic_exception: Final = HTTPException(
|
||||
status_code=response.status_code,
|
||||
detail=f"Upstream passthrough request failed with status {response.status_code}",
|
||||
detail=f"Upstream passthrough request failed with status {response.status_code}: {upstream_error_body}",
|
||||
)
|
||||
try:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
|
|
@ -892,6 +978,7 @@ async def _log_passthrough_upstream_failure(
|
|||
"pass_through_endpoint: post_call_failure_hook raised for upstream error",
|
||||
exc_info=True,
|
||||
)
|
||||
return relay_response
|
||||
|
||||
|
||||
async def _relay_reporting_failures(
|
||||
|
|
@ -1321,7 +1408,7 @@ async def pass_through_request(
|
|||
headers=response.headers,
|
||||
)
|
||||
|
||||
await _log_passthrough_upstream_failure(
|
||||
relay_response: Final = await _log_passthrough_upstream_failure(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
|
|
@ -1331,17 +1418,18 @@ async def pass_through_request(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
upstream_usage=upstream_usage,
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Call response headers hook for streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
headers=relay_response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
response=relay_response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
|
|
@ -1352,7 +1440,7 @@ async def pass_through_request(
|
|||
stream=_own_streamed_managed_ids(
|
||||
stream=_relay_reporting_failures(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
response=relay_response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
|
|
@ -1360,7 +1448,7 @@ async def pass_through_request(
|
|||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
),
|
||||
upstream_status=response.status_code,
|
||||
upstream_status=relay_response.status_code,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
parsed_body=_parsed_body,
|
||||
|
|
@ -1374,10 +1462,10 @@ async def pass_through_request(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=response.headers,
|
||||
upstream_headers=relay_response.headers,
|
||||
),
|
||||
headers=_response_headers,
|
||||
status_code=response.status_code,
|
||||
status_code=relay_response.status_code,
|
||||
)
|
||||
|
||||
if state_raw_body is not None:
|
||||
|
|
@ -1412,7 +1500,7 @@ async def pass_through_request(
|
|||
logging_obj.stream = True
|
||||
logging_obj.model_call_details["stream"] = True
|
||||
|
||||
await _log_passthrough_upstream_failure(
|
||||
detected_relay_response: Final = await _log_passthrough_upstream_failure(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
|
|
@ -1422,17 +1510,18 @@ async def pass_through_request(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
upstream_usage=upstream_usage,
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Call response headers hook for detected streaming pass-through
|
||||
_response_headers = HttpPassThroughEndpointHelpers.get_response_headers(
|
||||
headers=response.headers,
|
||||
headers=detected_relay_response.headers,
|
||||
litellm_call_id=litellm_call_id,
|
||||
)
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=_parsed_body or {},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
response=detected_relay_response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
|
|
@ -1443,7 +1532,7 @@ async def pass_through_request(
|
|||
stream=_own_streamed_managed_ids(
|
||||
stream=_relay_reporting_failures(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
response=detected_relay_response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
|
|
@ -1451,7 +1540,7 @@ async def pass_through_request(
|
|||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
),
|
||||
upstream_status=response.status_code,
|
||||
upstream_status=detected_relay_response.status_code,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=_build_passthrough_failure_request_payload(
|
||||
parsed_body=_parsed_body,
|
||||
|
|
@ -1465,10 +1554,10 @@ async def pass_through_request(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=response.headers,
|
||||
upstream_headers=detected_relay_response.headers,
|
||||
),
|
||||
headers=_response_headers,
|
||||
status_code=response.status_code,
|
||||
status_code=detected_relay_response.status_code,
|
||||
)
|
||||
|
||||
if not _should_buffer_passthrough_response(response):
|
||||
|
|
@ -1526,6 +1615,7 @@ async def pass_through_request(
|
|||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_payload=failure_request_payload,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
if response.status_code < 400 and response_body is not None and guardrails_to_run:
|
||||
|
|
@ -3435,7 +3525,7 @@ async def _filter_endpoints_by_team_allowed_routes(
|
|||
for endpoint in pass_through_endpoints
|
||||
if endpoint.path
|
||||
in cast( # cast-ok: guarded above; team metadata stores this key as a list of route paths
|
||||
"Sequence[str]", team_metadata.get("allowed_passthrough_routes")
|
||||
Sequence[str], team_metadata.get("allowed_passthrough_routes")
|
||||
)
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -43,7 +43,9 @@ class Wire:
|
|||
|
||||
|
||||
@contextmanager
|
||||
def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None) -> Generator[Wire, None, None]:
|
||||
def wire_server(
|
||||
respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None, port: int = 0
|
||||
) -> Generator[Wire, None, None]:
|
||||
"""Owned TCP peer; requests traverse the real HTTP client and serialization."""
|
||||
received: Final[SimpleQueue[Request]] = SimpleQueue()
|
||||
errors: Final[SimpleQueue[Exception]] = SimpleQueue()
|
||||
|
|
@ -114,7 +116,7 @@ def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None
|
|||
if tls is not None:
|
||||
self.socket = tls.wrap_socket(self.socket, server_side=True)
|
||||
|
||||
with OwnedHTTPServer(("127.0.0.1", 0), Handler) as server:
|
||||
with OwnedHTTPServer(("127.0.0.1", port), Handler) as server:
|
||||
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
|
||||
thread.start()
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,150 @@
|
|||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import signal
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
_GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}
|
||||
_NOT_FOUND_BODY: Final = json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"code": 404,
|
||||
"message": "models/nope-9 is not found for this scripted upstream",
|
||||
"status": "NOT_FOUND",
|
||||
}
|
||||
}
|
||||
).encode()
|
||||
_INTERNAL_BODY: Final = (
|
||||
'{"error":{"code":500,"message":"' + "chunked upstream failure body " * 200 + '","status":"INTERNAL"}}'
|
||||
).encode()
|
||||
_OK_CHUNKS: Final = tuple(f"data: ok-{index}\n\n".encode() for index in range(3))
|
||||
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
|
||||
|
||||
def _chaos_reply(request: Request) -> Reply:
|
||||
if "streamGenerateContent" in request.target:
|
||||
return Reply(status=500, chunks=tuple(_INTERNAL_BODY[i : i + 512] for i in range(0, len(_INTERNAL_BODY), 512)))
|
||||
if "healthy-model" in request.target:
|
||||
return Reply(status=200, chunks=_OK_CHUNKS, content_type="text/event-stream")
|
||||
return Reply(status=404, body=_NOT_FOUND_BODY)
|
||||
|
||||
|
||||
def _error_information(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
metadata: Final = rows[0]["metadata"]
|
||||
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
|
||||
return object_value(parsed["error_information"])
|
||||
|
||||
|
||||
def _single_spend_row(call_id: str) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert len(rows) == 1, call_id
|
||||
|
||||
|
||||
async def _fire_burst(
|
||||
base_url: str, key: str, count: int, *, tolerate_transport_errors: bool = False
|
||||
) -> tuple[httpx.Response, ...]:
|
||||
async def one(client: httpx.AsyncClient, index: int) -> httpx.Response:
|
||||
if index % 3 == 0:
|
||||
path: Final = "/gemini/v1beta/models/nope-9:generateContent"
|
||||
elif index % 3 == 1:
|
||||
path = "/gemini/v1beta/models/nope-9:streamGenerateContent?alt=sse"
|
||||
else:
|
||||
path = "/gemini/v1beta/models/healthy-model:streamGenerateContent?alt=sse"
|
||||
return await client.post(
|
||||
path,
|
||||
json=_GENERATE_CONTENT,
|
||||
headers={"Authorization": f"Bearer {key}", "x-goog-api-key": key},
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=30, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(one(client, index) for index in range(count)), return_exceptions=tolerate_transport_errors
|
||||
)
|
||||
for result in results:
|
||||
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
|
||||
return tuple(result for result in results if isinstance(result, httpx.Response))
|
||||
|
||||
|
||||
async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gateway: Gateway, tmp_path: Path) -> None:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = tmp_path / "chaos-outage.yaml"
|
||||
with wire_server(_chaos_reply) as wire:
|
||||
port: Final = int(wire.url.rsplit(":", 1)[1])
|
||||
config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"}
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
burst: Final = asyncio.create_task(_fire_burst(str(candidate.client.base_url), candidate.key, 30))
|
||||
await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 10, 30)
|
||||
with wire_server(_chaos_reply, port=port):
|
||||
responses: Final = await burst
|
||||
assert len(responses) == 30
|
||||
for response in responses:
|
||||
assert response.status_code in (200, 404, 500, 502), response.status_code
|
||||
assert "x-litellm-call-id" in response.headers, response.status_code
|
||||
assert len(_STARTED_WORKER.findall(owned.log.read_text())) >= 2
|
||||
for response in responses:
|
||||
_single_spend_row(response.headers["x-litellm-call-id"])
|
||||
if response.status_code == 404:
|
||||
error_information: Final = _error_information(response.headers["x-litellm-call-id"])
|
||||
assert "not found for this scripted upstream" in str(error_information["error_message"]), response.text
|
||||
elif response.status_code == 500:
|
||||
assert "chunked upstream failure body" in str(
|
||||
_error_information(response.headers["x-litellm-call-id"])["error_message"]
|
||||
), response.text
|
||||
|
||||
|
||||
async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gateway: Gateway, tmp_path: Path) -> None:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = tmp_path / "chaos-kill.yaml"
|
||||
with wire_server(_chaos_reply) as wire:
|
||||
config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"}
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
workers: Final = eventually(
|
||||
lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
|
||||
lambda pids: len(pids) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
burst: Final = asyncio.create_task(
|
||||
_fire_burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True)
|
||||
)
|
||||
await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 5, 30)
|
||||
psutil.Process(workers[0]).send_signal(signal.SIGKILL)
|
||||
responses: Final = await burst
|
||||
for response in responses:
|
||||
assert response.status_code in (200, 404, 500, 502), response.status_code
|
||||
follow_up: Final = candidate.request(
|
||||
"POST",
|
||||
"/gemini/v1beta/models/nope-9:generateContent",
|
||||
_GENERATE_CONTENT,
|
||||
headers={"x-goog-api-key": candidate.key},
|
||||
)
|
||||
assert follow_up.status_code == 404, follow_up.text
|
||||
assert follow_up.json() == json.loads(_NOT_FOUND_BODY), follow_up.text
|
||||
for response in responses:
|
||||
if "x-litellm-call-id" in response.headers:
|
||||
_single_spend_row(response.headers["x-litellm-call-id"])
|
||||
error_information: Final = _error_information(follow_up.headers["x-litellm-call-id"])
|
||||
assert "not found for this scripted upstream" in str(error_information["error_message"]), follow_up.text
|
||||
|
|
@ -0,0 +1,600 @@
|
|||
import gzip
|
||||
import json
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from openai import AsyncOpenAI, NotFoundError, OpenAI
|
||||
from pydantic import JsonValue
|
||||
|
||||
_UPSTREAM_ERROR: Final[dict[str, JsonValue]] = {
|
||||
"error": {
|
||||
"code": 404,
|
||||
"message": "Publisher Model `publishers/anthropic/models/claude-nope-9` was not found or your project does not have access to it. Please ensure you are using a valid model version.",
|
||||
"status": "NOT_FOUND",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_gemini_passthrough_upstream_error_body_reaches_proxy_log_and_spend_row(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=json.dumps(_UPSTREAM_ERROR).encode())
|
||||
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = tmp_path / "gemini-passthrough.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"}
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST",
|
||||
"/gemini/v1beta/models/claude-nope-9:generateContent",
|
||||
{"contents": [{"role": "user", "parts": [{"text": "hi"}]}]},
|
||||
headers={"x-goog-api-key": candidate.key},
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
assert response.json() == _UPSTREAM_ERROR, response.text
|
||||
try:
|
||||
eventually(
|
||||
lambda: owned.log.read_text(),
|
||||
lambda text: "was not found or your project" in text,
|
||||
seconds=30,
|
||||
)
|
||||
except AssertionError:
|
||||
pytest.fail(
|
||||
f"upstream 404 body never reached the proxy log after {response.status_code} passthrough; "
|
||||
f"log tail: {owned.log.read_text()[-2000:]}"
|
||||
)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
||||
(response.headers["x-litellm-call-id"],),
|
||||
),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
metadata: Final = rows[0]["metadata"]
|
||||
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
|
||||
error_information: Final = object_value(parsed["error_information"])
|
||||
assert "was not found or your project" in str(error_information["error_message"]), response.text
|
||||
assert error_information["error_code"] == "404", response.text
|
||||
|
||||
|
||||
_GEMINI_MODEL_PATH: Final = "/gemini/v1beta/models/claude-nope-9:generateContent"
|
||||
_GEMINI_STREAM_PATH: Final = "/gemini/v1beta/models/claude-nope-9:streamGenerateContent"
|
||||
_GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}
|
||||
_UPSTREAM_500_BODY: Final = (
|
||||
'{"error":{"code":500,"message":"' + "chunked upstream failure body " * 200 + '","status":"INTERNAL"}}'
|
||||
).encode()
|
||||
|
||||
|
||||
def _gemini_config(path: Path, wire_url: str) -> None:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["environment_variables"] = {"GEMINI_API_BASE": wire_url, "GEMINI_API_KEY": "scripted"}
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
|
||||
|
||||
def _gemini_headers(candidate: Gateway) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {candidate.key}", "x-goog-api-key": candidate.key}
|
||||
|
||||
|
||||
def _spend_error_information(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
metadata: Final = rows[0]["metadata"]
|
||||
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
|
||||
return object_value(parsed["error_information"])
|
||||
|
||||
|
||||
def _spend_status(call_id: str) -> str:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)),
|
||||
lambda values: len(values) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
return str(rows[0]["status"])
|
||||
|
||||
|
||||
def _upstream_warning(log: Path, needle: str = "pass_through_endpoint: upstream") -> str:
|
||||
text: Final = eventually(lambda: log.read_text(), lambda content: needle in content, seconds=30)
|
||||
return next(line for line in text.splitlines() if needle in line)
|
||||
|
||||
|
||||
def _upstream_warnings(log: Path, needle: str = "pass_through_endpoint: upstream") -> tuple[str, ...]:
|
||||
return tuple(line for line in log.read_text().splitlines() if needle in line)
|
||||
|
||||
|
||||
async def test_gemini_passthrough_async_client_404_body_reaches_proxy_log_and_spend_row(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=json.dumps(_UPSTREAM_ERROR).encode())
|
||||
|
||||
path: Final = tmp_path / "gemini-async.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
async with httpx.AsyncClient(
|
||||
base_url=str(candidate.client.base_url), timeout=15, trust_env=False
|
||||
) as async_client:
|
||||
response: Final = await async_client.post(
|
||||
_GEMINI_MODEL_PATH, json=_GENERATE_CONTENT, headers=_gemini_headers(candidate)
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
assert response.json() == _UPSTREAM_ERROR, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "was not found or your project" in warning, warning
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
assert "was not found or your project" in str(error_information["error_message"]), response.text
|
||||
assert error_information["error_code"] == "404", response.text
|
||||
|
||||
|
||||
def test_gemini_passthrough_streaming_500_relays_full_body_and_logs_bounded_preview(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
body: Final = _UPSTREAM_500_BODY
|
||||
assert len(body) == 6055
|
||||
chunks: Final = tuple(body[index * 512 : (index + 1) * 512] for index in range(11)) + (body[5632:],)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=500, chunks=chunks)
|
||||
|
||||
path: Final = tmp_path / "gemini-stream-500.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.client.stream(
|
||||
"POST",
|
||||
_GEMINI_STREAM_PATH,
|
||||
params={"alt": "sse"},
|
||||
json=_GENERATE_CONTENT,
|
||||
headers=_gemini_headers(candidate),
|
||||
) as response:
|
||||
assert response.status_code == 500, response.text
|
||||
streamed: Final = response.read()
|
||||
assert streamed == body
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert warning.endswith("... (truncated at 4096 chars)"), warning
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
error_message: Final = str(error_information["error_message"])
|
||||
assert error_message.endswith("... (truncated at 4096 chars)"), error_message
|
||||
assert error_information["error_code"] == "500", error_message
|
||||
|
||||
|
||||
def test_gemini_passthrough_success_logs_nothing_and_spend_row_is_success(gateway: Gateway, tmp_path: Path) -> None:
|
||||
upstream_ok: Final = {
|
||||
"candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}}],
|
||||
"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5},
|
||||
}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=200, body=json.dumps(upstream_ok).encode())
|
||||
|
||||
path: Final = tmp_path / "gemini-200.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == upstream_ok, response.text
|
||||
assert _spend_status(response.headers["x-litellm-call-id"]) == "success"
|
||||
assert not _upstream_warnings(owned.log), owned.log.read_text()[-2000:]
|
||||
|
||||
|
||||
def test_gemini_passthrough_streaming_200_relays_every_chunk(gateway: Gateway, tmp_path: Path) -> None:
|
||||
chunks: Final = tuple(f"data: chunk-{index}\n\n".encode() for index in range(5))
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=200, chunks=chunks, content_type="text/event-stream")
|
||||
|
||||
path: Final = tmp_path / "gemini-stream-200.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.client.stream(
|
||||
"POST",
|
||||
_GEMINI_STREAM_PATH,
|
||||
params={"alt": "sse"},
|
||||
json=_GENERATE_CONTENT,
|
||||
headers=_gemini_headers(candidate),
|
||||
) as response:
|
||||
assert response.status_code == 200
|
||||
streamed: Final = response.read()
|
||||
assert streamed == b"".join(chunks)
|
||||
assert not _upstream_warnings(owned.log), owned.log.read_text()[-2000:]
|
||||
|
||||
|
||||
def test_config_pass_through_route_logs_body_and_strips_query(gateway: Gateway, tmp_path: Path) -> None:
|
||||
upstream_error: Final = {"error": {"message": "max budget reached for this deployment"}}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=403, body=json.dumps(upstream_error).encode())
|
||||
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = tmp_path / "config-route.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
config["general_settings"]["pass_through_endpoints"] = [
|
||||
{
|
||||
"path": "/audit-pt",
|
||||
"target": f"{wire.url}/upstream?trace=secret-q",
|
||||
"include_subpath": True,
|
||||
"headers": {"Authorization": "Bearer scripted"},
|
||||
}
|
||||
]
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request("POST", "/audit-pt", _GENERATE_CONTENT)
|
||||
assert response.status_code == 403, response.text
|
||||
assert response.json() == upstream_error, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "max budget reached for this deployment" in warning, warning
|
||||
assert "?" not in warning and "secret-q" not in warning, warning
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
assert error_information["normalized_error"] == "500_UPSTREAM_PASSTHROUGH", response.text
|
||||
assert "max budget reached for this deployment" in str(error_information["error_message"]), response.text
|
||||
|
||||
|
||||
_OPENAI_UPSTREAM_404: Final[dict[str, JsonValue]] = {
|
||||
"error": {"message": "The model `nope-9` does not exist", "type": "invalid_request_error"}
|
||||
}
|
||||
|
||||
|
||||
def _openai_config(path: Path, wire_url: str) -> None:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["environment_variables"] = {"OPENAI_API_BASE": wire_url, "OPENAI_API_KEY": "scripted"}
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
|
||||
|
||||
def test_openai_passthrough_sdk_error_body_reaches_proxy_log_and_spend_row(gateway: Gateway, tmp_path: Path) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=json.dumps(_OPENAI_UPSTREAM_404).encode())
|
||||
|
||||
path: Final = tmp_path / "openai-404.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_openai_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with OpenAI(
|
||||
api_key=candidate.key,
|
||||
base_url=f"{str(candidate.client.base_url).rstrip('/')}/openai",
|
||||
max_retries=0,
|
||||
http_client=httpx.Client(timeout=15, trust_env=False),
|
||||
) as sdk:
|
||||
with pytest.raises(NotFoundError) as raised:
|
||||
sdk.chat.completions.create(model="nope-9", messages=[{"role": "user", "content": "hi"}])
|
||||
assert "does not exist" in str(raised.value), raised.value
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "does not exist" in warning, warning
|
||||
error_information: Final = _spend_error_information(raised.value.response.headers["x-litellm-call-id"])
|
||||
assert "does not exist" in str(error_information["error_message"])
|
||||
|
||||
|
||||
async def test_openai_passthrough_async_sdk_error_body_reaches_proxy_log_and_spend_row(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=json.dumps(_OPENAI_UPSTREAM_404).encode())
|
||||
|
||||
path: Final = tmp_path / "openai-async-404.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_openai_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
async with AsyncOpenAI(
|
||||
api_key=candidate.key,
|
||||
base_url=f"{str(candidate.client.base_url).rstrip('/')}/openai",
|
||||
max_retries=0,
|
||||
http_client=httpx.AsyncClient(timeout=15, trust_env=False),
|
||||
) as sdk:
|
||||
with pytest.raises(NotFoundError) as raised:
|
||||
await sdk.chat.completions.create(model="nope-9", messages=[{"role": "user", "content": "hi"}])
|
||||
assert "does not exist" in str(raised.value), raised.value
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "does not exist" in warning, warning
|
||||
error_information: Final = _spend_error_information(raised.value.response.headers["x-litellm-call-id"])
|
||||
assert "does not exist" in str(error_information["error_message"])
|
||||
|
||||
|
||||
def test_gemini_passthrough_control_characters_cannot_forge_log_lines(gateway: Gateway, tmp_path: Path) -> None:
|
||||
forged: Final = b'{"error": "line one"}\n2026-01-01 FAKE LOG LINE\x1b[31m\r' + b"x" * 4943 + b"\x00tail"
|
||||
assert len(forged) == 5000
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=502, body=forged, content_type="text/html")
|
||||
|
||||
path: Final = tmp_path / "gemini-forged.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 502, response.text
|
||||
assert response.content == forged, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "\n" not in warning and "\x1b" not in warning, warning
|
||||
assert "line one" in warning and "FAKE LOG LINE" in warning, warning
|
||||
assert warning.endswith("... (truncated at 4096 chars)"), warning
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
assert error_information["error_code"] == "502", response.text
|
||||
|
||||
|
||||
def test_gemini_passthrough_empty_error_body_still_logged_and_proxy_serves(gateway: Gateway, tmp_path: Path) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
if "claude-nope-9" in request.target:
|
||||
return Reply(status=404, body=b"")
|
||||
return Reply(status=200, body=b'{"candidates": [{"content": {"parts": [{"text": "ok"}]}}]}')
|
||||
|
||||
path: Final = tmp_path / "gemini-empty.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
assert response.content == b"", response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "returned 404" in warning, warning
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
assert error_information["error_code"] == "404", response.text
|
||||
follow_up: Final = candidate.request(
|
||||
"POST",
|
||||
"/gemini/v1beta/models/healthy-model:generateContent",
|
||||
_GENERATE_CONTENT,
|
||||
headers={"x-goog-api-key": candidate.key},
|
||||
)
|
||||
assert follow_up.status_code == 200, follow_up.text
|
||||
|
||||
|
||||
def test_gemini_passthrough_gzip_error_body_decoded_for_log_and_client(gateway: Gateway, tmp_path: Path) -> None:
|
||||
upstream_error: Final = {"error": {"message": "gzipped upstream says the model is gone"}}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(
|
||||
status=400,
|
||||
body=gzip.compress(json.dumps(upstream_error).encode()),
|
||||
headers={"content-encoding": "gzip"},
|
||||
)
|
||||
|
||||
path: Final = tmp_path / "gemini-gzip.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json() == upstream_error, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "gzipped upstream says the model is gone" in warning, warning
|
||||
|
||||
|
||||
def test_gemini_passthrough_streaming_gzip_error_body_decoded_for_log_and_client(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
upstream_error: Final = {"error": {"message": "streamed gzip upstream denies the deployment"}}
|
||||
compressed: Final = gzip.compress(json.dumps(upstream_error).encode())
|
||||
third: Final = len(compressed) // 3
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(
|
||||
status=403,
|
||||
chunks=(compressed[:third], compressed[third : 2 * third], compressed[2 * third :]),
|
||||
headers={"content-encoding": "gzip"},
|
||||
)
|
||||
|
||||
path: Final = tmp_path / "gemini-stream-gzip.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.client.stream(
|
||||
"POST",
|
||||
_GEMINI_STREAM_PATH,
|
||||
params={"alt": "sse"},
|
||||
json=_GENERATE_CONTENT,
|
||||
headers=_gemini_headers(candidate),
|
||||
) as response:
|
||||
assert response.status_code == 403
|
||||
streamed: Final = response.read()
|
||||
assert json.loads(streamed) == upstream_error, streamed
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "streamed gzip upstream denies the deployment" in warning, warning
|
||||
|
||||
|
||||
def test_gemini_passthrough_error_body_redacted_when_message_logging_off(gateway: Gateway, tmp_path: Path) -> None:
|
||||
upstream_error: Final = {"error": {"message": "sensitive upstream explanation"}}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=json.dumps(upstream_error).encode())
|
||||
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = tmp_path / "gemini-redacted.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"}
|
||||
config["litellm_settings"]["turn_off_message_logging"] = True
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
assert response.json() == upstream_error, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "redacted-by-litellm" in warning, warning
|
||||
assert "sensitive upstream explanation" not in warning, warning
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
error_message: Final = str(error_information["error_message"])
|
||||
assert "redacted-by-litellm" in error_message, error_message
|
||||
assert "sensitive upstream explanation" not in error_message, error_message
|
||||
|
||||
|
||||
def test_gemini_passthrough_exact_4096_byte_body_logged_without_marker(gateway: Gateway, tmp_path: Path) -> None:
|
||||
body: Final = b'{"error": "' + b"y" * 4083 + b'"}'
|
||||
assert len(body) == 4096
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=body)
|
||||
|
||||
path: Final = tmp_path / "gemini-exact.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert body[:512].decode() in warning, warning
|
||||
assert "(truncated at 4096 chars)" not in warning, warning
|
||||
|
||||
|
||||
def test_gemini_passthrough_4097_byte_body_truncated_with_marker(gateway: Gateway, tmp_path: Path) -> None:
|
||||
body: Final = b'{"error": "' + b"y" * 4084 + b'"}'
|
||||
assert len(body) == 4097
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=body)
|
||||
|
||||
path: Final = tmp_path / "gemini-over.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
response: Final = candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert body[:512].decode() in warning, warning
|
||||
assert warning.endswith("... (truncated at 4096 chars)"), warning
|
||||
|
||||
|
||||
def test_gemini_passthrough_one_byte_stream_chunks_reassembled_and_logged(gateway: Gateway, tmp_path: Path) -> None:
|
||||
body: Final = json.dumps(_UPSTREAM_ERROR).encode()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, chunks=tuple(bytes([byte]) for byte in body))
|
||||
|
||||
path: Final = tmp_path / "gemini-one-byte.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.client.stream(
|
||||
"POST",
|
||||
_GEMINI_STREAM_PATH,
|
||||
params={"alt": "sse"},
|
||||
json=_GENERATE_CONTENT,
|
||||
headers=_gemini_headers(candidate),
|
||||
) as response:
|
||||
assert response.status_code == 404
|
||||
streamed: Final = response.read()
|
||||
assert streamed == body
|
||||
warning: Final = _upstream_warning(owned.log)
|
||||
assert "was not found or your project" in warning, warning
|
||||
|
||||
|
||||
def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: Gateway, tmp_path: Path) -> None:
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(status=404, body=json.dumps(_UPSTREAM_ERROR).encode())
|
||||
|
||||
path: Final = tmp_path / "gemini-twice.yaml"
|
||||
with wire_server(respond) as wire:
|
||||
_gemini_config(path, wire.url)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
responses: Final = tuple(
|
||||
candidate.request(
|
||||
"POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key}
|
||||
)
|
||||
for _ in range(2)
|
||||
)
|
||||
call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses)
|
||||
assert len(set(call_ids)) == 2
|
||||
for response in responses:
|
||||
assert response.status_code == 404, response.text
|
||||
error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"])
|
||||
assert "was not found or your project" in str(error_information["error_message"]), response.text
|
||||
eventually(
|
||||
lambda: _upstream_warnings(owned.log, "returned 404"),
|
||||
lambda lines: len(lines) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
|
||||
|
||||
def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tmp_path: Path) -> None:
|
||||
path: Final = tmp_path / "budget.yaml"
|
||||
path.write_text(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
with candidate.scenario() as scenario:
|
||||
model: Final = scenario.model()
|
||||
key: Final = scenario.key(max_budget=0.000001)
|
||||
first: Final = candidate.chat(model, key=key)
|
||||
assert "id" in first, first
|
||||
rejected: Final = candidate.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": "over budget"}]},
|
||||
key=key,
|
||||
)
|
||||
assert rejected.status_code == 422 and "budget_exceeded" in rejected.text, rejected.text
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
||||
(digest,),
|
||||
),
|
||||
lambda values: any(
|
||||
"BUDGET_EXCEEDED"
|
||||
in str(
|
||||
object_value(
|
||||
json.loads(row["metadata"])
|
||||
if isinstance(row["metadata"], str)
|
||||
else object_value(row["metadata"])
|
||||
)["error_information"]
|
||||
)
|
||||
for row in values
|
||||
),
|
||||
seconds=70,
|
||||
)
|
||||
budget_rows: Final = tuple(
|
||||
row
|
||||
for row in rows
|
||||
if "BUDGET_EXCEEDED"
|
||||
in str(
|
||||
object_value(
|
||||
json.loads(row["metadata"])
|
||||
if isinstance(row["metadata"], str)
|
||||
else object_value(row["metadata"])
|
||||
)["error_information"]
|
||||
)
|
||||
)
|
||||
assert len(budget_rows) == 1, budget_rows
|
||||
|
|
@ -165,6 +165,18 @@ def test_variants_of_one_failure_share_a_normalized_error(messages: tuple[Except
|
|||
assert normalized == {expected}
|
||||
|
||||
|
||||
def test_normalize_error_passthrough_prefix_wins_over_upstream_body_text() -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
for detail in (
|
||||
'Upstream passthrough request failed with status 400: {"error": {"message": "no deployments available for this model"}}',
|
||||
'Upstream passthrough request failed with status 400: {"error": {"message": "max budget reached"}}',
|
||||
):
|
||||
exc = HTTPException(status_code=400, detail=detail)
|
||||
message = f"400: {detail}"
|
||||
assert normalize_error(exc, "400", message) == "500_UPSTREAM_PASSTHROUGH", message
|
||||
|
||||
|
||||
def test_router_no_healthy_deployment_wording_clusters_as_no_healthy_deployments() -> None:
|
||||
for message in (RouterErrors.no_healthy_deployments.value, "No healthy deployments found."):
|
||||
exc = litellm.BadRequestError(message, llm_provider="openai", model="gpt-4o")
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import asyncio
|
||||
import gzip
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import zlib
|
||||
from collections.abc import Callable
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from io import BytesIO
|
||||
|
|
@ -12,12 +14,14 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import Request, Response, UploadFile
|
||||
from fastapi import HTTPException, Request, Response, UploadFile
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import ValidationError
|
||||
from starlette.datastructures import FormData, Headers, QueryParams
|
||||
from starlette.datastructures import UploadFile as StarletteUploadFile
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
|
|
@ -27,6 +31,8 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
HttpPassThroughEndpointHelpers,
|
||||
InitPassThroughEndpointHelpers,
|
||||
_registered_pass_through_routes,
|
||||
_truncate_upstream_error_body,
|
||||
_with_trace_context,
|
||||
chat_completion_pass_through_endpoint,
|
||||
create_pass_through_route,
|
||||
initialize_pass_through_endpoints,
|
||||
|
|
@ -34,7 +40,6 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
|||
resolve_llm_passthrough_timeout,
|
||||
resolve_pass_through_request_timeout,
|
||||
websocket_passthrough_request,
|
||||
_with_trace_context,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
PassThroughEndpointLogging,
|
||||
|
|
@ -4126,6 +4131,705 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged(
|
|||
assert failure_call_kwargs["original_exception"].status_code == 403
|
||||
|
||||
|
||||
class _UpstreamErrorBodyStream(httpx.AsyncByteStream):
|
||||
def __init__(self, body: bytes) -> None:
|
||||
self._body: Final = body
|
||||
|
||||
async def __aiter__(self):
|
||||
yield self._body
|
||||
|
||||
|
||||
def _upstream_error_request() -> MagicMock:
|
||||
mock_request: Final = MagicMock(spec=Request)
|
||||
mock_request.method = "POST"
|
||||
mock_request.url = "http://test-proxy.com/mock-upstream/v1beta/models/claude-nope-9:generateContent"
|
||||
mock_request.body = AsyncMock(return_value=b'{"contents": []}')
|
||||
mock_request.headers = Headers({"content-type": "application/json"})
|
||||
mock_request.query_params = QueryParams({})
|
||||
return mock_request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_non_streaming_upstream_error_body_logged_and_in_failure_detail(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
upstream_body: Final = {
|
||||
"error": {
|
||||
"code": 404,
|
||||
"message": "Publisher Model `publishers/anthropic/models/claude-nope-9` was not found or your project does not have access",
|
||||
"status": "NOT_FOUND",
|
||||
}
|
||||
}
|
||||
upstream_content: Final = json.dumps(upstream_body).encode("utf-8")
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=404,
|
||||
headers={"content-type": "application/json"},
|
||||
content=upstream_content,
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:generateContent"),
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processing:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:generateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
warning_messages: Final = [record.getMessage() for record in caplog.records if record.levelno == logging.WARNING]
|
||||
upstream_warnings: Final = [
|
||||
message for message in warning_messages if "upstream" in message and "returned 404" in message
|
||||
]
|
||||
assert len(upstream_warnings) == 1, warning_messages
|
||||
assert "was not found or your project" in upstream_warnings[0]
|
||||
assert "/v1beta/models/claude-nope-9:generateContent" in upstream_warnings[0]
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.body == upstream_content
|
||||
|
||||
mock_proxy_logging.post_call_failure_hook.assert_called_once()
|
||||
failure_call_kwargs: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs
|
||||
original_exception: Final = failure_call_kwargs["original_exception"]
|
||||
assert isinstance(original_exception, HTTPException)
|
||||
assert original_exception.status_code == 404
|
||||
assert "was not found or your project" in original_exception.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_streaming_upstream_error_body_reaches_client_and_failure_detail():
|
||||
upstream_content: Final = (
|
||||
b'data: {"error": {"code": 403, "message": "stream access was not found or your project lacks"}}\n\n'
|
||||
)
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=403,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
stream=_UpstreamErrorBodyStream(upstream_content),
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"),
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
) as mock_success_handler:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.status_code == 403
|
||||
streamed_chunks: Final = [chunk async for chunk in response.body_iterator]
|
||||
streamed_bytes: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks
|
||||
)
|
||||
assert streamed_bytes == upstream_content
|
||||
|
||||
mock_proxy_logging.post_call_failure_hook.assert_called_once()
|
||||
original_exception: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"]
|
||||
assert "was not found or your project" in original_exception.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncate_upstream_error_body_caps_at_log_limit():
|
||||
short_body: Final = "x" * 4096
|
||||
assert _truncate_upstream_error_body(short_body) == short_body
|
||||
|
||||
long_body: Final = "a" * 5000
|
||||
truncated: Final = _truncate_upstream_error_body(long_body)
|
||||
assert truncated == f"{'a' * 4096}... (truncated at 4096 chars)"
|
||||
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=500,
|
||||
headers={"content-type": "text/plain"},
|
||||
content=long_body.encode("utf-8"),
|
||||
request=httpx.Request("POST", "http://target-api.com/api/big-error"),
|
||||
)
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processing:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/api/big-error",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
detail: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail
|
||||
assert detail == f"Upstream passthrough request failed with status 500: {'a' * 4096}... (truncated at 4096 chars)"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_upstream_error_log_strips_provider_key_from_url():
|
||||
upstream_content: Final = b'{"error": "denied"}'
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=404,
|
||||
headers={"content-type": "application/json"},
|
||||
content=upstream_content,
|
||||
request=httpx.Request(
|
||||
"POST",
|
||||
"http://target-api.com/v1beta/models/claude-nope-9:generateContent?key=AIzaSySecretProviderKey123",
|
||||
),
|
||||
)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning") as mock_warning:
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processing:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:generateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
upstream_warnings: Final = [
|
||||
call
|
||||
for call in mock_warning.call_args_list
|
||||
if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s"
|
||||
]
|
||||
assert len(upstream_warnings) == 1, mock_warning.call_args_list
|
||||
logged_url: Final = str(upstream_warnings[0].args[2])
|
||||
assert "/v1beta/models/claude-nope-9:generateContent" in logged_url
|
||||
assert "AIzaSySecretProviderKey123" not in logged_url
|
||||
assert "key=" not in logged_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("turn_off_message_logging", [True, False])
|
||||
async def test_passthrough_upstream_error_body_redacted_when_message_logging_off(
|
||||
turn_off_message_logging: bool,
|
||||
):
|
||||
upstream_content: Final = b'{"error": {"message": "upstream body says the project was not found"}}'
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=404,
|
||||
headers={"content-type": "application/json"},
|
||||
content=upstream_content,
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:generateContent"),
|
||||
)
|
||||
user_api_key_dict: Final = MagicMock()
|
||||
user_api_key_dict.metadata = {
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "prometheus",
|
||||
"callback_type": "success_and_failure",
|
||||
"callback_vars": {"turn_off_message_logging": turn_off_message_logging},
|
||||
}
|
||||
]
|
||||
}
|
||||
user_api_key_dict.team_metadata = None
|
||||
user_api_key_dict.team_id = None
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning") as mock_warning:
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processing:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:generateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.body == upstream_content
|
||||
|
||||
upstream_warnings: Final = [
|
||||
call
|
||||
for call in mock_warning.call_args_list
|
||||
if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s"
|
||||
]
|
||||
assert len(upstream_warnings) == 1, mock_warning.call_args_list
|
||||
logged_body: Final = str(upstream_warnings[0].args[4])
|
||||
|
||||
mock_proxy_logging.post_call_failure_hook.assert_called_once()
|
||||
detail: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail
|
||||
|
||||
if turn_off_message_logging:
|
||||
assert logged_body == "redacted-by-litellm"
|
||||
assert "upstream body says the project was not found" not in logged_body
|
||||
assert detail == "Upstream passthrough request failed with status 404: redacted-by-litellm"
|
||||
else:
|
||||
assert "upstream body says the project was not found" in logged_body
|
||||
assert detail == f"Upstream passthrough request failed with status 404: {upstream_content.decode()}"
|
||||
|
||||
|
||||
class _ChunkedUpstreamErrorBodyStream(httpx.AsyncByteStream):
|
||||
def __init__(self, chunks: tuple[bytes, ...]) -> None:
|
||||
self._chunks: Final = chunks
|
||||
self.served: int = 0
|
||||
|
||||
async def __aiter__(self):
|
||||
for chunk in self._chunks:
|
||||
self.served += 1
|
||||
yield chunk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_streaming_upstream_error_reads_only_preview_and_relays_full_body():
|
||||
chunk_size: Final = 1024
|
||||
chunks: Final = tuple(b"x" * chunk_size for _ in range(10))
|
||||
upstream_content: Final = b"".join(chunks)
|
||||
body_stream: Final = _ChunkedUpstreamErrorBodyStream(chunks)
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=500,
|
||||
headers={"content-type": "text/plain"},
|
||||
stream=body_stream,
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"),
|
||||
)
|
||||
|
||||
served_at_warning: list[int] = []
|
||||
real_warning: Final = verbose_proxy_logger.warning
|
||||
|
||||
def _recording_warning(*args, **kwargs):
|
||||
if args and args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s":
|
||||
served_at_warning.append(body_stream.served)
|
||||
return real_warning(*args, **kwargs)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning):
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
) as mock_success_handler:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.status_code == 500
|
||||
streamed_chunks: Final = [chunk async for chunk in response.body_iterator]
|
||||
streamed_bytes: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks
|
||||
)
|
||||
assert streamed_bytes == upstream_content
|
||||
|
||||
assert served_at_warning == [5], (
|
||||
"each raw chunk is yielded as-is; five 1024-byte chunks are the first point the preview budget is exceeded"
|
||||
)
|
||||
expected_body: Final = f"{'x' * 4096}... (truncated at 4096 chars)"
|
||||
assert (
|
||||
mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail
|
||||
== f"Upstream passthrough request failed with status 500: {expected_body}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_streaming_upstream_error_single_large_chunk_stays_bounded():
|
||||
first_chunk: Final = b"x" * 65536
|
||||
second_chunk: Final = b'{"error": "tail"}'
|
||||
upstream_content: Final = first_chunk + second_chunk
|
||||
body_stream: Final = _ChunkedUpstreamErrorBodyStream((first_chunk, second_chunk))
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=500,
|
||||
headers={"content-type": "text/plain"},
|
||||
stream=body_stream,
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"),
|
||||
)
|
||||
|
||||
served_at_warning: list[int] = []
|
||||
real_warning: Final = verbose_proxy_logger.warning
|
||||
|
||||
def _recording_warning(*args, **kwargs):
|
||||
if args and args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s":
|
||||
served_at_warning.append(body_stream.served)
|
||||
return real_warning(*args, **kwargs)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning):
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
) as mock_success_handler:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.status_code == 500
|
||||
streamed_chunks: Final = [chunk async for chunk in response.body_iterator]
|
||||
streamed_bytes: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks
|
||||
)
|
||||
assert streamed_bytes == upstream_content
|
||||
|
||||
assert served_at_warning == [1], (
|
||||
"the rechunked preview is served from the first raw chunk; the second must not be pulled before the warning"
|
||||
)
|
||||
expected_body: Final = f"{'x' * 4096}... (truncated at 4096 chars)"
|
||||
assert (
|
||||
mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail
|
||||
== f"Upstream passthrough request failed with status 500: {expected_body}"
|
||||
)
|
||||
|
||||
|
||||
class _UpstreamErrorBodyStreamDropping(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'{"error": "half'
|
||||
raise httpx.ReadError("peer reset")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_streaming_upstream_error_body_read_failure_keeps_status_and_partial_body():
|
||||
"""
|
||||
Regression: a 502 whose upstream dies while the error preview is being read
|
||||
must still reach the client with status 502 and the bytes already received;
|
||||
the read failure must not escape as a ProxyException 500.
|
||||
"""
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=502,
|
||||
headers={"content-type": "application/json"},
|
||||
stream=_UpstreamErrorBodyStreamDropping(),
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"),
|
||||
)
|
||||
|
||||
recorded_warnings: list[tuple] = []
|
||||
real_warning: Final = verbose_proxy_logger.warning
|
||||
|
||||
def _recording_warning(*args, **kwargs):
|
||||
if args and str(args[0]).startswith("pass_through_endpoint: upstream"):
|
||||
recorded_warnings.append(args)
|
||||
return real_warning(*args, **kwargs)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning):
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
) as mock_success_handler:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.status_code == 502
|
||||
streamed_chunks: Final = [chunk async for chunk in response.body_iterator]
|
||||
streamed_bytes: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks
|
||||
)
|
||||
assert streamed_bytes == b'{"error": "half'
|
||||
await upstream_response.aclose()
|
||||
|
||||
rendered: Final = [str(args[0]) for args in recorded_warnings]
|
||||
formats: Final = [args[0] for args in recorded_warnings]
|
||||
assert any(
|
||||
fmt == "pass_through_endpoint: upstream %s %s returned %s: %s" and '{"error": "half' in str(args[4])
|
||||
for args, fmt in zip(recorded_warnings, formats)
|
||||
), rendered
|
||||
assert any(
|
||||
fmt == "pass_through_endpoint: upstream error body read failed after %d bytes: %s"
|
||||
and args[1] == 15
|
||||
and args[2] == "ReadError"
|
||||
for args, fmt in zip(recorded_warnings, formats)
|
||||
), rendered
|
||||
|
||||
|
||||
class _UpstreamErrorGzipStreamDropping(httpx.AsyncByteStream):
|
||||
def __init__(self, flushed_prefix: bytes) -> None:
|
||||
self._flushed_prefix: Final = flushed_prefix
|
||||
|
||||
async def __aiter__(self):
|
||||
yield self._flushed_prefix
|
||||
raise httpx.ReadError("peer reset")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_streaming_upstream_error_gzip_read_failure_relays_decoded_partial():
|
||||
"""
|
||||
Regression: a mid-read failure on a gzip upstream must relay the decoded
|
||||
plaintext, not the compressed bytes; the relay strips content-encoding so
|
||||
raw compressed bytes would reach the client as garbage.
|
||||
"""
|
||||
plaintext: Final = b'{"error": "half'
|
||||
compressor: Final = zlib.compressobj(level=6, wbits=31)
|
||||
flushed_prefix: Final = compressor.compress(plaintext) + compressor.flush(zlib.Z_SYNC_FLUSH)
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=502,
|
||||
headers={"content-type": "application/json", "content-encoding": "gzip"},
|
||||
stream=_UpstreamErrorGzipStreamDropping(flushed_prefix),
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"),
|
||||
)
|
||||
|
||||
recorded_warnings: list[tuple] = []
|
||||
real_warning: Final = verbose_proxy_logger.warning
|
||||
|
||||
def _recording_warning(*args, **kwargs):
|
||||
if args and str(args[0]).startswith("pass_through_endpoint: upstream"):
|
||||
recorded_warnings.append(args)
|
||||
return real_warning(*args, **kwargs)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning):
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
) as mock_success_handler:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
assert response.status_code == 502
|
||||
assert "content-encoding" not in response.headers
|
||||
streamed_chunks: Final = [chunk async for chunk in response.body_iterator]
|
||||
streamed_bytes: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks
|
||||
)
|
||||
assert streamed_bytes == plaintext
|
||||
await upstream_response.aclose()
|
||||
|
||||
rendered: Final = [str(args[0]) for args in recorded_warnings]
|
||||
assert any(
|
||||
args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s" and plaintext.decode() in str(args[4])
|
||||
for args in recorded_warnings
|
||||
), rendered
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_streaming_upstream_error_gzip_body_decoded_for_log_and_client():
|
||||
upstream_content: Final = b'{"error": {"message": "gzipped upstream says the project was not found"}}'
|
||||
compressed: Final = gzip.compress(upstream_content)
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=502,
|
||||
headers={"content-type": "text/event-stream", "content-encoding": "gzip"},
|
||||
stream=_ChunkedUpstreamErrorBodyStream((compressed[:10], compressed[10:])),
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"),
|
||||
)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning") as mock_warning:
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler"
|
||||
) as mock_success_handler:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_success_handler.return_value = None
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
response: Final = await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingResponse)
|
||||
streamed_chunks: Final = [chunk async for chunk in response.body_iterator]
|
||||
streamed_bytes: Final = b"".join(
|
||||
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks
|
||||
)
|
||||
assert streamed_bytes == upstream_content
|
||||
|
||||
upstream_warnings: Final = [
|
||||
call
|
||||
for call in mock_warning.call_args_list
|
||||
if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s"
|
||||
]
|
||||
assert len(upstream_warnings) == 1, mock_warning.call_args_list
|
||||
logged_body: Final = str(upstream_warnings[0].args[4])
|
||||
assert "gzipped upstream says the project was not found" in logged_body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_request_upstream_error_body_sanitized_against_log_forging():
|
||||
upstream_content: Final = b'{"error": "line one"}\n2026-01-01 FAKE LOG LINE\x1b[31m'
|
||||
upstream_response: Final = httpx.Response(
|
||||
status_code=404,
|
||||
headers={"content-type": "application/json"},
|
||||
content=upstream_content,
|
||||
request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:generateContent"),
|
||||
)
|
||||
|
||||
with patch.object(verbose_proxy_logger, "warning") as mock_warning:
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
|
||||
) as mock_get_client:
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
|
||||
) as mock_processing:
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
|
||||
async_client: Final = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
mock_get_client.return_value = MagicMock(client=async_client)
|
||||
|
||||
await pass_through_request(
|
||||
request=_upstream_error_request(),
|
||||
target="http://target-api.com/v1beta/models/claude-nope-9:generateContent",
|
||||
custom_headers={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
)
|
||||
|
||||
upstream_warnings: Final = [
|
||||
call
|
||||
for call in mock_warning.call_args_list
|
||||
if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s"
|
||||
]
|
||||
assert len(upstream_warnings) == 1, mock_warning.call_args_list
|
||||
logged_body: Final = str(upstream_warnings[0].args[4])
|
||||
assert logged_body == '{"error": "line one"} 2026-01-01 FAKE LOG LINE [31m'
|
||||
assert "\n" not in logged_body
|
||||
assert "\x1b" not in logged_body
|
||||
|
||||
detail: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail
|
||||
assert (
|
||||
detail
|
||||
== 'Upstream passthrough request failed with status 404: {"error": "line one"} 2026-01-01 FAKE LOG LINE [31m'
|
||||
)
|
||||
|
||||
|
||||
class _UpstreamDroppingMidStream(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "hi"}}]}\n\n'
|
||||
|
|
@ -4287,7 +4991,9 @@ async def test_pass_through_request_claims_the_budget_reservation_only_when_its_
|
|||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
mock_processing.get_custom_headers.return_value = {}
|
||||
mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close())
|
||||
mock_worker.ensure_initialized_and_enqueue = MagicMock(
|
||||
side_effect=lambda async_coroutine: async_coroutine.close()
|
||||
)
|
||||
async_client = MagicMock()
|
||||
async_client.build_request = MagicMock(return_value=MagicMock())
|
||||
async_client.send = AsyncMock(return_value=upstream_response)
|
||||
|
|
@ -5304,9 +6010,7 @@ async def test_websocket_passthrough_propagates_active_trace_context(
|
|||
mock_proxy_logging.post_call_success_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_worker = MagicMock()
|
||||
mock_worker.ensure_initialized_and_enqueue = MagicMock(
|
||||
side_effect=lambda async_coroutine: async_coroutine.close()
|
||||
)
|
||||
mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue