diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index b86cc954b0e..3c08f16fa80 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,8 +1,11 @@ from collections.abc import Mapping from typing import Final +import httpx from pydantic import TypeAdapter, ValidationError +from litellm.constants import PROVIDER_REQUEST_ID_HEADERS + _HEADERS: Final = TypeAdapter(Mapping[str, str]) _REQUEST_ID_HEADERS: Final = ("request-id", "x-request-id", "llm_provider-request-id", "llm_provider-x-request-id") @@ -16,6 +19,16 @@ def get_provider_request_id(headers: object) -> str | None: return next((normalized[name] for name in _REQUEST_ID_HEADERS if normalized.get(name)), None) +def provider_request_id_headers(headers: Mapping[str, object] | None) -> httpx.Headers | None: + """The provider request id headers in `headers` whose values are ASCII strings, or None when there are none.""" + request_id_headers: Final = tuple( + (name, value) + for name, value in (headers or {}).items() + if name.lower() in PROVIDER_REQUEST_ID_HEADERS and isinstance(value, str) and value.isascii() + ) + return httpx.Headers(request_id_headers) if request_id_headers else None + + def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index a207bceb418..6ad400661be 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data from litellm.llms.anthropic.chat.handler import ( @@ -448,6 +449,7 @@ class _EventStreamTally: f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, " f"{self.bytes_received} bytes received" ) + request_id_headers: Final = provider_request_id_headers(response_headers) if not self.events: return BedrockError( status_code=502, @@ -455,6 +457,7 @@ class _EventStreamTally: "Bedrock answered the stream with HTTP 200 but its body decoded to no events " f"({detail}, first bytes={self.head!r})" ), + headers=request_id_headers, ) if not self.recognized_events: return BedrockError( @@ -464,10 +467,12 @@ class _EventStreamTally: f"event type (event types={sorted(self.unrecognized_event_types)}, {detail}, " f"first payload={self.unrecognized_head!r})" ), + headers=request_id_headers, ) return BedrockError( status_code=502, message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})", + headers=request_id_headers, ) @@ -837,11 +842,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) @@ -858,11 +864,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) async for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: _data = json.loads(message) yield self.chunk_parser(chunk_data=_data) @@ -882,27 +889,34 @@ class AWSEventStreamDecoder: ) return response_stream_shape - def _decode_event(self, event: "EventStreamMessage", tally: _EventStreamTally) -> str | None: + def _decode_event( + self, event: "EventStreamMessage", tally: _EventStreamTally, request_id_headers: httpx.Headers | None + ) -> str | None: response_stream_shape: Final = self._response_stream_shape() response_dict: Final = bedrock_event_stream_response(event) tally.add_event(event, response_dict["headers"]) - return self._parse_message_from_response(response_dict, response_stream_shape) + return self._parse_message_from_response(response_dict, response_stream_shape, request_id_headers) def _parse_message_from_event(self, event: "EventStreamMessage") -> str | None: response_stream_shape: Final = self._response_stream_shape() return self._parse_message_from_response(bedrock_event_stream_response(event), response_stream_shape) def _parse_message_from_response( - self, response_dict: BedrockEventStreamResponseDict, response_stream_shape: "Shape" + self, + response_dict: BedrockEventStreamResponseDict, + response_stream_shape: "Shape", + request_id_headers: httpx.Headers | None = None, ) -> str | None: parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - raise build_bedrock_stream_error(response_dict, response_stream_shape) + raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers) event_type: Final = bedrock_event_stream_header(response_dict["headers"], ":event-type") event_error_status: Final = bedrock_stream_event_error_status(event_type) if event_type is not None and event_error_status is not None: - raise build_bedrock_stream_event_error(event_type, event_error_status, response_dict["body"]) + raise build_bedrock_stream_event_error( + event_type, event_error_status, response_dict["body"], headers=request_id_headers + ) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index e79865ef6bf..5862b05a619 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1913,13 +1913,18 @@ def bedrock_event_stream_header(headers: Mapping[str, object], name: str) -> str return value if isinstance(value, str) else None -def build_bedrock_stream_event_error(event_type: str, status_code: int, body: bytes) -> BedrockError: - return BedrockError(status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}") +def build_bedrock_stream_event_error( + event_type: str, status_code: int, body: bytes, headers: httpx.Headers | None = None +) -> BedrockError: + return BedrockError( + status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}", headers=headers + ) def build_bedrock_stream_error( response_dict: BedrockEventStreamResponseDict, response_stream_shape: Shape | None, + headers: httpx.Headers | None = None, ) -> BedrockError: """Build a BedrockError for a non-200 event-stream error event. @@ -1929,12 +1934,14 @@ def build_bedrock_stream_error( """ exception_type: Final = bedrock_event_stream_header(response_dict["headers"], ":exception-type") if exception_type is None: - return BedrockError(status_code=response_dict["status_code"], message=response_dict["body"].decode()) + return BedrockError( + status_code=response_dict["status_code"], message=response_dict["body"].decode(), headers=headers + ) member: Final = _structure_members(response_stream_shape).get(exception_type) modeled_status: Final = None if member is None else _modeled_error_status(member) status_code: Final = response_dict["status_code"] if modeled_status is None else modeled_status - return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"]) + return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"], headers=headers) class BedrockEventStreamDecoderBase: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7cf68332c07..89c32aed1fb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) @@ -262,7 +263,9 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: return _status_code_for_error_fields(error_type, error_code), message -def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: +def _map_stream_error_to_exception( + error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None +) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException error_message, error_type, error_code = _error_event_fields(error_obj) @@ -272,6 +275,7 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro status_code=status_code, message=f"Error code: {status_code} - {{'error': {error_body}}}", body=error_body, + headers=headers, ) try: return litellm.exception_type( @@ -667,7 +671,12 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: - return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "") + return _map_stream_error_to_exception( + error_obj, + self.model or "", + self.custom_llm_provider or "", + headers=provider_request_id_headers(self._raw_response_headers), + ) def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) @@ -683,11 +692,17 @@ class BaseResponsesAPIStreamingIterator: mapped_exception: Final = self._map_error_event_exception(error_obj) if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception + mapped_status: Final = getattr(mapped_exception, "status_code", None) raise MidStreamFallbackError( message=str(mapped_exception), model=self.model or "", llm_provider=self.custom_llm_provider or "", original_exception=mapped_exception, + response=httpx.Response( + status_code=int(mapped_status) if mapped_status is not None else 503, + headers=provider_request_id_headers(self._raw_response_headers), + request=httpx.Request(method="POST", url=f"https://{self.custom_llm_provider or ''}.com/v1/"), + ), generated_content=self._generated_content, is_pre_first_chunk=not self._yielded_first_chunk, ) diff --git a/tests/integration/observability/test_bedrock_error_request_id.py b/tests/integration/observability/test_bedrock_error_request_id.py index 230619a2863..436dcc572d5 100644 --- a/tests/integration/observability/test_bedrock_error_request_id.py +++ b/tests/integration/observability/test_bedrock_error_request_id.py @@ -1,13 +1,208 @@ +import asyncio +import base64 import json import uuid -from typing import Final +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from types import MappingProxyType +from typing import Final, Literal, NamedTuple +from urllib.parse import unquote -from integration._support.client import Gateway, eventually, object_value +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value from integration._support.database import read_rows +from integration._support.responses_stream import ( + created, + delta, + error_event, + failed, + frame, + healthy_stream, + openai_error_event, + spend_row, +) +from integration._support.upstream import _aws_event_frame, aws_event_stream_frame from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id _MODEL: Final = "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" +_BEDROCK_RESPONSES_MODEL: Final = "bedrock/us.openai.gpt-5.6-sol" _TOKEN: Final = "synthetic-bedrock-bearer" +_SHAPES: Final = ("nested-error", "top-level-error", "response.failed", "error-before-output") +_ERROR_CODES: Final = ("server_error", "rate_limit_exceeded", "invalid_request_error", "content_policy_violation") +_INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" +_BEDROCK_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BEDROCK_STREAM_TARGETS: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream", + f"/model/{_INVOKE_MODEL_ID}/invoke-with-response-stream", + ) +) + +BedrockFailure = Literal["exception-event", "exception-message", "no-known-event", "truncated-frame"] + + +class _FailureShape(NamedTuple): + bedrock_status: int + converse_class: str + after_output: tuple[bool, ...] + + +BedrockEndpoint = Literal["chat", "responses", "messages", "messages-invoke"] +_BEDROCK_ENDPOINTS: Final[tuple[BedrockEndpoint, ...]] = ("chat", "responses", "messages", "messages-invoke") +_BEDROCK_FAILURES: Final[Mapping[BedrockFailure, _FailureShape]] = MappingProxyType( + { + "exception-event": _FailureShape(500, "ServiceUnavailableError", (False, True)), + "exception-message": _FailureShape(429, "RateLimitError", (False, True)), + "no-known-event": _FailureShape(502, "BadGatewayError", (False,)), + "truncated-frame": _FailureShape(502, "BadGatewayError", (True,)), + } +) +_BEDROCK_CELLS: Final = tuple( + pytest.param( + endpoint, failure, after_output, id=f"{endpoint}-{failure}-{'after' if after_output else 'before'}-output" + ) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _stream(identity: str, shape: str, error_code: str, message: str) -> tuple[bytes, ...]: + opening: Final = (frame(created(identity)), frame(delta(identity, "partial output"))) + nested: Final = frame(error_event({"type": error_code, "code": error_code, "message": message})) + match shape: + case "nested-error": + return (*opening, nested) + case "top-level-error": + return (*opening, frame(openai_error_event(error_code, message))) + case "response.failed": + return (*opening, frame(failed(identity, error_code, message))) + case "error-before-output": + return (nested,) + raise AssertionError(shape) + + +def _bedrock_frame(event_type: str, payload: dict[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _bedrock_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +def _bedrock_output(invoke: bool, text: str = "partial output") -> tuple[bytes, ...]: + if not invoke: + return ( + _bedrock_frame("messageStart", {"role": "assistant"}), + _bedrock_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}), + ) + started: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "content": [], + "model": _INVOKE_MODEL_ID, + "stop_reason": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + return ( + _invoke_chunk({"type": "message_start", "message": started}), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + + +def _bedrock_failure_frame(failure: BedrockFailure, message: str) -> bytes: + match failure: + case "exception-event": + return _bedrock_frame("internalServerException", {"message": message}) + case "exception-message": + headers: Final = { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + } + return aws_event_stream_frame(headers, json.dumps({"message": message}).encode()) + case "no-known-event": + return _bedrock_frame("somethingBedrockAddedLater", {"message": message}) + case "truncated-frame": + return _bedrock_frame("contentBlockDelta", {"delta": {"text": message}, "contentBlockIndex": 0})[:-7] + + +def _bedrock_failure_text(failure: BedrockFailure, message: str) -> str: + return "undecoded bytes after" if failure == "truncated-frame" else message + + +def _bedrock_stream_request( + endpoint: BedrockEndpoint, model: str, prompt: str, *, stream: bool = True +) -> tuple[str, dict[str, JsonValue]]: + extra: Final[dict[str, JsonValue]] = {"stream": stream, "num_retries": 0, "cache": {"no-cache": True}} + turn: Final[list[JsonValue]] = [{"role": "user", "content": prompt}] + match endpoint: + case "chat": + return "/v1/chat/completions", {"model": model, "messages": turn, **extra} + case "responses": + return "/v1/responses", {"model": model, "input": prompt, **extra} + case "messages" | "messages-invoke": + return "/v1/messages", {"model": model, "messages": turn, "max_tokens": 16, **extra} + + +def _failure_information(call_id: str) -> dict[str, JsonValue]: + row: Final = spend_row(call_id) + assert row["status"] == "failure", row + metadata: Final = row["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + return object_value(parsed["error_information"]) + + +def _streamed_response(gateway: Gateway, model: str, prompt: str) -> httpx.Response: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": prompt, "stream": True, "num_retries": 0, "cache": {"no-cache": True}}, + ) + assert response.status_code == 200, response.text + return response + + +async def _consume_sdk_stream(asynchronous: bool, api_base: str, identity: str) -> None: + if asynchronous: + async_stream: Final = await litellm.aresponses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + async for _ in async_stream: + pass + return + tuple( + litellm.responses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + ) def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gateway: Gateway) -> None: @@ -54,3 +249,556 @@ def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gate parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) error_information: Final = object_value(parsed["error_information"]) assert error_information["error_provider_request_id"] == amzn_request_id, error_information + + +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize("error_code", _ERROR_CODES) +def test_responses_stream_failure_log_retains_provider_request_id( + gateway: Gateway, shape: str, error_code: str +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, error_code, message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert ("partial output" in response.text) is (shape != "error-before-output"), response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == provider_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("shape", _SHAPES) +def test_responses_stream_failure_without_provider_request_id_logs_none(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply(content_type="text/event-stream", chunks=_stream(identity, shape, "server_error", message)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert "llm_provider-x-amzn-requestid" not in response.headers, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_class"] == "InternalServerError", information + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("shape", ("nested-error", "response.failed")) +def test_bedrock_responses_stream_server_error_logs_amzn_request_id(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": amzn_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_BEDROCK_RESPONSES_MODEL, + api_key=_TOKEN, + api_base=wire.url, + aws_region_name="us-east-1", + num_retries=0, + ) + response: Final = _streamed_response(gateway, model, identity) + assert "partial output" in response.text and message in response.text, response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["llm_provider"] == "bedrock", information + assert information["error_provider_request_id"] == amzn_request_id, information + + +def test_responses_stream_failure_log_reads_openai_x_request_id(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = f"req_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-request-id": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, "top-level-error", "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +async def test_responses_sdk_stream_error_retains_provider_request_id(asynchronous: bool, event_type: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + shape: Final = "nested-error" if event_type == "error" else "response.failed" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-amzn-RequestId": provider_request_id, "retry-after": "7"}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", identity), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_sdk_stream(asynchronous, wire.url, identity) + assert raised.value.status_code == 500 + assert raised.value.response.status_code == 500 + assert str(raised.value.response.request.url) == "https://openai.com/v1/" + assert raised.value.generated_content == "partial output" + assert dict(raised.value.response.headers) == {"x-amzn-requestid": provider_request_id} + assert _get_provider_request_id(raised.value) == provider_request_id + assert raised.value.original_exception is not None + assert _get_provider_request_id(raised.value.original_exception) == provider_request_id + assert len(wire.drain()) == 1 + + +def _concurrent_reply(request: Request, failing: tuple[str, ...], recovery: str) -> Reply: + prompt: Final = json.loads(request.body)["input"] + assert prompt in (*failing, recovery) + chunks: Final = ( + healthy_stream(f"resp_{prompt}", "recovered") + if prompt == recovery + else _stream(f"resp_{prompt}", "nested-error", "server_error", prompt) + ) + return Reply(content_type="text/event-stream", headers={"x-amzn-requestid": prompt}, chunks=chunks) + + +def test_concurrent_responses_failures_keep_distinct_request_ids_and_recover(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + prompts: Final = tuple(f"{identity}-{index}" for index in range(4)) + recovery: Final = f"{identity}-recovery" + + with ( + wire_server(lambda request: _concurrent_reply(request, prompts, recovery)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + + def call(prompt: str) -> str: + response: Final = _streamed_response(gateway, model, prompt) + assert (prompt if prompt != recovery else "recovered") in response.text, response.text + return response.headers["x-litellm-call-id"] + + with ThreadPoolExecutor(max_workers=4) as workers: + call_ids: Final = tuple(workers.map(call, prompts)) + recovered_call_id: Final = call(recovery) + assert len(wire.drain()) == 5 + for prompt, call_id in zip(prompts, call_ids, strict=True): + assert _failure_information(call_id)["error_provider_request_id"] == prompt + assert spend_row(recovered_call_id)["status"] == "success" + + +def _bedrock_logged_error(endpoint: BedrockEndpoint, failure: BedrockFailure) -> tuple[str, str]: + shape: Final = _BEDROCK_FAILURES[failure] + if endpoint == "messages-invoke": + return "BedrockError", str(shape.bedrock_status) + return shape.converse_class, "503" if failure == "exception-event" else str(shape.bedrock_status) + + +def _bedrock_client_status(endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool) -> int: + logged_code: Final = int(_bedrock_logged_error(endpoint, failure)[1]) + if after_output or endpoint in ("responses", "messages"): + return 200 + if endpoint == "messages-invoke" and failure == "no-known-event": + return 200 + return logged_code + + +@pytest.mark.parametrize(("endpoint", "failure", "after_output"), _BEDROCK_CELLS) +def test_bedrock_event_stream_failure_log_retains_amzn_request_id( + gateway: Gateway, endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + invoke: Final = endpoint == "messages-invoke" + frames: Final = (*(_bedrock_output(invoke) if after_output else ()), _bedrock_failure_frame(failure, message)) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if invoke else _MODEL, + api_base=wire.url, + num_retries=0, + **_BEDROCK_AWS, + ) + path, body = _bedrock_stream_request(endpoint, model, message) + with gateway.client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + status: Final = response.status_code + response_headers: Final = dict(response.headers) + assert status == _bedrock_client_status(endpoint, failure, after_output), text + assert _bedrock_failure_text(failure, message) in text and ("partial output" in text) is after_output, text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(endpoint, failure) + assert information["error_provider_request_id"] == amzn_request_id, information + if status != 200 and not invoke: + assert response_headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, response_headers + assert "llm_provider-retry-after" not in response_headers, response_headers + + +def _bedrock_model(scenario: Scenario, endpoint: BedrockEndpoint, api_base: str) -> str: + return scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if endpoint == "messages-invoke" else _MODEL, + api_base=api_base, + num_retries=0, + **_BEDROCK_AWS, + ) + + +def _bedrock_stream( + gateway: Gateway, endpoint: BedrockEndpoint, model: str, prompt: str +) -> tuple[int, str, dict[str, str]]: + path, body = _bedrock_stream_request(endpoint, model, prompt) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + return response.status_code, text, dict(response.headers) + + +@pytest.mark.parametrize("after_output", (False, True), ids=("before-output", "after-output")) +@pytest.mark.parametrize("endpoint", _BEDROCK_ENDPOINTS) +def test_bedrock_event_stream_failure_without_amzn_request_id_logs_none( + gateway: Gateway, endpoint: BedrockEndpoint, after_output: bool +) -> None: + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = ( + *(_bedrock_output(endpoint == "messages-invoke") if after_output else ()), + _bedrock_failure_frame("exception-event", message), + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"retry-after": "7"}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + status, text, response_headers = _bedrock_stream( + gateway, endpoint, _bedrock_model(scenario, endpoint, wire.url), message + ) + assert status == _bedrock_client_status(endpoint, "exception-event", after_output), text + assert message in text, text + assert not any("amzn" in name for name in response_headers), response_headers + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error( + endpoint, "exception-event" + ) + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("endpoint", ("responses", "messages", "messages-invoke")) +def test_bedrock_non_stream_500_logs_amzn_request_id_on_every_endpoint( + gateway: Gateway, endpoint: BedrockEndpoint +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + targets: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse", + f"/model/{_INVOKE_MODEL_ID}/invoke", + ) + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in targets, request.target + return Reply( + status=500, + headers={"x-amzn-RequestId": amzn_request_id}, + body=b'{"message":"synthetic bedrock failure"}', + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + path, body = _bedrock_stream_request( + endpoint, _bedrock_model(scenario, endpoint, wire.url), amzn_request_id, stream=False + ) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 503, response.text + if endpoint == "responses": + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == amzn_request_id, information + + +@pytest.mark.parametrize("amzn_request_id", ("", "a" * 5000, "dup-1, dup-1"), ids=("empty", "5kb", "joined-duplicate")) +def test_bedrock_stream_failure_logs_odd_request_ids_like_the_non_stream_error( + gateway: Gateway, amzn_request_id: str +) -> None: + stream_body: Final = b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "odd id"))) + + def respond(request: Request) -> Reply: + if unquote(request.target) in _BEDROCK_STREAM_TARGETS: + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=stream_body) + return Reply(status=500, headers={"x-amzn-RequestId": amzn_request_id}, body=b'{"message":"odd id"}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_model(scenario, "chat", wire.url) + status, text, stream_headers = _bedrock_stream(gateway, "chat", model, "odd id") + path, body = _bedrock_stream_request("chat", model, "odd id", stream=False) + non_stream: Final = gateway.request("POST", path, body) + assert status == 200 and "odd id" in text, text + assert non_stream.status_code == 503, non_stream.text + assert len(wire.drain()) == 2 + streamed: Final = _failure_information(stream_headers["x-litellm-call-id"]) + unstreamed: Final = _failure_information(non_stream.headers["x-litellm-call-id"]) + assert streamed["error_provider_request_id"] == unstreamed["error_provider_request_id"], (streamed, unstreamed) + + +_BURST_CALLS: Final = tuple( + (endpoint, failure, after_output) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _healthy_bedrock_stream(invoke: bool) -> bytes: + closing: Final = ( + ( + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}} + ), + _invoke_chunk({"type": "message_stop"}), + ) + if invoke + else ( + _bedrock_frame("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_frame("messageStop", {"stopReason": "end_turn"}), + _bedrock_frame("metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}}), + ) + ) + return b"".join((*_bedrock_output(invoke, "recovered"), *closing)) + + +def test_concurrent_bedrock_stream_failures_on_every_endpoint_keep_distinct_request_ids(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + failing: Final = {f"{identity}-{index:02d}": call for index, call in enumerate(_BURST_CALLS)} + recoveries: Final = {f"{identity}-r{index}": endpoint for index, endpoint in enumerate(_BEDROCK_ENDPOINTS)} + + def respond(request: Request) -> Reply: + (token,) = (token for token in (*failing, *recoveries) if token in request.body.decode()) + if token in recoveries: + body = _healthy_bedrock_stream(recoveries[token] == "messages-invoke") + else: + endpoint, failure, after_output = failing[token] + opening = _bedrock_output(endpoint == "messages-invoke") if after_output else () + body = b"".join((*opening, _bedrock_failure_frame(failure, token))) + return Reply(headers={"x-amzn-RequestId": token, "retry-after": "7"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + models: Final = {endpoint: _bedrock_model(scenario, endpoint, wire.url) for endpoint in _BEDROCK_ENDPOINTS} + + def call(token: str, endpoint: BedrockEndpoint) -> tuple[int, str, dict[str, str]]: + return _bedrock_stream(gateway, endpoint, models[endpoint], token) + + with ThreadPoolExecutor(max_workers=8) as workers: + outcomes: Final = dict( + zip(failing, workers.map(call, failing, (endpoint for endpoint, _, _ in failing.values())), strict=True) + ) + recovered: Final = {token: call(token, endpoint) for token, endpoint in recoveries.items()} + assert len(wire.drain()) == len(failing) + len(recoveries) + for token, (endpoint, failure, after_output) in failing.items(): + status, text, response_headers = outcomes[token] + assert status == _bedrock_client_status(endpoint, failure, after_output), (token, text) + assert _bedrock_failure_text(failure, token) in text, (token, text) + information = _failure_information(response_headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == token, (token, information) + for token, (status, text, response_headers) in recovered.items(): + assert status == 200 and "recovered" in text, (token, text) + assert spend_row(response_headers["x-litellm-call-id"])["status"] == "success" + + +def _sdk_stream_text( + client_kind: str, asynchronous: bool, gateway: Gateway, model: str, prompt: str +) -> tuple[str, str]: + extra_body: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}} + turn: Final = [{"role": "user", "content": prompt}] + + def open_stream(client: openai.OpenAI | openai.AsyncOpenAI | anthropic.Anthropic | anthropic.AsyncAnthropic): + match client_kind: + case "openai-chat": + return client.chat.completions.with_raw_response.create( + model=model, messages=turn, stream=True, extra_body=extra_body + ) + case "openai-responses": + return client.responses.with_raw_response.create( + model=model, input=prompt, stream=True, extra_body=extra_body + ) + case "anthropic-messages": + return client.messages.with_raw_response.create( + model=model, messages=turn, max_tokens=16, stream=True, extra_body=extra_body + ) + raise AssertionError(client_kind) + + seen: Final[list[str]] = [] + base_url: Final = str(gateway.client.base_url).rstrip("/") + if not asynchronous: + with httpx.Client(timeout=15, trust_env=False) as transport: + client = ( + anthropic.Anthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.OpenAI(api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport) + ) + raw = open_stream(client) + try: + seen.extend(repr(event) for event in raw.parse()) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"], "\n".join(seen) + + async def run() -> str: + async with httpx.AsyncClient(timeout=15, trust_env=False) as transport: + client = ( + anthropic.AsyncAnthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.AsyncOpenAI( + api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport + ) + ) + raw = await open_stream(client) + try: + async for event in raw.parse(): + seen.append(repr(event)) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"] + + return asyncio.run(run()), "\n".join(seen) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("client_kind", ("openai-chat", "openai-responses", "anthropic-messages")) +def test_sdk_clients_streaming_bedrock_failure_through_proxy_log_amzn_request_id( + gateway: Gateway, client_kind: str, asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + call_id, seen = _sdk_stream_text( + client_kind, asynchronous, gateway, _bedrock_model(scenario, "chat", wire.url), message + ) + assert "partial output" in seen and message in seen, seen + assert len(wire.drain()) == 1 + information: Final = _failure_information(call_id) + assert information["error_provider_request_id"] == amzn_request_id, information + + +async def _consume_bedrock_sdk_stream(api: Literal["chat", "responses"], asynchronous: bool, api_base: str) -> str: + common: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "stream": True, + "api_base": api_base, + "num_retries": 0, + **_BEDROCK_AWS, + } + turn: Final[list[JsonValue]] = [{"role": "user", "content": "probe"}] + match api, asynchronous: + case "chat", True: + return "".join([repr(chunk) async for chunk in await litellm.acompletion(messages=turn, **common)]) + case "chat", False: + return "".join(repr(chunk) for chunk in litellm.completion(messages=turn, **common)) + case "responses", True: + return "".join([repr(event) async for event in await litellm.aresponses(input="probe", **common)]) + case "responses", False: + return "".join(repr(event) for event in litellm.responses(input="probe", **common)) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("api", ("chat", "responses")) +async def test_bedrock_sdk_stream_exception_carries_amzn_request_id( + api: Literal["chat", "responses"], asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, asynchronous, wire.url) + assert len(wire.drain()) == 1 + provider_error: Final = raised.value.original_exception + assert isinstance(provider_error, litellm.ServiceUnavailableError), provider_error + assert raised.value.generated_content == "partial output" + assert message in str(provider_error) + assert _get_provider_request_id(provider_error) == amzn_request_id + assert "retry-after" not in provider_error.response.headers, provider_error.response.headers + + +@pytest.mark.parametrize("api", ("chat", "responses")) +@pytest.mark.parametrize("healthy", (True, False), ids=("healthy", "failing")) +async def test_sync_bedrock_sdk_stream_with_non_ascii_request_id_header_keeps_streaming( + api: Literal["chat", "responses"], healthy: bool +) -> None: + body: Final = ( + _healthy_bedrock_stream(False) + if healthy + else b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "non-ascii id failure"))) + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": "req-\u00fc-1"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire: + if healthy: + assert "recovered" in await _consume_bedrock_sdk_stream(api, False, wire.url) + else: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, False, wire.url) + assert "non-ascii id failure" in str(raised.value.original_exception) + assert _get_provider_request_id(raised.value.original_exception) is None + assert len(wire.drain()) == 1 diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py index f52e837136a..1eccf5010ec 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( _get_llm_provider_headers, get_response_headers, get_provider_request_id, + provider_request_id_headers, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -89,6 +90,31 @@ def test_invalid_provider_request_ids_remain_absent(headers: object) -> None: assert get_provider_request_id(headers) is None +@pytest.mark.parametrize( + ("headers", "expected"), + ( + ( + {"X-Amzn-RequestId": "amzn-1", "retry-after": "7", "content-type": "text/event-stream"}, + {"x-amzn-requestid": "amzn-1"}, + ), + ( + {"x-request-id": "req_1", "cf-ray": "ray-1", "x-ratelimit-remaining-requests": "0"}, + {"x-request-id": "req_1", "cf-ray": "ray-1"}, + ), + ({"retry-after": "7"}, None), + ({"x-request-id": 7, "request-id": None, "x-amzn-requestid": "amzn-2"}, {"x-amzn-requestid": "amzn-2"}), + ({"x-amzn-requestid": "amzn-\u00fc", "x-request-id": "req_2"}, {"x-request-id": "req_2"}), + ({}, None), + (None, None), + ), +) +def test_provider_request_id_headers_keeps_only_request_id_headers( + headers: dict[str, object] | None, expected: dict[str, str] | None +) -> None: + result = provider_request_id_headers(headers) + assert (None if result is None else dict(result)) == expected + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 0609f122949..069438e25c0 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -16,6 +16,7 @@ import litellm from litellm import AmazonInvokeConfig from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AmazonOpenAICompatibleStreamDecoder, @@ -864,14 +865,16 @@ async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers( def _event_stream_frame(event_type: str, payload: bytes) -> bytes: + return _typed_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _typed_event_stream_frame(frame_headers: Mapping[str, str], payload: bytes) -> bytes: def header(name: str, value: str) -> bytes: return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() - headers: Final = ( - header(":event-type", event_type) - + header(":content-type", "application/json") - + header(":message-type", "event") - ) + headers: Final = b"".join(header(name, value) for name, value in frame_headers.items()) prelude: Final = struct.pack(">II", 12 + len(headers) + len(payload) + 4, len(headers)) body: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + headers + payload return body + struct.pack(">I", binascii.crc32(body)) @@ -1367,6 +1370,87 @@ def test_iter_bytes_keeps_a_stream_whose_unknown_event_sits_beside_known_frames( assert "stop" in finish_reasons +_STREAM_FAILURE_BODIES: Final = ( + pytest.param( + _converse_event_frame("internalServerException", {"message": _UPSTREAM_REJECTION}), 500, id="exception-event" + ), + pytest.param( + _typed_event_stream_frame( + { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": _UPSTREAM_REJECTION}).encode(), + ), + 429, + id="exception-message", + ), + pytest.param(_unknown_event_frame(), 502, id="no-known-event"), + pytest.param(b"", 502, id="no-events"), + pytest.param(_CONVERSE_CONTENT_FRAMES[0] + _truncated_frame(), 502, id="undecoded-tail"), +) + + +def _stream_headers(request_id: str | None) -> httpx.Headers: + request_id_header: Final = {} if request_id is None else {"x-amzn-RequestId": request_id} + return httpx.Headers( + {"content-type": "application/vnd.amazon.eventstream", "retry-after": "7", **request_id_header} + ) + + +def _assert_carries_only_the_request_id(error: BedrockError, status: int, request_id: str | None) -> None: + assert error.status_code == status + assert _get_provider_request_id(error) == request_id + if request_id is None: + assert error.headers is None + return + assert dict(error.headers or {}) == {"x-amzn-requestid": request_id} + assert dict(error.response.headers) == {"x-amzn-requestid": request_id} + + +@pytest.mark.parametrize("request_id", ("req-stream-1", None)) +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +def test_iter_bytes_stream_error_carries_only_the_provider_request_id_header( + body: bytes, status: int, request_id: str | None +) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([body]), response_headers=_stream_headers(request_id))) + + _assert_carries_only_the_request_id(exc_info.value, status, request_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +async def test_aiter_bytes_stream_error_carries_only_the_provider_request_id_header(body: bytes, status: int) -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield body + + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_stream_headers("req-stream-2"))] + + _assert_carries_only_the_request_id(exc_info.value, status, "req-stream-2") + + +@pytest.mark.parametrize("request_id", ("req-\u00fc-1", "req-\udcfc-1"), ids=("latin-1", "surrogate-escaped")) +def test_iter_bytes_streams_and_fails_normally_with_a_non_ascii_provider_request_id(request_id: str) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + headers: Final = {"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": request_id} + + chunks: Final = list(decoder.iter_bytes(iter(_CONVERSE_CONTENT_FRAMES), response_headers=headers)) + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([_unknown_event_frame()]), response_headers=headers)) + + texts: Final = [chunk.choices[0].delta.content for chunk in chunks if isinstance(chunk, ModelResponseStream)] + assert "".join(text or "" for text in texts) == "hi" + assert exc_info.value.status_code == 502 + assert exc_info.value.headers is None + + def _assert_exception_event_surfaced_with_its_modeled_status(error: BaseException, event_type: str) -> None: assert not isinstance(error, litellm.BadGatewayError) assert getattr(error, "status_code", None) == _modeled_status(event_type) diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index c37e57a43ca..001be7e1fbd 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -20,14 +20,18 @@ Pydantic ValidationError (previously typed as Optional[str]). """ import json +from datetime import datetime from importlib import import_module +from typing import Final from unittest.mock import Mock, patch +import httpx import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, @@ -64,6 +68,52 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj) +@pytest.mark.parametrize("provider", ("openai", "azure", "bedrock")) +@pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None)) +@pytest.mark.parametrize( + "code,status", (("server_error", 500), ("rate_limit_exceeded", 429), ("invalid_request_error", 400)) +) +def test_stream_error_preserves_provider_request_id(provider: str, header: str | None, code: str, status: int) -> None: + stream_headers: Final = {"retry-after": "7", "content-type": "text/event-stream"} + response: Final = httpx.Response( + 200, headers={**stream_headers, header: "provider-request-id"} if header else stream_headers + ) + logging: Final = LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=True, + call_type="responses", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + iterator: Final = BaseResponsesAPIStreamingIterator( + response=response, + model="test-model", + responses_api_provider_config=None, + logging_obj=logging, + custom_llm_provider=provider, + ) + expected_id: Final = "provider-request-id" if header else None + expected_headers: Final = {header.lower(): "provider-request-id"} if header else {} + mapped: Final = iterator._map_error_event_exception({"type": code, "code": code, "message": "failure"}) + assert _get_provider_request_id(mapped) == expected_id + assert dict(getattr(mapped, "litellm_response_headers", None) or {}) == expected_headers + expected_status: Final = 503 if provider == "bedrock" and status == 500 else status + with pytest.raises((litellm.BadRequestError, MidStreamFallbackError)) as raised: + iterator._maybe_raise_for_error_event(_make_error_chunk(code, code, "failure")) + assert raised.value.status_code == expected_status + assert _get_provider_request_id(raised.value) == expected_id + assert response.status_code == 200 + if status in (429, 500): + assert isinstance(raised.value, MidStreamFallbackError) + assert dict(raised.value.response.headers) == expected_headers + assert raised.value.response.status_code == expected_status + assert str(raised.value.response.request.url) == f"https://{provider}.com/v1/" + assert raised.value.generated_content == "" + assert raised.value.is_pre_first_chunk is True + + def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback(): iterator = _make_iterator() chunk = _make_error_chunk("server_error", "internal_error", "something went wrong")