diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d5bbf180e97..7c176a36a7a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -650,6 +650,7 @@ class Logging(LiteLLMLoggingBaseClass): # enqueue closure here instead of firing it immediately. self._defer_async_logging: bool = False self._enqueue_deferred_logging: Callable[[], None] | None = None + self._async_success_scheduled: bool = False self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None: @@ -2065,6 +2066,15 @@ class Logging(LiteLLMLoggingBaseClass): except Exception: return True + def claim_async_success_log(self) -> bool: + """One async success log per request. The innermost @client wrapper always exits first, + since the outer one is awaiting it, so it claims the log here and the outer wrapper gets + False and schedules nothing.""" + if self._async_success_scheduled: + return False + self._async_success_scheduled = True + return True + def has_run_logging( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], diff --git a/litellm/litellm_core_utils/realtime_errors.py b/litellm/litellm_core_utils/realtime_errors.py index 3c064728a66..7d5f8ca23ff 100644 --- a/litellm/litellm_core_utils/realtime_errors.py +++ b/litellm/litellm_core_utils/realtime_errors.py @@ -9,13 +9,20 @@ frame itself fail, which is how a loud failure turns back into a silent one. """ import json -from typing import Final +from types import MappingProxyType +from typing import Final, Protocol from litellm.types.realtime import RealtimeErrorDetail, RealtimeErrorEvent WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 +class _ClientWebSocket(Protocol): + async def send_text(self, data: str) -> None: ... + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + def realtime_error_event(message: str, error_type: str) -> str: detail: Final[RealtimeErrorDetail] = {"type": error_type, "message": message} event: Final[RealtimeErrorEvent] = {"type": "error", "error": detail} @@ -37,3 +44,31 @@ def client_close_code(upstream_code: int) -> int: if upstream_code in EXTERNAL_CLOSE_CODES or 3000 <= upstream_code < 5000: return upstream_code return int(CloseCode.INTERNAL_ERROR) + + +def upstream_handshake_close_code(status_code: int) -> int: + from websockets.frames import CloseCode + + refusal_codes: Final = MappingProxyType( + { + 401: int(CloseCode.POLICY_VIOLATION), + 403: int(CloseCode.POLICY_VIOLATION), + 429: int(CloseCode.TRY_AGAIN_LATER), + } + ) + return refusal_codes.get(status_code, int(CloseCode.INTERNAL_ERROR)) + + +async def close_after_upstream_handshake_refusal(websocket: _ClientWebSocket, status_code: int) -> None: + message: Final = f"Upstream realtime handshake rejected with HTTP {status_code}" + try: + await websocket.send_text(realtime_error_event(message, error_type="server_error")) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + pass + try: + await websocket.close( + code=upstream_handshake_close_code(status_code), + reason=websocket_close_reason(message, fallback="Upstream handshake rejected"), + ) + except Exception: # noqa: BLE001 # the client may have closed the socket already; closing twice is not an error + pass diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 146915dd6fd..df74975ad0f 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -8,11 +8,15 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Any, Final, Protocol, cast -from litellm._logging import _redact_string, verbose_proxy_logger +from litellm._logging import verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ....litellm_core_utils.realtime_errors import ( + close_after_upstream_handshake_refusal, + realtime_error_event, +) from ....litellm_core_utils.realtime_streaming import ( RealTimeStreaming, ScopedWebSocket, @@ -49,7 +53,9 @@ def azure_realtime_protocol_for_client( class _ProxyClientWebSocket(Protocol): - """Client-facing websocket handle: this path only closes it after a failed handshake.""" + """Client-facing websocket handle: this path only writes to it after a failed handshake.""" + + async def send_text(self, data: str) -> None: ... async def close(self, code: int = ..., reason: str | None = ...) -> None: ... @@ -181,7 +187,16 @@ class AzureOpenAIRealtime(AzureChatCompletion): ) await realtime_streaming.bidirectional_forward() - except websockets.exceptions.InvalidStatusCode as e: - await websocket.close(code=e.status_code, reason=_redact_string(str(e))) + except websockets.exceptions.InvalidStatus as e: + verbose_proxy_logger.exception("Error in AzureOpenAIRealtime.async_realtime") + await close_after_upstream_handshake_refusal(websocket, e.response.status_code) except Exception: verbose_proxy_logger.exception("Error in AzureOpenAIRealtime.async_realtime") + try: + await websocket.send_text(realtime_error_event("Internal server error", error_type="server_error")) + except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below + pass + try: + await websocket.close(code=1011, reason="Internal server error") + except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error + pass diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 49a332e62bb..8e85f10781e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -44,7 +44,11 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( ) from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields -from litellm.litellm_core_utils.realtime_errors import realtime_error_event, websocket_close_reason +from litellm.litellm_core_utils.realtime_errors import ( + close_after_upstream_handshake_refusal, + realtime_error_event, + websocket_close_reason, +) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -6372,9 +6376,9 @@ class BaseLLMHTTPHandler: await realtime_streaming.bidirectional_forward() - except websockets.exceptions.InvalidStatusCode as e: + except websockets.exceptions.InvalidStatus as e: verbose_logger.exception("Error connecting to backend: %s", e) - await websocket.close(code=e.status_code, reason=_redact_string(str(e))) + await close_after_upstream_handshake_refusal(websocket, e.response.status_code) except Exception as e: verbose_logger.exception("Error connecting to backend: %s", e) redacted_error: Final = _redact_string(str(e)) @@ -6787,9 +6791,9 @@ class BaseLLMHTTPHandler: ) return await streaming.bidirectional_forward() - except websockets.exceptions.InvalidStatusCode as e: + except websockets.exceptions.InvalidStatus as e: verbose_logger.exception("Error connecting to responses WS backend: %s", e) - await websocket.close(code=e.status_code, reason=_redact_string(str(e))) + await close_after_upstream_handshake_refusal(websocket, e.response.status_code) except Exception as e: verbose_logger.exception("Error in responses WS: %s", e) try: diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index e3ecbac1a53..bdc3a6c7908 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -12,6 +12,7 @@ from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ....litellm_core_utils.realtime_errors import close_after_upstream_handshake_refusal from ....litellm_core_utils.realtime_streaming import ( RealtimeEventNormalizer, RealTimeStreaming, @@ -175,8 +176,8 @@ class OpenAIRealtime(OpenAIChatCompletion): ) await realtime_streaming.bidirectional_forward() - except websockets.exceptions.InvalidStatusCode as e: - await websocket.close(code=e.status_code, reason=_redact_string(str(e))) + except websockets.exceptions.InvalidStatus as e: + await close_after_upstream_handshake_refusal(websocket, e.response.status_code) except Exception as e: try: await websocket.close(code=1011, reason=_redact_string(f"Internal server error: {e}")) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 63dbab23b54..c168c6f210f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -302,6 +302,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.realtime_errors import ( + close_after_upstream_handshake_refusal, realtime_error_event, websocket_close_reason, ) @@ -12385,9 +12386,9 @@ async def realtime_websocket_endpoint( user_model=user_model, ) await llm_call - except websockets.exceptions.InvalidStatusCode as e: + except websockets.exceptions.InvalidStatus as e: verbose_proxy_logger.exception("Invalid status code") - await websocket.close(code=e.status_code, reason="Invalid status code") + await close_after_upstream_handshake_refusal(websocket, e.response.status_code) except Exception as e: verbose_proxy_logger.exception("Internal server error") redacted_error: Final = _redact_string(str(e)) diff --git a/litellm/utils.py b/litellm/utils.py index b724313641f..9eed0e8345f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1268,11 +1268,14 @@ def _schedule_async_success_logging( Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses bridge) each exit through here with the same logging object and their own shape of the same - response. The immediate path already logs one request once, since the first task marks - ``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same - first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log - reads usage from, and a later wrapper never swaps in its client-shaped translation. + response. The innermost wrapper always exits first, since the outer one is awaiting it, so it + claims the async success log here, synchronously, and the outer wrapper returns without + enqueuing: one request queues one handler, whatever the handler tasks do later. The deferred + slot keeps the same innermost-wins rule: the provider-shaped result is the one the spend log + reads usage from, and the outer wrapper never swaps in its client-shaped translation. """ + if not logging_obj.claim_async_success_log(): + return def _enqueue_async_logging() -> None: asyncio.create_task( diff --git a/pyproject.toml b/pyproject.toml index bb7cd4e2901..d01deecb465 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.103.5" +version = "1.103.6" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -342,7 +342,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.103.5" +version = "1.103.6" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index 7c4a9cc4af9..7da7ceac4d3 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -286,7 +286,7 @@ def function_call_item(events: tuple[ReceivedEvent, ...]) -> OutputItem | None: # ---- session + client -------------------------------------------------- -def _as_text(message: str | bytes) -> str: +def as_text(message: str | bytes) -> str: return message.decode("utf-8") if isinstance(message, bytes) else message @@ -304,7 +304,7 @@ class RealtimeSession: collected: list[ReceivedEvent] = [] while time.monotonic() < deadline: try: - text = _as_text( + text = as_text( self.connection.recv(timeout=deadline - time.monotonic()) ) except TimeoutError: diff --git a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py index f99fa8d86b3..d7870b26497 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py @@ -13,8 +13,9 @@ failure. See REALTIME_COVERAGE_MATRIX.md. """ import pytest +from lifecycle import ResourceManager +from models import LiteLLMParamsBody from pydantic import BaseModel - from realtime_client import ( PROVIDERS, ConversationItemCreate, @@ -27,14 +28,17 @@ from realtime_client import ( RealtimeProvider, ResponseCreate, ResponseDone, + ServerEnvelope, SessionConfig, SessionUpdate, + as_text, function_call_item, parse_last, realtime_model, transcript, user_message, ) +from websockets.exceptions import ConnectionClosedError pytestmark = pytest.mark.e2e @@ -147,3 +151,39 @@ def test_tool_call_round_trip( second = session.collect_until("response.done", timeout=60) assert "72" in transcript(second), "follow-up did not use the tool result" + + +_REFUSED_UPSTREAMS = ( + RealtimeProvider( + "azure-bad-key", + "azure-realtime-refused", + LiteLLMParamsBody( + model="azure/gpt-realtime", + api_key="invalid-e2e-key", + api_version="2025-08-28", + realtime_protocol="GA", + ), + ), +) + + +@pytest.mark.parametrize("provider", _REFUSED_UPSTREAMS, ids=[p.id for p in _REFUSED_UPSTREAMS]) +def test_upstream_handshake_refusal_is_an_error_event_and_policy_close( + client: RealtimeClient, + resources: ResourceManager, + scoped_key: str, + provider: RealtimeProvider, +) -> None: + model_name, model_id = client.provision(provider) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + with client.connect(key=scoped_key, model=model_name) as session: + first = ServerEnvelope.model_validate_json( + as_text(session.connection.recv(timeout=15)) + ) + assert first.type == "error", first + with pytest.raises(ConnectionClosedError) as closed: + session.connection.recv(timeout=15) + + assert closed.value.rcvd is not None + assert closed.value.rcvd.code == 1008, closed.value diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index a4d4790ec8b..95b6063c6f6 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -1629,6 +1629,9 @@ ], "tests/integration/providers/test_responses_bridge_incomplete.py::test_chat_stream_over_responses_deployment_folds_message_and_two_function_calls_into_one_choice": [ "other.provider_wire.responses_bridge.stream_message_and_two_calls_fold_numbered_0_1" + ], + "tests/integration/spend/test_responses_bridge_spend_once.py::test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_once": [ + "quota_management.spend_tracking.responses_bridge.logs_and_charges_once_above_offload_threshold" ] }, "browser": { diff --git a/tests/integration/spend/test_responses_bridge_spend_once.py b/tests/integration/spend/test_responses_bridge_spend_once.py new file mode 100644 index 00000000000..4390c68ad33 --- /dev/null +++ b/tests/integration/spend/test_responses_bridge_spend_once.py @@ -0,0 +1,112 @@ +import json +from base64 import b64decode +from collections.abc import Callable +from hashlib import sha256 +from typing import Final +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +PROMPT_TOKENS: Final = 30 +COMPLETION_TOKENS: Final = 5 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE + + +def _responses_reply(prompt: str, response_id: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method == "GET" and request.target == "/v1/models": + return Reply( + body=json.dumps({"object": "list", "data": [{"id": "gpt-6.1-sol", "object": "model"}]}).encode() + ) + assert request.method == "POST", request.method + assert request.target == "/v1/responses", request.target + body: Final = object_value(json.loads(request.body)) + encoded_input: Final = json.dumps(body["input"]) + assert prompt in encoded_input, encoded_input + assert body.get("stream") is not True, body + response: Final = { + "id": response_id, + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": body["model"], + "output": [ + { + "type": "message", + "id": f"msg_{uuid4().hex}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": { + "input_tokens": PROMPT_TOKENS, + "output_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + } + return Reply(body=json.dumps(response).encode()) + + return respond + + +def _rows_for_key(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT litellm_call_id, request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +@pytest.mark.timeout(120) +@pytest.mark.covers("quota_management.spend_tracking.responses_bridge.logs_and_charges_once_above_offload_threshold") +def test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_once(gateway: Gateway) -> None: + prompt_prefix: Final = f"spend once {uuid4().hex[:8]} " + long_prompt: Final = (prompt_prefix + "lorem ipsum " * 30_000)[:300_000] + upstream_id: Final = f"resp_{uuid4().hex}" + with ( + wire_server(_responses_reply(long_prompt, upstream_id)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/responses/gpt-6.1-sol", + api_base=f"{wire.url}/v1", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ) + key: Final = scenario.key(models=[model]) + digest: Final = sha256(key.encode()).hexdigest() + + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": long_prompt}], "stream": False}, + key=key, + ) + assert response.status_code == 200, response.text + assert "ok" in response.text, response.text + assert response.json()["id"] == upstream_id, response.text + call_id: Final = response.headers["x-litellm-call-id"] + + rows: Final = eventually( + lambda: _rows_for_key(key), + lambda values: any(row["litellm_call_id"] == call_id for row in values), + seconds=70, + ) + call_rows: Final = [row for row in rows if row["litellm_call_id"] == call_id] + assert len(call_rows) == 1, rows + assert float(str(call_rows[0]["spend"])) == pytest.approx(EXPECTED_SPEND), rows + logged_id: Final = b64decode(str(call_rows[0]["request_id"]).removeprefix("resp_")).decode() + assert logged_id.startswith("litellm:") and f";response_id:{upstream_id}" in logged_id, rows + + key_spend: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda values: len(values) == 1 and float(str(values[0]["spend"])) >= EXPECTED_SPEND, + seconds=70, + ) + assert float(str(key_spend[0]["spend"])) == pytest.approx(EXPECTED_SPEND), key_spend diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 792a7bbaa89..fa4884ca47e 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4884,7 +4884,7 @@ def _interactions_logging_obj(stream: bool, call_type: str = "acreate"): messages=[], stream=stream, call_type=call_type, - start_time=time.time(), + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc).timestamp(), litellm_call_id="interactions-call-id", function_id="interactions-fn-id", ) @@ -7427,3 +7427,61 @@ class TestAzurePTUSpilloverCost: finally: litellm.model_cost.pop(custom_model_id, None) self._unregister_models() + + +@pytest.mark.asyncio +async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch): + """ + The create's in_progress success handler can still be parked on awaited work when + the settlement arrives. The completed result is the only event carrying usage and + cost, so it has to reach the success callbacks whatever the one-success-log-per-request + dedupe does, or the settlement is marked billed with no spend row behind it. + """ + from litellm.litellm_core_utils import litellm_logging + from litellm.types.interactions import InteractionsAPIResponse + + class CountingLogger(CustomLogger): + def __init__(self): + super().__init__() + self.logged_results: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_results.append(response_obj) + + counting_logger: Final = CountingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger]) + + in_progress_parked: Final = asyncio.Event() + release_in_progress: Final = asyncio.Event() + real_truncate: Final = litellm_logging.truncate_base64_in_messages_async + calls: list[int] = [] + + async def parked_then_real_truncate(messages): + calls.append(1) + if len(calls) == 1: + in_progress_parked.set() + await release_in_progress.wait() + return await real_truncate(messages) + + monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", parked_then_real_truncate) + + logging_obj: Final = _interactions_logging_obj(stream=False) + in_progress: Final = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + completed: Final = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + first: Final = asyncio.create_task(logging_obj.async_success_handler(result=in_progress)) + await in_progress_parked.wait() + completion: Final = asyncio.create_task(logging_obj.async_log_background_interaction_completion(result=completed)) + await asyncio.sleep(0) + await asyncio.sleep(0) + release_in_progress.set() + await first + await completion + + assert counting_logger.logged_results == [completed, in_progress], counting_logger.logged_results diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py index 1d2cf905f4e..572a4443ad6 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_errors.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_errors.py @@ -1,13 +1,18 @@ import json +from typing import cast import pytest +from starlette.websockets import WebSocketDisconnect from litellm.litellm_core_utils.realtime_errors import ( WEBSOCKET_CLOSE_REASON_MAX_BYTES, client_close_code, + close_after_upstream_handshake_refusal, realtime_error_event, + upstream_handshake_close_code, websocket_close_reason, ) +from litellm.types.realtime import RealtimeErrorEvent def test_realtime_error_event_shape(): @@ -52,3 +57,81 @@ def test_websocket_close_reason_truncates_multibyte_message_by_bytes(): ) def test_client_close_code_only_forwards_codes_a_server_may_send(upstream_code, expected): assert client_close_code(upstream_code) == expected + + +@pytest.mark.parametrize( + ("status_code", "expected"), + [(401, 1008), (403, 1008), (429, 1013), (500, 1011)], +) +def test_upstream_handshake_close_code_maps_http_status_to_close_code(status_code: int, expected: int): + assert upstream_handshake_close_code(status_code) == expected + + +class _RecordingWebSocket: + def __init__(self) -> None: + self.sent: list[str] = [] + self.closed: tuple[int, str | None] | None = None + + async def send_text(self, data: str) -> None: + self.sent.append(data) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + self.closed = (code, reason) + + +@pytest.mark.asyncio +async def test_close_after_upstream_handshake_refusal_sends_error_event_then_policy_close(): + websocket = _RecordingWebSocket() + + await close_after_upstream_handshake_refusal(websocket, 401) + + assert len(websocket.sent) == 1 + event = cast(RealtimeErrorEvent, json.loads(websocket.sent[0])) + assert event["type"] == "error" + assert event["error"]["type"] == "server_error" + assert "401" in event["error"]["message"] + assert websocket.closed is not None + assert websocket.closed[0] == 1008 + assert websocket.closed[1] + + +@pytest.mark.asyncio +async def test_close_after_upstream_handshake_refusal_still_closes_when_send_fails(): + class _DeadWebSocket(_RecordingWebSocket): + async def send_text(self, data: str) -> None: + raise RuntimeError("socket gone") + + websocket = _DeadWebSocket() + + await close_after_upstream_handshake_refusal(websocket, 500) + + assert websocket.closed == (1011, "Upstream realtime handshake rejected with HTTP 500") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "close_error", + [ + RuntimeError('Cannot call "send" once a close message has been sent.'), + OSError("client disconnected"), + WebSocketDisconnect(code=1006), + ], +) +async def test_close_after_upstream_handshake_refusal_returns_when_client_already_closed(close_error: Exception): + class _ClosedWebSocket(_RecordingWebSocket): + def __init__(self) -> None: + super().__init__() + self.close_attempts: list[int] = [] + + async def send_text(self, data: str) -> None: + raise OSError("client disconnected") + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + self.close_attempts.append(code) + raise close_error + + websocket = _ClosedWebSocket() + + await close_after_upstream_handshake_refusal(websocket, 401) + + assert websocket.close_attempts == [1008] diff --git a/tests/test_litellm/llms/azure/realtime/__init__.py b/tests/test_litellm/llms/azure/realtime/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm/llms/azure/realtime/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm/llms/azure/realtime/test_handler.py b/tests/test_litellm/llms/azure/realtime/test_handler.py new file mode 100644 index 00000000000..e9d24b459d8 --- /dev/null +++ b/tests/test_litellm/llms/azure/realtime/test_handler.py @@ -0,0 +1,86 @@ +import json +from typing import cast +from unittest.mock import MagicMock, patch + +import pytest + + +class _RecordingClientWebSocket: + scope: dict[str, list[tuple[bytes, bytes]]] = {"headers": []} + + def __init__(self) -> None: + self.sent: list[str] = [] + self.closed: list[tuple[int, str | None]] = [] + + async def send_text(self, data: str) -> None: + self.sent.append(data) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + self.closed.append((code, reason)) + + +@pytest.mark.asyncio +async def test_async_realtime_upstream_handshake_refusal_sends_error_event_then_policy_close(): + from websockets.datastructures import Headers + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response + + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + from litellm.types.realtime import RealtimeErrorEvent + + handler = AzureOpenAIRealtime() + model = "gpt-realtime" + + dummy_websocket = _RecordingClientWebSocket() + dummy_logging_obj = MagicMock() + + refused = InvalidStatus(Response(401, "Unauthorized", Headers())) + + with patch("websockets.connect", side_effect=refused): + await handler.async_realtime( # pyright: ignore[reportUnknownMemberType] # handler's websocket param is a Protocol here but the mock connect type is incomplete + model=model, + websocket=dummy_websocket, + logging_obj=dummy_logging_obj, + api_base="https://example.openai.azure.com", + api_key="bad-key", + api_version="2025-08-28", + query_params={"model": model}, + ) + + assert len(dummy_websocket.sent) == 1 + event = cast(RealtimeErrorEvent, json.loads(dummy_websocket.sent[0])) + assert event["type"] == "error" + assert event["error"]["type"] == "server_error" + assert "401" in event["error"]["message"] + assert dummy_websocket.closed and dummy_websocket.closed[0][0] == 1008 + + +@pytest.mark.asyncio +async def test_async_realtime_unexpected_error_sends_error_event_then_internal_close(): + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + from litellm.types.realtime import RealtimeErrorEvent + + handler = AzureOpenAIRealtime() + model = "gpt-realtime" + + dummy_websocket = _RecordingClientWebSocket() + dummy_logging_obj = MagicMock() + + with patch("websockets.connect", side_effect=OSError("connection reset")): + await handler.async_realtime( # pyright: ignore[reportUnknownMemberType] # same as above + model=model, + websocket=dummy_websocket, + logging_obj=dummy_logging_obj, + api_base="https://example.openai.azure.com", + api_key="bad-key", + api_version="2025-08-28", + query_params={"model": model}, + ) + + assert len(dummy_websocket.sent) == 1 + event = cast(RealtimeErrorEvent, json.loads(dummy_websocket.sent[0])) + assert event["type"] == "error" + assert event["error"]["type"] == "server_error" + assert event["error"]["message"] == "Internal server error" + assert "connection reset" not in dummy_websocket.sent[0] + assert dummy_websocket.closed and dummy_websocket.closed[0] == (1011, "Internal server error") diff --git a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py index 4221954d787..f7a88b5ba63 100644 --- a/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py +++ b/tests/test_litellm/llms/openai/realtime/test_openai_realtime_handler.py @@ -416,3 +416,52 @@ async def test_async_realtime_ws_url_has_no_ssl(): # Verify ssl is None for ws:// URLs (the fix for issue #19222) assert called_kwargs["ssl"] is None + + +@pytest.mark.asyncio +async def test_async_realtime_upstream_handshake_refusal_sends_error_event_then_policy_close(): + from typing import cast + + from websockets.datastructures import Headers + from websockets.exceptions import InvalidStatus + from websockets.http11 import Response + + from litellm.llms.openai.realtime.handler import OpenAIRealtime + from litellm.types.realtime import RealtimeErrorEvent + + handler = OpenAIRealtime() + model = "gpt-realtime" + + sent: list[str] = [] + closed: list[tuple[int, str | None]] = [] + + class RecordingClientWebSocket: + scope: dict[str, list[tuple[bytes, bytes]]] = {"headers": []} + + async def send_text(self, data: str) -> None: + sent.append(data) + + async def close(self, code: int = 1000, reason: str | None = None) -> None: + closed.append((code, reason)) + + dummy_websocket = RecordingClientWebSocket() + dummy_logging_obj = MagicMock() + + refused = InvalidStatus(Response(401, "Unauthorized", Headers())) + + with patch("websockets.connect", side_effect=refused): + await handler.async_realtime( # pyright: ignore[reportUnknownMemberType] # handler's websocket param is Any + model=model, + websocket=dummy_websocket, + logging_obj=dummy_logging_obj, + api_base="https://api.openai.com/", + api_key="bad-key", + query_params={"model": model}, + ) + + assert len(sent) == 1 + event = cast(RealtimeErrorEvent, json.loads(sent[0])) + assert event["type"] == "error" + assert event["error"]["type"] == "server_error" + assert "401" in event["error"]["message"] + assert closed and closed[0][0] == 1008 diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index 07d1ec5f523..a5128a87b0d 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -2,7 +2,7 @@ Tests for _redact_string usage in error/logging paths. Covers actual execution of redaction in: -- WebSocket close reasons in realtime handlers (openai, azure, bedrock) +- WebSocket close reasons in realtime handlers (openai, bedrock) - Gemini RAG ingestion x-goog-api-key header usage - Traceback redaction pattern used in proxy streaming - Router fallback-failure traceback redaction @@ -72,25 +72,6 @@ class TestOpenAIRealtimeRedaction: api_key="test-key", ) - @pytest.mark.asyncio - async def test_invalid_status_code_redacts_reason(self): - import websockets.exceptions - - from litellm.llms.openai.realtime.handler import OpenAIRealtime - - handler = OpenAIRealtime() - exc = websockets.exceptions.InvalidStatusCode(403, None) - exc.status_code = 403 - - kwargs = self._call_kwargs() - mock_ws = kwargs["websocket"] - p1, p2, p3 = self._make_patches(handler) - with p1, p2, p3, patch("websockets.connect", side_effect=exc): - await handler.async_realtime(**kwargs) - - mock_ws.close.assert_called_once() - assert mock_ws.close.call_args[1]["code"] == 403 - @pytest.mark.asyncio async def test_generic_exception_redacts_reason(self): from litellm.llms.openai.realtime.handler import OpenAIRealtime @@ -111,41 +92,6 @@ class TestOpenAIRealtimeRedaction: assert "sk-1234567890abcdefghij" not in mock_ws.close.call_args[1]["reason"] -class TestAzureRealtimeRedaction: - """Test that Azure realtime handler redacts secrets in websocket close reasons.""" - - @pytest.mark.asyncio - async def test_invalid_status_code_redacts_reason(self): - import websockets.exceptions - - from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime - - handler = AzureOpenAIRealtime() - mock_ws = AsyncMock() - exc = websockets.exceptions.InvalidStatusCode(403, None) - exc.status_code = 403 - - with ( - patch.object( - handler, - "_construct_url", - return_value="wss://test.openai.azure.com/openai/realtime", - ), - patch("websockets.connect", side_effect=exc), - ): - await handler.async_realtime( - model="gpt-4", - websocket=mock_ws, - logging_obj=MagicMock(), - api_base="https://test.openai.azure.com/", - api_key="test-key", - api_version="2024-10-01-preview", - ) - - mock_ws.close.assert_called_once() - assert mock_ws.close.call_args[1]["code"] == 403 - - class TestBedrockRealtimeRedaction: """Test that _redact_string produces safe close reasons for Bedrock-style errors.""" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3336ad6d33a..da9f9dbeda6 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -5792,3 +5792,67 @@ def test_calculate_max_parallel_requests_precedence( ) == expected ) + + +@pytest.mark.asyncio +async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None: + """Chat over the Responses bridge exits two @client wrappers with one logging object. Issue + #44500: both exits enqueued a success handler, and on a large prompt the first one yielded to + the worker-thread base64 offload before it marked ``has_logged_async_success``, so the second + passed the check too and the request was logged and billed twice. The schedule step claims the + log for the object synchronously, so only the inner provider-shaped result is ever logged.""" + from litellm.litellm_core_utils import litellm_logging + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.utils import _dispatch_success_logging + + class CountingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_results: list[object] = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + if kwargs["litellm_call_id"] == "bridge-call-id": + self.logged_results.append(response_obj) + + counting_logger: Final = CountingLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger]) + real_truncate: Final = litellm_logging.truncate_base64_in_messages_async + + async def yielding_truncate(messages): + await asyncio.sleep(0) + return await real_truncate(messages) + + monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", yielding_truncate) + + messages: Final = [{"role": "user", "content": "hello"}] + logging_obj: Final = Logging( + model="gpt-5.6-luna", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp(), + litellm_call_id="bridge-call-id", + function_id="bridge-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="gpt-5.6-luna", custom_llm_provider="openai", input=messages + ) + inner_result: Final = ModelResponse(id="inner") + outer_result: Final = ModelResponse(id="outer") + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + for result in (inner_result, outer_result): + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=now, + end_time=now, + is_completion_with_fallbacks=False, + is_litellm_internal_call=False, + ) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + + assert len(counting_logger.logged_results) == 1, counting_logger.logged_results + assert counting_logger.logged_results[0] is inner_result diff --git a/uv.lock b/uv.lock index 02994dfe53c..e2d5225415d 100644 --- a/uv.lock +++ b/uv.lock @@ -4509,7 +4509,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.103.5" +version = "1.103.6" source = { editable = "." } dependencies = [ { name = "aiohttp" },